diff --git a/source/Halibut.Tests/Support/Streams/SegmentedReadStream.cs b/source/Halibut.Tests/Support/Streams/SegmentedReadStream.cs
new file mode 100644
index 00000000..19a58aa0
--- /dev/null
+++ b/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
+{
+ ///
+ /// 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.
+ ///
+ public class SegmentedReadStream : Stream
+ {
+ readonly Queue segments;
+ byte[] currentSegment = Array.Empty();
+ int currentSegmentOffset;
+
+ public SegmentedReadStream(params byte[][] segments)
+ {
+ this.segments = new Queue(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 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();
+ }
+ }
+}
diff --git a/source/Halibut.Tests/Transport/Protocol/MessageSerializerFixture.cs b/source/Halibut.Tests/Transport/Protocol/MessageSerializerFixture.cs
index 3ac48bed..89728f89 100644
--- a/source/Halibut.Tests/Transport/Protocol/MessageSerializerFixture.cs
+++ b/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(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 CreateBytesFromMessage(object message)
{
var writingSerializer = new MessageSerializerBuilder(new LogFactory()).Build();
diff --git a/source/Halibut/Transport/Protocol/MessageSerializer.cs b/source/Halibut/Transport/Protocol/MessageSerializer.cs
index f9de4444..ad6908ee 100644
--- a/source/Halibut/Transport/Protocol/MessageSerializer.cs
+++ b/source/Halibut/Transport/Protocol/MessageSerializer.cs
@@ -129,6 +129,8 @@ public async Task> WriteMessageAsync(Stream stream,
{
var (messageEnvelope, dataStreams) = DeserializeMessageAndDataStreams(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> WriteMessageAsync(Stream stream,
}
}
+ ///
+ /// Ensures the DeflateStream has consumed the whole compressed message from the underlying stream.
+ ///
+ ///
+ /// 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.
+ ///
+ 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 MessageEnvelope, IReadOnlyList DataStreams) DeserializeMessageAndDataStreams(JsonReader reader)
{
var streamCapturingSerializer = CreateStreamCapturingSerializer();