wx-win-agent/node-agent/WxAgent.Core/RemoteControlClient.cs

using System.Net.Http.Headers;
using System.Net.Security;
using System.Security.Cryptography.X509Certificates;
using System.Text;
using System.Text.Json;

namespace WxAgent.Core;

public sealed class RemoteClientException(string code, int statusCode, string message, string correlationId)
    : Exception(message)
{
    public string Code { get; } = code;
    public int StatusCode { get; } = statusCode;
    public string CorrelationId { get; } = correlationId;
}

public sealed class RemoteControlClient : IDisposable
{
    private readonly HttpClient _http;
    private readonly bool _ownsHttp;
    private readonly RemoteAgentOptions _options;
    private readonly Uri? _baseAddress;
    private readonly X509Certificate2? _clientCertificate;
    private readonly X509Certificate2? _serverCaCertificate;
    private bool _authenticated;

    public RemoteControlClient(RemoteAgentOptions options, HttpClient? httpClient = null)
    {
        _options = options ?? throw new ArgumentNullException(nameof(options));
        _ownsHttp = httpClient is null;
        if (httpClient is null)
        {
            _clientCertificate = LoadClientCertificate(options);
            _serverCaCertificate = LoadServerCaCertificate(options);
            _http = CreateHttpClient(_clientCertificate, _serverCaCertificate);
        }
        else
        {
            _http = httpClient;
        }
        if (Uri.TryCreate(options.AuthAddress, UriKind.Absolute, out var address))
            _baseAddress = new Uri($"{address.Scheme}://{address.Authority}/", UriKind.Absolute);
        AuthState = options.IsConfigured ? RemoteAuthState.NotConfigured : RemoteAuthState.NotConfigured;
    }

    public RemoteAuthState AuthState { get; private set; }
    public string? LastErrorCode { get; private set; }
    public DateTimeOffset? LastSuccessAt { get; private set; }

    public async Task<RemoteNodeRegistrationResponse> RegisterAsync(
        RemoteNodeRegistration registration,
        CancellationToken cancellationToken = default)
    {
        EnsureConfigured();
        if (!string.Equals(registration.NodeId, _options.NodeId, StringComparison.Ordinal))
            throw new WxAgentException(WxAgentErrorCode.InvalidArgument, "Registration nodeId does not match the configured nodeId.");
        AuthState = RemoteAuthState.Authenticating;
        try
        {
            var response = await SendAsync<RemoteNodeRegistrationResponse>(HttpMethod.Post, "/v1/nodes/register", registration, false, cancellationToken);
            _authenticated = response.Authenticated;
            AuthState = response.Authenticated ? RemoteAuthState.Authenticated : RemoteAuthState.AuthenticationFailed;
            LastSuccessAt = DateTimeOffset.UtcNow;
            LastErrorCode = null;
            return response;
        }
        catch (Exception exception) when (exception is RemoteClientException or HttpRequestException or TaskCanceledException)
        {
            _authenticated = false;
            AuthState = RemoteAuthState.AuthenticationFailed;
            LastErrorCode = exception is RemoteClientException remote ? remote.Code : "ConnectionFailed";
            throw;
        }
    }

    public Task<RemoteHeartbeatResponse> HeartbeatAsync(RemoteHeartbeat heartbeat, CancellationToken cancellationToken = default) =>
        SendAuthenticatedAsync<RemoteHeartbeatResponse>(HttpMethod.Post,
            $"/v1/nodes/{Escape(heartbeat.NodeId)}/heartbeat", heartbeat, cancellationToken);

    public async Task<IReadOnlyList<RemoteTaskEnvelope>> PollTasksAsync(
        string accountId,
        CancellationToken cancellationToken = default,
        int waitSeconds = 0)
    {
        RemoteAgentOptions.ValidateIdentifier(accountId, "accountId", 200);
        if (waitSeconds is < 0 or > 30)
            throw new WxAgentException(WxAgentErrorCode.InvalidArgument, "waitSeconds must be between 0 and 30.");
        var query = $"/v1/nodes/{Escape(_options.NodeId!)}/tasks?account_id={Uri.EscapeDataString(accountId)}";
        if (waitSeconds > 0) query += $"&wait_seconds={waitSeconds}";
        var batch = await SendAuthenticatedAsync<RemoteTaskBatch>(HttpMethod.Get, query, null, cancellationToken);
        return batch.Tasks ?? [];
    }

    public Task<RemoteTaskEnvelope> AcknowledgeTaskAsync(RemoteTaskEnvelope task, CancellationToken cancellationToken = default) =>
        SendAuthenticatedAsync<RemoteTaskEnvelope>(HttpMethod.Post,
            $"/v1/nodes/{Escape(task.NodeId)}/tasks/{Escape(task.TaskId)}/ack", new
            {
                task_id = task.TaskId,
                account_id = task.AccountId,
                lease_generation = task.LeaseGeneration
            }, cancellationToken);

    public Task<RemoteTaskEnvelope> StartTaskAsync(RemoteTaskEnvelope task, CancellationToken cancellationToken = default) =>
        SendAuthenticatedAsync<RemoteTaskEnvelope>(HttpMethod.Post,
            $"/v1/nodes/{Escape(task.NodeId)}/tasks/{Escape(task.TaskId)}/start", new
            {
                task_id = task.TaskId,
                account_id = task.AccountId,
                lease_generation = task.LeaseGeneration
            }, cancellationToken);

    public Task<RemoteTaskEnvelope> RenewTaskAsync(RemoteTaskEnvelope task, CancellationToken cancellationToken = default) =>
        SendAuthenticatedAsync<RemoteTaskEnvelope>(HttpMethod.Post,
            $"/v1/nodes/{Escape(task.NodeId)}/tasks/{Escape(task.TaskId)}/renew", new
            {
                task_id = task.TaskId,
                account_id = task.AccountId,
                lease_generation = task.LeaseGeneration
            }, cancellationToken);

    public Task<RemoteTaskEnvelope> SendTaskResultAsync(
        RemoteTaskResult result,
        ReportingConfig? reportingConfig = null,
        string? chatId = null,
        ReportingChatType? chatType = null,
        IReadOnlyList<RemoteReportingScope>? chatScopes = null,
        CancellationToken cancellationToken = default)
    {
        var filtered = chatScopes is null
            ? ReportingAuthorization.FilterTaskResult(reportingConfig, result, chatId, chatType, out _)
            : ReportingAuthorization.FilterTaskResultForChats(reportingConfig, result, chatScopes, out _);
        return SendAuthenticatedAsync<RemoteTaskEnvelope>(HttpMethod.Post,
            $"/v1/nodes/{Escape(_options.NodeId!)}/tasks/{Escape(result.TaskId)}/result", filtered, cancellationToken,
            RemoteProtocol.MaxTaskResultBytes);
    }

    public async Task<RemoteEventReceipt> SubmitEventAsync(
        ReportingConfig reportingConfig,
        RemoteMessageEvent messageEvent,
        CancellationToken cancellationToken = default)
    {
        var filtered = ReportingAuthorization.FilterEvent(reportingConfig, messageEvent, out var decision);
        if (filtered is null)
            return new RemoteEventReceipt(false, false, null, decision.Reason);
        return await SendAuthenticatedAsync<RemoteEventReceipt>(HttpMethod.Post,
            $"/v1/nodes/{Escape(messageEvent.NodeId)}/events", filtered, cancellationToken);
    }

    public async Task<int> FlushEventsAsync(
        RemoteEventQueue queue,
        ReportingConfig reportingConfig,
        CancellationToken cancellationToken = default)
    {
        EnsureAuthenticated();
        var sent = 0;
        foreach (var messageEvent in queue.PrepareForSend(reportingConfig))
        {
            var receipt = await SubmitEventAsync(reportingConfig, messageEvent, cancellationToken);
            if (!receipt.Accepted)
                continue;
            if (queue.MarkSent(messageEvent)) sent++;
        }
        return sent;
    }

    public Task<RemoteTaskEnvelope> AcknowledgeCancellationAsync(
        RemoteTaskEnvelope task,
        RemoteTaskStatus status,
        string? errorCode,
        string? message,
        CancellationToken cancellationToken = default) =>
        SendTaskResultAsync(new RemoteTaskResult(task.TaskId, task.AccountId, task.LeaseGeneration,
            status, errorCode, message, false, null, NewCorrelationId()), cancellationToken: cancellationToken);

    private Task<T> SendAuthenticatedAsync<T>(HttpMethod method, string path, object? body, CancellationToken cancellationToken,
        int maxBodyBytes = RemoteProtocol.MaxTaskPayloadBytes)
    {
        EnsureAuthenticated();
        return SendAsync<T>(method, path, body, true, cancellationToken, maxBodyBytes);
    }

    private async Task<T> SendAsync<T>(HttpMethod method, string path, object? body, bool authenticated,
        CancellationToken cancellationToken, int maxBodyBytes = RemoteProtocol.MaxTaskPayloadBytes)
    {
        if (_baseAddress is null)
            throw new WxAgentException(WxAgentErrorCode.InvalidArgument, "authAddress is not a valid absolute URL.");
        using var request = new HttpRequestMessage(method, new Uri(_baseAddress, path.TrimStart('/')));
        request.Headers.TryAddWithoutValidation("X-Correlation-Id", NewCorrelationId());
        if (authenticated || method == HttpMethod.Post && path == "/v1/nodes/register")
            request.Headers.Authorization = new AuthenticationHeaderValue("Bearer", _options.GetToken());
        if (body is not null)
        {
            var bytes = JsonSerializer.SerializeToUtf8Bytes(body, RemoteJson.Options);
            if (bytes.Length > maxBodyBytes)
                throw new WxAgentException(WxAgentErrorCode.InvalidArgument, "Remote request payload is too large.");
            request.Content = new ByteArrayContent(bytes);
            request.Content.Headers.ContentType = new MediaTypeHeaderValue("application/json") { CharSet = "utf-8" };
        }

        using var response = await _http.SendAsync(request, HttpCompletionOption.ResponseHeadersRead, cancellationToken);
        var responseBytes = await response.Content.ReadAsByteArrayAsync(cancellationToken);
        var correlationId = response.Headers.TryGetValues("X-Correlation-Id", out var values)
            ? values.FirstOrDefault() ?? request.Headers.GetValues("X-Correlation-Id").First()
            : request.Headers.GetValues("X-Correlation-Id").First();
        if (!response.IsSuccessStatusCode)
        {
            var error = TryDeserializeError(responseBytes, correlationId);
            if (response.StatusCode == System.Net.HttpStatusCode.Unauthorized)
            {
                _authenticated = false;
                AuthState = RemoteAuthState.AuthenticationFailed;
            }
            LastErrorCode = error.Code;
            throw new RemoteClientException(error.Code, (int)response.StatusCode, error.Message, error.CorrelationId);
        }

        try
        {
            var result = JsonSerializer.Deserialize<T>(responseBytes, RemoteJson.Options);
            if (result is null) throw new JsonException("The control plane returned an empty response.");
            LastSuccessAt = DateTimeOffset.UtcNow;
            return result;
        }
        catch (JsonException exception)
        {
            throw new RemoteClientException("InvalidResponse", (int)response.StatusCode,
                "The control plane returned an invalid response.", correlationId) { Source = exception.Source };
        }
    }

    private void EnsureConfigured()
    {
        if (!_options.IsConfigured)
        {
            AuthState = RemoteAuthState.NotConfigured;
            throw new WxAgentException(WxAgentErrorCode.InvalidArgument, "Remote authAddress, token or tokenFile and nodeId must all be configured.");
        }
        _options.Validate();
    }

    private void EnsureAuthenticated()
    {
        EnsureConfigured();
        if (!_authenticated || AuthState != RemoteAuthState.Authenticated)
            throw new WxAgentException(WxAgentErrorCode.PermissionMismatch, "The node must authenticate before remote operations.");
    }

    private static RemoteApiError TryDeserializeError(byte[] bytes, string correlationId)
    {
        try
        {
            using var document = JsonDocument.Parse(bytes);
            var root = document.RootElement;
            if (root.TryGetProperty("error", out var error))
            {
                var code = error.TryGetProperty("code", out var codeValue) ? codeValue.GetString() : null;
                var message = error.TryGetProperty("message", out var messageValue) ? messageValue.GetString() : null;
                return new RemoteApiError(code ?? "RemoteRequestFailed", message ?? "Remote request failed.", correlationId);
            }
        }
        catch (JsonException) { }
        return new RemoteApiError("RemoteRequestFailed", "Remote request failed.", correlationId);
    }

    private static X509Certificate2? LoadClientCertificate(RemoteAgentOptions options)
    {
        if (string.IsNullOrWhiteSpace(options.ClientCertificateFile)) return null;
        return X509Certificate2.CreateFromPemFile(options.ClientCertificateFile, options.ClientCertificateKeyFile!);
    }

    private static X509Certificate2? LoadServerCaCertificate(RemoteAgentOptions options) =>
        string.IsNullOrWhiteSpace(options.ServerCaFile) ? null : new X509Certificate2(options.ServerCaFile);

    private static HttpClient CreateHttpClient(X509Certificate2? clientCertificate, X509Certificate2? serverCaCertificate)
    {
        var handler = new HttpClientHandler();
        if (clientCertificate is not null) handler.ClientCertificates.Add(clientCertificate);
        if (serverCaCertificate is not null)
        {
            handler.ServerCertificateCustomValidationCallback = (_, certificate, _, errors) =>
            {
                if (certificate is null || (errors & (SslPolicyErrors.RemoteCertificateNameMismatch | SslPolicyErrors.RemoteCertificateNotAvailable)) != 0)
                    return false;
                using var chain = new X509Chain();
                chain.ChainPolicy.TrustMode = X509ChainTrustMode.CustomRootTrust;
                chain.ChainPolicy.CustomTrustStore.Add(serverCaCertificate);
                chain.ChainPolicy.RevocationMode = X509RevocationMode.NoCheck;
                return chain.Build(new X509Certificate2(certificate));
            };
        }
        return new HttpClient(handler) { Timeout = TimeSpan.FromSeconds(15) };
    }

    private static string Escape(string value) => Uri.EscapeDataString(value);
    private static string NewCorrelationId() => Guid.NewGuid().ToString("N");

    public void Dispose()
    {
        if (_ownsHttp) _http.Dispose();
        _clientCertificate?.Dispose();
        _serverCaCertificate?.Dispose();
    }
}