Support Client Migration (Client Side)

src/core/binding.c

@@ -560,46 +560,101 @@
 BOOLEAN
 QuicBindingAddSourceConnectionID(
     _In_ QUIC_BINDING* Binding,
-    _In_ QUIC_CID_HASH_ENTRY* SourceCid
+    _In_ QUIC_CID_SLIST_ENTRY* SourceCid
     )
 {
     return QuicLookupAddLocalCid(&Binding->Lookup, SourceCid, NULL);
 }
 
+_IRQL_requires_max_(DISPATCH_LEVEL)
+BOOLEAN
+QuicBindingAddAllSourceConnectionIDs(
+    _In_ QUIC_BINDING* Binding,
+    _In_ QUIC_CONNECTION* Connection
+    )
+{
+    for (CXPLAT_SLIST_ENTRY* Link = Connection->SourceCids.Next;
+        Link != NULL;
+        Link = Link->Next) {
+
+        QUIC_CID_SLIST_ENTRY* Entry =
+            CXPLAT_CONTAINING_RECORD(
+                Link,
+                QUIC_CID_SLIST_ENTRY,
+                Link);
+        if (!QuicBindingAddSourceConnectionID(Binding, Entry)) {
+            QuicBindingRemoveAllSourceConnectionIDs(Binding, Connection);
+            return FALSE;
+        }
+    }
+
+    return TRUE;
+}
+
 _IRQL_requires_max_(DISPATCH_LEVEL)
 void
 QuicBindingRemoveSourceConnectionID(
     _In_ QUIC_BINDING* Binding,
-    _In_ QUIC_CID_HASH_ENTRY* SourceCid,
-    _In_ CXPLAT_SLIST_ENTRY** Entry
+    _In_ QUIC_CID_HASH_ENTRY* SourceCid
     )
 {
-    QuicLookupRemoveLocalCid(&Binding->Lookup, SourceCid, Entry);
+    QuicLookupRemoveLocalCid(&Binding->Lookup, SourceCid);
 }
 
 _IRQL_requires_max_(DISPATCH_LEVEL)
 void
-QuicBindingRemoveConnection(
+QuicBindingRemoveAllSourceConnectionIDs(
     _In_ QUIC_BINDING* Binding,
     _In_ QUIC_CONNECTION* Connection
     )
 {
-    if (Connection->RemoteHashEntry != NULL) {
-        QuicLookupRemoveRemoteHash(&Binding->Lookup, Connection->RemoteHashEntry);
+    CXPLAT_SLIST_ENTRY EntriesToFree = {0};
+
+    for (CXPLAT_SLIST_ENTRY* Link = Connection->SourceCids.Next;
+        Link != NULL;
+        Link = Link->Next) {
+
+        QUIC_CID_SLIST_ENTRY* Entry =
+            CXPLAT_CONTAINING_RECORD(
+                Link,
+                QUIC_CID_SLIST_ENTRY,
+                Link);
+
+        CXPLAT_SLIST_ENTRY** Link1 = &Entry->HashEntries.Next;
+        while (*Link1 != NULL) {
+            QUIC_CID_HASH_ENTRY* Entry1 = 
+                CXPLAT_CONTAINING_RECORD(
+                    *Link1,
+                    QUIC_CID_HASH_ENTRY,
+                    Link);
+            if (Entry1->Binding == Binding) {
+                QuicBindingRemoveSourceConnectionID(Binding, Entry1);
+                *Link1 = (*Link1)->Next;
+                CxPlatListPushEntry(&EntriesToFree, &Entry1->Link);
+            } else {
+                Link1 = &(*Link1)->Next;
+            }
+        }
+    }
+
+    while (EntriesToFree.Next != NULL) {
+        QUIC_CID_HASH_ENTRY* Entry = 
+            CXPLAT_CONTAINING_RECORD(
+                CxPlatListPopEntry(&EntriesToFree),
+                QUIC_CID_HASH_ENTRY,
+                Link);
+        CXPLAT_FREE(Entry, QUIC_POOL_CIDHASH);
     }
-    QuicLookupRemoveLocalCids(&Binding->Lookup, Connection);
 }
 
 _IRQL_requires_max_(DISPATCH_LEVEL)
 void
-QuicBindingMoveSourceConnectionIDs(
-    _In_ QUIC_BINDING* BindingSrc,
-    _In_ QUIC_BINDING* BindingDest,
-    _In_ QUIC_CONNECTION* Connection
+QuicBindingRemoveRemoteHash(
+    _In_ QUIC_BINDING* Binding,
+    _In_ QUIC_REMOTE_HASH_ENTRY* RemoteHashEntry
     )
 {
-    QuicLookupMoveLocalConnectionIDs(
-        &BindingSrc->Lookup, &BindingDest->Lookup, Connection);
+    QuicLookupRemoveRemoteHash(&Binding->Lookup, RemoteHashEntry);
 }
 
 _IRQL_requires_max_(DISPATCH_LEVEL)
@@ -1280,10 +1335,10 @@
 
     BOOLEAN BindingRefAdded = FALSE;
     CXPLAT_DBG_ASSERT(NewConnection->SourceCids.Next != NULL);
-    QUIC_CID_HASH_ENTRY* SourceCid =
+    QUIC_CID_SLIST_ENTRY* SourceCid =
         CXPLAT_CONTAINING_RECORD(
             NewConnection->SourceCids.Next,
-            QUIC_CID_HASH_ENTRY,
+            QUIC_CID_SLIST_ENTRY,
             Link);
 
     QuicConnAddRef(NewConnection, QUIC_CONN_REF_LOOKUP_RESULT);
@@ -1318,6 +1373,8 @@
         }
         goto Exit;
     }
+    CXPLAT_DBG_ASSERT(NewConnection->RemoteHashEntry != NULL);
+    NewConnection->RemoteHashEntry->Binding = Binding;
 
     QuicWorkerQueueConnection(NewConnection->Worker, NewConnection);
 
@@ -1351,7 +1408,7 @@
 
     } else {
         NewConnection->SourceCids.Next = NULL;
-        CXPLAT_FREE(SourceCid, QUIC_POOL_CIDHASH);
+        CXPLAT_FREE(SourceCid, QUIC_POOL_CIDSLIST);
         QuicConnRelease(NewConnection, QUIC_CONN_REF_LOOKUP_RESULT);
 #pragma prefast(suppress:6001, "SAL doesn't understand ref counts")
         QuicConnRelease(NewConnection, QUIC_CONN_REF_HANDLE_OWNER);

src/core/binding.h

@@ -373,7 +373,18 @@ _IRQL_requires_max_(DISPATCH_LEVEL)
 BOOLEAN
 QuicBindingAddSourceConnectionID(
     _In_ QUIC_BINDING* Binding,
-    _In_ QUIC_CID_HASH_ENTRY* SourceCid
+    _In_ QUIC_CID_SLIST_ENTRY* SourceCid
+    );
+
+//
+// Attempts to insert all the connection's new source CIDs into the binding's
+// lookup table.
+//
+_IRQL_requires_max_(DISPATCH_LEVEL)
+BOOLEAN
+QuicBindingAddAllSourceConnectionIDs(
+    _In_ QUIC_BINDING* Binding,
+    _In_ QUIC_CONNECTION* Connection
     );
 
 //
@@ -383,30 +394,24 @@ _IRQL_requires_max_(DISPATCH_LEVEL)
 void
 QuicBindingRemoveSourceConnectionID(
     _In_ QUIC_BINDING* Binding,
-    _In_ QUIC_CID_HASH_ENTRY* SourceCid,
-    _In_ CXPLAT_SLIST_ENTRY** Entry
+    _In_ QUIC_CID_HASH_ENTRY* SourceCid
     );
 
 //
-// Removes all the connection's source CIDs from the binding's lookup table.
+// Removes all the source CIDs from the binding's lookup table.
 //
 _IRQL_requires_max_(DISPATCH_LEVEL)
 void
-QuicBindingRemoveConnection(
+QuicBindingRemoveAllSourceConnectionIDs(
     _In_ QUIC_BINDING* Binding,
     _In_ QUIC_CONNECTION* Connection
     );
 
-//
-// Moves all the connections source CIDs from the one binding's lookup table to
-// another.
-//
 _IRQL_requires_max_(DISPATCH_LEVEL)
 void
-QuicBindingMoveSourceConnectionIDs(
-    _In_ QUIC_BINDING* BindingSrc,
-    _In_ QUIC_BINDING* BindingDest,
-    _In_ QUIC_CONNECTION* Connection
+QuicBindingRemoveRemoteHash(
+    _In_ QUIC_BINDING* Binding,
+    _In_ QUIC_REMOTE_HASH_ENTRY* RemoteHashEntry
     );
 
 //
@@ -516,3 +521,16 @@ QuicRetryTokenDecrypt(
     CxPlatDispatchLockRelease(&MsQuicLib.StatelessRetryKeysLock);
     return QUIC_SUCCEEDED(Status);
 }
+
+//
+// Helper to get the owning QUIC_BINDING for the lookup module.
+//
+inline
+_Ret_notnull_
+QUIC_BINDING*
+QuicLookupGetBinding(
+    _In_ QUIC_LOOKUP* Lookup
+    )
+{
+    return CXPLAT_CONTAINING_RECORD(Lookup, QUIC_BINDING, Lookup);
+}

src/core/cid.h

@@ -163,12 +163,22 @@ typedef struct QUIC_CID_LIST_ENTRY {
 #define QUIC_CID_VALIDATE_NULL(Conn, Cid) UNREFERENCED_PARAMETER(Cid)
 #endif
 
+typedef struct QUIC_CID_SLIST_ENTRY {
+
+    CXPLAT_SLIST_ENTRY Link;
+    QUIC_CONNECTION* Connection;
+    CXPLAT_SLIST_ENTRY HashEntries;
+    QUIC_CID CID;
+
+} QUIC_CID_SLIST_ENTRY;
+
 typedef struct QUIC_CID_HASH_ENTRY {
 
     CXPLAT_HASHTABLE_ENTRY Entry;
     CXPLAT_SLIST_ENTRY Link;
     QUIC_CONNECTION* Connection;
-    QUIC_CID CID;
+    QUIC_BINDING* Binding;
+    QUIC_CID_SLIST_ENTRY* CID;
 
 } QUIC_CID_HASH_ENTRY;
 
@@ -178,18 +188,19 @@ typedef struct QUIC_CID_HASH_ENTRY {
 //
 inline
 _Success_(return != NULL)
-QUIC_CID_HASH_ENTRY*
+QUIC_CID_SLIST_ENTRY*
 QuicCidNewNullSource(
     _In_ QUIC_CONNECTION* Connection
     )
 {
-    QUIC_CID_HASH_ENTRY* Entry =
-        (QUIC_CID_HASH_ENTRY*)CXPLAT_ALLOC_NONPAGED(
-            sizeof(QUIC_CID_HASH_ENTRY),
-            QUIC_POOL_CIDHASH);
+    QUIC_CID_SLIST_ENTRY* Entry =
+        (QUIC_CID_SLIST_ENTRY*)CXPLAT_ALLOC_NONPAGED(
+            sizeof(QUIC_CID_SLIST_ENTRY),
+            QUIC_POOL_CIDSLIST);
 
     if (Entry != NULL) {
         Entry->Connection = Connection;
+        Entry->HashEntries.Next = NULL;
         CxPlatZeroMemory(&Entry->CID, sizeof(Entry->CID));
     }
 
@@ -201,23 +212,24 @@ QuicCidNewNullSource(
 //
 inline
 _Success_(return != NULL)
-QUIC_CID_HASH_ENTRY*
+QUIC_CID_SLIST_ENTRY*
 QuicCidNewSource(
     _In_ QUIC_CONNECTION* Connection,
     _In_ uint8_t Length,
     _In_reads_(Length)
         const uint8_t* const Data
     )
 {
-    QUIC_CID_HASH_ENTRY* Entry =
-        (QUIC_CID_HASH_ENTRY*)
+    QUIC_CID_SLIST_ENTRY* Entry =
+        (QUIC_CID_SLIST_ENTRY*)
         CXPLAT_ALLOC_NONPAGED(
-            sizeof(QUIC_CID_HASH_ENTRY) +
+            sizeof(QUIC_CID_SLIST_ENTRY) +
             Length,
-            QUIC_POOL_CIDHASH);
+            QUIC_POOL_CIDSLIST);
 
     if (Entry != NULL) {
         Entry->Connection = Connection;
+        Entry->HashEntries.Next = NULL;
         CxPlatZeroMemory(&Entry->CID, sizeof(Entry->CID));
         Entry->CID.Length = Length;
         if (Length != 0) {