Add a mode to the receive buffer to handle external buffers

docs/api/StreamOpen.md

@@ -35,6 +35,7 @@ Value | Meaning
 **QUIC_STREAM_OPEN_FLAG_UNIDIRECTIONAL**<br>1 | Opens a unidirectional stream.
 **QUIC_STREAM_OPEN_FLAG_0_RTT**<br>2 | Indicates that the stream may be sent in 0-RTT.
 **QUIC_STREAM_OPEN_FLAG_DELAY_ID_FC_UPDATES**<br>4 | Indicates stream ID flow control limit updates for the connection should be delayed to StreamClose.
+**QUIC_STREAM_OPEN_FLAG_EXTERNAL_BUFFERS**<br>5 | No buffer will be allocated for the stream, external buffers must be provided using StreamProvideReceiveBuffers.
 
 `Handler`
 

src/core/api.c

@@ -658,6 +658,15 @@
         goto Error;
     }
 
+    if (!!(Flags & QUIC_STREAM_OPEN_FLAG_EXTERNAL_BUFFERS) &&
+        Connection->Settings.StreamMultiReceiveEnabled) {
+        //
+        // External buffers are not supported with multi-receive.
+        //
+        Status = QUIC_STATUS_INVALID_PARAMETER;
+        goto Error;
+    }
+
     Status = QuicStreamInitialize(Connection, FALSE, Flags, (QUIC_STREAM**)NewStream);
     if (QUIC_FAILED(Status)) {
         goto Error;
@@ -1351,6 +1360,148 @@
         "[ api] Exit");
 }
 
+_IRQL_requires_max_(DISPATCH_LEVEL)
+QUIC_STATUS
+QUIC_API
+MsQuicStreamProvideReceiveBuffers(
+    _In_ _Pre_defensive_ HQUIC Handle,
+    _In_ uint32_t BufferCount,
+    _In_reads_(BufferCount) const QUIC_BUFFER* Buffers
+    )
+{
+    QUIC_STATUS Status;
+    QUIC_OPERATION* Oper;
+    CXPLAT_LIST_ENTRY ChunkList;
+    CxPlatListInitializeHead(&ChunkList);
+
+    QuicTraceEvent(
+        ApiEnter,
+        "[ api] Enter %u (%p).",
+        QUIC_TRACE_API_STREAM_PROVIDE_RECEIVE_BUFFERS,
+        Handle);
+
+    if (!IS_STREAM_HANDLE(Handle) || Buffers == NULL || BufferCount == 0) {
+        Status = QUIC_STATUS_INVALID_PARAMETER;
+        goto Error;
+    }
+
+    for (uint32_t i = 0; i < BufferCount; ++i) {
+        if (Buffers[i].Length == 0) {
+            Status = QUIC_STATUS_INVALID_PARAMETER;
+            goto Error;
+        }
+    }
+
+#pragma prefast(suppress: __WARNING_25024, "Pointer cast already validated.")
+    QUIC_STREAM* Stream = (QUIC_STREAM*)Handle;
+
+    CXPLAT_TEL_ASSERT(!Stream->Flags.HandleClosed);
+    CXPLAT_TEL_ASSERT(!Stream->Flags.Freed);
+
+    QUIC_CONNECTION* Connection = Stream->Connection;
+    QUIC_CONN_VERIFY(Connection, !Connection->State.Freed);
+
+    //
+    // Execute this API call inline if called on the worker thread.
+    //
+    BOOLEAN IsWorkerThread = Connection->WorkerThreadID == CxPlatCurThreadID();
+    BOOLEAN IsAlreadyInline = Connection->State.InlineApiExecution;
+
+    if (!Stream->Flags.UseExternalRecvBuffers) {
+        if (Stream->Flags.PeerStreamStartEventActive) {
+            CXPLAT_DBG_ASSERT(IsWorkerThread);
+            //
+            // We are inline from the callback indicating a peer opened a stream.
+            // No data was received yet so we can setup external buffers.
+            //
+            Connection->State.InlineApiExecution = TRUE;
+            QuicStreamSwitchToExternalBuffers(Stream);
+            Connection->State.InlineApiExecution = IsAlreadyInline;
+        } else {
+            //
+            // External buffers can't be provided after the stream has been
+            // started using internal buffers.
+            //
+            Status = QUIC_STATUS_INVALID_STATE;
+            goto Error;
+        }
+    }
+
+    //
+    // Allocate a chunk for each buffer, linking them together.
+    // The allocation is done here to make the worker thread task failure free.
+    //
+    for (uint32_t i = 0; i < BufferCount; ++i) {
+        QUIC_RECV_CHUNK* Chunk = CXPLAT_ALLOC_NONPAGED(sizeof(QUIC_RECV_CHUNK), QUIC_POOL_RECVBUF);
+        if (Chunk == NULL) {
+            QuicTraceEvent(
+                AllocFailure,
+                "Allocation of '%s' failed. (%llu bytes)",
+                "provided_chunk",
+                0);
+            Status = QUIC_STATUS_OUT_OF_MEMORY;
+            goto Error;
+        }
+        QuicRecvChunkInitialize(Chunk, Buffers[i].Length, Buffers[i].Buffer);
+        CxPlatListInsertTail(&ChunkList, &Chunk->Link);
+    }
+
+    if (IsWorkerThread) {
+        //
+        // Execute this API call inline if called on the worker thread.
+        //
+        Connection->State.InlineApiExecution = TRUE;
+        Status = QuicStreamProvideRecvBuffers(Stream, &ChunkList);
+        Connection->State.InlineApiExecution = IsAlreadyInline;
+    } else {
+        //
+        // Queue the operation to insert the chunks in the recv buffer, without waiting for the result.
+        //
+        Oper = QuicOperationAlloc(Connection->Worker, QUIC_OPER_TYPE_API_CALL);
+        if (Oper == NULL) {
+            Status = QUIC_STATUS_OUT_OF_MEMORY;
+            QuicTraceEvent(
+                AllocFailure,
+                "Allocation of '%s' failed. (%llu bytes)",
+                "STRM_PROVIDE_RECV_BUFFERS, operation",
+                0);
+            goto Error;
+        }
+        Oper->API_CALL.Context->Type = QUIC_API_TYPE_STRM_PROVIDE_RECV_BUFFERS;
+        Oper->API_CALL.Context->STRM_PROVIDE_RECV_BUFFERS.Stream = Stream;
+        CxPlatListInitializeHead(&Oper->API_CALL.Context->STRM_PROVIDE_RECV_BUFFERS.Chunks);
+        CxPlatListMoveItems(&ChunkList, &Oper->API_CALL.Context->STRM_PROVIDE_RECV_BUFFERS.Chunks);
+
+        //
+        // Async stream operations need to hold a ref on the stream so that the
+        // stream isn't freed before the operation can be processed. The ref is
+        // released after the operation is processed.
+        //
+        QuicStreamAddRef(Stream, QUIC_STREAM_REF_OPERATION);
+
+        //
+        // Queue the operation but don't wait for the completion.
+        //
+        QuicConnQueueOper(Connection, Oper);
+    }
+
+    Status = QUIC_STATUS_SUCCESS;
+
+Error:
+    // Cleanup allocated chunks if the operation failed.
+    while (!CxPlatListIsEmpty(&ChunkList)) {
+        QUIC_RECV_CHUNK* Chunk = CXPLAT_CONTAINING_RECORD(CxPlatListRemoveHead(&ChunkList), QUIC_RECV_CHUNK, Link);
+        CXPLAT_FREE(Chunk, QUIC_POOL_RECVBUF);
+    }
+
+    QuicTraceEvent(
+        ApiExitStatus,
+        "[ api] Exit %u",
+        Status);
+
+    return Status;
+}
+
 _IRQL_requires_max_(PASSIVE_LEVEL)
 QUIC_STATUS
 QUIC_API

src/core/api.h

@@ -216,6 +216,15 @@ MsQuicStreamReceiveComplete(
     _In_ uint64_t BufferLength
     );
 
+_IRQL_requires_max_(DISPATCH_LEVEL)
+QUIC_STATUS
+QUIC_API
+MsQuicStreamProvideReceiveBuffers(
+    _In_ _Pre_defensive_ HQUIC Handle,
+    _In_ uint32_t BufferCount,
+    _In_reads_(BufferCount) const QUIC_BUFFER *Buffers
+    );
+
 _IRQL_requires_max_(DISPATCH_LEVEL)
 QUIC_STATUS
 QUIC_API

src/core/connection.c

@@ -7508,6 +7508,25 @@
                 ApiCtx->STRM_RECV_SET_ENABLED.IsEnabled);
         break;
 
+    case QUIC_API_TYPE_STRM_PROVIDE_RECV_BUFFERS:
+        Status =
+            QuicStreamProvideRecvBuffers(
+                ApiCtx->STRM_PROVIDE_RECV_BUFFERS.Stream,
+                &ApiCtx->STRM_PROVIDE_RECV_BUFFERS.Chunks);
+
+        if (Status != QUIC_STATUS_SUCCESS) {
+            //
+            // If we cannot accept the app provided buffers at this point, we need to abort
+            // the connection: otherwise, we break the contract with the app about writting
+            // data to the provided buffers in order.
+            //
+            QuicConnFatalError(
+                ApiCtx->STRM_PROVIDE_RECV_BUFFERS.Stream->Connection,
+                Status,
+                "Failed to accept app provided receive buffers");
+        }
+        break;
+
     case QUIC_API_TYPE_SET_PARAM:
         Status =
             QuicLibrarySetParam(

src/core/library.c

@@ -1798,6 +1798,7 @@ MsQuicOpenVersion(
     Api->StreamSend = MsQuicStreamSend;
     Api->StreamReceiveComplete = MsQuicStreamReceiveComplete;
     Api->StreamReceiveSetEnabled = MsQuicStreamReceiveSetEnabled;
+    Api->StreamProvideReceiveBuffers = MsQuicStreamProvideReceiveBuffers;
 
     Api->DatagramSend = MsQuicDatagramSend;
 

src/core/operation.c

@@ -119,6 +119,8 @@ QuicOperationFree(
             }
         } else if (ApiCtx->Type == QUIC_API_TYPE_STRM_RECV_SET_ENABLED) {
             QuicStreamRelease(ApiCtx->STRM_RECV_SET_ENABLED.Stream, QUIC_STREAM_REF_OPERATION);
+        } else if (ApiCtx->Type == QUIC_API_TYPE_STRM_PROVIDE_RECV_BUFFERS) {
+            QuicStreamRelease(ApiCtx->STRM_PROVIDE_RECV_BUFFERS.Stream, QUIC_STREAM_REF_OPERATION);
         }
         CxPlatPoolFree(&Worker->ApiContextPool, ApiCtx);
     } else if (Oper->Type == QUIC_OPER_TYPE_FLUSH_STREAM_RECV) {

src/core/operation.h

@@ -56,6 +56,7 @@ typedef enum QUIC_API_TYPE {
     QUIC_API_TYPE_STRM_SEND,
     QUIC_API_TYPE_STRM_RECV_COMPLETE,
     QUIC_API_TYPE_STRM_RECV_SET_ENABLED,
+    QUIC_API_TYPE_STRM_PROVIDE_RECV_BUFFERS,
 
     QUIC_API_TYPE_SET_PARAM,
     QUIC_API_TYPE_GET_PARAM,
@@ -154,6 +155,10 @@ typedef struct QUIC_API_CONTEXT {
             QUIC_STREAM* Stream;
             BOOLEAN IsEnabled;
         } STRM_RECV_SET_ENABLED;
+        struct {
+            QUIC_STREAM* Stream;
+            CXPLAT_LIST_ENTRY /* QUIC_RECV_CHUNK */ Chunks;
+        } STRM_PROVIDE_RECV_BUFFERS;
 
         struct {
             HQUIC Handle;