Refactor migration test to use QUIC_MIGRATION_TYPE and improve path m…

…anagement logic
microsoft · Jan 10, 2025 · 64a35a7 · 64a35a7
1 parent 1cd0b78
commit 64a35a7
Show file tree

Hide file tree

Showing 8 changed files with 279 additions and 28 deletions.
diff --git a/src/core/connection.c b/src/core/connection.c
@@ -6531,11 +6531,8 @@ QuicConnRemoveLocalAddress(
 
     QUIC_PATH* Path = &Connection->Paths[PathIndex];
 
-    if (Path->IsActive && Connection->State.Started) {
-        return QUIC_STATUS_INVALID_STATE;
-    }
-
-    if (Path->DestCid != NULL) {
+    if (Path->DestCid != NULL &&
+        Connection->State.Started && Connection->State.HandshakeConfirmed) {
         QuicConnRetireCid(Connection, Path->DestCid);
     }
 
@@ -6546,9 +6543,65 @@ QuicConnRemoveLocalAddress(
     }
 
     if (Connection->PathsCount == 1) {
-        CXPLAT_DBG_ASSERT(!Connection->State.Started);
-        Connection->State.LocalAddressSet = FALSE;
+        if (!Connection->State.Started) {
+            Connection->State.LocalAddressSet = FALSE;
+        } else {
+            QuicTraceEvent(
+                ConnError,
+                "[conn][%p] ERROR, %s.",
+                Connection,
+                "Last Local Address Removed!");
+            QuicConnSilentlyAbort(Connection);
+            return QUIC_STATUS_ABORTED;
+        }
     } else {
+        if (Path->IsActive) {
+            CXPLAT_DBG_ASSERT(PathIndex == 0);
+            if (!Connection->State.Started) {
+                CXPLAT_DBG_ASSERT(Path->DestCid != NULL);
+                CXPLAT_DBG_ASSERT(!Path->DestCid->CID.Retired);
+#if DEBUG
+                QUIC_CID_CLEAR_PATH(Path->DestCid);
+#endif
+                // Move the dest CID to the new active path.
+                QUIC_CID_LIST_ENTRY* DestCid = Path->DestCid;
+                Path->DestCid = NULL;
+                QUIC_PATH* NewActivePath = &Connection->Paths[1];
+                NewActivePath->DestCid = DestCid;
+                QUIC_CID_SET_PATH(Connection, NewActivePath->DestCid, NewActivePath);
+
+                QuicPathSetActive(Connection, NewActivePath);
+                PathIndex = 1; // The removing path is now at index 1.
+            } else if (!Connection->State.HandshakeConfirmed) {
+                QuicTraceEvent(
+                    ConnError,
+                    "[conn][%p] ERROR, %s.",
+                    Connection,
+                    "Active Local Address Removed during Handshake!");
+                QuicConnSilentlyAbort(Connection);
+                return QUIC_STATUS_ABORTED;
+            } else {
+                uint8_t NewActivePathIndex = Connection->PathsCount;
+                for (uint8_t i = 0; i < Connection->PathsCount; ++i) {
+                    if (i != PathIndex && Connection->Paths[i].DestCid != NULL) {
+                        NewActivePathIndex = i;
+                        break;
+                    }
+                }
+                if (NewActivePathIndex == Connection->PathsCount) {
+                    QuicTraceEvent(
+                        ConnError,
+                        "[conn][%p] ERROR, %s.",
+                        Connection,
+                        "No Active Local Address Remaining!");
+                    QuicConnSilentlyAbort(Connection);
+                    return QUIC_STATUS_ABORTED;
+                }
+                QUIC_PATH* NewActivePath = &Connection->Paths[NewActivePathIndex];
+                QuicPathSetActive(Connection, NewActivePath);
+                PathIndex = NewActivePathIndex; // The removing path is now at the new active index.
+            }
+        }
         QuicPathRemove(Connection, PathIndex);
     }
 

diff --git a/src/core/path.c b/src/core/path.c
@@ -231,6 +231,11 @@ QuicConnGetPathForPacket(
         return &Connection->Paths[i];
     }
 
+    if (!QuicConnIsServer(Connection)) {
+        // Client doesn't create a new path.
+        return NULL;
+    }
+
     if (Connection->PathsCount == QUIC_MAX_PATH_COUNT) {
         //
         // See if any old paths share the same remote address, and is just a rebind.
@@ -329,8 +334,8 @@ QuicPathSetActive(
     if (!UdpPortChangeOnly) {
         QuicCongestionControlReset(&Connection->CongestionControl, FALSE);
     }
-    CXPLAT_DBG_ASSERT(Path->DestCid != NULL);
-    CXPLAT_DBG_ASSERT(!Path->DestCid->CID.Retired);
+    CXPLAT_DBG_ASSERT(Connection->Paths[0].DestCid != NULL);
+    CXPLAT_DBG_ASSERT(!Connection->Paths[0].DestCid->CID.Retired);
 }
 
 _IRQL_requires_max_(PASSIVE_LEVEL)

diff --git a/src/test/MsQuicTests.h b/src/test/MsQuicTests.h
@@ -372,11 +372,17 @@ QuicTestProbePath(
     _In_ uint32_t DropPacketCount
     );
 
+typedef enum QUIC_MIGRATION_TYPE {
+    MigrateWithProbe,
+    MigrateWithoutProbe,
+    DeleteAndMigrate,
+} QUIC_MIGRATION_TYPE;
+
 void
 QuicTestMigration(
     _In_ int Family,
     _In_ BOOLEAN ShareBinding,
-    _In_ BOOLEAN PathProbe
+    _In_ QUIC_MIGRATION_TYPE Type
     );
 
 void
@@ -1368,7 +1374,7 @@ typedef struct {
 typedef struct {
     int Family;
     BOOLEAN ShareBinding;
-    BOOLEAN Smooth;
+    QUIC_MIGRATION_TYPE Type;
 } QUIC_RUN_MIGRATION_PARAMS;
 
 #define IOCTL_QUIC_RUN_MIGRATION \

diff --git a/src/test/bin/quic_gtest.cpp b/src/test/bin/quic_gtest.cpp
@@ -1692,11 +1692,11 @@ TEST_P(WithMigrationArgs, Migration) {
         QUIC_RUN_MIGRATION_PARAMS Params = {
             GetParam().Family,
             GetParam().ShareBinding,
-            GetParam().Smooth
+            GetParam().Type
         };
         ASSERT_TRUE(DriverClient.Run(IOCTL_QUIC_RUN_MIGRATION, Params));
     } else {
-        QuicTestMigration(GetParam().Family, GetParam().ShareBinding, GetParam().Smooth);
+        QuicTestMigration(GetParam().Family, GetParam().ShareBinding, GetParam().Type);
     }
 }
 

diff --git a/src/test/bin/quic_gtest.h b/src/test/bin/quic_gtest.h
@@ -961,21 +961,21 @@ class WithProbePathArgs : public testing::Test,
 struct MigrationArgs {
     int Family;
     BOOLEAN ShareBinding;
-    BOOLEAN Smooth;
+    QUIC_MIGRATION_TYPE Type;
     static ::std::vector<MigrationArgs> Generate() {
         ::std::vector<MigrationArgs> list;
         for (int Family : { 4, 6 })
         for (BOOLEAN ShareBinding : { TRUE, FALSE })
-        for (BOOLEAN Smooth : { TRUE, FALSE })
-            list.push_back({ Family, ShareBinding, Smooth });
+        for (QUIC_MIGRATION_TYPE Type : { MigrateWithProbe, MigrateWithoutProbe, DeleteAndMigrate })
+            list.push_back({ Family, ShareBinding, Type });
         return list;
     }
 };
 
 std::ostream& operator << (std::ostream& o, const MigrationArgs& args) {
     return o << (args.Family == 4 ? "v4" : "v6") << "/"
         << (args.ShareBinding ? "ShareBinding" : "not ShareBinding") << "/"
-        << (args.Smooth ? "Smooth" : "not Smooth");
+        << (args.Type ? (args.Type == MigrateWithoutProbe ? "Migrate without Probe" : "Delete and Migrate") : "Migrate with Probe");
 }
 
 class WithMigrationArgs : public testing::Test,

diff --git a/src/test/bin/winkernel/control.cpp b/src/test/bin/winkernel/control.cpp
@@ -975,7 +975,7 @@ QuicTestCtlEvtIoDeviceControl(
             QuicTestMigration(
                 Params->MigrationParams.Family,
                 Params->MigrationParams.ShareBinding,
-                Params->MigrationParams.Smooth));
+                Params->MigrationParams.Type));
         break;
 
     case IOCTL_QUIC_RUN_NAT_PORT_REBIND: