desktop/WirelessTextSyncer.Windows/Services/WebSocketServerService.cs

using System.Net;
using System.Net.NetworkInformation;
using System.Net.Sockets;
using System.Text.Json;
using Fleck;
using WirelessTextSyncer.Windows.Models;

namespace WirelessTextSyncer.Windows.Services;

public sealed class WebSocketServerService : IDisposable
{
    private static readonly byte[] HeartbeatPayload = "wts-heartbeat"u8.ToArray();
    private static readonly JsonSerializerOptions JsonOptions = new(JsonSerializerDefaults.Web)
    {
        PropertyNameCaseInsensitive = true
    };

    private readonly TimeSpan heartbeatInterval = TimeSpan.FromSeconds(5);
    private readonly SyncMessageHandler messageHandler;
    private readonly string deviceName;
    private readonly List<IWebSocketConnection> clients = [];
    private readonly object clientsLock = new();
    private WebSocketServer? server;
    private System.Threading.Timer? heartbeatTimer;

    public WebSocketServerService(SyncMessageHandler messageHandler, string? deviceName = null)
    {
        this.messageHandler = messageHandler;
        this.deviceName = string.IsNullOrWhiteSpace(deviceName) ? Environment.MachineName : deviceName;
    }

    public int Port { get; private set; } = 8181;

    public string LocalIpAddress { get; private set; } = "127.0.0.1";

    public bool HasClient
    {
        get
        {
            lock (clientsLock)
            {
                clients.RemoveAll(client => !client.IsAvailable);
                return clients.Count > 0;
            }
        }
    }

    public event EventHandler? StatusChanged;

    public string BuildServiceInfoJson()
    {
        return JsonSerializer.Serialize(
            new ServiceInfoMessage(
                DiscoveryResponderService.ServiceResponseType,
                1,
                deviceName,
                LocalIpAddress,
                Port),
            JsonOptions);
    }

    public void Start(int port)
    {
        Port = port;
        LocalIpAddress = ResolveLocalIpAddress();
        AppLogger.Info($"Starting WebSocket server on 0.0.0.0:{Port}. Local IP: {LocalIpAddress}");

        server = new WebSocketServer($"ws://0.0.0.0:{Port}");
        server.Start(socket =>
        {
            socket.OnOpen = () =>
            {
                IWebSocketConnection[] replacedClients;
                lock (clientsLock)
                {
                    replacedClients = clients
                        .Where(client => client.ConnectionInfo.ClientIpAddress == socket.ConnectionInfo.ClientIpAddress)
                        .ToArray();
                    clients.RemoveAll(client => !client.IsAvailable || replacedClients.Contains(client));
                    clients.Add(socket);
                }

                foreach (var replacedClient in replacedClients)
                {
                    replacedClient.Close();
                }

                AppLogger.Info($"Client connected: {socket.ConnectionInfo.ClientIpAddress}:{socket.ConnectionInfo.ClientPort}");
                _ = SendServiceInfoAsync(socket);
                StatusChanged?.Invoke(this, EventArgs.Empty);
            };
            socket.OnClose = () =>
            {
                lock (clientsLock)
                {
                    clients.Remove(socket);
                }

                AppLogger.Info($"Client disconnected: {socket.ConnectionInfo.ClientIpAddress}:{socket.ConnectionInfo.ClientPort}");
                StatusChanged?.Invoke(this, EventArgs.Empty);
            };
            socket.OnError = exception =>
            {
                AppLogger.Error("WebSocket connection error.", exception);
                RemoveClient(socket, "connection error");
            };
            socket.OnMessage = HandleRawMessage;
        });
        heartbeatTimer = new System.Threading.Timer(_ => ProbeClients(), null, heartbeatInterval, heartbeatInterval);
    }

    public void Dispose()
    {
        heartbeatTimer?.Dispose();
        IWebSocketConnection[] currentClients;
        lock (clientsLock)
        {
            currentClients = clients.ToArray();
            clients.Clear();
        }

        foreach (var client in currentClients)
        {
            client.Close();
        }

        server?.Dispose();
    }

    private void ProbeClients()
    {
        IWebSocketConnection[] currentClients;
        var removedUnavailableClients = false;
        lock (clientsLock)
        {
            var countBefore = clients.Count;
            clients.RemoveAll(client => !client.IsAvailable);
            removedUnavailableClients = clients.Count != countBefore;
            currentClients = clients.ToArray();
        }

        if (removedUnavailableClients)
        {
            AppLogger.Info("Removed unavailable WebSocket clients during heartbeat.");
            StatusChanged?.Invoke(this, EventArgs.Empty);
        }

        foreach (var client in currentClients)
        {
            _ = SendHeartbeatAsync(client);
        }
    }

    private async Task SendHeartbeatAsync(IWebSocketConnection client)
    {
        try
        {
            if (!client.IsAvailable)
            {
                RemoveClient(client, "heartbeat unavailable");
                return;
            }

            await client.SendPing(HeartbeatPayload);
        }
        catch (Exception exception)
        {
            AppLogger.Error("WebSocket heartbeat failed.", exception);
            RemoveClient(client, "heartbeat failure");
        }
    }

    private async Task SendServiceInfoAsync(IWebSocketConnection client)
    {
        try
        {
            await client.Send(BuildServiceInfoJson());
        }
        catch (Exception exception)
        {
            AppLogger.Error("Failed to send WebSocket service info.", exception);
            RemoveClient(client, "service info failure");
        }
    }

    private void RemoveClient(IWebSocketConnection client, string reason)
    {
        var removed = false;
        lock (clientsLock)
        {
            removed = clients.Remove(client);
        }

        if (!removed)
        {
            return;
        }

        AppLogger.Info($"Removed WebSocket client after {reason}: {client.ConnectionInfo.ClientIpAddress}:{client.ConnectionInfo.ClientPort}");
        StatusChanged?.Invoke(this, EventArgs.Empty);
    }

    private void HandleRawMessage(string rawMessage)
    {
        try
        {
            AppLogger.Info($"Received message: {rawMessage}");
            var message = JsonSerializer.Deserialize<SyncMessage>(rawMessage, JsonOptions);
            if (message is null)
            {
                AppLogger.Info("Ignored empty sync message.");
                return;
            }

            messageHandler.Handle(message);
        }
        catch (JsonException exception)
        {
            AppLogger.Error($"Invalid sync message: {rawMessage}", exception);
        }
        catch (Exception exception)
        {
            AppLogger.Error($"Failed to handle sync message: {rawMessage}", exception);
        }
    }

    private static string ResolveLocalIpAddress()
    {
        var address = NetworkInterface.GetAllNetworkInterfaces()
            .Where(networkInterface =>
                networkInterface.OperationalStatus == OperationalStatus.Up
                && networkInterface.NetworkInterfaceType != NetworkInterfaceType.Loopback)
            .Select(networkInterface => networkInterface.GetIPProperties())
            .Where(properties => properties.GatewayAddresses.Any(
                gateway => gateway.Address.AddressFamily == AddressFamily.InterNetwork
                    && !IPAddress.Any.Equals(gateway.Address)))
            .SelectMany(properties => properties.UnicastAddresses)
            .Select(unicast => unicast.Address)
            .FirstOrDefault(address =>
                address.AddressFamily == AddressFamily.InterNetwork
                && !IPAddress.IsLoopback(address)
                && !address.ToString().StartsWith("169.254.", StringComparison.Ordinal));

        return address?.ToString() ?? "127.0.0.1";
    }

    private sealed record ServiceInfoMessage(string Type, int Version, string Name, string Host, int Port);
}