Add ConnectionOpenInPartition (#4987)

nibanks · web-flow · commit a65ad5ea4f61 · 2025-04-11T15:16:03.000-04:00
diff --git a/docs/api/ConnectionOpen.md b/docs/api/ConnectionOpen.md
@@ -52,6 +52,7 @@ Once the connection has been shutdown, it must be cleaned up with a call to [Con
 
 # See Also
 
+[ConnectionOpenInPartition](ConnectionOpenInPartition.md)<br>
 [ConnectionClose](ConnectionClose.md)<br>
 [ConnectionShutdown](ConnectionShutdown.md)<br>
 [ConnectionStart](ConnectionStart.md)<br>
diff --git a/docs/api/ConnectionOpenInPartition.md b/docs/api/ConnectionOpenInPartition.md
@@ -0,0 +1,61 @@
+ConnectionOpenInPartition function
+======
+
+Creates a new connection in a specific partition.
+
+# Syntax
+
+```C
+typedef
+_IRQL_requires_max_(DISPATCH_LEVEL)
+QUIC_STATUS
+(QUIC_API * QUIC_CONNECTION_OPEN_IN_PARTITION_FN)(
+    _In_ _Pre_defensive_ HQUIC Registration,
+    _In_ uint16_t PartitionIndex,
+    _In_ _Pre_defensive_ QUIC_CONNECTION_CALLBACK_HANDLER Handler,
+    _In_opt_ void* Context,
+    _Outptr_ _At_(*Connection, __drv_allocatesMem(Mem)) _Pre_defensive_
+        HQUIC* Connection
+    );
+```
+
+# Parameters
+
+`Registration`
+
+The valid handle to an open registration object.
+
+`PartitionIndex`
+
+An index into the global partition set.
+
+`Handler`
+
+A pointer to the app's callback handler to be invoked for all connection events.
+
+`Context`
+
+The app context pointer (possibly null) to be associated with the connection object.
+
+`Connection`
+
+On success, returns a handle to the newly opened connection object.
+
+# Return Value
+
+The function returns a [QUIC_STATUS](QUIC_STATUS.md). The app may use `QUIC_FAILED` or `QUIC_SUCCEEDED` to determine if the function failed or succeeded.
+
+# Remarks
+
+See [ConnectionOpen](ConnectionOpen.md#Remarks) for more remarks.
+
+This function is the same as `ConnectionOpen` with the exception that this puts the connection in an explicit partition instead of inferring it based on the calling thread's current processor.
+
+# See Also
+
+[ConnectionOpen](ConnectionOpen.md)<br>
+[ConnectionClose](ConnectionClose.md)<br>
+[ConnectionShutdown](ConnectionShutdown.md)<br>
+[ConnectionStart](ConnectionStart.md)<br>
+[QUIC_CONNECTION_CALLBACK](QUIC_CONNECTION_CALLBACK.md)<br>
+[QUIC_CONNECTION_EVENT](QUIC_CONNECTION_EVENT.md)<br>
diff --git a/src/core/api.c b/src/core/api.c
@@ -40,6 +40,27 @@ MsQuicConnectionOpen(
     _Outptr_ _At_(*NewConnection, __drv_allocatesMem(Mem)) _Pre_defensive_
         HQUIC *NewConnection
     )
+{
+    return
+        MsQuicConnectionOpenInPartition(
+            RegistrationHandle,
+            QuicLibraryGetCurrentPartition()->Index,
+            Handler,
+            Context,
+            NewConnection);
+}
+
+_IRQL_requires_max_(DISPATCH_LEVEL)
+QUIC_STATUS
+QUIC_API
+MsQuicConnectionOpenInPartition(
+    _In_ _Pre_defensive_ HQUIC RegistrationHandle,
+    _In_ uint16_t PartitionIndex,
+    _In_ _Pre_defensive_ QUIC_CONNECTION_CALLBACK_HANDLER Handler,
+    _In_opt_ void* Context,
+    _Outptr_ _At_(*NewConnection, __drv_allocatesMem(Mem)) _Pre_defensive_
+        HQUIC *NewConnection
+    )
 {
     QUIC_STATUS Status;
     QUIC_REGISTRATION* Registration;
@@ -52,6 +73,7 @@ MsQuicConnectionOpen(
         RegistrationHandle);
 
     if (!IS_REGISTRATION_HANDLE(RegistrationHandle) ||
+        PartitionIndex >= MsQuicLib.PartitionCount ||
         NewConnection == NULL ||
         Handler == NULL) {
         Status = QUIC_STATUS_INVALID_PARAMETER;
@@ -68,7 +90,7 @@ MsQuicConnectionOpen(
     Status =
         QuicConnAlloc(
             Registration,
-            QuicLibraryGetCurrentPartition(),
+            &MsQuicLib.Partitions[PartitionIndex],
             NULL,
             NULL,
             &Connection);
diff --git a/src/core/api.h b/src/core/api.h
@@ -111,6 +111,18 @@ MsQuicConnectionOpen(
         HQUIC *Connection
     );
 
+_IRQL_requires_max_(DISPATCH_LEVEL)
+QUIC_STATUS
+QUIC_API
+MsQuicConnectionOpenInPartition(
+    _In_ _Pre_defensive_ HQUIC Registration,
+    _In_ uint16_t PartitionIndex,
+    _In_ _Pre_defensive_ QUIC_CONNECTION_CALLBACK_HANDLER Handler,
+    _In_opt_ void* Context,
+    _Outptr_ _At_(*Connection, __drv_allocatesMem(Mem)) _Pre_defensive_
+        HQUIC *Connection
+    );
+
 _IRQL_requires_max_(PASSIVE_LEVEL)
 void
 QUIC_API
diff --git a/src/core/library.c b/src/core/library.c
@@ -1787,6 +1787,7 @@ MsQuicOpenVersion(
     Api->ListenerStop = MsQuicListenerStop;
 
     Api->ConnectionOpen = MsQuicConnectionOpen;
+    Api->ConnectionOpenInPartition = MsQuicConnectionOpenInPartition;
     Api->ConnectionClose = MsQuicConnectionClose;
     Api->ConnectionShutdown = MsQuicConnectionShutdown;
     Api->ConnectionStart = MsQuicConnectionStart;
diff --git a/src/cs/lib/msquic_generated.cs b/src/cs/lib/msquic_generated.cs
@@ -3388,6 +3388,9 @@ internal unsafe partial struct QUIC_API_TABLE
         [NativeTypeName("QUIC_CONNECTION_COMP_CERT_FN")]
         internal delegate* unmanaged[Cdecl]<QUIC_HANDLE*, byte, QUIC_TLS_ALERT_CODES, int> ConnectionCertificateValidationComplete;
 
+        [NativeTypeName("QUIC_CONNECTION_OPEN_IN_PARTITION_FN")]
+        internal delegate* unmanaged[Cdecl]<QUIC_HANDLE*, ushort, delegate* unmanaged[Cdecl]<QUIC_HANDLE*, void*, QUIC_CONNECTION_EVENT*, int>, void*, QUIC_HANDLE**, int> ConnectionOpenInPartition;
+
         [NativeTypeName("QUIC_STREAM_PROVIDE_RECEIVE_BUFFERS_FN")]
         internal delegate* unmanaged[Cdecl]<QUIC_HANDLE*, uint, QUIC_BUFFER*, int> StreamProvideReceiveBuffers;
 
diff --git a/src/inc/msquic.h b/src/inc/msquic.h
@@ -1316,6 +1316,17 @@ QUIC_STATUS
     _Outptr_ _At_(*Connection, __drv_allocatesMem(Mem)) _Pre_defensive_
         HQUIC* Connection
     );
+typedef
+_IRQL_requires_max_(DISPATCH_LEVEL)
+QUIC_STATUS
+(QUIC_API * QUIC_CONNECTION_OPEN_IN_PARTITION_FN)(
+    _In_ _Pre_defensive_ HQUIC Registration,
+    _In_ uint16_t PartitionIndex,
+    _In_ _Pre_defensive_ QUIC_CONNECTION_CALLBACK_HANDLER Handler,
+    _In_opt_ void* Context,
+    _Outptr_ _At_(*Connection, __drv_allocatesMem(Mem)) _Pre_defensive_
+        HQUIC* Connection
+    );
 
 //
 // Closes an existing connection.
@@ -1714,6 +1725,9 @@ typedef struct QUIC_API_TABLE {
     QUIC_CONNECTION_COMP_RESUMPTION_FN  ConnectionResumptionTicketValidationComplete; // Available from v2.2
     QUIC_CONNECTION_COMP_CERT_FN        ConnectionCertificateValidationComplete;      // Available from v2.2
 
+    QUIC_CONNECTION_OPEN_IN_PARTITION_FN
+                                        ConnectionOpenInPartition;   // Available from v2.5
+
 #ifdef QUIC_API_ENABLE_PREVIEW_FEATURES
     QUIC_STREAM_PROVIDE_RECEIVE_BUFFERS_FN
                                         StreamProvideReceiveBuffers; // Available from v2.5
diff --git a/src/inc/msquic.hpp b/src/inc/msquic.hpp
@@ -1024,6 +1024,29 @@ struct MsQuicConnection {
         }
     }
 
+    MsQuicConnection(
+        _In_ const MsQuicRegistration& Registration,
+        _In_ uint16_t PartitionIndex,
+        _In_ MsQuicCleanUpMode CleanUpMode = CleanUpManual,
+        _In_ MsQuicConnectionCallback* Callback = NoOpCallback,
+        _In_ void* Context = nullptr
+        ) noexcept : CleanUpMode(CleanUpMode), Callback(Callback), Context(Context) {
+        if (!Registration.IsValid()) {
+            InitStatus = Registration.GetInitStatus();
+            return;
+        }
+        if (QUIC_FAILED(
+            InitStatus =
+                MsQuic->ConnectionOpenInPartition(
+                    Registration,
+                    PartitionIndex,
+                    (QUIC_CONNECTION_CALLBACK_HANDLER)MsQuicCallback,
+                    this,
+                    &Handle))) {
+            Handle = nullptr;
+        }
+    }
+
     MsQuicConnection(
         _In_ HQUIC ConnectionHandle,
         _In_ MsQuicCleanUpMode CleanUpMode,
diff --git a/src/rs/ffi/linux_bindings.rs b/src/rs/ffi/linux_bindings.rs
@@ -5697,6 +5697,15 @@ pub type QUIC_CONNECTION_OPEN_FN = ::std::option::Option<
         Connection: *mut HQUIC,
     ) -> ::std::os::raw::c_uint,
 >;
+pub type QUIC_CONNECTION_OPEN_IN_PARTITION_FN = ::std::option::Option<
+    unsafe extern "C" fn(
+        Registration: HQUIC,
+        PartitionIndex: u16,
+        Handler: QUIC_CONNECTION_CALLBACK_HANDLER,
+        Context: *mut ::std::os::raw::c_void,
+        Connection: *mut HQUIC,
+    ) -> ::std::os::raw::c_uint,
+>;
 pub type QUIC_CONNECTION_CLOSE_FN = ::std::option::Option<unsafe extern "C" fn(Connection: HQUIC)>;
 pub type QUIC_CONNECTION_SHUTDOWN_FN = ::std::option::Option<
     unsafe extern "C" fn(
@@ -6367,12 +6376,13 @@ pub struct QUIC_API_TABLE {
     pub DatagramSend: QUIC_DATAGRAM_SEND_FN,
     pub ConnectionResumptionTicketValidationComplete: QUIC_CONNECTION_COMP_RESUMPTION_FN,
     pub ConnectionCertificateValidationComplete: QUIC_CONNECTION_COMP_CERT_FN,
+    pub ConnectionOpenInPartition: QUIC_CONNECTION_OPEN_IN_PARTITION_FN,
     pub StreamProvideReceiveBuffers: QUIC_STREAM_PROVIDE_RECEIVE_BUFFERS_FN,
     pub ConnectionPoolCreate: QUIC_CONN_POOL_CREATE_FN,
 }
 #[allow(clippy::unnecessary_operation, clippy::identity_op)]
 const _: () = {
-    ["Size of QUIC_API_TABLE"][::std::mem::size_of::<QUIC_API_TABLE>() - 264usize];
+    ["Size of QUIC_API_TABLE"][::std::mem::size_of::<QUIC_API_TABLE>() - 272usize];
     ["Alignment of QUIC_API_TABLE"][::std::mem::align_of::<QUIC_API_TABLE>() - 8usize];
     ["Offset of field: QUIC_API_TABLE::SetContext"]
         [::std::mem::offset_of!(QUIC_API_TABLE, SetContext) - 0usize];
@@ -6440,10 +6450,12 @@ const _: () = {
         QUIC_API_TABLE,
         ConnectionCertificateValidationComplete
     ) - 240usize];
+    ["Offset of field: QUIC_API_TABLE::ConnectionOpenInPartition"]
+        [::std::mem::offset_of!(QUIC_API_TABLE, ConnectionOpenInPartition) - 248usize];
     ["Offset of field: QUIC_API_TABLE::StreamProvideReceiveBuffers"]
-        [::std::mem::offset_of!(QUIC_API_TABLE, StreamProvideReceiveBuffers) - 248usize];
+        [::std::mem::offset_of!(QUIC_API_TABLE, StreamProvideReceiveBuffers) - 256usize];
     ["Offset of field: QUIC_API_TABLE::ConnectionPoolCreate"]
-        [::std::mem::offset_of!(QUIC_API_TABLE, ConnectionPoolCreate) - 256usize];
+        [::std::mem::offset_of!(QUIC_API_TABLE, ConnectionPoolCreate) - 264usize];
 };
 pub const QUIC_STATUS_SUCCESS: QUIC_STATUS = 0;
 pub const QUIC_STATUS_PENDING: QUIC_STATUS = 4294967294;
diff --git a/src/rs/ffi/win_bindings.rs b/src/rs/ffi/win_bindings.rs
@@ -5726,6 +5726,15 @@ pub type QUIC_CONNECTION_OPEN_FN = ::std::option::Option<
         Connection: *mut HQUIC,
     ) -> HRESULT,
 >;
+pub type QUIC_CONNECTION_OPEN_IN_PARTITION_FN = ::std::option::Option<
+    unsafe extern "C" fn(
+        Registration: HQUIC,
+        PartitionIndex: u16,
+        Handler: QUIC_CONNECTION_CALLBACK_HANDLER,
+        Context: *mut ::std::os::raw::c_void,
+        Connection: *mut HQUIC,
+    ) -> HRESULT,
+>;
 pub type QUIC_CONNECTION_CLOSE_FN = ::std::option::Option<unsafe extern "C" fn(Connection: HQUIC)>;
 pub type QUIC_CONNECTION_SHUTDOWN_FN = ::std::option::Option<
     unsafe extern "C" fn(
@@ -6388,12 +6397,13 @@ pub struct QUIC_API_TABLE {
     pub DatagramSend: QUIC_DATAGRAM_SEND_FN,
     pub ConnectionResumptionTicketValidationComplete: QUIC_CONNECTION_COMP_RESUMPTION_FN,
     pub ConnectionCertificateValidationComplete: QUIC_CONNECTION_COMP_CERT_FN,
+    pub ConnectionOpenInPartition: QUIC_CONNECTION_OPEN_IN_PARTITION_FN,
     pub StreamProvideReceiveBuffers: QUIC_STREAM_PROVIDE_RECEIVE_BUFFERS_FN,
     pub ConnectionPoolCreate: QUIC_CONN_POOL_CREATE_FN,
 }
 #[allow(clippy::unnecessary_operation, clippy::identity_op)]
 const _: () = {
-    ["Size of QUIC_API_TABLE"][::std::mem::size_of::<QUIC_API_TABLE>() - 264usize];
+    ["Size of QUIC_API_TABLE"][::std::mem::size_of::<QUIC_API_TABLE>() - 272usize];
     ["Alignment of QUIC_API_TABLE"][::std::mem::align_of::<QUIC_API_TABLE>() - 8usize];
     ["Offset of field: QUIC_API_TABLE::SetContext"]
         [::std::mem::offset_of!(QUIC_API_TABLE, SetContext) - 0usize];
@@ -6461,10 +6471,12 @@ const _: () = {
         QUIC_API_TABLE,
         ConnectionCertificateValidationComplete
     ) - 240usize];
+    ["Offset of field: QUIC_API_TABLE::ConnectionOpenInPartition"]
+        [::std::mem::offset_of!(QUIC_API_TABLE, ConnectionOpenInPartition) - 248usize];
     ["Offset of field: QUIC_API_TABLE::StreamProvideReceiveBuffers"]
-        [::std::mem::offset_of!(QUIC_API_TABLE, StreamProvideReceiveBuffers) - 248usize];
+        [::std::mem::offset_of!(QUIC_API_TABLE, StreamProvideReceiveBuffers) - 256usize];
     ["Offset of field: QUIC_API_TABLE::ConnectionPoolCreate"]
-        [::std::mem::offset_of!(QUIC_API_TABLE, ConnectionPoolCreate) - 256usize];
+        [::std::mem::offset_of!(QUIC_API_TABLE, ConnectionPoolCreate) - 264usize];
 };
 pub const QUIC_STATUS_SUCCESS: QUIC_STATUS = 0;
 pub const QUIC_STATUS_PENDING: QUIC_STATUS = 459749;
diff --git a/src/test/lib/ApiTest.cpp b/src/test/lib/ApiTest.cpp
@@ -692,6 +692,22 @@ void QuicTestValidateConnection()
                 &Connection.Handle));
     }
 
+    //
+    // Invalid partition index.
+    //
+    {
+        TestScopeLogger logScope("Invalid partition index");
+        ConnectionScope Connection;
+        TEST_QUIC_STATUS(
+            QUIC_STATUS_INVALID_PARAMETER,
+            MsQuic->ConnectionOpenInPartition(
+                Registration,
+                UINT16_MAX,
+                DummyConnectionCallback,
+                nullptr,
+                &Connection.Handle));
+    }
+
     //
     // Null connection parameter.
     //
