Fix stray compressed bytes before PROCEED on large messages

source/Halibut.Tests/Support/Streams/SegmentedReadStream.cs

@@ -0,0 +1,67 @@
+using System;
+using System.Collections.Generic;
+using System.IO;
+using System.Linq;
+using System.Threading;
+using System.Threading.Tasks;
+
+namespace Halibut.Tests.Support.Streams
+{
+    /// <summary>
+    /// A read-only stream that never returns data from more than one segment in a single read.
+    ///
+    /// Useful for mimicking data arriving over the network in separate writes (e.g. separate TLS records),
+    /// where a read can return before the data from a later write has arrived.
+    /// </summary>
+    public class SegmentedReadStream : Stream
+    {
+        readonly Queue<byte[]> segments;
+        byte[] currentSegment = Array.Empty<byte>();
+        int currentSegmentOffset;
+
+        public SegmentedReadStream(params byte[][] segments)
+        {
+            this.segments = new Queue<byte[]>(segments.Where(s => s.Length > 0));
+        }
+
+        public override int Read(byte[] buffer, int offset, int count)
+        {
+            if (currentSegmentOffset == currentSegment.Length)
+            {
+                if (segments.Count == 0)
+                {
+                    return 0;
+                }
+
+                currentSegment = segments.Dequeue();
+                currentSegmentOffset = 0;
+            }
+
+            var bytesToRead = Math.Min(count, currentSegment.Length - currentSegmentOffset);
+            Array.Copy(currentSegment, currentSegmentOffset, buffer, offset, bytesToRead);
+            currentSegmentOffset += bytesToRead;
+            return bytesToRead;
+        }
+
+        public override Task<int> ReadAsync(byte[] buffer, int offset, int count, CancellationToken cancellationToken)
+        {
+            return Task.FromResult(Read(buffer, offset, count));
+        }
+
+        public override void Flush() => throw new NotSupportedException();
+        public override long Seek(long offset, SeekOrigin origin) => throw new NotSupportedException();
+        public override void SetLength(long value) => throw new NotSupportedException();
+        public override void Write(byte[] buffer, int offset, int count) => throw new NotSupportedException();
+
+        public override bool CanRead => true;
+        public override bool CanSeek => false;
+        public override bool CanWrite => false;
+        public override long Length => throw new NotSupportedException();
+
+        public override long Position
+        {
+            get => throw new NotSupportedException();
+            set => throw new NotSupportedException();
+        }
+    }
+}

source/Halibut.Tests/Transport/Protocol/MessageSerializerFixture.cs

@@ -1,12 +1,14 @@
 using System;
 using System.IO;
 using System.IO.Compression;
+using System.Linq;
 using System.Text;
 using System.Threading;
 using System.Threading.Tasks;
 using FluentAssertions;
 using Halibut.Diagnostics;
 using Halibut.Tests.Support;
+using Halibut.Tests.Support.Streams;
 using Halibut.Tests.Util;
 using Halibut.Transport.Protocol;
 using Halibut.Transport.Streams;
@@ -285,6 +287,54 @@ public async Task ReadMessage_Rewindable_ShouldNotObserveTrailingDataWhenReading
             readMessage.DecompressedBytesReadIntoMemory.Should().Be(expectedDecompressedBytesReadIntoMemory);
         }
 
+        [Test]
+        [TestCaseSource(typeof(MessageSerializerTestCaseSource))]
+        public async Task ReadMessage_Rewindable_WhenTheEndOfTheCompressedMessageArrivesInALaterRead_ShouldNotLeaveCompressedBytesBeforeTrailingData(MessageSerializerTestCase testCase)
+        {
+            // The BSON document can be fully decoded before the DeflateStream has read the final byte(s) of the
+            // compressed stream (which only hold the end-of-block code). If those bytes arrive in a later read, e.g.
+            // because the sender wrote them separately, they must still be consumed as part of the message, and must
+            // not be seen as the start of the next control message (e.g. "\u001fPROCEED").
+            const string trailingData = "PROCEED\r\n";
+
+            var sut = new MessageSerializerBuilder(new LogFactory())
+                .WithAsyncMemoryLimits(testCase.AsyncMemoryLimit, testCase.AsyncMemoryLimit)
+                .Build();
+
+            var message = RandomAsciiString(new Random(2142), 100_000);
+            var compressedMessage = await CreateBytesFromMessage(message);
+            var trailingBytes = Encoding.UTF8.GetBytes(trailingData);
+
+            for (var bytesInLaterRead = 1; bytesInLaterRead <= 2; bytesInLaterRead++)
+            {
+                var firstSegment = compressedMessage.Take(compressedMessage.Length - bytesInLaterRead).ToArray();
+                var laterSegment = compressedMessage.Skip(compressedMessage.Length - bytesInLaterRead).Concat(trailingBytes).ToArray();
+
+                await using (var stream = new RewindableBufferStream(new SegmentedReadStream(firstSegment, laterSegment)))
+                {
+                    var readMessage = await ReadMessage<string>(sut, stream);
+                    readMessage.Should().Be(message);
+
+                    using (var reader = new StreamReader(stream, new UTF8Encoding(false)))
+                    {
+                        var trailingResult = await reader.ReadToEndAsync();
+                        trailingResult.Should().Be(trailingData, $"the final {bytesInLaterRead} byte(s) of the compressed message arrived in a later read");
+                    }
+                }
+            }
+        }
+
+        static string RandomAsciiString(Random random, int length)
+        {
+            var chars = new char[length];
+            for (var i = 0; i < length; i++)
+            {
+                chars[i] = (char)random.Next(32, 127);
+            }
+
+            return new string(chars);
+        }
+
         static async Task<byte[]> CreateBytesFromMessage(object message)
         {
             var writingSerializer = new MessageSerializerBuilder(new LogFactory()).Build();

source/Halibut/Transport/Protocol/MessageSerializer.cs

@@ -129,6 +129,8 @@ public async Task<IReadOnlyList<DataStream>> WriteMessageAsync<T>(Stream stream,
                 {
                     var (messageEnvelope, dataStreams) = DeserializeMessageAndDataStreams<T>(bson);
 
+                    await ConsumeEndOfCompressedStreamAsync(zip, cancellationToken);
+
                     // Find the unused bytes in the DeflateStream input buffer
                     if (deflateReflector.TryGetAvailableInputBufferSize(zip, out var unusedBytesCount))
                     {
@@ -150,6 +152,31 @@ public async Task<IReadOnlyList<DataStream>> WriteMessageAsync<T>(Stream stream,
             }
         }
 
+        /// <summary>
+        /// Ensures the DeflateStream has consumed the whole compressed message from the underlying stream.
+        /// </summary>
+        /// <remarks>
+        /// The BSON reader can finish before the DeflateStream has read the end of the compressed stream, as the last
+        /// byte or two may only hold the end-of-block code. If those bytes had not been received when the BSON reader
+        /// finished, they would be left on the underlying stream and read as part of the next control message
+        /// (e.g. "\u001fPROCEED").
+        ///
+        /// The DeflateStream can only report the end of the decompressed data (a read returning 0) once it has decoded
+        /// the end-of-block code. So this read pulls any remaining compressed bytes from the underlying stream, waiting
+        /// for them if needed, even though it produces no decompressed bytes. If the compressed stream has already been
+        /// fully consumed, it returns 0 without reading from the underlying stream.
+        ///
+        /// Decompressed bytes after the BSON document are not part of the protocol, so they are treated as an error.
+        /// </remarks>
+        static async Task ConsumeEndOfCompressedStreamAsync(DeflateStream zip, CancellationToken cancellationToken)
+        {
+            var buffer = new byte[1];
+            if (await zip.ReadAsync(buffer, 0, buffer.Length, cancellationToken) != 0)
+            {
+                throw new ProtocolException("Unexpected data found in the compressed message after the end of the message.");
+            }
+        }
+
         (MessageEnvelope<T> MessageEnvelope, IReadOnlyList<DataStream> DataStreams) DeserializeMessageAndDataStreams<T>(JsonReader reader)
         {
             var streamCapturingSerializer = CreateStreamCapturingSerializer();