Don't do full drain on chunks containing future buffered data

src/core/recv_buffer.c

@@ -830,7 +830,8 @@ QuicRecvBufferPartialDrain(
 }
 
 //
-// Handles draining the entire first chunk (and possibly more). Return the new
+// Handles draining the entire first chunk (and possibly more). This function
+// expects the chunk to not contain more (unread) data. Return the new
 // drain length.
 //
 _IRQL_requires_max_(DISPATCH_LEVEL)
@@ -907,9 +908,16 @@ QuicRecvBufferDrain(
     }
 
     do {
-        if ((uint64_t)RecvBuffer->ReadLength > DrainLength) {
+        BOOLEAN PartialDrain = (uint64_t)RecvBuffer->ReadLength > DrainLength;
+        if (PartialDrain ||
+            //
+            // If there are 2 or more written ranges, it means that there may be
+            // more data later in the chunk that couldn't be read because there is a gap.
+            // Reuse the partial drain logic to preserve data after the gap.
+            //
+            QuicRangeSize(&RecvBuffer->WrittenRanges) > 1) {
             QuicRecvBufferPartialDrain(RecvBuffer, DrainLength);
-            return FALSE;
+            return !PartialDrain;
         }
 
         DrainLength = QuicRecvBufferFullDrain(RecvBuffer, DrainLength);

src/core/unittest/RecvBufferTest.cpp

@@ -625,6 +625,74 @@ TEST_P(WithMode, RecvCrypto) // Note - Values based on a previous failing MsQuic
     }
 }
 
+TEST_P(WithMode, DrainFrontChunkWithPendingGap)
+{
+    RecvBuffer RecvBuf;
+    ASSERT_EQ(QUIC_STATUS_SUCCESS, RecvBuf.Initialize(GetParam(), false));
+    uint64_t InOutWriteLength = DEF_TEST_BUFFER_LENGTH;
+    BOOLEAN NewDataReady = FALSE;
+
+    // place data at some future offset
+    ASSERT_EQ(
+        QUIC_STATUS_SUCCESS,
+        RecvBuf.Write(
+            2,
+            1,
+            &InOutWriteLength,
+            &NewDataReady));
+    ASSERT_FALSE(RecvBuf.HasUnreadData());
+
+    // place data to the front and drain
+    InOutWriteLength = DEF_TEST_BUFFER_LENGTH;
+    ASSERT_EQ(
+        QUIC_STATUS_SUCCESS,
+        RecvBuf.Write(
+            0,
+            1,
+            &InOutWriteLength,
+            &NewDataReady));
+    ASSERT_TRUE(NewDataReady);
+
+    uint64_t ReadOffset;
+    QUIC_BUFFER ReadBuffers[3];
+    uint32_t BufferCount = ARRAYSIZE(ReadBuffers);
+    RecvBuf.Read(&ReadOffset, &BufferCount, ReadBuffers);
+    ASSERT_EQ(1ul, BufferCount);
+    ASSERT_EQ(1ul, ReadBuffers[0].Length);
+    ASSERT_TRUE(RecvBuf.Drain(1));
+
+    // insert missing chunk and drain the rest
+    InOutWriteLength = DEF_TEST_BUFFER_LENGTH;
+    ASSERT_EQ(
+        QUIC_STATUS_SUCCESS,
+        RecvBuf.Write(
+            1,
+            1,
+            &InOutWriteLength,
+            &NewDataReady));
+    ASSERT_TRUE(NewDataReady);
+
+    BufferCount = ARRAYSIZE(ReadBuffers);
+    RecvBuf.Read(&ReadOffset, &BufferCount, ReadBuffers);
+    uint64_t TotalRead = 0;
+    for (uint32_t i = 0; i < BufferCount; ++i) {
+        TotalRead += ReadBuffers[i].Length;
+    }
+    ASSERT_EQ(2ul, TotalRead);
+    ASSERT_FALSE(RecvBuf.Drain(1)); // more data left in buffer
+
+    if (GetParam() != QUIC_RECV_BUF_MODE_MULTIPLE) {
+        BufferCount = ARRAYSIZE(ReadBuffers);
+        RecvBuf.Read(&ReadOffset, &BufferCount, ReadBuffers);
+        TotalRead = 0;
+        for (uint32_t i = 0; i < BufferCount; ++i) {
+            TotalRead += ReadBuffers[i].Length;
+        }
+        ASSERT_EQ(1ul, TotalRead);
+    }
+    ASSERT_TRUE(RecvBuf.Drain(1));
+}
+
 INSTANTIATE_TEST_SUITE_P(
     RecvBufferTest,
     WithMode,