MoShiro/backend/shared/tests/MoShiro.Protocol.Tests/RawChunkCodecTests.cs

using System.Buffers.Binary;
using System.IO.Compression;
using Fleet.Core;
using MoShiro.Protocol;
using MoShiro.TestSupport;

namespace MoShiro.Protocol.Tests;

public sealed class RawChunkCodecTests
{
    private static readonly NodeId s_agent = new("agt_test");

    private static RawChunk SampleChunk()
    {
        RawChunk chunk = ChunkFactory.Echo(
            s_agent, "udp_echo", 1_700_000_000_000, 6, 100,
            i => new EchoPacket(LostUp: i == 2, LostDown: i == 4, 20_000 + (i * 100), 21_000));
        chunk.Header.ChunkId = new ChunkId(Guid.Parse("01890000-0000-7000-8000-000000000001"));
        return chunk;
    }

    [Fact]
    public void RoundTrip_PreservesEveryValue_IncludingMissingOnes()
    {
        RawChunk original = SampleChunk();
        RawChunk decoded = RawChunkCodec.Decode(RawChunkCodec.Encode(original)).Value;

        Assert.Equal(original.Header.ChunkId, decoded.Header.ChunkId);
        Assert.Equal("udp_echo", decoded.Header.Kind);
        Assert.Equal(original.Header.Clock.ProbeOffsetUs, decoded.Header.Clock.ProbeOffsetUs);
        RawStream before = original.Find("msg")!;
        RawStream after = decoded.Find("msg")!;
        Assert.Equal(before.Count, after.Count);
        foreach (RawColumn column in before.Columns)
        {
            RawColumn other = after.Find(column.Name)!;
            Assert.Equal(column.Role, other.Role);
            for (int row = 0; row < before.Count; row++)
            {
                Assert.Equal(column.IsValid(row), other.IsValid(row));
                if (column.IsValid(row))
                {
                    Assert.Equal(column.Values[row], other.Values[row]);
                }
            }
        }

        Assert.False(after.Find("t2")!.IsValid(2));
        Assert.False(after.Find("t4")!.IsValid(4));
    }

    [Fact]
    public void RoundTrip_HandlesTypeExtremes_AndNegativeDeltas()
    {
        var spec = new StreamSpec("x", [
            Col.U8("a"), Col.U16("b"), Col.U32("c"), Col.U64("d"), Col.Seq("e"),
            new ColumnSpec("f", ColumnType.I32), new ColumnSpec("g", ColumnType.I64, ColumnEnc.Delta, true),
        ]);
        var builder = new StreamBuilder(spec)
            .Add(byte.MaxValue, ushort.MaxValue, uint.MaxValue, long.MaxValue, 5, int.MinValue, long.MinValue / 2)
            .Add(0, 0, 0, 0, 2, int.MaxValue, null)
            .Add(1, 1, 1, 1, 9, -1, long.MaxValue / 2);
        var header = new RawChunkHeader { Kind = "x", AgentId = s_agent };
        RawChunk decoded = RawChunkCodec.Decode(RawChunkCodec.Encode(new RawChunk(header, [builder.Build()]))).Value;
        RawStream s = decoded.Find("x")!;

        Assert.Equal(uint.MaxValue, s.Find("c")!.Values[0]);
        Assert.Equal(long.MaxValue, s.Find("d")!.Values[0]);
        Assert.Equal(new long[] { 5, 2, 9 }, s.Find("e")!.Values);
        Assert.Equal(int.MinValue, s.Find("f")!.Values[0]);
        Assert.Equal(long.MinValue / 2, s.Find("g")!.Values[0]);
        Assert.False(s.Find("g")!.IsValid(1));
        Assert.Equal(long.MaxValue / 2, s.Find("g")!.Values[2]);
    }

    [Fact]
    public void GoldenBody_LayoutIsStable()
    {
        // 布局变化（列编码、头部字段、字节序）会让这个测试失败：必须同时升级 container_version 并更新黄金字节。
        RawChunk chunk = SampleChunk();
        chunk.Header.Clock = new ChunkClock { AgentMonoUs = 1000, UtcMs = 1_700_000_000_000, ProbeOffsetUs = 500 };
        byte[] body = RawChunkCodec.EncodeBody(chunk);
        uint headerLength = BinaryPrimitives.ReadUInt32LittleEndian(body);
        string columns = Convert.ToHexStringLower(body.AsSpan(4 + (int)headerLength));
        string header = System.Text.Encoding.UTF8.GetString(body.AsSpan(4, (int)headerLength));

        Assert.StartsWith("{\"chunk_id\":\"01890000-0000-7000-8000-000000000001\",\"container_version\":1,\"kind\":\"udp_echo\"", header);
        Assert.Contains("\"streams\":[{\"name\":\"msg\",\"count\":6,\"columns\":[{\"name\":\"seq\",\"type\":\"u32\",\"enc\":\"delta\",\"nullable\":false},", header);
        Assert.Equal(GoldenColumns, columns);
    }

    private const string GoldenColumns = "00020202020240404040404080fbf3b407c09a0cc09a0cc09a0cc09a0cc09a0c3bc0e0d8b907889c0c90b818889c0c889c0c3ba4e1d8b907889c0c90b818889c0c889c0c2bb4fcf8b407889c0c90b81890b818";

    [Fact]
    public void Decode_RejectsBadMagicVersionCrcAndTruncation()
    {
        byte[] good = RawChunkCodec.Encode(SampleChunk());

        byte[] badMagic = (byte[])good.Clone();
        badMagic[0] = (byte)'X';
        byte[] badVersion = (byte[])good.Clone();
        badVersion[4] = 99;
        byte[] badCrc = (byte[])good.Clone();
        badCrc[^1] ^= 0xFF;

        Assert.Equal(MoShiroErrors.RawCorrupt, RawChunkCodec.Decode(badMagic).Error.Code);
        Assert.Contains("version", RawChunkCodec.Decode(badVersion).Error.Message);
        Assert.Contains("crc", RawChunkCodec.Decode(badCrc).Error.Message);
        Assert.False(RawChunkCodec.Decode(good.AsSpan(0, good.Length / 2)).IsSuccess);
        Assert.False(RawChunkCodec.Decode([]).IsSuccess);
        Assert.False(RawChunkCodec.Decode(good.AsSpan(0, 8)).IsSuccess);
    }

    [Fact]
    public void Decode_RejectsDecompressionBombs()
    {
        byte[] zeros = new byte[ProtocolLimits.MaxDecodedChunkBytes + 1024];
        using var compressed = new MemoryStream();
        using (var brotli = new BrotliStream(compressed, CompressionLevel.Fastest, leaveOpen: true))
        {
            brotli.Write(zeros);
        }

        byte[] payload = compressed.ToArray();
        Assert.True(payload.Length < 1_000_000, "the bomb must be small on the wire");
        Assert.False(RawChunkCodec.Decode(Wrap(payload)).IsSuccess);
    }

    [Fact]
    public void Decode_RejectsHeaderLengthBeyondData_AndTrailingBytes()
    {
        byte[] body = new byte[16];
        BinaryPrimitives.WriteUInt32LittleEndian(body, 1_000_000);
        Assert.Contains("header length", RawChunkCodec.Decode(Wrap(Brotli(body))).Error.Message);

        byte[] valid = RawChunkCodec.EncodeBody(SampleChunk());
        byte[] withTrailing = [.. valid, 1, 2, 3];
        Assert.Contains("trailing", RawChunkCodec.Decode(Wrap(Brotli(withTrailing))).Error.Message);
    }

    [Fact]
    public void Decode_RejectsOutOfRangeValues_AndTooManyRecords()
    {
        // u8 列里写入 300
        var spec = new StreamSpec("x", [new ColumnSpec("a", ColumnType.I32)]);
        RawChunk chunk = new(new RawChunkHeader { Kind = "x", AgentId = s_agent }, [new StreamBuilder(spec).Add(300).Build()]);
        byte[] body = RawChunkCodec.EncodeBody(chunk);
        string json = System.Text.Encoding.UTF8.GetString(body.AsSpan(4, (int)BinaryPrimitives.ReadUInt32LittleEndian(body)));
        string patched = json.Replace("\"type\":\"i32\"", "\"type\":\"u8\"");
        byte[] patchedJson = System.Text.Encoding.UTF8.GetBytes(patched);
        byte[] rebuilt = new byte[4 + patchedJson.Length + (body.Length - 4 - json.Length)];
        BinaryPrimitives.WriteUInt32LittleEndian(rebuilt, (uint)patchedJson.Length);
        patchedJson.CopyTo(rebuilt, 4);
        body.AsSpan(4 + json.Length).CopyTo(rebuilt.AsSpan(4 + patchedJson.Length));
        Assert.Contains("out of range", RawChunkCodec.Decode(Wrap(Brotli(rebuilt))).Error.Message);

        string tooMany = json.Replace("\"count\":1", $"\"count\":{ProtocolLimits.MaxRecordsPerChunk + 1}");
        byte[] tooManyJson = System.Text.Encoding.UTF8.GetBytes(tooMany);
        byte[] big = new byte[4 + tooManyJson.Length];
        BinaryPrimitives.WriteUInt32LittleEndian(big, (uint)tooManyJson.Length);
        tooManyJson.CopyTo(big, 4);
        Assert.Contains("invalid size", RawChunkCodec.Decode(Wrap(Brotli(big))).Error.Message);
    }

    [Fact]
    public void Encode_IsSmallerThanTheNaiveRepresentation()
    {
        // 10 Hz 的 udp_echo 一个 10 秒块应当远小于每行 6×8 字节
        RawChunk chunk = ChunkFactory.UdpEchoClean(s_agent, 1_700_000_000_000);
        Assert.True(RawChunkCodec.Encode(chunk).Length < 100 * 6 * 8 / 4);
    }

    [Fact]
    public void Crc32c_MatchesTheStandardCheckValue()
    {
        Assert.Equal(0xE3069283u, Crc32c.Compute("123456789"u8));
    }

    private static byte[] Brotli(byte[] data)
    {
        using var output = new MemoryStream();
        using (var brotli = new BrotliStream(output, CompressionLevel.Fastest, leaveOpen: true))
        {
            brotli.Write(data);
        }

        return output.ToArray();
    }

    private static byte[] Wrap(byte[] payload)
    {
        byte[] result = new byte[9 + payload.Length];
        "MSRC"u8.CopyTo(result);
        result[4] = ProtocolVersion.Container;
        BinaryPrimitives.WriteUInt32LittleEndian(result.AsSpan(5), Crc32c.Compute(payload));
        payload.CopyTo(result, 9);
        return result;
    }
}