Add Load Balancing Tests

src/core/binding.c

@@ -125,14 +125,42 @@ QuicBindingInitialize(
     }
 #endif
 
-    Status =
-        CxPlatSocketCreateUdp(
-            MsQuicLib.Datapath,
-            LocalAddress,
-            RemoteAddress,
-            Binding,
-            0,
-            &Binding->Socket);
+#if QUIC_TEST_DATAPATH_HOOKS_ENABLED
+    QUIC_TEST_DATAPATH_HOOKS* Hooks = MsQuicLib.TestDatapathHooks;
+    if (Hooks != NULL) {
+        QUIC_ADDR RemoteAddressCopy;
+        if (RemoteAddress != NULL) {
+            RemoteAddressCopy = *RemoteAddress;
+        }
+        QUIC_ADDR LocalAddressCopy;
+        if (LocalAddress != NULL) {
+            LocalAddressCopy = *LocalAddress;
+        }
+        Hooks->Create(
+            RemoteAddress != NULL ? &RemoteAddressCopy : NULL,
+            LocalAddress != NULL ? &LocalAddressCopy : NULL);
+
+        Status =
+            CxPlatSocketCreateUdp(
+                MsQuicLib.Datapath,
+                LocalAddress != NULL ? &LocalAddressCopy : NULL,
+                RemoteAddress != NULL ? &RemoteAddressCopy : NULL,
+                Binding,
+                0,
+                &Binding->Socket);
+    } else {
+#endif
+        Status =
+            CxPlatSocketCreateUdp(
+                MsQuicLib.Datapath,
+                LocalAddress,
+                RemoteAddress,
+                Binding,
+                0,
+                &Binding->Socket);
+#if QUIC_TEST_DATAPATH_HOOKS_ENABLED
+    }
+#endif
 
 #ifdef QUIC_COMPARTMENT_ID
     if (RevertCompartmentId) {
@@ -151,8 +179,8 @@ QuicBindingInitialize(
     }
 
     QUIC_ADDR DatapathLocalAddr, DatapathRemoteAddr;
-    CxPlatSocketGetLocalAddress(Binding->Socket, &DatapathLocalAddr);
-    CxPlatSocketGetRemoteAddress(Binding->Socket, &DatapathRemoteAddr);
+    QuicBindingGetLocalAddress(Binding, &DatapathLocalAddr);
+    QuicBindingGetRemoteAddress(Binding, &DatapathRemoteAddr);
     QuicTraceEvent(
         BindingCreated,
         "[bind][%p] Created, Udp=%p LocalAddr=%!ADDR! RemoteAddr=%!ADDR!",
@@ -248,8 +276,8 @@ QuicBindingTraceRundown(
     // TODO - Trace datapath binding
 
     QUIC_ADDR DatapathLocalAddr, DatapathRemoteAddr;
-    CxPlatSocketGetLocalAddress(Binding->Socket, &DatapathLocalAddr);
-    CxPlatSocketGetRemoteAddress(Binding->Socket, &DatapathRemoteAddr);
+    QuicBindingGetLocalAddress(Binding, &DatapathLocalAddr);
+    QuicBindingGetRemoteAddress(Binding, &DatapathRemoteAddr);
     QuicTraceEvent(
         BindingRundown,
         "[bind][%p] Rundown, Udp=%p LocalAddr=%!ADDR! RemoteAddr=%!ADDR!",
@@ -270,6 +298,38 @@ QuicBindingTraceRundown(
     CxPlatDispatchRwLockReleaseShared(&Binding->RwLock);
 }
 
+_IRQL_requires_max_(DISPATCH_LEVEL)
+void
+QuicBindingGetLocalAddress(
+    _In_ QUIC_BINDING* Binding,
+    _Out_ QUIC_ADDR* Address
+    )
+{
+    CxPlatSocketGetLocalAddress(Binding->Socket, Address);
+#if QUIC_TEST_DATAPATH_HOOKS_ENABLED
+    QUIC_TEST_DATAPATH_HOOKS* Hooks = MsQuicLib.TestDatapathHooks;
+    if (Hooks != NULL) {
+        Hooks->GetLocalAddress(Address);
+    }
+#endif
+}
+
+_IRQL_requires_max_(DISPATCH_LEVEL)
+void
+QuicBindingGetRemoteAddress(
+    _In_ QUIC_BINDING* Binding,
+    _Out_ QUIC_ADDR* Address
+    )
+{
+    CxPlatSocketGetRemoteAddress(Binding->Socket, Address);
+#if QUIC_TEST_DATAPATH_HOOKS_ENABLED
+    QUIC_TEST_DATAPATH_HOOKS* Hooks = MsQuicLib.TestDatapathHooks;
+    if (Hooks != NULL) {
+        Hooks->GetRemoteAddress(Address);
+    }
+#endif
+}
+
 //
 // Returns TRUE if there are any registered listeners on this binding.
 //

src/core/binding.h

@@ -289,6 +289,26 @@ QuicBindingTraceRundown(
     _In_ QUIC_BINDING* Binding
     );
 
+//
+// Queries the local IP address of the binding.
+//
+_IRQL_requires_max_(DISPATCH_LEVEL)
+void
+QuicBindingGetLocalAddress(
+    _In_ QUIC_BINDING* Binding,
+    _Out_ QUIC_ADDR* Address
+    );
+
+//
+// Queries the remote IP address of the binding.
+//
+_IRQL_requires_max_(DISPATCH_LEVEL)
+void
+QuicBindingGetRemoteAddress(
+    _In_ QUIC_BINDING* Binding,
+    _Out_ QUIC_ADDR* Address
+    );
+
 //
 // Looks up the listener based on the ALPN list. Optionally, outputs the
 // first ALPN that matches.

src/core/connection.c

@@ -1868,9 +1868,7 @@ QuicConnStart(
     }
 
     Connection->State.LocalAddressSet = TRUE;
-    CxPlatSocketGetLocalAddress(
-        Path->Binding->Socket,
-        &Path->LocalAddress);
+    QuicBindingGetLocalAddress(Path->Binding, &Path->LocalAddress);
     QuicTraceEvent(
         ConnLocalAddrAdded,
         "[conn][%p] New Local IP: %!ADDR!",
@@ -5675,8 +5673,8 @@ QuicConnParamSet(
                 Connection,
                 CASTED_CLOG_BYTEARRAY(sizeof(Connection->Paths[0].LocalAddress), &Connection->Paths[0].LocalAddress));
 
-            CxPlatSocketGetLocalAddress(
-                Connection->Paths[0].Binding->Socket,
+            QuicBindingGetLocalAddress(
+                Connection->Paths[0].Binding,
                 &Connection->Paths[0].LocalAddress);
 
             QuicTraceEvent(

src/core/library.c

@@ -1390,7 +1390,7 @@ QuicLibraryLookupBinding(
 #endif
 
         QUIC_ADDR BindingLocalAddr;
-        CxPlatSocketGetLocalAddress(Binding->Socket, &BindingLocalAddr);
+        QuicBindingGetLocalAddress(Binding, &BindingLocalAddr);
 
         if (!QuicAddrCompare(LocalAddress, &BindingLocalAddr)) {
             continue;
@@ -1402,7 +1402,7 @@ QuicLibraryLookupBinding(
             }
 
             QUIC_ADDR BindingRemoteAddr;
-            CxPlatSocketGetRemoteAddress(Binding->Socket, &BindingRemoteAddr);
+            QuicBindingGetRemoteAddress(Binding, &BindingRemoteAddr);
             if (!QuicAddrCompare(RemoteAddress, &BindingRemoteAddr)) {
                 continue;
             }
@@ -1507,7 +1507,7 @@ QuicLibraryGetBinding(
         goto Exit;
     }
 
-    CxPlatSocketGetLocalAddress((*NewBinding)->Socket, &NewLocalAddress);
+    QuicBindingGetLocalAddress(*NewBinding, &NewLocalAddress);
 
     CxPlatDispatchLockAcquire(&MsQuicLib.DatapathLock);
 

src/core/listener.c

@@ -293,9 +293,7 @@ MsQuicListenerStart(
     }
 
     if (PortUnspecified) {
-        CxPlatSocketGetLocalAddress(
-            Listener->Binding->Socket,
-            &BindingLocalAddress);
+        QuicBindingGetLocalAddress(Listener->Binding, &BindingLocalAddress);
         QuicAddrSetPort(
             &Listener->LocalAddress,
             QuicAddrGetPort(&BindingLocalAddress));