Fix client hang when server closes during handshake

Sources/SwiftNetwork/QUIC/QUICConnection.swift

@@ -1833,6 +1833,10 @@ public final class QUICConnection: ManyToManyApplicationStreamProtocol,
             )
             if !continueProcessing {
                 frame.finalize(success: true)
+                // close(with:) only records the error while deferClosing is set, so close here
+                if closeError != nil {
+                    close(in: &eventContext)
+                }
                 break
             }
 

Sources/SwiftNetwork/QUIC/SendItems.swift

@@ -1503,6 +1503,16 @@ extension FrameApplicationClose: SendableItem {
             )
             return
         }
+        if pendingItems.packetNumberSpace != .applicationData {
+            // RFC 9000 §10.2.3: Initial and Handshake packets must not reveal application state, so
+            // send a CONNECTION_CLOSE with APPLICATION_ERROR and no reason instead
+            let errorCode = UInt64(QUICTransportError.QUICTransportErrorCode.applicationError.rawValue)
+            try FrameConnectionClose.write(frame: &frame, stats: &stats, errorCode: errorCode, frameType: nil)
+            shorthandFrames?.append(
+                FrameConnectionClose.toShorthandLogEntry(errorCode: errorCode, frameType: nil, reason: "")
+            )
+            return
+        }
         let errorCode = UInt64(error.code)
         let errorReason = error.reason
         try write(

Tests/SwiftNetworkTests/QUICTestHarness.swift

@@ -158,6 +158,7 @@ class QUICTestHarness {
         serverLinkDelay: NetworkDuration = .zero,
         clientDrops: DatagramDrops? = nil,
         serverDrops: DatagramDrops? = nil,
+        waitForServerConnected: Bool = true,  // false: return once the client alone is connected
         timeout: TimeInterval = 5.0,
         clientOptions: ProtocolOptions<QUICProtocol> = QUICProtocol.options(),
         serverOptions: ProtocolOptions<QUICProtocol> = QUICProtocol.options(),
@@ -408,14 +409,17 @@ class QUICTestHarness {
             serverHarness.start { connected in
                 if connected {
                     serverConnected = true
-                    if expectHandshakeError == nil {
+                    if expectHandshakeError == nil, waitForServerConnected {
                         handshakeExpectation.fulfill()  // server transitions to connected last, wait for it
                     }
                 }
             }
             clientHarness.start { connected in
                 if connected {
                     clientConnected = true
+                    if expectHandshakeError == nil, !waitForServerConnected {
+                        handshakeExpectation.fulfill()
+                    }
                 }
             }
         }
@@ -435,9 +439,9 @@ class QUICTestHarness {
         }
 
         XCTAssertTrue(clientConnected, "QUIC client failed to become connected")
-        XCTAssertTrue(serverConnected, "QUIC server failed to become connected")
+        XCTAssertTrue(serverConnected || !waitForServerConnected, "QUIC server failed to become connected")
         XCTAssertNotNil(state, "Result cannot be nil here for the tests to proceed")
-        guard clientConnected, serverConnected else {
+        guard clientConnected, serverConnected || !waitForServerConnected else {
             XCTFail("This test cannot continue without both client and server being connected")
             throw NetworkError.posix(EINVAL)
         }
@@ -1026,6 +1030,7 @@ class QUICTestHarness {
         serverLinkDelay: NetworkDuration = .zero,
         clientDrops: DatagramDrops? = nil,
         serverDrops: DatagramDrops? = nil,
+        waitForServerConnected: Bool = true,
         clientReadChunkSize: Int = Int.max,
         timeout: TimeInterval = 5.0,
         applicationError: UInt64? = nil,
@@ -1062,6 +1067,7 @@ class QUICTestHarness {
                 serverLinkDelay: serverLinkDelay,
                 clientDrops: clientDrops,
                 serverDrops: serverDrops,
+                waitForServerConnected: waitForServerConnected,
                 timeout: timeout,
                 clientOptions: clientOptions,
                 serverOptions: serverOptions,

Tests/SwiftNetworkTests/SwiftNetworkQUICUngracefulCloseTests.swift

@@ -129,6 +129,63 @@ final class SwiftNetworkQUICUngracefulCloseTests: NetTestCase {
         )
     }
 
+    func testQUICServerApplicationCloseBeforeHandshakeComplete() throws {
+        // Drop all packets to the server after the first. The client completes the handshake, but the
+        // server never sees the client's Finished, so it closes while its key state is still .handshake
+        QUICTestHarness().runQUICTest(
+            clientDrops: .init(1...Int.max),
+            waitForServerConnected: false,
+            afterHandshake: { harness in
+                let expectation = XCTestExpectation(description: "Wait for client disconnect")
+                harness.context.async {
+                    harness.state?.clientHarness.waitForDisconnected {
+                        // RFC 9000 §10.2.3: APPLICATION_CLOSE must not be sent in a Handshake packet
+                        XCTAssertNotEqual(
+                            harness.state?.clientInstance.closeFrameType,
+                            .applicationClose,
+                            "Server sent APPLICATION_CLOSE in a Handshake packet"
+                        )
+                        expectation.fulfill()
+                    }
+                    harness.state?.serverHarness.stop(error: .init(quicApplicationError: 10, reason: "test"))
+                }
+                // Well below the 30 s idle timeout the client falls back to if it loses the close
+                self.wait(for: [expectation], timeout: 5.0)
+            }
+        )
+    }
+
+    func testQUICClientClosesOnFrameNotAllowedDuringHandshake() throws {
+        // Same setup as above, so the server still sends Handshake packets. RFC 9000 §12.4: a frame
+        // that is not permitted in a Handshake packet is a PROTOCOL_VIOLATION for the client
+        QUICTestHarness().runQUICTest(
+            clientDrops: .init(1...Int.max),
+            waitForServerConnected: false,
+            afterHandshake: { harness in
+                let expectation = XCTestExpectation(description: "Wait for the client to handle the packet")
+                harness.context.async {
+                    if let server = harness.state?.serverInstance {
+                        server.fromExternal { eventContext in
+                            // HANDSHAKE_DONE is only allowed in 1-RTT packets
+                            server.withPendingItems(for: .handshake) { $0.handshakeDone = true }
+                            server.sendFrames(in: &eventContext)
+                        }
+                    }
+                    // Queued behind the delivery of that packet. A later check would pass even without
+                    // closing on the violation, because the next server packet also closes the client
+                    harness.context.async {
+                        XCTAssertTrue(
+                            harness.state?.clientHarness.receivedDisconnected ?? false,
+                            "Client did not close on a frame not allowed in a Handshake packet"
+                        )
+                        expectation.fulfill()
+                    }
+                }
+                self.wait(for: [expectation], timeout: 5.0)
+            }
+        )
+    }
+
     func testQUICStatelessResetTokenWithSCID() throws {
         // This test seeds the stateless reset token and the SCID on the server.
         // Then sends a stateless reset packet to the client with the seeded token and verifies the connection closes.