Lumi/companion/src/Lumi.Companion.App/UpdateService.cs

using System.Diagnostics;
using System.IO.Compression;
using System.Net.Http.Headers;
using System.Security.Cryptography;
using System.Text.Json;
using System.Text.Json.Serialization;
using Lumi.Companion.Protocol;

namespace Lumi.Companion.App;

public sealed class UpdateService(HttpClient http, CompanionPaths paths)
{
    public async Task<CompanionUpdate?> CheckAsync(DeviceCredential credential, string currentVersion, CancellationToken cancellationToken = default)
    {
        var host = new Uri(credential.Host);
        var endpoint = new Uri(host, $"/plugins/lumi_transcription/api/companion/update?current_version={Uri.EscapeDataString(currentVersion)}");
        using var request = new HttpRequestMessage(HttpMethod.Get, endpoint);
        request.Headers.Authorization = new AuthenticationHeaderValue("LumiDevice", $"{credential.DeviceId}.{credential.DeviceSecret}");
        using var response = await http.SendAsync(request, HttpCompletionOption.ResponseHeadersRead, cancellationToken);
        response.EnsureSuccessStatusCode();
        var update = await JsonSerializer.DeserializeAsync<CompanionUpdate>(await response.Content.ReadAsStreamAsync(cancellationToken), ProtocolV1.JsonOptions, cancellationToken);
        return update is { Ok: true, UpdateAvailable: true } ? update : null;
    }

    public async Task<StagedUpdate> StageAsync(CompanionUpdate update, CancellationToken cancellationToken = default)
    {
        if (update.Artifact.Bytes is <= 0 or > 300 * 1024 * 1024 || !IsSha256(update.Artifact.Sha256)) throw new InvalidDataException("The update manifest is invalid.");
        if (!Uri.TryCreate(update.Artifact.Url, UriKind.Absolute, out var artifactUri) || artifactUri.Scheme != Uri.UriSchemeHttps) throw new InvalidDataException("The update manifest does not use a secure artifact URL.");
        var safeVersion = SafeVersion(update.Version);
        if (string.IsNullOrWhiteSpace(safeVersion)) throw new InvalidDataException("The update manifest has an invalid version.");
        Directory.CreateDirectory(paths.UpdatesDirectory);
        var archivePath = Path.Combine(paths.UpdatesDirectory, $"companion-{safeVersion}.zip.partial");
        var stageRoot = Path.Combine(paths.UpdatesDirectory, safeVersion);
        try
        {
            using var response = await http.GetAsync(update.Artifact.Url, HttpCompletionOption.ResponseHeadersRead, cancellationToken);
            response.EnsureSuccessStatusCode();
            await using (var source = await response.Content.ReadAsStreamAsync(cancellationToken))
            await using (var target = new FileStream(archivePath, FileMode.Create, FileAccess.Write, FileShare.None, 81920, true))
            {
                var buffer = new byte[81920];
                long total = 0;
                while (true)
                {
                    var read = await source.ReadAsync(buffer, cancellationToken);
                    if (read == 0) break;
                    total = checked(total + read);
                    if (total > update.Artifact.Bytes) throw new InvalidDataException("The Companion update download exceeded its declared size.");
                    await target.WriteAsync(buffer.AsMemory(0, read), cancellationToken);
                }
            }
            if (new FileInfo(archivePath).Length != update.Artifact.Bytes || !HashFile(archivePath).Equals(update.Artifact.Sha256, StringComparison.OrdinalIgnoreCase))
                throw new InvalidDataException("The downloaded Companion update did not match its trusted checksum.");
            if (Directory.Exists(stageRoot)) Directory.Delete(stageRoot, true);
            Directory.CreateDirectory(stageRoot);
            using var archive = ZipFile.OpenRead(archivePath);
            if (archive.Entries.Count is 0 or > 256) throw new InvalidDataException("The Companion update contains an invalid number of files.");
            long extractedBytes = 0;
            foreach (var entry in archive.Entries)
            {
                var relative = entry.FullName.Replace('\\', '/').TrimStart('/');
                if (string.IsNullOrEmpty(relative) || relative.EndsWith('/')) continue;
                if (relative.Split('/').Any(segment => segment is "" or "." or "..")) throw new InvalidDataException("The Companion update contains an unsafe path.");
                extractedBytes = checked(extractedBytes + entry.Length);
                if (extractedBytes > 500 * 1024 * 1024) throw new InvalidDataException("The Companion update expands beyond its safety limit.");
                var destination = Path.GetFullPath(Path.Combine(stageRoot, relative.Replace('/', Path.DirectorySeparatorChar)));
                if (!destination.StartsWith(stageRoot + Path.DirectorySeparatorChar, StringComparison.OrdinalIgnoreCase)) throw new InvalidDataException("The Companion update contains an unsafe path.");
                Directory.CreateDirectory(Path.GetDirectoryName(destination)!);
                entry.ExtractToFile(destination, true);
            }
            var staged = Path.GetFullPath(Path.Combine(stageRoot, update.Artifact.Entrypoint.Replace('/', Path.DirectorySeparatorChar)));
            if (!File.Exists(staged) || !staged.StartsWith(stageRoot + Path.DirectorySeparatorChar, StringComparison.OrdinalIgnoreCase))
                throw new InvalidDataException("The Companion update is missing its application entrypoint.");
            return new StagedUpdate(stageRoot, staged, HashFile(staged));
        }
        finally { try { File.Delete(archivePath); } catch { } }
    }

    public void LaunchApplier(StagedUpdate update)
    {
        if (!OperatingSystem.IsWindows() || string.IsNullOrWhiteSpace(Environment.ProcessPath)) throw new PlatformNotSupportedException("In-place updates require the packaged Windows application.");
        var target = Path.GetFullPath(Environment.ProcessPath);
        var helper = Path.Combine(paths.UpdatesDirectory, $"Lumi.Companion.UpdateHelper.{Environment.ProcessId}.exe");
        File.Copy(target, helper, true);
        var start = new ProcessStartInfo(helper) { UseShellExecute = false, CreateNoWindow = true };
        foreach (var argument in new[] { "--apply-update", Environment.ProcessId.ToString(), update.StageRoot, update.ExecutablePath, target, update.Sha256, helper }) start.ArgumentList.Add(argument);
        Process.Start(start)?.Dispose();
    }

    internal static string HashFile(string path) { using var stream = File.OpenRead(path); return Convert.ToHexString(SHA256.HashData(stream)).ToLowerInvariant(); }
    private static bool IsSha256(string value) => value.Length == 64 && value.All(Uri.IsHexDigit);
    private static string SafeVersion(string value) => string.Concat(value.Where(character => char.IsLetterOrDigit(character) || character is '.' or '-')).Trim('.');
}

public sealed record CompanionUpdate(
    [property: JsonPropertyName("ok")] bool Ok,
    [property: JsonPropertyName("version")] string Version,
    [property: JsonPropertyName("update_available")] bool UpdateAvailable,
    [property: JsonPropertyName("artifact")] CompanionUpdateArtifact Artifact,
    [property: JsonPropertyName("signed")] bool Signed,
    [property: JsonPropertyName("release_notes")] string ReleaseNotes);
public sealed record CompanionUpdateArtifact(
    [property: JsonPropertyName("url")] string Url,
    [property: JsonPropertyName("sha256")] string Sha256,
    [property: JsonPropertyName("bytes")] long Bytes,
    [property: JsonPropertyName("entrypoint")] string Entrypoint);
public sealed record StagedUpdate(string StageRoot, string ExecutablePath, string Sha256);