voice-cat/tests/VoiceCat.Tests/FramingTests.cs

using System.Buffers;
using System.IO.Pipelines;
using Google.Protobuf;
using VoiceCat.Protocol;
using Voicecat.V1;

namespace VoiceCat.Tests;

public class FramingTests
{
    [Theory]
    [InlineData(0)]
    [InlineData(1)]
    [InlineData(65536)]
    [InlineData(ControlFraming.MaxPayloadLength)]
    public void PayloadRoundTrips(int size)
    {
        byte[] payload = Enumerable.Range(0, size).Select(i => (byte)i).ToArray();
        var output = new ArrayBufferWriter<byte>();
        ControlFraming.WriteFrame(output, payload);
        var input = new ReadOnlySequence<byte>(output.WrittenMemory);
        Assert.True(ControlFraming.TryReadFrame(ref input, out var actual));
        Assert.Equal(payload, actual.ToArray());
        Assert.True(input.IsEmpty);
    }

    [Fact]
    public void IncompleteFramesDoNotConsumeInput()
    {
        byte[] frame = [0, 0, 0, 3, 1, 2, 3];
        for (int size = 0; size < frame.Length; size++)
        {
            var input = new ReadOnlySequence<byte>(frame.AsMemory(0, size));
            Assert.False(ControlFraming.TryReadFrame(ref input, out _));
            Assert.Equal(size, input.Length);
        }
    }

    [Fact]
    public void SegmentsAndBatchedFramesAreHandled()
    {
        byte[] bytes = [0, 0, 0, 3, 1, 2, 3, 0, 0, 0, 0];
        var first = new Segment(bytes.AsMemory(0, 1));
        var last = first;
        for (int i = 1; i < bytes.Length; i++) last = last.Append(bytes.AsMemory(i, 1));
        var input = new ReadOnlySequence<byte>(first, 0, last, last.Memory.Length);
        Assert.True(ControlFraming.TryReadFrame(ref input, out var payload));
        Assert.Equal(new byte[] { 1, 2, 3 }, payload.ToArray());
        Assert.True(ControlFraming.TryReadFrame(ref input, out payload));
        Assert.True(payload.IsEmpty);
        Assert.True(input.IsEmpty);
    }

    [Fact]
    public void OversizedLengthsAreRejectedImmediately()
    {
        var input = new ReadOnlySequence<byte>(new byte[] { 1, 0, 0, 1 });
        Assert.Throws<InvalidDataException>(() => ControlFraming.TryReadFrame(ref input, out _));
        Assert.Throws<ArgumentOutOfRangeException>(() => ControlFraming.WriteFrame(new ArrayBufferWriter<byte>(), new byte[ControlFraming.MaxPayloadLength + 1]));
    }

    [Fact]
    public async Task EnvelopesRoundTripThroughPipe()
    {
        var expected = new Envelope { RequestId = 42, ClientHello = new() { ProtoVersion = 1, ClientName = "test-client", ClientVersion = "0.0.1" } };
        expected.ClientHello.Features.Add("text");
        var pipe = new Pipe();
        ControlFraming.WriteEnvelope(pipe.Writer, expected);
        ControlFraming.WriteEnvelope(pipe.Writer, new());
        await pipe.Writer.CompleteAsync();
        var actual = new List<Envelope>();
        await foreach (var envelope in ControlFraming.ReadEnvelopesAsync(pipe.Reader)) actual.Add(envelope);
        Assert.Equal(new[] { expected, new Envelope() }, actual);
        await pipe.Reader.CompleteAsync();
    }

    [Theory]
    [InlineData(new byte[] { 0 })]
    [InlineData(new byte[] { 0, 0, 0, 2, 1 })]
    public async Task TruncatedEndOfStreamIsRejected(byte[] bytes)
    {
        var pipe = new Pipe();
        pipe.Writer.Write(bytes);
        await pipe.Writer.CompleteAsync();
        await Assert.ThrowsAsync<InvalidDataException>(async () =>
        {
            await foreach (var _ in ControlFraming.ReadEnvelopesAsync(pipe.Reader)) { }
        });
        await pipe.Reader.CompleteAsync();
    }

    [Fact]
    public async Task InvalidProtobufIsRejected()
    {
        var pipe = new Pipe();
        ControlFraming.WriteFrame(pipe.Writer, new byte[] { 0xff });
        await pipe.Writer.CompleteAsync();
        await Assert.ThrowsAsync<InvalidProtocolBufferException>(async () =>
        {
            await foreach (var _ in ControlFraming.ReadEnvelopesAsync(pipe.Reader)) { }
        });
        await pipe.Reader.CompleteAsync();
    }

    [Fact]
    public async Task ReadCanBeCanceled()
    {
        var pipe = new Pipe();
        using var cancellation = new CancellationTokenSource();
        await using var enumerator = ControlFraming.ReadEnvelopesAsync(pipe.Reader, cancellation.Token).GetAsyncEnumerator();
        var pending = enumerator.MoveNextAsync().AsTask();
        cancellation.Cancel();
        await Assert.ThrowsAnyAsync<OperationCanceledException>(() => pending);
        await pipe.Writer.CompleteAsync();
        await pipe.Reader.CompleteAsync();
    }

    [Fact]
    public void UnknownFieldsSurviveParsing()
    {
        byte[] bytes = [8, 42, 0xa0, 6, 7];
        Assert.Equal(bytes, Envelope.Parser.ParseFrom(bytes).ToByteArray());
    }

    [Fact]
    public async Task FragmentedLargeEnvelopeMakesProgressUnderBackpressure()
    {
        var envelope = new Envelope { ClientHello = new() { ClientName = new string('a', 200000) } };
        var framed = new ArrayBufferWriter<byte>();
        ControlFraming.WriteEnvelope(framed, envelope);
        var pipe = new Pipe(new PipeOptions(pauseWriterThreshold: 32, resumeWriterThreshold: 16));
        using var timeout = new CancellationTokenSource(TimeSpan.FromSeconds(10));
        async Task Produce()
        {
            for (int offset = 0; offset < framed.WrittenCount; offset += 7)
                await pipe.Writer.WriteAsync(framed.WrittenMemory.Slice(offset, Math.Min(7, framed.WrittenCount - offset)), timeout.Token);
            await pipe.Writer.CompleteAsync();
        }
        var producer = Produce();
        var actual = new List<Envelope>();
        await foreach (var item in ControlFraming.ReadEnvelopesAsync(pipe.Reader, timeout.Token)) actual.Add(item);
        await producer;
        Assert.Equal(new[] { envelope }, actual);
        await pipe.Reader.CompleteAsync();
    }

    [Fact]
    public async Task StoppingEnumerationLeavesFollowingFramesAvailable()
    {
        var pipe = new Pipe();
        ControlFraming.WriteEnvelope(pipe.Writer, new() { RequestId = 1 });
        ControlFraming.WriteEnvelope(pipe.Writer, new() { RequestId = 2 });
        await pipe.Writer.FlushAsync();
        await using (var first = ControlFraming.ReadEnvelopesAsync(pipe.Reader).GetAsyncEnumerator())
        {
            Assert.True(await first.MoveNextAsync());
            Assert.Equal(1UL, first.Current.RequestId);
        }
        using var timeout = new CancellationTokenSource(TimeSpan.FromSeconds(5));
        await using (var second = ControlFraming.ReadEnvelopesAsync(pipe.Reader, timeout.Token).GetAsyncEnumerator())
        {
            Assert.True(await second.MoveNextAsync());
            Assert.Equal(2UL, second.Current.RequestId);
        }
        await pipe.Writer.CompleteAsync();
        await pipe.Reader.CompleteAsync();
    }

    private sealed class Segment : ReadOnlySequenceSegment<byte>
    {
        public Segment(ReadOnlyMemory<byte> memory) => Memory = memory;

        public Segment Append(ReadOnlyMemory<byte> memory)
        {
            var segment = new Segment(memory) { RunningIndex = RunningIndex + Memory.Length };
            Next = segment;
            return segment;
        }

    }
}