MoShiro/backend/fleet/src/Fleet.Link/Transports.cs

using System.Buffers.Binary;
using System.Net.WebSockets;
using Fleet.Rpc;

namespace Fleet.Link;

/// <summary>WebSocket 传输（节点 → 服务端）：文本帧是 JSON-RPC，二进制帧是发件箱。</summary>
public sealed class WsTransport : IRpcTransport
{
    private readonly WebSocket _socket;
    private readonly int _maxFrameBytes;
    private readonly SemaphoreSlim _sendLock = new(1, 1);

    public WsTransport(WebSocket socket, int maxFrameBytes)
    {
        _socket = socket;
        _maxFrameBytes = maxFrameBytes;
    }

    public async ValueTask SendAsync(RpcFrame frame, CancellationToken ct)
    {
        await _sendLock.WaitAsync(ct);
        try
        {
            await _socket.SendAsync(frame.Data, frame.Kind == FrameKind.Text ? WebSocketMessageType.Text : WebSocketMessageType.Binary, endOfMessage: true, ct);
        }
        finally
        {
            _sendLock.Release();
        }
    }

    public async ValueTask<RpcFrame?> ReceiveAsync(CancellationToken ct)
    {
        var message = new MemoryStream();
        byte[] buffer = new byte[16 * 1024];
        while (true)
        {
            ValueWebSocketReceiveResult result = await _socket.ReceiveAsync(buffer.AsMemory(), ct);
            if (result.MessageType == WebSocketMessageType.Close)
            {
                if (_socket.State == WebSocketState.CloseReceived)
                {
                    // 按协议确认对端发起的关闭
                    try
                    {
                        await _socket.CloseOutputAsync(WebSocketCloseStatus.NormalClosure, "ack", CancellationToken.None);
                    }
                    catch (WebSocketException)
                    {
                        // 对端已经断开
                    }
                }

                return null;
            }

            message.Write(buffer, 0, result.Count);
            if (message.Length > _maxFrameBytes)
            {
                await CloseAsync("frame too large", CancellationToken.None);
                return null;
            }

            if (result.EndOfMessage)
            {
                return new RpcFrame(result.MessageType == WebSocketMessageType.Text ? FrameKind.Text : FrameKind.Binary, message.ToArray());
            }
        }
    }

    public async ValueTask CloseAsync(string reason, CancellationToken ct)
    {
        if (_socket.State is WebSocketState.Open or WebSocketState.CloseReceived)
        {
            try
            {
                await _socket.CloseOutputAsync(WebSocketCloseStatus.NormalClosure, reason, ct);
            }
            catch (WebSocketException)
            {
                // 对端已经断开
            }
        }
    }

    public ValueTask DisposeAsync()
    {
        _socket.Dispose();
        _sendLock.Dispose();
        return ValueTask.CompletedTask;
    }
}

/// <summary>
/// TLS 帧传输（Probe → Agent）：<c>u32 LE 正文长度 | u8 类型(0 文本, 1 二进制) | 正文</c>。
/// 比 WebSocket 轻得多，Agent 不需要内置 HTTP 服务端。
/// </summary>
public sealed class TlsFrameTransport : IRpcTransport
{
    private const int HeaderBytes = 5;

    private readonly Stream _stream;
    private readonly int _maxFrameBytes;
    private readonly SemaphoreSlim _sendLock = new(1, 1);

    public TlsFrameTransport(Stream stream, int maxFrameBytes)
    {
        _stream = stream;
        _maxFrameBytes = maxFrameBytes;
    }

    public async ValueTask SendAsync(RpcFrame frame, CancellationToken ct)
    {
        if (frame.Data.Length > _maxFrameBytes)
        {
            throw new InvalidOperationException("frame exceeds the maximum size");
        }

        byte[] header = new byte[HeaderBytes];
        BinaryPrimitives.WriteUInt32LittleEndian(header, (uint)frame.Data.Length);
        header[4] = (byte)frame.Kind;
        await _sendLock.WaitAsync(ct);
        try
        {
            await _stream.WriteAsync(header, ct);
            await _stream.WriteAsync(frame.Data, ct);
            await _stream.FlushAsync(ct);
        }
        finally
        {
            _sendLock.Release();
        }
    }

    public async ValueTask<RpcFrame?> ReceiveAsync(CancellationToken ct)
    {
        byte[] header = new byte[HeaderBytes];
        try
        {
            if (!await ReadExactlyOrEofAsync(header, ct))
            {
                return null;
            }

            uint length = BinaryPrimitives.ReadUInt32LittleEndian(header);
            if (length > _maxFrameBytes || header[4] > (byte)FrameKind.Binary)
            {
                await _stream.DisposeAsync(); // 非法帧：直接断开
                return null;
            }

            byte[] body = new byte[length];
            if (!await ReadExactlyOrEofAsync(body, ct))
            {
                return null;
            }

            return new RpcFrame((FrameKind)header[4], body);
        }
        catch (IOException)
        {
            return null;
        }
        catch (ObjectDisposedException)
        {
            return null;
        }
    }

    public async ValueTask CloseAsync(string reason, CancellationToken ct)
    {
        try
        {
            await _stream.FlushAsync(ct);
        }
        catch (Exception ex) when (ex is IOException or ObjectDisposedException)
        {
            // 对端已经断开
        }

        await _stream.DisposeAsync();
    }

    public async ValueTask DisposeAsync()
    {
        await _stream.DisposeAsync();
        _sendLock.Dispose();
    }

    private async ValueTask<bool> ReadExactlyOrEofAsync(byte[] buffer, CancellationToken ct)
    {
        int read = 0;
        while (read < buffer.Length)
        {
            int n = await _stream.ReadAsync(buffer.AsMemory(read), ct);
            if (n == 0)
            {
                return false;
            }

            read += n;
        }

        return true;
    }
}