NetworkTool/windows/NetworkTool.Client/Services/DiscoveryService.cs

using System.Net;
using System.Net.Sockets;
using System.Runtime.InteropServices;
using System.Text;
using System.Text.Json;
using System.Text.Json.Serialization;
using NetworkTool.Client.Models;

namespace NetworkTool.Client.Services;

public sealed class DiscoveryService
{
    private const int DiscoveryPort = 50000;
    private static readonly TimeSpan DiscoveryTimeout = TimeSpan.FromSeconds(5);
    private static readonly TimeSpan DiscoveryRetryInterval = TimeSpan.FromSeconds(1);

    [DllImport("iphlpapi.dll", ExactSpelling = true)]
    private static extern int SendARP(int destinationIp, int sourceIp, byte[] macAddress, ref int physicalAddressLength);

    public async Task<IReadOnlyList<DiscoveredDevice>> DiscoverManyAsync(
        string localIPv4,
        Action<DiscoveredDevice>? onDeviceDiscovered = null,
        CancellationToken cancellationToken = default)
    {
        if (string.IsNullOrWhiteSpace(localIPv4))
        {
            return [];
        }

        using var client = new UdpClient(AddressFamily.InterNetwork)
        {
            EnableBroadcast = true,
        };

        client.Client.Bind(new IPEndPoint(IPAddress.Parse(localIPv4), 0));
        var request = new
        {
            protocol_version = 1,
            message_type = "discover",
            request_id = Guid.NewGuid().ToString(),
            client_name = Environment.MachineName,
        };

        var payload = JsonSerializer.SerializeToUtf8Bytes(request);
        var broadcastEndPoint = new IPEndPoint(IPAddress.Broadcast, DiscoveryPort);
        await client.SendAsync(payload, payload.Length, broadcastEndPoint);

        using var timeoutCts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
        timeoutCts.CancelAfter(DiscoveryTimeout);

        var devicesByIp = new Dictionary<string, DiscoveredDevice>(StringComparer.OrdinalIgnoreCase);
        var nextRetryAt = DateTimeOffset.UtcNow + DiscoveryRetryInterval;
        while (!timeoutCts.IsCancellationRequested)
        {
            try
            {
                var retryDelay = nextRetryAt - DateTimeOffset.UtcNow;
                if (retryDelay <= TimeSpan.Zero)
                {
                    await client.SendAsync(payload, payload.Length, broadcastEndPoint);
                    nextRetryAt = DateTimeOffset.UtcNow + DiscoveryRetryInterval;
                    continue;
                }

                using var receiveCts = CancellationTokenSource.CreateLinkedTokenSource(timeoutCts.Token);
                var receiveTask = client.ReceiveAsync(receiveCts.Token).AsTask();
                var retryDelayTask = Task.Delay(retryDelay, timeoutCts.Token);
                var completedTask = await Task.WhenAny(receiveTask, retryDelayTask);

                if (completedTask == retryDelayTask)
                {
                    receiveCts.Cancel();
                    try
                    {
                        await receiveTask;
                    }
                    catch (OperationCanceledException)
                    {
                    }

                    if (timeoutCts.IsCancellationRequested)
                    {
                        break;
                    }

                    await client.SendAsync(payload, payload.Length, broadcastEndPoint);
                    nextRetryAt = DateTimeOffset.UtcNow + DiscoveryRetryInterval;
                    continue;
                }

                var result = await receiveTask;
                var response = JsonSerializer.Deserialize<DiscoveryResponse>(result.Buffer);
                if (response is null
                    || response.MessageType != "discover_response"
                    || !TryGetDiscoveredLinkLocalIp(response.Lan2Ip, result.RemoteEndPoint.Address, out var discoveredIp))
                {
                    continue;
                }

                var mac = response.Mac;
                if (string.IsNullOrWhiteSpace(mac))
                {
                    mac = ResolveMacAddress(result.RemoteEndPoint.Address);
                }

                var device = new DiscoveredDevice
                {
                    DeviceId = response.DeviceId ?? string.Empty,
                    Hostname = response.Hostname ?? string.Empty,
                    ServerVersion = response.ServerVersion ?? string.Empty,
                    Mac = mac ?? string.Empty,
                    Lan2Ip = discoveredIp,
                    HttpPort = response.HttpPort,
                    AuthRequired = response.AuthRequired,
                };

                devicesByIp[discoveredIp] = device;
                onDeviceDiscovered?.Invoke(device);
            }
            catch (OperationCanceledException)
            {
                break;
            }
        }

        return devicesByIp.Values.OrderBy(device => device.Lan2Ip).ToList();
    }

    private static bool TryGetDiscoveredLinkLocalIp(string? lan2Ip, IPAddress remoteAddress, out string discoveredIp)
    {
        if (IsLinkLocalIPv4(lan2Ip))
        {
            discoveredIp = lan2Ip!.Trim();
            return true;
        }

        if (IsLinkLocalIPv4(remoteAddress))
        {
            discoveredIp = remoteAddress.ToString();
            return true;
        }

        discoveredIp = string.Empty;
        return false;
    }

    private static bool IsLinkLocalIPv4(string? ipAddress)
    {
        return IPAddress.TryParse(ipAddress, out var parsed) && IsLinkLocalIPv4(parsed);
    }

    private static bool IsLinkLocalIPv4(IPAddress ipAddress)
    {
        var bytes = ipAddress.GetAddressBytes();
        return ipAddress.AddressFamily == AddressFamily.InterNetwork && bytes[0] == 169 && bytes[1] == 254;
    }

    private static string ResolveMacAddress(IPAddress ipAddress)
    {
        if (ipAddress.AddressFamily != AddressFamily.InterNetwork)
        {
            return string.Empty;
        }

        var macAddress = new byte[6];
        var macAddressLength = macAddress.Length;
        var destinationIp = BitConverter.ToInt32(ipAddress.GetAddressBytes(), 0);
        if (SendARP(destinationIp, 0, macAddress, ref macAddressLength) != 0 || macAddressLength <= 0)
        {
            return string.Empty;
        }

        return string.Join(":", macAddress.Take(macAddressLength).Select(value => value.ToString("X2")));
    }

    private sealed class DiscoveryResponse
    {
        [JsonPropertyName("message_type")]
        public string? MessageType { get; set; }

        [JsonPropertyName("device_id")]
        public string? DeviceId { get; set; }

        [JsonPropertyName("hostname")]
        public string? Hostname { get; set; }

        [JsonPropertyName("server_version")]
        public string? ServerVersion { get; set; }

        [JsonPropertyName("mac")]
        public string? Mac { get; set; }

        [JsonPropertyName("lan2_ip")]
        public string? Lan2Ip { get; set; }

        [JsonPropertyName("http_port")]
        public int HttpPort { get; set; }

        [JsonPropertyName("auth_required")]
        public bool AuthRequired { get; set; }
    }
}