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

using Fleet.Identity;
using Fleet.Rpc;

namespace Fleet.Link;

/// <summary>
/// 认证阶段是明文，认证通过之后换成应用层加密（<see cref="LinkCipher"/>）：此后每帧都是加密的二进制帧，明文帧一律拒绝。
/// 认证阶段服务器的拒绝回复可能是明文（还没有协商好密钥），所以第一帧加密的回复到达之前允许一次明文文本帧。
/// </summary>
internal sealed class SecureTransport : IRpcTransport
{
    private readonly IRpcTransport _inner;
    private LinkCipher? _cipher;
    private bool _allowPlaintext = true;

    public SecureTransport(IRpcTransport inner)
    {
        _inner = inner;
    }

    public bool Secure => _cipher is not null;

    public void Enable(LinkCipher cipher) => _cipher = cipher;

    public ValueTask SendAsync(RpcFrame frame, CancellationToken ct) =>
        _cipher is null ? _inner.SendAsync(frame, ct) : _inner.SendAsync(new RpcFrame(FrameKind.Binary, _cipher.Seal((byte)frame.Kind, frame.Data.Span)), ct);

    public async ValueTask<RpcFrame?> ReceiveAsync(CancellationToken ct)
    {
        RpcFrame? frame = await _inner.ReceiveAsync(ct);
        if (frame is not { } received || _cipher is null)
        {
            return frame;
        }

        if (received.Kind == FrameKind.Binary && _cipher.TryOpen(received.Data.Span, out byte kind, out byte[] data))
        {
            _allowPlaintext = false;
            return new RpcFrame((FrameKind)kind, data);
        }

        if (_allowPlaintext && received.Kind == FrameKind.Text)
        {
            return received; // 认证被拒绝时服务器还没有协商好密钥，回复是明文
        }

        throw new IOException("a frame on the secure link failed to authenticate");
    }

    public ValueTask CloseAsync(string reason, CancellationToken ct) => _inner.CloseAsync(reason, ct);

    public async ValueTask DisposeAsync()
    {
        _cipher?.Dispose();
        await _inner.DisposeAsync();
    }
}