MoShiro/backend/shared/MoShiro.Protocol/Wire/TcpWire.cs

using System.Buffers;
using System.Buffers.Binary;
using System.Security.Cryptography;
using Fleet.Core;

namespace MoShiro.Protocol;

/// <summary>
/// 数据面的 TCP 连接（Fleet 之外的线格式）。连接都由 Probe 发起；TLS 是可选的（ClientHello 以 0x16 开头，
/// 明文首帧以 "MN" 开头，所以同一个端口可以区分）。首帧：
/// <c>"MN"(2) version(1) session_id(8) stream_id(4) nonce(16) mac(16)</c>，认证失败直接关闭连接；
/// 之后是 <c>VarInt 长度 | 正文</c> 的消息，回显类消息的正文是 <c>seq(u32) send_ts_us(i64) | padding</c>。
/// </summary>
public static class TcpWire
{
    public const int AuthBytes = 2 + 1 + 8 + 4 + 16 + 16;
    public const int NonceBytes = 16;
    public const int MinEchoBody = 12;
    public const int MaxMessageBytes = 1 << 20;

    public static void WriteAuth(Span<byte> destination, long sessionId, uint streamId, ReadOnlySpan<byte> nonce, ReadOnlySpan<byte> key)
    {
        BinaryPrimitives.WriteUInt16LittleEndian(destination, MeasurementDatagram.Magic);
        destination[2] = MeasurementDatagram.Version;
        BinaryPrimitives.WriteInt64LittleEndian(destination[3..], sessionId);
        BinaryPrimitives.WriteUInt32LittleEndian(destination[11..], streamId);
        nonce[..NonceBytes].CopyTo(destination[15..]);
        Span<byte> full = stackalloc byte[32];
        HMACSHA256.HashData(key, destination[..(AuthBytes - 16)], full);
        full[..16].CopyTo(destination[(AuthBytes - 16)..]);
    }

    /// <summary>解析首帧（MAC 要在查到会话密钥之后用 <see cref="VerifyAuth"/> 校验）。</summary>
    public static bool TryParseAuth(ReadOnlySpan<byte> frame, out long sessionId, out uint streamId, out ReadOnlySpan<byte> nonce)
    {
        sessionId = 0;
        streamId = 0;
        nonce = default;
        if (frame.Length != AuthBytes || BinaryPrimitives.ReadUInt16LittleEndian(frame) != MeasurementDatagram.Magic || frame[2] != MeasurementDatagram.Version)
        {
            return false;
        }

        sessionId = BinaryPrimitives.ReadInt64LittleEndian(frame[3..]);
        streamId = BinaryPrimitives.ReadUInt32LittleEndian(frame[11..]);
        nonce = frame.Slice(15, NonceBytes);
        return true;
    }

    public static bool VerifyAuth(ReadOnlySpan<byte> frame, ReadOnlySpan<byte> key)
    {
        Span<byte> full = stackalloc byte[32];
        HMACSHA256.HashData(key, frame[..(AuthBytes - 16)], full);
        return CryptographicOperations.FixedTimeEquals(full[..16], frame[(AuthBytes - 16)..]);
    }

    /// <summary>回显消息正文：<c>seq | send_ts_us | padding</c>，总长度 <paramref name="size"/>。</summary>
    public static void WriteEchoBody(Span<byte> destination, uint seq, long sendTsUs)
    {
        destination.Clear();
        BinaryPrimitives.WriteUInt32LittleEndian(destination, seq);
        BinaryPrimitives.WriteInt64LittleEndian(destination[4..], sendTsUs);
    }

    public static bool TryReadEchoBody(ReadOnlySpan<byte> body, out uint seq, out long sendTsUs)
    {
        seq = 0;
        sendTsUs = 0;
        if (body.Length < MinEchoBody)
        {
            return false;
        }

        seq = BinaryPrimitives.ReadUInt32LittleEndian(body);
        sendTsUs = BinaryPrimitives.ReadInt64LittleEndian(body[4..]);
        return true;
    }

    /// <summary>读取一条 <c>VarInt 长度 | 正文</c> 消息；连接结束或长度非法时返回 null。</summary>
    public static async ValueTask<byte[]?> ReadMessageAsync(Stream stream, CancellationToken ct)
    {
        ulong length = 0;
        int shift = 0;
        byte[] one = new byte[1];
        for (int i = 0; i < VarInt.MaxBytes; i++)
        {
            if (await stream.ReadAsync(one, ct) == 0)
            {
                return null;
            }

            length |= (ulong)(one[0] & 0x7F) << shift;
            if ((one[0] & 0x80) == 0)
            {
                break;
            }

            shift += 7;
            if (i == VarInt.MaxBytes - 1)
            {
                return null;
            }
        }

        if (length > MaxMessageBytes)
        {
            return null;
        }

        byte[] body = new byte[length];
        int read = 0;
        while (read < body.Length)
        {
            int n = await stream.ReadAsync(body.AsMemory(read), ct);
            if (n == 0)
            {
                return null;
            }

            read += n;
        }

        return body;
    }

    public static async ValueTask WriteMessageAsync(Stream stream, ReadOnlyMemory<byte> body, CancellationToken ct)
    {
        byte[] frame = ArrayPool<byte>.Shared.Rent(checked(VarInt.MaxBytes + body.Length));
        try
        {
            int n = VarInt.Write(frame, (ulong)body.Length);
            body.CopyTo(frame.AsMemory(n));
            // TCP_NODELAY is enabled: two writes otherwise create a tiny
            // length-only segment before every echo body (and two TLS records).
            await stream.WriteAsync(frame.AsMemory(0, n + body.Length), ct);
        }
        finally
        {
            ArrayPool<byte>.Shared.Return(frame);
        }
    }
}