duplicati/Duplicati/Library/RemoteControl/KeepRemoteConnection.cs

// Copyright (C) 2026, The Duplicati Team
// https://duplicati.com, hello@duplicati.com
// 
// Permission is hereby granted, free of charge, to any person obtaining a 
// copy of this software and associated documentation files (the "Software"), 
// to deal in the Software without restriction, including without limitation 
// the rights to use, copy, modify, merge, publish, distribute, sublicense, 
// and/or sell copies of the Software, and to permit persons to whom the 
// Software is furnished to do so, subject to the following conditions:
// 
// The above copyright notice and this permission notice shall be included in 
// all copies or substantial portions of the Software.
// 
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS 
// OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING 
// FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER 
// DEALINGS IN THE SOFTWARE.

using System.Diagnostics;
using System.Net;
using System.Net.Http.Headers;
using System.Security.Cryptography;
using System.Text.Json;
using CoCoL;
using Duplicati.Library.AutoUpdater;
using Duplicati.Library.Logging;
using Duplicati.Library.Utility;
using Uri = System.Uri;

namespace Duplicati.Library.RemoteControl;

/// <summary>
/// Support class for keeping a connection to a remote server
/// </summary>
public class KeepRemoteConnection : IDisposable
{
    /// <summary>
    /// The protocol version to use
    /// </summary>
    private const int PROTOCOL_VERSION = 1;
    /// <summary>
    /// The log tag for messages from this class
    /// </summary>
    private static readonly string LogTag = Log.LogTagFromType<KeepRemoteConnection>();

    /// <summary>
    /// The interval between reconnect attempts
    /// </summary>
    private static readonly TimeSpan ReconnectInterval = TimeSpan.FromSeconds(30);

    /// <summary>
    /// The interval between heartbeats
    /// </summary>
    private static readonly TimeSpan HeartbeatInterval = TimeSpan.FromSeconds(60);

    /// <summary>
    /// The time between reconnect attempts if no response is received
    /// </summary>
    private static readonly TimeSpan NoResponseTimeout = HeartbeatInterval * 2;

    /// <summary>
    /// The minimum time between certificate refreshes
    /// </summary>
    private static readonly TimeSpan MinimumCertificateRefreshInterval = TimeSpan.FromMinutes(5);

    /// <summary>
    /// The interval between certificate refreshes
    /// </summary>
    private static readonly TimeSpan CertificateRefreshInterval = TimeSpan.FromDays(7);

    /// <summary>
    /// The client key to use for signing messages
    /// </summary>
    private static readonly RSA ClientKey = RSA.Create(2048);

    /// <summary>
    /// The client ID to use for identifying the client
    /// </summary>
    private static readonly string ClientId = Guid.NewGuid().ToString();

    /// <summary>
    /// The JSON options to use for deserialization
    /// </summary>
    internal static readonly JsonSerializerOptions JsonOptions = new JsonSerializerOptions
    {
        PropertyNamingPolicy = JsonNamingPolicy.CamelCase,
        PropertyNameCaseInsensitive = true
    };

    /// <summary>
    /// The stats the connection can be in
    /// </summary>
    public enum ConnectionState
    {
        /// <summary>
        /// The connection is not established
        /// </summary>
        NotConnected,
        /// <summary>
        /// We received a welcome message
        /// </summary>
        WelcomeReceived,
        /// <summary>
        /// The connection is authenticated
        /// </summary>
        Authenticated,
        /// <summary>
        /// The connection is in an error state
        /// </summary>
        Error
    }

    /// <summary>
    /// The websocket client
    /// </summary>
    private readonly Websocket.Client.WebsocketClient _client;
    /// <summary>
    /// The cancellation token source
    /// </summary>
    private readonly CancellationTokenSource _cancellationTokenSource;
    /// <summary>
    /// The current state of the connection
    /// </summary>
    private ConnectionState _state = ConnectionState.NotConnected;

    /// <summary>
    /// Event raised when the connection state changes
    /// </summary>
    public event EventHandler<ConnectionState>? StateChanged;

    /// <summary>
    /// Sets the connection state and raises the StateChanged event if it changed
    /// </summary>
    /// <param name="newState">The new state</param>
    private void SetState(ConnectionState newState)
    {
        if (_state == newState)
            return;

        _state = newState;
        StateChanged?.Invoke(this, newState);
    }
    /// <summary>
    /// The task that runs the connection
    /// </summary>
    private Task _runnerTask;
    /// <summary>
    /// The currently negotiated server certificate
    /// </summary>
    private MiniServerCertificate? _serverCertificate;
    /// <summary>
    /// The public key of the server
    /// </summary>
    private RSA? _serverPublicKey;

    /// <summary>
    /// The callback to call when connecting
    /// </summary>
    private readonly Func<Dictionary<string, string?>, Task<Dictionary<string, string?>>> _onConnect;
    /// <summary>
    /// The callback to call when rekeying
    /// </summary>
    private readonly Func<ClaimedClientData, Task> _onReKey;
    /// <summary>
    /// The callback to call when a control message is received
    /// </summary>
    private readonly Func<ControlMessage, Task> _onControl;
    /// <summary>
    /// The callback to call when a command message is received
    /// </summary>
    private readonly Func<CommandMessage, Task> _onMessage;

    /// <summary>
    /// The current JWT token
    /// </summary>
    private string _token;
    /// <summary>
    /// The server URL
    /// </summary>
    private string _serverUrl;
    /// <summary>
    /// The certificate URL
    /// </summary>
    private string _certificateUrl;
    /// <summary>
    /// The server keys
    /// </summary>
    private IEnumerable<MiniServerCertificate> _serverKeys;
    /// <summary>
    /// The last time a message was received
    /// </summary>
    private DateTimeOffset _lastMessageReceived = DateTimeOffset.MinValue;
    /// <summary>
    /// The time to refresh the settings by, or null if keeping a persistent connection
    /// </summary>
    private DateTimeOffset _refreshSettingsBy = DateTimeOffset.MinValue;

    /// <summary>
    /// Creates a new connection to the remote server
    /// </summary>
    /// <param name="serverUrl">The url to use</param>
    /// <param name="JWT">The JWT token to use</param>
    /// <param name="certificateUrl">The certificate url to use</param>
    /// <param name="serverKeys">The server keys to use</param>
    /// <param name="refreshSettingsBy">The time to refresh the settings by, or null if keeping a persistent connection</param>
    /// <param name="forceConnect">If the connection should be force enabled, ignoring re-connect delays</param>
    /// <param name="cancellationToken">The token to cancel the connection</param>
    /// <param name="onConnect">The callback to call when connecting</param>
    /// <param name="onReKey">The callback to call when rekeying</param>
    /// <param name="onControl">The callback to call when a control message is received</param>
    /// <param name="onMessage">The callback to call when a command message is received</param>
    private KeepRemoteConnection(
        string serverUrl,
        string JWT,
        string certificateUrl,
        IEnumerable<MiniServerCertificate> serverKeys,
        DateTimeOffset? refreshSettingsBy,
        bool forceConnect,
        CancellationToken cancellationToken,
        Func<Dictionary<string, string?>, Task<Dictionary<string, string?>>> onConnect,
        Func<ClaimedClientData, Task> onReKey,
        Func<ControlMessage, Task> onControl,
        Func<CommandMessage, Task> onMessage)
    {
        _serverUrl = serverUrl;
        _certificateUrl = certificateUrl;
        _token = JWT;
        _serverKeys = serverKeys;
        _cancellationTokenSource = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
        _onConnect = onConnect;
        _onReKey = onReKey;
        _onControl = onControl;
        _onMessage = onMessage;
        _refreshSettingsBy = refreshSettingsBy ?? DateTimeOffset.MinValue;

        _client = new Websocket.Client.WebsocketClient(new Uri(serverUrl));
        _runnerTask = RunMainLoopAsync(forceConnect);
    }

    /// <summary>
    /// Runs the inner loop of the connection
    /// </summary>
    /// <param name="forceConnect">If the connection should be force enabled, ignoring re-connect delays</param>
    private async Task RunMainLoopAsync(bool forceConnect)
    {
        _client.ReconnectTimeout = NoResponseTimeout;
        _client.LostReconnectTimeout = ReconnectInterval;
        _client.IsReconnectionEnabled = true;

        // Set up the periodic refreshers
        using var reconnectHelper = new PeriodicRefresher(
            System.Threading.Timeout.InfiniteTimeSpan,
            ReconnectInterval,
            async token =>
            {
                await _client.Start();
            },
            _cancellationTokenSource.Token);

        using var heartbeatHelper = new PeriodicRefresher(
            HeartbeatInterval,
            TimeSpan.FromSeconds(1),
            _ =>
            {
                // Reconnect if we have disconnected (but not if auto-reconnect is disabled)
                if (!_client.IsRunning && !IsAutoReconnectDisabled())
                {
                    reconnectHelper.Signal();
                }
                // If we do not get any response from the server, we should reconnect
                else if ((_state == ConnectionState.Authenticated || _state == ConnectionState.WelcomeReceived) && _lastMessageReceived.Add(NoResponseTimeout) < DateTimeOffset.Now)
                {
                    SetState(ConnectionState.Error);
                    Log.WriteMessage(LogMessageType.Warning, LogTag, "WebsocketDisconnect", "No response from server");
                    _client.Stop(System.Net.WebSockets.WebSocketCloseStatus.NormalClosure, "No response").FireAndForget();
                }

                SendEnvelope(new EnvelopedMessage()
                {
                    From = ClientId,
                    To = "server",
                    Type = "ping",
                    ErrorMessage = null,
                    Payload = null,
                    MessageId = Guid.NewGuid().ToString()
                });

                return Task.CompletedTask;
            },
            _cancellationTokenSource.Token);

        using var certificateRefreshHelper = new PeriodicRefresher(
            CertificateRefreshInterval,
            MinimumCertificateRefreshInterval,
            RefreshCertificatesAsync,
            _cancellationTokenSource.Token);

        _client.DisconnectionHappened.Subscribe(info =>
        {
            SetState(ConnectionState.NotConnected);
            _serverCertificate = null;
            _serverPublicKey = null;
            if (!IsAutoReconnectDisabled())
            {
                reconnectHelper.Signal();
                Log.WriteMessage(LogMessageType.Warning, LogTag, "WebsocketDisconnect", "Disconnected from the server");
            }
            else
            {
                Log.WriteMessage(LogMessageType.Information, LogTag, "WebsocketDisconnect", "Disconnected from the server. Auto-reconnect disabled until {0}", _refreshSettingsBy.ToLocalTime());
            }
        });

        _client.MessageReceived.Subscribe(msg => OnMessageAsync(msg, certificateRefreshHelper, reconnectHelper).Await());

        // Start the connection
        if (forceConnect || !IsAutoReconnectDisabled())
            reconnectHelper.Signal();

        var t = await Task.WhenAny(
            heartbeatHelper.RunLoopAsync(),
            reconnectHelper.RunLoopAsync(),
            certificateRefreshHelper.RunLoopAsync()
        );

        await _cancellationTokenSource.CancelAsync();

        // Re-throw any exceptions
        await t;
    }

    private async Task OnMessageAsync(Websocket.Client.ResponseMessage msg, PeriodicRefresher certificateRefreshHelper, PeriodicRefresher reconnectHelper)
    {
        // Ignore messages if we are in an error state
        if (_state == ConnectionState.Error)
            return;

        _lastMessageReceived = DateTimeOffset.Now;
        if (_state == ConnectionState.NotConnected)
            Log.WriteMessage(LogMessageType.Verbose, LogTag, "WebsocketMessage", "Received message from server: {0}", msg);
        else // Encrypted messages are not logged, as the content has no meaning before being decrypted
            Log.WriteMessage(LogMessageType.Verbose, LogTag, "WebsocketMessage", "Received encrypted message from server");

        try
        {
            if (string.IsNullOrWhiteSpace(msg.Text))
                throw new ProtocolViolationException("Empty message");

            if (_serverCertificate == null || _serverPublicKey == null || _state == ConnectionState.NotConnected)
            {
                // Should be safe from replay, as the response is encrypted with the server public key
                // So even a replay attack would not let the attacker know the client's token
                var welcomeEnvelope = EnvelopedMessage.ForceParse(msg.Text);
                if (welcomeEnvelope.GetMessageType() != MessageType.Welcome)
                    throw new ProtocolViolationException("Expected welcome message");
                if (string.IsNullOrWhiteSpace(welcomeEnvelope.Payload))
                    throw new ProtocolViolationException("No payload in welcome message");

                var welcomeMessage = welcomeEnvelope.GetPayload<WelcomeMessage>()
                    ?? throw new ProtocolViolationException("Invalid welcome message");

                if (string.IsNullOrWhiteSpace(welcomeMessage.PublicKeyHash))
                    throw new ProtocolViolationException("No public key hash in welcome message");
                _serverCertificate = _serverKeys.FirstOrDefault(x => x.PublicKeyHash == welcomeMessage.PublicKeyHash && x.Expiry > DateTimeOffset.Now);

                if (_serverCertificate == null)
                {
                    certificateRefreshHelper.Signal();
                    throw new ProtocolViolationException("No valid server certificate");
                }

                try
                {
                    var tmp = RSA.Create();
                    tmp.ImportFromPem(_serverCertificate.PublicKey);
                    _serverPublicKey = tmp;
                }
                catch
                {
                    certificateRefreshHelper.Signal();
                    throw new ProtocolViolationException("Invalid server certificate");
                }

                SetState(ConnectionState.WelcomeReceived);

                // Prepare basic metadata and allow additional metadata to be added
                var metadata = await _onConnect(new Dictionary<string, string?>() {
                        { "client-version", UpdaterManager.SelfVersion?.Version ?? "0.0.0.0" },
                        { "client-id", ClientId },
                        { "client-uptime", (DateTime.Now - Process.GetCurrentProcess().StartTime).ToString() },
                        { "machine-name", DataFolderManager.MachineName },
                        { "machine-id", DataFolderManager.MachineID },
                        { "install-id", DataFolderManager.InstallID },
                        { "machine-os", UpdaterManager.OperatingSystemName },
                        { "package-id", UpdaterManager.PackageTypeId },
                        { "update-channel", UpdaterManager.CurrentChannel.ToString() }
                    });

                SendEnvelope(
                    welcomeEnvelope.RespondWith(
                        new AuthMessage(
                            _token,
                            ClientKey.ExportRSAPublicKeyPem(),
                            UpdaterManager.SelfVersion?.Version ?? "0.0.0.0",
                            PROTOCOL_VERSION,
                            metadata
                        ),
                        "auth"
                    ),
                    force: true);
                return;
            }

            if (_serverCertificate == null || _serverPublicKey == null || _serverCertificate.HasExpired())
            {
                certificateRefreshHelper.Signal();
                throw new ProtocolViolationException("No valid server certificate");
            }

            var envelope = TransportHelper.ParseFromEncryptedMessage(msg.Text, ClientKey);
            if (_state == ConnectionState.WelcomeReceived)
            {
                if (envelope.GetMessageType() != MessageType.Auth)
                    throw new ProtocolViolationException("Expected welcome message");

                var authMessage = envelope.GetPayload<AuthResultMessage>();
                if (!authMessage.Accepted ?? false)
                    throw new ProtocolViolationException("Authentication failed");

                SetState(ConnectionState.Authenticated);

                if ((authMessage.WillReplaceToken ?? false) && authMessage.NewToken != null)
                {
                    _token = authMessage.NewToken;
                    await InvokeReKeyAsync();
                }
            }
            else if (_state == ConnectionState.Authenticated)
            {
                Log.WriteVerboseMessage(LogTag, "WebsocketMessage", "Processing message of type {0}", envelope.GetMessageType());
                switch (envelope.GetMessageType())
                {
                    case MessageType.Pong:
                        break;

                    case MessageType.Command:
                        await _onMessage(new CommandMessage(
                            envelope.GetPayload<CommandRequestMessage>(),
                            response => SendEnvelope(envelope.RespondWith(response))
                        ));
                        break;

                    case MessageType.Control:
                        await _onControl(new ControlMessage(
                            envelope.GetPayload<ControlRequestMessage>(),
                            response => SendEnvelope(envelope.RespondWith(response)),
                            refreshSettingsBy =>
                            {
                                _refreshSettingsBy = DateTimeOffset.FromUnixTimeMilliseconds(Math.Max(refreshSettingsBy.ToUnixTimeMilliseconds(), DateTimeOffset.UtcNow.AddSeconds(30).ToUnixTimeMilliseconds()));
                                _client.Stop(System.Net.WebSockets.WebSocketCloseStatus.NormalClosure, "Disconnect requested").FireAndForget();
                            }
                        ));
                        break;

                    default:
                        throw new ProtocolViolationException("Unexpected message");
                }
            }
            else
            {
                throw new ProtocolViolationException("Unexpected message");
            }
        }
        catch (Exception ex)
        {
            SetState(ConnectionState.Error);
            Log.WriteMessage(LogMessageType.Error, LogTag, "WebsocketMessage", ex, "Failed to process message: {0}", msg);

            await _client.Stop(System.Net.WebSockets.WebSocketCloseStatus.NormalClosure, "Error");
            reconnectHelper.Signal();
        }

    }

    /// <summary>
    /// Helper method to invoke the rekey callback
    /// </summary>
    /// <returns>An awaitable task</returns>
    private Task InvokeReKeyAsync()
        => _onReKey(new ClaimedClientData(_token, _serverUrl, _certificateUrl, _serverKeys, null, null));

    /// <summary>
    /// Creates a new connection to the remote server
    /// </summary>
    /// <param name="serverUrl">The url to use</param>
    /// <param name="JWT">The JWT to use</param>
    /// <param name="certificateUrl">The certificate url to use</param>
    /// <param name="serverKeys">The server keys to use</param>
    /// <param name="refreshSettingsBy">The time to refresh settings by</param>
    /// <param name="forceConnect">If the connection should be force enabled, ignoring re-connect delays</param>
    /// <param name="cancellationToken">The token to cancel the connection</param>
    /// <param name="onConnect">The callback to call when connecting</param>
    /// <param name="onReKey">The callback to call when rekeying</param>
    /// <param name="onControl">The callback to call when a control message is received</param>
    /// <param name="onMessage">The callback to call when a command message is received</param>
    /// <returns></returns>
    public static Task StartAsync(
        string serverUrl,
        string JWT,
        string certificateUrl,
        IEnumerable<MiniServerCertificate> serverKeys,
        DateTimeOffset? refreshSettingsBy,
        bool forceConnect,
        CancellationToken cancellationToken,
        Func<Dictionary<string, string?>, Task<Dictionary<string, string?>>> onConnect,
        Func<ClaimedClientData, Task> onReKey,
        Func<ControlMessage, Task> onControl,
        Func<CommandMessage, Task> onMessage)
        => Task.Run(async () =>
        {
            using var connection = new KeepRemoteConnection(serverUrl, JWT, certificateUrl, serverKeys, refreshSettingsBy, forceConnect, cancellationToken, onConnect, onReKey, onControl, onMessage);
            await connection._runnerTask;
        });

    /// <summary>
    /// Gets the task representing the connection
    /// </summary>
    /// <returns>The task</returns>
    public Task RunAsync()
        => _runnerTask;

    /// <summary>
    /// Stops the connection
    /// </summary>
    /// <returns>An awaitable task</returns>
    public async Task StopAsync()
    {
        await _cancellationTokenSource.CancelAsync();
        await _runnerTask;
    }

    /// <summary>
    /// Sends an enveloped message to the remote server
    /// </summary>
    /// <param name="envelope">The envelope to send</param>
    /// <returns>True if the message was sent</returns>
    private bool SendEnvelope(EnvelopedMessage envelope, bool force = true)
    {
        if ((_state != ConnectionState.Authenticated && !force) || _serverPublicKey == null)
            return false;

        _client.Send(TransportHelper.CreateEncryptedMessage(envelope with { From = ClientId }, _serverPublicKey));
        return true;
    }

    /// <summary>
    /// Sends a new command to the server
    /// </summary>
    /// <param name="message">The message to send</param>
    /// <returns>True if the message was sent</returns>
    public bool SendCommand(CommandRequestMessage message)
    {
        if (_state != ConnectionState.Authenticated || _serverPublicKey == null)
            return false;

        _client.Send(TransportHelper.CreateEncryptedMessage(new EnvelopedMessage()
        {
            From = ClientId,
            To = "server",
            Type = "command",
            MessageId = Guid.NewGuid().ToString(),
            Payload = JsonSerializer.Serialize(message, options: JsonOptions),
            ErrorMessage = null
        }, _serverPublicKey));

        return true;
    }

    /// <summary>
    /// The current state of the connection
    /// </summary>
    public ConnectionState State => _state;

    /// <summary>
    /// Checks if automatic reconnection should be disabled based on the RefreshSettingsBy timestamp.
    /// </summary>
    /// <returns>True if automatic reconnect is disabled; otherwise, false.</returns>
    public bool IsAutoReconnectDisabled()
        => DateTimeOffset.UtcNow < _refreshSettingsBy;

    /// <summary>
    /// Creates a new connection to the remote server
    /// </summary>
    /// <param name="serverUrl">The url to use</param>
    /// <param name="JWT">The JWT token to use</param>
    /// <param name="certificateUrl">The certificate url to use</param>
    /// <param name="serverKeys">The server keys to use</param>
    /// <param name="refreshSettingsBy">The timestamp to disable automatic reconnect</param>
    /// <param name="forceConnect">If the connection should be force enabled, ignoring re-connect delays</param>
    /// <param name="cancellationToken">The cancellation token to use</param>
    /// <param name="onConnect">The callback to call when connecting</param>
    /// <param name="onReKey">The callback to call when rekeying</param>
    /// <param name="onControl">The callback to call when a control message is received</param>
    /// <param name="onMessage">The callback to call when a message is received</param>
    /// <returns>The connection object</returns>
    public static KeepRemoteConnection CreateRemoteListener(
        string serverUrl,
        string JWT,
        string certificateUrl,
        IEnumerable<MiniServerCertificate> serverKeys,
        DateTimeOffset? refreshSettingsBy,
        bool forceConnect,
        CancellationToken cancellationToken,
        Func<Dictionary<string, string?>, Task<Dictionary<string, string?>>> onConnect,
        Func<ClaimedClientData, Task> onReKey,
        Func<ControlMessage, Task> onControl,
        Func<CommandMessage, Task> onMessage)
        => new KeepRemoteConnection(serverUrl, JWT, certificateUrl, serverKeys, refreshSettingsBy, forceConnect, cancellationToken, onConnect, onReKey, onControl, onMessage);

    /// <summary>
    /// Requests a certificate refresh
    /// </summary>
    /// <param name="cancelToken">The cancellation token</param>
    /// <returns>An awaitable task</returns>
    private async Task RefreshCertificatesAsync(CancellationToken cancelToken)
    {
        using var client = HttpClientHelper.CreateClient(); // We won't set infiniteTimeout and keep the default 100s timeout
        var response = await client.GetAsync(_certificateUrl, cancelToken);
        if (response.IsSuccessStatusCode)
        {
            await using var stream = await response.Content.ReadAsStreamAsync(cancelToken);
            var serverKeys = await JsonSerializer.DeserializeAsync<IEnumerable<MiniServerCertificate>>(stream, options: RegisterForRemote.JsonOptions, cancellationToken: cancelToken);
            if (serverKeys != null && serverKeys.Any())
            {
                _serverKeys = serverKeys
                    .Where(x => !x.HasExpired() && !string.IsNullOrWhiteSpace(x.PublicKeyHash) && !string.IsNullOrWhiteSpace(x.PublicKey))
                    .ToList();

                await InvokeReKeyAsync();
            }
        }
    }

    /// </inheritdoc>
    public void Dispose()
    {
        _cancellationTokenSource.Cancel();
        _client.Dispose();
        _cancellationTokenSource.Dispose();
    }

    /// <summary>
    /// A wrapper for allowing external code to handle a command message
    /// </summary>
    public sealed class CommandMessage
    {
        /// <summary>
        /// The callback method that will receive the response
        /// </summary>
        private readonly Func<CommandResponseMessage, bool> _respondCommand;
        /// <summary>
        /// The command request message
        /// </summary>
        public CommandRequestMessage CommandRequestMessage { get; }

        /// <summary>
        /// Creates a new command message
        /// </summary>
        /// <param name="commandRequestMessage">The command request message</param>
        /// <param name="respondCommand">The callback method that will receive the response</param>
        public CommandMessage(CommandRequestMessage commandRequestMessage, Func<CommandResponseMessage, bool> respondCommand)
        {
            CommandRequestMessage = commandRequestMessage;
            _respondCommand = respondCommand;
        }

        /// <summary>
        /// Responds to the command message
        /// </summary>
        /// <param name="response">The response to send</param>
        /// <returns>True if the response was sent</returns>
        public bool Respond(CommandResponseMessage response)
            => _respondCommand(response);

        /// <summary>
        /// Handles the command message with a configured http client.
        /// The client must be configured with the correct base address and authorization headers.
        /// </summary>
        /// <param name="client">The pre-configured http client</param>
        /// <returns>An awaitable task</returns>
        public async Task HandleAsync(HttpClient client)
        {
            try
            {
                Log.WriteVerboseMessage(LogTag, "WebsocketCommand", "Handling command {0} {1}", CommandRequestMessage.Method, CommandRequestMessage.Path);

                var request = new HttpRequestMessage(new HttpMethod(CommandRequestMessage.Method), CommandRequestMessage.Path);
                if (!string.IsNullOrWhiteSpace(CommandRequestMessage.Body))
                    request.Content = new ByteArrayContent(Convert.FromBase64String(CommandRequestMessage.Body));
                if (CommandRequestMessage.Headers != null)
                {
                    foreach (var header in CommandRequestMessage.Headers)
                    {
                        if (header.Key == "Content-Type")
                        {
                            if (request.Content != null)
                                request.Content.Headers.ContentType = new MediaTypeHeaderValue(header.Value);
                        }
                        else
                            request.Headers.Add(header.Key, header.Value);
                    }
                }

                var response = await client.SendAsync(request);
                var responseBody = await response.Content.ReadAsByteArrayAsync();
                var responseHeaders = response.Headers.ToDictionary(x => x.Key, x => x.Value.First());

                Respond(new CommandResponseMessage((int)response.StatusCode, responseBody == null ? null : Convert.ToBase64String(responseBody), responseHeaders));
            }
            catch (Exception ex)
            {
                Respond(new CommandResponseMessage(500, ex.Message, null));
            }
        }
    }

    /// <summary>
    /// A wrapper for allowing external code to handle a control message
    /// </summary>
    public sealed class ControlMessage
    {
        /// <summary>
        /// The callback method that will receive the response
        /// </summary>
        private readonly Func<ControlResponseMessage, bool> _respondCommand;
        /// <summary>
        /// The callback method to request disconnect after responding
        /// </summary>
        private readonly Action<DateTimeOffset>? _requestDisconnect;
        /// <summary>
        /// The command request message
        /// </summary>
        public ControlRequestMessage ControlRequestMessage { get; }

        /// <summary>
        /// Creates a new command message
        /// </summary>
        /// <param name="controlRequestMessage">The command request message</param>
        /// <param name="respondCommand">The callback method that will receive the response</param>
        /// <param name="requestDisconnect">Optional callback to request disconnect after responding</param>
        public ControlMessage(ControlRequestMessage controlRequestMessage, Func<ControlResponseMessage, bool> respondCommand, Action<DateTimeOffset>? requestDisconnect)
        {
            ControlRequestMessage = controlRequestMessage;
            _respondCommand = respondCommand;
            _requestDisconnect = requestDisconnect;
        }

        /// <summary>
        /// Responds to the command message
        /// </summary>
        /// <param name="response">The response to send</param>
        /// <returns>True if the response was sent</returns>
        public bool Respond(ControlResponseMessage response)
            => _respondCommand(response);

        /// <summary>
        /// Requests the connection to be closed after responding.
        /// This is used when the server indicates the client should disconnect
        /// </summary>
        /// <param name="refreshSettingsBy">The time to refresh settings by</param>
        public void RequestDisconnect(DateTimeOffset refreshSettingsBy)
            => _requestDisconnect?.Invoke(refreshSettingsBy);

        /// <summary>
        /// Creates a control message for initial settings from ClaimedClientData.
        /// This is used to apply settings that are returned when the machine is claimed.
        /// </summary>
        /// <param name="settings">The settings dictionary from ClaimedClientData.Settings</param>
        /// <returns>A control message with the settings as parameters</returns>
        public static ControlMessage CreateSettingsControlMessage(Dictionary<string, string?> settings)
            => new(
                new ControlRequestMessage(ControlRequestMessage.UpdateSettingsCommand, settings),
                _ => true, // No-op response handler since there's no websocket connection for this
                _ => { } // No-op disconnect handler since there's no websocket connection for this
            );
    }
}