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();