Merge pull request #50 from apple/kkuk/teardown-path-after-migration

Tear down old path after migration

Author: kkuk24

Sources/SwiftNetwork/QUIC/Migration.swift

@@ -106,6 +106,7 @@ struct Migration: ~Copyable {
             return
         }
 
+        let oldPath = connection.currentPath
         connection.log.notice("Migrating to path \(path.identifier)")
         connection.currentPath = path
         path.spinValue = connection.initialSpinValue
@@ -128,6 +129,10 @@ struct Migration: ~Copyable {
         }
         // TODO: Handle preferred address migration
 
+        // Remove the path we just migrated away from.
+        if let oldPath, oldPath != path {
+            connection.tearDownMigratedPath(oldPath)
+        }
     }
 
     func probingPathCount(_ connection: QUICConnection) -> Int {
@@ -197,16 +202,7 @@ extension QUICConnection {
             }
             break
         case .unavailable:
-            if path.isOpenForSending, let dcid = path.dcid,
-                let sequence = remoteCIDs.retire(connectionID: dcid)
-            {
-                withPendingItems(
-                    for: .applicationData,
-                    block: {
-                        $0.addRetireConnectionID(FrameRetireConnectionID(sequence: sequence))
-                    }
-                )
-            }
+            retireOutboundCID(forPathGoingAway: path)
             path.changeState(to: .routeUnavailable)
             break
         }
@@ -225,5 +221,35 @@ extension QUICConnection {
             sendFrames(on: path)
         }
     }
+
+    // Retires a path's outbound CID and queues a RETIRE_CONNECTION_ID frame for it.
+    func retireOutboundCID(forPathGoingAway path: QUICPath) {
+        guard path.isOpenForSending, !path.hasPreAssignedCIDs, let dcid = path.dcid,
+            let sequence = remoteCIDs.retire(connectionID: dcid)
+        else {
+            return
+        }
+        withPendingItems(for: .applicationData) {
+            $0.addRetireConnectionID(FrameRetireConnectionID(sequence: sequence))
+        }
+    }
+
+    // Removes a path we migrated away from.
+    func tearDownMigratedPath(_ oldPath: QUICPath) {
+        guard oldPath !== currentPath else {
+            log.fault("Refusing to tear down the current path \(oldPath.identifier)")
+            return
+        }
+        log.notice("Tearing down old path \(oldPath.identifier) after migration")
+
+        retireOutboundCID(forPathGoingAway: oldPath)
+
+        if oldPath.state.isValidStateChange(to: .routeUnavailable) {
+            oldPath.changeState(to: .routeUnavailable)
+        }
+        oldPath.tearDownLowerStack()
+        multiplexingPaths.removeValue(forKey: oldPath.identifier)
+        sendFrames()
+    }
 }
 #endif

Sources/SwiftNetwork/QUIC/QUICPath.swift

@@ -623,6 +623,10 @@ public final class QUICPath: MultiplexingDatagramPath<QUICConnection>, Equatable
             parentProtocol.migration.migrate(to: self, connection: parentProtocol)
         }
     }
+
+    func tearDownLowerStack() {
+        try? lower.invokeDetach(self.reference)
+    }
 }
 
 // Congestion Control access

Tests/QUICTests/MigrationTests.swift

@@ -0,0 +1,114 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the Swift open source project
+//
+// Copyright (c) 2026 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0
+//
+// See LICENSE.txt for license information
+// See CONTRIBUTORS.txt for the list of Swift project authors
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+//===----------------------------------------------------------------------===//
+
+#if !NETWORK_NO_SWIFT_QUIC
+
+import XCTest
+
+#if canImport(SwiftNetwork)
+@_spi(Essentials) @_spi(ProtocolProvider) @testable import SwiftNetwork
+#elseif canImport(Network)
+@_spi(Essentials) @_spi(ProtocolProvider) @testable import Network
+#endif
+
+@available(Network 0.1.0, *)
+let migrationTestsLogPrefixer: LogPrefixer = LogPrefixer("[MigrationTests]")
+
+@available(Network 0.1.0, *)
+final class MigrationTests: XCTestCase {
+    var connection = QUICConnection(context: .implicitContext)
+
+    static let oldCID = QUICConnectionID([0xA1, 0xA2, 0xA3, 0xA4])!
+    static let newCID = QUICConnectionID([0xB1, 0xB2, 0xB3, 0xB4])!
+
+    override func setUp() {
+        let expectation = XCTestExpectation()
+        connection.context.async {
+            try? self.connection.setup(remote: nil, local: nil, parameters: nil, path: nil)
+            self.connection.recovery = Recovery(logPrefixer: migrationTestsLogPrefixer)
+            self.connection.recovery.connection = self.connection
+            expectation.fulfill()
+        }
+        wait(for: [expectation], timeout: 5.0)
+    }
+
+    override func tearDown() {
+        self.connection.currentPath = nil
+    }
+
+    // Builds a path that is open for sending, backed by a lower harness, with its DCID
+    // registered in `remoteCIDs` so it can be retired. `validated` drives it to the
+    // validated state so `migrate(to:)` will accept it.
+    private func makePath(dcid: QUICConnectionID, sequenceNumber: UInt64, validated: Bool) -> QUICPath {
+        let lower = DatagramLowerHarness(identifier: "\(sequenceNumber)", context: .implicitContext)
+        lower.connect()
+        var path = QUICPath(parent: connection)
+        path.set(interface: nil, priority: 1, isInitial: true)  // -> .routeEstablished
+        path.assignDCID(dcid)  // -> .cidAssigned (open for sending)
+        if validated {
+            path.changeState(to: .probing)
+            path.changeState(to: .validated)
+        }
+        try? path.attachLowerProtocol(lower.reference, remote: nil, local: nil, parameters: nil, path: nil)
+        try? connection.remoteCIDs.insert(
+            sequenceNumber: sequenceNumber,
+            connectionID: dcid,
+            token: QUICStatelessResetToken(Array(repeating: UInt8(sequenceNumber & 0xff), count: 16))!
+        )
+        return path
+    }
+
+    func testMigrationRemovesOldPathAndRetiresItsCID() {
+        let expectation = XCTestExpectation()
+        connection.context.async {
+            let oldPath = self.makePath(dcid: Self.oldCID, sequenceNumber: 1, validated: false)
+            let newPath = self.makePath(dcid: Self.newCID, sequenceNumber: 2, validated: true)
+
+            self.connection.currentPath = oldPath
+            self.connection.multiplexingPaths[oldPath.identifier] = oldPath
+            self.connection.multiplexingPaths[newPath.identifier] = newPath
+            let oldPathID = oldPath.identifier
+
+            self.connection.migration.migrate(to: newPath, connection: self.connection)
+
+            // The path we migrated away from is dropped from the connection and its
+            // remote CID is retired.
+            XCTAssertNil(
+                self.connection.multiplexingPaths[oldPathID],
+                "Old path not removed"
+            )
+            XCTAssertNil(
+                self.connection.remoteCIDs.retire(connectionID: Self.oldCID),
+                "Old path CID not retired"
+            )
+
+            // Only the new path remains, and it is now the current path.
+            XCTAssertEqual(
+                self.connection.multiplexingPaths.count,
+                1,
+                "Unexpected path count"
+            )
+            XCTAssertEqual(
+                self.connection.currentPath?.identifier,
+                newPath.identifier,
+                "Current path not switched"
+            )
+
+            expectation.fulfill()
+        }
+        wait(for: [expectation], timeout: 5.0)
+    }
+}
+
+#endif

Tests/SwiftNetworkTests/SwiftNetworkQUICStackTests.swift

@@ -452,6 +452,17 @@ final class SwiftNetworkQUICStackTests: NetTestCase {
                     }
                 }
 
+                // Currently we only keep the current path object and remove all the
+                // other path objects we are migrating away from, so the number of
+                // multiplexing paths should be one.
+                if case .quic(let clientConnection) = clientQUICReference.reference {
+                    XCTAssertEqual(
+                        clientConnection.multiplexingPaths.count,
+                        1,
+                        "Path leaked after migration \(pathIndex)"
+                    )
+                }
+
                 if let dataToSend {
                     Logger.test.info("Writing data to send on path \(pathIndex)")
                     _ = clientUpperHarness?.write(dataToSend)
@@ -538,6 +549,17 @@ final class SwiftNetworkQUICStackTests: NetTestCase {
             dataToSend: Array("Hello World!".utf8)
         )
     }
+
+    func testQUICStackMigrationDoesNotLeakPaths() {
+        let ipv4Client = Endpoint(address: IPv4Address(SwiftNetworkQUICStackTests.localIPv4Address)!, port: 1234)
+        let ipv4Server = Endpoint(address: IPv4Address(SwiftNetworkQUICStackTests.localIPv4Address)!, port: 8080)
+        quicStackHandshake(
+            clientEndpoint: ipv4Client,
+            serverEndpoint: ipv4Server,
+            migrateCount: 4,
+            dataToSend: Array("Hello World!".utf8)
+        )
+    }
 }
 #endif
 #endif