Add Slow Receiver Test (#1576)

src/inc/msquic.hpp

@@ -813,6 +813,22 @@ struct MsQuicStream {
         return MsQuic->StreamSend(Handle, Buffers, BufferCount, Flags, ClientSendContext);
     }
 
+    _IRQL_requires_max_(DISPATCH_LEVEL)
+    QUIC_STATUS
+    ReceiveComplete(
+        _In_ uint64_t BufferLength
+        ) noexcept {
+        return MsQuic->StreamReceiveComplete(Handle, BufferLength);
+    }
+
+    _IRQL_requires_max_(DISPATCH_LEVEL)
+    QUIC_STATUS
+    ReceiveSetEnabled(
+        _In_ bool IsEnabled = true
+        ) noexcept {
+        return MsQuic->StreamReceiveSetEnabled(Handle, IsEnabled ? TRUE : FALSE);
+    }
+
     QUIC_STATUS GetInitStatus() const noexcept { return InitStatus; }
     bool IsValid() const { return QUIC_SUCCEEDED(InitStatus); }
     MsQuicStream(MsQuicStream& other) = delete;

src/test/MsQuicTests.h

@@ -351,6 +351,10 @@ QuicTestAbortReceive(
     _In_ QUIC_ABORT_RECEIVE_TYPE Type
     );
 
+void
+QuicTestSlowReceive(
+    );
+
 //
 // QuicDrill tests
 //
@@ -805,4 +809,7 @@ typedef struct {
 #define IOCTL_QUIC_RUN_KEY_UPDATE_RANDOM_LOSS \
     QUIC_CTL_CODE(64, METHOD_BUFFERED, FILE_WRITE_DATA)
 
-#define QUIC_MAX_IOCTL_FUNC_CODE 64
+#define IOCTL_QUIC_RUN_SLOW_RECEIVE \
+    QUIC_CTL_CODE(65, METHOD_BUFFERED, FILE_WRITE_DATA)
+
+#define QUIC_MAX_IOCTL_FUNC_CODE 65

src/test/bin/quic_gtest.cpp

@@ -1318,6 +1318,15 @@ TEST(Misc, AbortIncompleteReceive) {
     }
 }
 
+TEST(Misc, SlowReceive) {
+    TestLogger Logger("SlowReceive");
+    if (TestingKernelMode) {
+        ASSERT_TRUE(DriverClient.Run(IOCTL_QUIC_RUN_SLOW_RECEIVE));
+    } else {
+        QuicTestSlowReceive();
+    }
+}
+
 TEST(Drill, VarIntEncoder) {
     TestLogger Logger("QuicDrillTestVarIntEncoder");
     if (TestingKernelMode) {

src/test/bin/winkernel/control.cpp

@@ -433,7 +433,8 @@ size_t QUIC_IOCTL_BUFFER_SIZES[] =
     sizeof(QUIC_RUN_CRED_VALIDATION),
     sizeof(QUIC_RUN_CRED_VALIDATION),
     sizeof(QUIC_ABORT_RECEIVE_TYPE),
-    sizeof(QUIC_RUN_KEY_UPDATE_RANDOM_LOSS_PARAMS)
+    sizeof(QUIC_RUN_KEY_UPDATE_RANDOM_LOSS_PARAMS),
+    0
 };
 
 CXPLAT_STATIC_ASSERT(
@@ -1029,6 +1030,11 @@ QuicTestCtlEvtIoDeviceControl(
                 Params->KeyUpdateRandomLossParams.Family,
                 Params->KeyUpdateRandomLossParams.RandomLossPercentage))
         break;
+
+    case IOCTL_QUIC_RUN_SLOW_RECEIVE:
+        QuicTestCtlRun(QuicTestSlowReceive());
+        break;
+
     default:
         Status = STATUS_NOT_IMPLEMENTED;
         break;

src/test/lib/DataTest.cpp

@@ -2149,3 +2149,93 @@ QuicTestAbortReceive(
     TEST_QUIC_SUCCEEDED(RecvContext.ServerStream->Shutdown(1));
     TEST_TRUE(RecvContext.ServerStreamShutdown.WaitTimeout(TestWaitTimeout));
 }
+
+struct SlowRecvTestContext {
+    CxPlatEvent ServerStreamRecv;
+    CxPlatEvent ServerStreamShutdown;
+    MsQuicStream* ServerStream {nullptr};
+    bool ServerStreamHasShutdown {false};
+
+    static QUIC_STATUS StreamCallback(_In_ MsQuicStream* Stream, _In_opt_ void* Context, _Inout_ QUIC_STREAM_EVENT* Event) {
+        auto TestContext = (SlowRecvTestContext*)Context;
+        if (Event->Type == QUIC_STREAM_EVENT_RECEIVE) {
+            TestContext->ServerStreamRecv.Set();
+            return QUIC_STATUS_PENDING;
+        } else if (Event->Type == QUIC_STREAM_EVENT_SHUTDOWN_COMPLETE) {
+            TestContext->ServerStreamHasShutdown = true;
+            TestContext->ServerStreamShutdown.Set();
+            Stream->ConnectionShutdown(1);
+        }
+        return QUIC_STATUS_SUCCESS;
+    }
+
+    static QUIC_STATUS ConnCallback(_In_ MsQuicConnection*, _In_opt_ void* Context, _Inout_ QUIC_CONNECTION_EVENT* Event) {
+        auto TestContext = (SlowRecvTestContext*)Context;
+        if (Event->Type == QUIC_CONNECTION_EVENT_PEER_STREAM_STARTED) {
+            TestContext->ServerStream = new MsQuicStream(Event->PEER_STREAM_STARTED.Stream, CleanUpAutoDelete, StreamCallback, Context);
+        }
+        return QUIC_STATUS_SUCCESS;
+    }
+};
+
+void
+QuicTestSlowReceive(
+    void
+    )
+{
+    MsQuicRegistration Registration;
+    TEST_QUIC_SUCCEEDED(Registration.GetInitStatus());
+
+    MsQuicConfiguration ServerConfiguration(Registration, "MsQuicTest", MsQuicSettings().SetPeerUnidiStreamCount(1), ServerSelfSignedCredConfig);
+    TEST_QUIC_SUCCEEDED(ServerConfiguration.GetInitStatus());
+
+    MsQuicConfiguration ClientConfiguration(Registration, "MsQuicTest", MsQuicCredentialConfig());
+    TEST_QUIC_SUCCEEDED(ClientConfiguration.GetInitStatus());
+
+    SlowRecvTestContext Context;
+    MsQuicAutoAcceptListener Listener(Registration, ServerConfiguration, SlowRecvTestContext::ConnCallback, &Context);
+    TEST_QUIC_SUCCEEDED(Listener.GetInitStatus());
+    TEST_QUIC_SUCCEEDED(Listener.Start("MsQuicTest"));
+    QuicAddr ServerLocalAddr;
+    TEST_QUIC_SUCCEEDED(Listener.GetLocalAddr(ServerLocalAddr));
+
+    MsQuicConnection Connection(Registration);
+    TEST_QUIC_SUCCEEDED(Connection.GetInitStatus());
+    TEST_QUIC_SUCCEEDED(Connection.StartLocalhost(ClientConfiguration, ServerLocalAddr));
+
+    MsQuicStream Stream(Connection, QUIC_STREAM_OPEN_FLAG_UNIDIRECTIONAL);
+    TEST_QUIC_SUCCEEDED(Stream.GetInitStatus());
+
+    //
+    // Open a stream, send some data and a FIN.
+    //
+    uint8_t RawBuffer[100];
+    QUIC_BUFFER Buffer { sizeof(RawBuffer), RawBuffer };
+    TEST_QUIC_SUCCEEDED(Stream.Send(&Buffer, 1, QUIC_SEND_FLAG_START | QUIC_SEND_FLAG_FIN));
+
+    //
+    // Wait for the first received data on the server side. The handler always
+    // returns pending, so make sure that pending is respected (no shutdown).
+    //
+    TEST_TRUE(Context.ServerStreamRecv.WaitTimeout(TestWaitTimeout));
+    CxPlatSleep(50);
+    TEST_FALSE(Context.ServerStreamHasShutdown);
+
+    //
+    // Complete the receive and drain only the first half of the data, and then
+    // repeat the steps above to make sure we get another receive and it doesn't
+    // shutdown the stream.
+    //
+    TEST_QUIC_SUCCEEDED(Context.ServerStream->ReceiveComplete(50));
+    TEST_QUIC_SUCCEEDED(Context.ServerStream->ReceiveSetEnabled()); // Need to reenable because the partial receive completion pauses additional events.
+    TEST_TRUE(Context.ServerStreamRecv.WaitTimeout(TestWaitTimeout));
+    CxPlatSleep(50);
+    TEST_FALSE(Context.ServerStreamHasShutdown);
+
+    //
+    // Receive the rest of the data and make sure the shutdown is then delivered.
+    //
+    TEST_QUIC_SUCCEEDED(Context.ServerStream->ReceiveComplete(50));
+    TEST_TRUE(Context.ServerStreamShutdown.WaitTimeout(TestWaitTimeout));
+    TEST_TRUE(Context.ServerStreamHasShutdown);
+}