Don't gate PTO packets on address validation

Sources/SwiftNetwork/QUIC/Recovery.swift

@@ -363,7 +363,11 @@ struct Recovery: ~Copyable, PrefixedLoggable, NonCopyableTimerUser {
         ) {
             if sentEntry.lostTime == .zero && sentEntry.packet.isInFlightEligible {
                 if sentEntry.packet.isAckEliciting {
-                    ackElicitingPacketsInFlight -= 1
+                    if ackElicitingPacketsInFlight > 0 {
+                        ackElicitingPacketsInFlight -= 1
+                    } else {
+                        log.fault("Cannot decrement ackElicitingPacketsInFlight below zero")
+                    }
                     let number = sentEntry.packet.number
                     log.datapath(
                         "Ack eliciting packet \(number) acked, decrementing ackElicitingPacketsInFlight to: \(ackElicitingPacketsInFlight)"
@@ -656,14 +660,16 @@ struct Recovery: ~Copyable, PrefixedLoggable, NonCopyableTimerUser {
             _ packets: consuming NetworkUniqueDeque<SentPacketRecord>,
             connection: QUICConnection
         ) -> Bool {
-            var packets = packets
-            guard !packets.isEmpty else {
+            if packets.isEmpty {
                 return false
             }
+
+            var packets = packets
             while !packets.isEmpty {
-                let packet = packets.remove(at: 0)
+                let packet = packets.removeFirst()
                 sentPacket(packet, time: connection.now, connection: connection)
             }
+
             return true
         }
     }
@@ -910,7 +916,11 @@ struct Recovery: ~Copyable, PrefixedLoggable, NonCopyableTimerUser {
                     // and set the lost time.
                     entry.lostTime = timeNow
                     if entry.packet.isAckEliciting {
-                        ackElicitingPacketsInFlight -= 1
+                        if ackElicitingPacketsInFlight > 0 {
+                            ackElicitingPacketsInFlight -= 1
+                        } else {
+                            connection.log.fault("Cannot decrement ackElicitingPacketsInFlight below zero")
+                        }
                     }
                     lostPackets.append(entry.packet.identifier)
                     Recovery.logAckElicitingPacketsInFlight(
@@ -1114,42 +1124,41 @@ struct Recovery: ~Copyable, PrefixedLoggable, NonCopyableTimerUser {
 
     mutating func sendPTO(connection: QUICConnection, path: QUICPath) {
         var sentPTO = false
-        let (_, pnSpace) = getEarliestTime(
-            earliestTimeType: EarliestTimeType.lastSentAckElicitingTime,
-            connection: connection
-        )
-        var hasAckEliciting = false
-        connection.withPendingItems(for: pnSpace) { pendingItems in
-            hasAckEliciting = pendingItems.hasAckElicitingPendingItems
-        }
-        let peerCompletedValidation = peerCompletedValidation(connection: connection)
-        var shouldClearTimer = false
+
         var discardInitialRecoveryState = false
-        var hadAckElicitingInFlight = false
         applyToAllInnerStatesMutable { innerState, packetNumberSpace in
             let ackElicitingPacketsInFlight = innerState.ackElicitingPacketsInFlight
-            if ackElicitingPacketsInFlight == 0 {
+            guard ackElicitingPacketsInFlight > 0 else {
                 return
             }
-            hadAckElicitingInFlight = true
+
             connection.log.datapath(
                 "PTO \(path.recoveryState.PTOCount) (\(packetNumberSpace)) fired on path \(path.identifier) with \(ackElicitingPacketsInFlight) ack-eliciting packets in flight"
             )
+
+            let hasAckEliciting = connection.withPendingItems(for: packetNumberSpace) {
+                $0.hasAckElicitingPendingItems
+            }
+
             if hasAckEliciting {
                 connection.log.datapath("Sending next frames with new data as PTOs")
-                sentPTO = true
                 let packets = connection.sendFramesFromRecovery(
                     on: path,
                     ignoreCongestionWindow: true,
                     discardInitialRecoveryState: &discardInitialRecoveryState
                 )
-                if !innerState.recordSentPackets(packets, connection: connection) {
+                // Only a recorded packet counts as a probe; the pending items may write no payload.
+                if innerState.recordSentPackets(packets, connection: connection) {
+                    sentPTO = true
+                } else {
                     connection.log.datapath(
                         "Unable to force send PTOs, likely flow-controlled or unavailable"
                     )
                 }
-            } else if ackElicitingPacketsInFlight > 0 {
+
+            } else {
                 connection.log.datapath("Retransmitting two tail-packets as PTO")
+
                 var addedPackets = 0
                 let packetCount = innerState.outstandingPackets.count
                 for i in 0..<packetCount {
@@ -1175,6 +1184,7 @@ struct Recovery: ~Copyable, PrefixedLoggable, NonCopyableTimerUser {
                             discardInitialRecoveryState: &discardInitialRecoveryState
                         )
                     }
+
                     if let packets, innerState.recordSentPackets(packets, connection: connection) {
                         sentPTO = true
                         addedPackets += 1
@@ -1190,47 +1200,43 @@ struct Recovery: ~Copyable, PrefixedLoggable, NonCopyableTimerUser {
         }
 
         if !sentPTO {
+            if totalAckElicitingPacketsInFlight == 0, peerCompletedValidation(connection: connection) {
+                // Nothing is in flight to probe for, so the state must have changed between arming the
+                // timer and it firing (e.g. the handshake completed and cleared the in-flight packets).
+                // `resetTimer` cancels the timer for this state once we return.
+                connection.log.fault("PTO fired after validation")
+                return
+            }
+
+            // Send an ack-eliciting probe because either:
+            // - packets are in flight, RFC 9002 Section 6.2.4 requires a probe
+            // - nothing is in flight and the peer has not validated our address, RFC 9002 Section 6.2.2.1
+            // requires an anti-deadlock packet to unblock the server.
+            // The PING is padded when it goes out in an initial packet.
             connection.log.datapath("Sending a PING as PTO")
-            if peerCompletedValidation {
-                if hadAckElicitingInFlight {
-                    // A PTO fired with ack-eliciting packets still in flight, but the probe could
-                    // not be sent because the retransmit was flow-controlled or otherwise
-                    // unavailable. Under heavy loss this is a legitimate loss-recovery outcome
-                    // rather than an anomaly, so log it at error level instead of faulting. The
-                    // timer is still cleared below so the PTO does not immediately re-arm and spin.
-                    connection.log.error("PTO fired after validation but could not send probe, likely flow-controlled")
-                } else {
-                    // A PTO fired after validation with nothing ack-eliciting in flight. The PTO
-                    // timer should not have been armed in that state, so log a fault.
-                    connection.log.fault("PTO fired after validation")
-                }
-                shouldClearTimer = true
-            } else {
-                // Anti deadlock PING frame (i.e PADDED PING). The PING will be padded when we send an initial packet.
-                let pnSpace =
-                    !connection.receivedHandshakePacket
-                    ? PacketNumberSpace.initial : PacketNumberSpace.applicationData
-                connection.withPendingItems(for: pnSpace) { item in
-                    item.ping = true
-                }
-                sentPTO = true
-                withMutableInnerState(packetNumberSpace: pnSpace) { innerState in
-                    let packets = connection.sendFramesFromRecovery(
-                        on: path,
-                        ignoreCongestionWindow: true,
-                        discardInitialRecoveryState: &discardInitialRecoveryState
+            let packetNumberSpace =
+                !connection.receivedHandshakePacket
+                ? PacketNumberSpace.initial : PacketNumberSpace.applicationData
+            connection.withPendingItems(for: packetNumberSpace) { item in
+                item.ping = true
+            }
+
+            sentPTO = true
+
+            withMutableInnerState(packetNumberSpace: packetNumberSpace) { innerState in
+                let packets = connection.sendFramesFromRecovery(
+                    on: path,
+                    ignoreCongestionWindow: true,
+                    discardInitialRecoveryState: &discardInitialRecoveryState
+                )
+                if !innerState.recordSentPackets(packets, connection: connection) {
+                    connection.log.datapath(
+                        "Unable to force send PTOs, likely flow-controlled or unavailable"
                     )
-                    if !innerState.recordSentPackets(packets, connection: connection) {
-                        connection.log.datapath(
-                            "Unable to force send PTOs, likely flow-controlled or unavailable"
-                        )
-                    }
                 }
             }
         }
-        if shouldClearTimer {
-            setTimer(delay: .zero, connection: connection)
-        }
+
         if discardInitialRecoveryState {
             self.resetPNSpace(packetNumberSpace: .initial, connection: connection)
             connection.withCurrentPath {
@@ -1329,13 +1335,8 @@ struct Recovery: ~Copyable, PrefixedLoggable, NonCopyableTimerUser {
     }
 
     mutating func resetTimer(connection: QUICConnection) {
-        var ackElicitingPacketsInFlight = 0
         // if there are ack eliciting packets on any of the innerStates, the L4S error should not be emitted
-        applyToAllInnerStatesImmutable { innerState, _ in
-            ackElicitingPacketsInFlight += innerState.ackElicitingPacketsInFlight
-        }
-
-        if ackElicitingPacketsInFlight == 0 && peerCompletedValidation(connection: connection) {
+        if totalAckElicitingPacketsInFlight == 0 && peerCompletedValidation(connection: connection) {
             log.datapath("No ack eliciting packets in flight, cancelling timer")
             setTimer(delay: .zero, connection: connection)
             connection.withCurrentPath { path in
@@ -1417,6 +1418,16 @@ struct Recovery: ~Copyable, PrefixedLoggable, NonCopyableTimerUser {
         return hasOutstandingPackets
     }
 
+    // The PTO timer is shared across packet number spaces, so decisions about whether there is
+    // anything left to probe for consider every space.
+    var totalAckElicitingPacketsInFlight: Int {
+        var totalAckElicitingPacketsInFlight = 0
+        applyToAllInnerStatesImmutable { innerState, _ in
+            totalAckElicitingPacketsInFlight += innerState.ackElicitingPacketsInFlight
+        }
+        return totalAckElicitingPacketsInFlight
+    }
+
     mutating func resetPNSpace(
         packetNumberSpace: PacketNumberSpace,
         connection: QUICConnection

Tests/QUICTests/RecoveryTests.swift

@@ -653,6 +653,120 @@ final class RecoveryTests: XCTestCase {
         XCTAssertEqual(path.recoveryState.PTOCount, 3)
         XCTAssertGreaterThan(connection.recovery.computedTimeout, .milliseconds(7900))
     }
+
+    // A PTO with ack-eliciting data in flight must emit a probe, even when the only outstanding
+    // packet carries STREAM data for a now-closed flow (so it can't be rebuilt) and the connection
+    // is validated; otherwise `sendPTO` sends nothing and the connection makes no progress until the
+    // idle timeout closes it.
+    func testValidatedPTOProbesWhenTailRetransmitProducesNothing() {
+        // Register a flow and close it, so its STREAM data can never be rebuilt for retransmission.
+        let stream = QUICStreamInstance(parent: connection, inbound: true)
+        stream.setup(streamID: QUICStreamID(0), logPrefixer: recoveryTestsLogPrefixer)
+        connection.multiplexedFlows[stream.identifier] = stream
+        stream.closed = true
+        XCTAssertFalse(stream.isOpen)
+
+        // A single ack-eliciting application-data packet is outstanding, carrying only that flow's
+        // STREAM data. This keeps ackElicitingPacketsInFlight > 0 (so the PTO stays armed) while
+        // being unrebuildable once the flow is closed.
+        var packet = SentPacketRecord()
+        packet.identifier = .init(space: .applicationData, number: 0)
+        packet.isInFlightEligible = true
+        packet.isAckEliciting = true
+        packet.totalLength = 20 + 96
+        packet.sentPath = connection.currentPath?.identifier ?? .none
+        packet.transmittedItems.sentStreams.append(
+            TransmittedItems.SentStream(
+                flowID: stream.identifier,
+                streamID: QUICStreamID(0),
+                offset: 0,
+                length: 32,
+                isFinal: true
+            )
+        )
+        XCTAssertTrue(packet.transmittedItems.hasRetransmissibleItems)
+        sentPacket(packet, connection: connection)
+
+        connection.recovery.withImmutableInnerState(packetNumberSpace: .applicationData) { innerState in
+            XCTAssertEqual(innerState.ackElicitingPacketsInFlight, 1)
+        }
+
+        // Establish (address-validated) connection: peerCompletedValidation must be true. This is
+        // the condition under which the anti-deadlock PING is incorrectly skipped.
+        connection.recovery.received1RTTAck = true
+        XCTAssertTrue(connection.recovery.peerCompletedValidation(connection: connection))
+        XCTAssertEqual(path.recoveryState.PTOCount, 0)
+
+        // Fire the PTO with no new ack-eliciting data pending.
+        let expectation = XCTestExpectation()
+        self.connection.context.async {
+            self.connection.withCurrentPath { path in
+                self.connection.recovery.sendPTO(connection: self.connection, path: path)
+            }
+            expectation.fulfill()
+        }
+        wait(for: [expectation], timeout: 5.0)
+
+        // A probe must have been recorded, taking the packets in flight to two, and the PTO counted.
+        XCTAssertEqual(
+            connection.recovery.totalAckElicitingPacketsInFlight,
+            2,
+            "PTO produced no probe for a closed-flow tail packet"
+        )
+        XCTAssertEqual(path.recoveryState.PTOCount, 1)
+    }
+
+    // `sendPTO` must emit a probe; otherwise the PTO makes no progress. A pending item whose flow was
+    // torn down writes no payload, so ensure the probe only counts once the packet is recorded.
+    func testPTOProbesWhenNewDataProducesNothing() {
+        // A stream queued for service whose flow has since been torn down: it is absent from
+        // `multiplexedFlows`, so writing it produces no payload.
+        let unregisteredStream = QUICStreamInstance(parent: connection, inbound: true)
+        unregisteredStream.setup(streamID: QUICStreamID(0), logPrefixer: recoveryTestsLogPrefixer)
+        XCTAssertNil(connection.flow(for: unregisteredStream.identifier))
+        connection.withPendingItems(for: .initial) { pendingItems in
+            pendingItems.streamsToService.append(unregisteredStream.identifier)
+            pendingItems.stream = true
+        }
+
+        // Recovery still sees new ack-eliciting data, so the PTO sends that rather than retransmitting.
+        let hasPendingAckEliciting = connection.withPendingItems(for: .initial) {
+            $0.hasAckElicitingPendingItems
+        }
+        XCTAssertTrue(hasPendingAckEliciting)
+
+        // One ack-eliciting packet outstanding, so the PTO is armed and the per-space loop runs.
+        var packet = SentPacketRecord()
+        packet.identifier = .init(space: .initial, number: 0)
+        packet.isInFlightEligible = true
+        packet.isAckEliciting = true
+        packet.totalLength = 20 + 96
+        packet.sentPath = connection.currentPath?.identifier ?? .none
+
+        sentPacket(packet, connection: connection)
+
+        connection.recovery.withImmutableInnerState(packetNumberSpace: .initial) { innerState in
+            XCTAssertEqual(innerState.ackElicitingPacketsInFlight, 1)
+        }
+
+        let expectation = XCTestExpectation()
+        self.connection.context.async {
+            self.connection.withCurrentPath { path in
+                self.connection.recovery.sendPTO(connection: self.connection, path: path)
+            }
+            expectation.fulfill()
+        }
+        wait(for: [expectation], timeout: 5.0)
+
+        // Ensure a probe has been recorded, taking the packets in flight to two.
+        connection.recovery.withImmutableInnerState(packetNumberSpace: .initial) { innerState in
+            XCTAssertEqual(
+                innerState.ackElicitingPacketsInFlight,
+                2,
+                "PTO reported a probe but no packet was sent"
+            )
+        }
+    }
 }
 
 #endif