fix(quic): RFC 9002 §6.1.2 loss-detection timer + PSK rejection signal

Round 9 — the two largest items remaining from the audit. * RFC 9002 §6.1.2 timer-driven loss detection. [detectAndRemoveLost] now returns a [Result] data class carrying both the list of lost packets and the absolute monotonic deadline at which the next earliest sub-largest in-flight packet will cross the time threshold. The parser stores that deadline on each [LevelState.nextLossTimeMs] per encryption level, and the driver's send loop now sleeps for `min(ptoDeadline, minNextLossTimeAcrossLevels) - now`. On expiry, the driver distinguishes: * Loss-timer wake → run [detectAndRemoveLost] across all levels; declare time-threshold-lost packets and dispatch their tokens. No probe budget, no exponential backoff. * PTO wake → existing [handlePtoFired] path (probe + backoff). Pre-fix tail loss waited for the full PTO (often 5x the time threshold) before retransmitting because we had no event between ACK arrivals to fire loss detection. Now `9/8 * max_rtt` is the ceiling. * TLS PSK rejection: instead of the prior generic [QuicCodecException] ("server rejected PSK; full-handshake fallback not implemented"), raise a typed [PskRejectedException] subclass. Application reconnect logic can selectively catch this and retry the handshake without cached resumption state — a path that's correct by construction (fresh ClientHello, no PSK history, no transcript-rebuild complexity). [QuicCodecException] is now `open` so the subclass can extend it. In-place fallback (clear early secret, rebuild transcript without PSK extension, replay derivation) remains deferred — the subtle transcript-hash discrepancies it could introduce would be much harder to debug than a hard failure that the application intentionally turns into a retry. All 269 :quic:jvmTest tests pass. BUILD SUCCESSFUL in 47s. https://claude.ai/code/session_01AhGvbMV8uPRse3TmAGaddM
2026-05-09 00:23:37 +00:00
parent 90f59a3351
commit 28c1a355da
8 changed files with 194 additions and 39 deletions
@@ -287,7 +287,7 @@ class QuicReader(
    }
 }

-class QuicCodecException(
+open class QuicCodecException(
    message: String,
    cause: Throwable? = null,
 ) : RuntimeException(message, cause)
@@ -85,6 +85,21 @@ class LevelState {
     */
    var largestAckedSentTimeMs: Long? = null

+    /**
+     * RFC 9002 §6.1.2 timer-driven loss-detection deadline for this
+     * encryption level. Absolute monotonic time at which the next
+     * earliest in-flight sub-largest packet will cross the time
+     * threshold. The driver's send loop uses
+     * `min(nextLossTimeMs across levels, ptoDeadline)` as its wakeup
+     * deadline so tail-loss recovery doesn't wait for the next ACK
+     * or PTO. Null when no sub-largest packets are in flight (or
+     * before the first ACK arrives).
+     *
+     * Updated by [com.vitorpamplona.quic.connection.QuicConnectionParser]
+     * each time a fresh ACK runs `detectAndRemoveLost`.
+     */
+    var nextLossTimeMs: Long? = null
+
    /**
     * RFC 9001 §4.9: latches true once [discardKeys] runs. Used by
     * the writer / parser to short-circuit operations on a discarded
@@ -131,6 +146,7 @@ class LevelState {
        sentPackets.clear()
        largestAckedPn = null
        largestAckedSentTimeMs = null
+        nextLossTimeMs = null
        keysDiscarded = true
    }

@@ -163,6 +179,7 @@ class LevelState {
        sentPackets.clear()
        largestAckedPn = null
        largestAckedSentTimeMs = null
+        nextLossTimeMs = null
        keysDiscarded = false
        this.sendProtection = sendProtection
        this.receiveProtection = receiveProtection
@@ -204,6 +221,7 @@ class LevelState {
        sentPackets.clear()
        largestAckedPn = null
        largestAckedSentTimeMs = null
+        nextLossTimeMs = null
        keysDiscarded = false
        this.sendProtection = sendProtection
        this.receiveProtection = receiveProtection
@@ -268,18 +268,78 @@ class QuicConnectionDriver(
                    .coerceAtLeast(1L)
            val backoff = (1L shl connection.consecutivePtoCount.coerceAtMost(6))
            val ptoMillis = (ptoBaseMs * backoff).coerceAtMost(60_000L)
-            // Suspend until either: a wakeup arrives, or the PTO timer expires.
+            // RFC 9002 §6.1.2 timer-driven loss detection: take the
+            // earliest of (PTO deadline, next-loss-time across levels).
+            // Without this, tail loss waits for the PTO instead of the
+            // shorter `9/8 * max_rtt` time threshold — visible to the
+            // user as a recovery delay equal to the PTO minus that
+            // threshold (often ~5x worse than necessary on lossy
+            // links). A null nextLossTime contributes nothing.
+            val nowForLossTimer = nowMillis()
+            val nextLossDelta =
+                listOfNotNull(
+                    connection.initial.nextLossTimeMs,
+                    connection.handshake.nextLossTimeMs,
+                    connection.application.nextLossTimeMs,
+                ).minOrNull()?.let { (it - nowForLossTimer).coerceAtLeast(0L) }
+            val sleepMillis =
+                if (nextLossDelta != null && nextLossDelta < ptoMillis) nextLossDelta else ptoMillis
+            // Suspend until either: a wakeup arrives, or the timer expires.
            val woke =
-                withTimeoutOrNull(ptoMillis) {
+                withTimeoutOrNull(sleepMillis) {
                    sendWakeup.receive()
                    Unit
                }
            if (woke == null) {
-                handlePtoFired(connection)
+                // Distinguish loss-timer expiry from PTO expiry. A
+                // loss-timer wake just runs `detectAndRemoveLost`
+                // across the levels — newly time-threshold-lost
+                // packets get their tokens re-queued for retransmit on
+                // the next drain. A PTO wake additionally bumps the
+                // consecutive-PTO counter and arms the probe budget.
+                val pickedLossTimer = nextLossDelta != null && nextLossDelta < ptoMillis
+                if (pickedLossTimer) {
+                    handleLossTimerFired(connection)
+                } else {
+                    handlePtoFired(connection)
+                }
            }
        }
    }

+    /**
+     * RFC 9002 §6.1.2 timer-driven loss detection. Re-runs
+     * `detectAndRemoveLost` across each encryption level; any newly
+     * time-threshold-lost packets dispatch their tokens to the
+     * pending retransmit queues, and the next drain emits them.
+     * Cheaper than [handlePtoFired] because no probe budget /
+     * exponential backoff is needed — the peer ACKs that fix the
+     * loss state arrive on their normal cadence.
+     */
+    private fun handleLossTimerFired(conn: QuicConnection) {
+        val nowMs = conn.nowMillis()
+        runLossDetectionForLevel(conn, conn.initial, EncryptionLevel.INITIAL, nowMs)
+        runLossDetectionForLevel(conn, conn.handshake, EncryptionLevel.HANDSHAKE, nowMs)
+        runLossDetectionForLevel(conn, conn.application, EncryptionLevel.APPLICATION, nowMs)
+    }
+
+    private fun runLossDetectionForLevel(
+        conn: QuicConnection,
+        state: LevelState,
+        level: EncryptionLevel,
+        nowMs: Long,
+    ) {
+        val largest = state.largestAckedPn ?: return
+        val result = conn.lossDetection.detectAndRemoveLost(state.sentPackets, largest, nowMs)
+        for (lostPacket in result.lost) {
+            conn.onTokensLost(lostPacket.tokens)
+        }
+        if (result.lost.isNotEmpty()) {
+            conn.qlogObserver.onLossDetected(level, result.lost.map { it.packetNumber })
+        }
+        state.nextLossTimeMs = result.nextLossTimeMs
+    }
+
    /**
     * Cleanly tear down the driver. Runs on [parentScope] so the caller (which
     * may itself live inside the driver's [scope]) isn't cancelled before its
@@ -549,7 +549,7 @@ private fun dispatchFrames(
                    conn.consecutivePtoCount = 0
                }
                state.largestAckedPn?.let { largestAckedPn ->
-                    val lost =
+                    val result =
                        conn.lossDetection.detectAndRemoveLost(
                            sentPackets = state.sentPackets,
                            largestAckedPn = largestAckedPn,
@@ -558,12 +558,18 @@ private fun dispatchFrames(
                    // Step 6: dispatch each lost packet's tokens to the
                    // matching pending* field. The supersede check (lost
                    // value still == advertised) lives inside onTokensLost.
-                    for (lostPacket in lost) {
+                    for (lostPacket in result.lost) {
                        conn.onTokensLost(lostPacket.tokens)
                    }
-                    if (lost.isNotEmpty()) {
-                        conn.qlogObserver.onLossDetected(level, lost.map { it.packetNumber })
+                    if (result.lost.isNotEmpty()) {
+                        conn.qlogObserver.onLossDetected(level, result.lost.map { it.packetNumber })
                    }
+                    // RFC 9002 §6.1.2 timer-driven loss detection. Record
+                    // the earliest pending time-threshold deadline so the
+                    // driver's send loop can wake at that instant rather
+                    // than waiting for the next inbound ACK or PTO. Set
+                    // null when no sub-largest packets remain in flight.
+                    state.nextLossTimeMs = result.nextLossTimeMs
                }
            }

@@ -134,18 +134,30 @@ class QuicLossDetection {
     *       [largestAckedPn], OR
     *     - it was sent more than [lossDelayMs] ago.
     *
-     * Returns the list of lost packets in arbitrary order. Caller
-     * dispatches their [SentPacket.tokens] (step 6).
+     * Returns the list of lost packets PLUS the absolute monotonic
+     * time at which the next time-threshold loss will fire — if any
+     * surviving in-flight packet predates [largestAckedPn]. The
+     * caller schedules a timer at that instant so tail-loss recovery
+     * doesn't have to wait for the PTO. RFC 9002 §6.1.2 (loss-
+     * detection timer): without this, a quiet link's last few packets
+     * sit in `sentPackets` until either the PTO fires or fresh ACKs
+     * arrive. Returning [Result.nextLossTimeMs] = null means there's
+     * nothing left for the time-threshold to fire on.
     */
    fun detectAndRemoveLost(
        sentPackets: MutableMap<Long, SentPacket>,
        largestAckedPn: Long,
        nowMs: Long,
-    ): List<SentPacket> {
-        if (sentPackets.isEmpty()) return emptyList()
+    ): Result {
+        if (sentPackets.isEmpty()) return Result.EMPTY
        val lossDelay = lossDelayMs()
        val lossThresholdSentMs = nowMs - lossDelay
        val lost = mutableListOf<SentPacket>()
+        // Track the EARLIEST sentAtMillis among surviving sub-largest
+        // packets. The next loss-detection timer fires at
+        // `earliestSent + lossDelay` — the earliest moment any of
+        // them crosses the time threshold.
+        var earliestSentSubLargestMs = Long.MAX_VALUE
        val it = sentPackets.entries.iterator()
        while (it.hasNext()) {
            val (pn, pkt) = it.next()
@@ -155,9 +167,29 @@ class QuicLossDetection {
            if (packetThresholdLost || timeThresholdLost) {
                lost += pkt
                it.remove()
+            } else if (pkt.sentAtMillis < earliestSentSubLargestMs) {
+                earliestSentSubLargestMs = pkt.sentAtMillis
            }
        }
-        return lost
+        val nextLossTimeMs =
+            if (earliestSentSubLargestMs == Long.MAX_VALUE) null else earliestSentSubLargestMs + lossDelay
+        return Result(lost, nextLossTimeMs)
+    }
+
+    /**
+     * Outcome of [detectAndRemoveLost]. [lost] is the list of packets
+     * declared lost on this call (for retransmit dispatch);
+     * [nextLossTimeMs] is the absolute monotonic time at which the
+     * next earliest surviving sub-largest packet will cross the
+     * time threshold, or null if none remain.
+     */
+    data class Result(
+        val lost: List<SentPacket>,
+        val nextLossTimeMs: Long?,
+    ) {
+        companion object {
+            internal val EMPTY = Result(emptyList(), null)
+        }
    }

    /**
@@ -358,15 +358,32 @@ class TlsClient(
                val pskExt = sh.extensions.firstOrNull { it.type == TlsConstants.EXT_PRE_SHARED_KEY }
                if (resumption != null) {
                    if (pskExt == null) {
-                        // We offered PSK but server picked full-handshake.
-                        // Plumbing for the fallback path (clear early secret,
-                        // re-run binder-less ClientHello transcript) is real
-                        // work; for now hard-fail. In production we'd want
-                        // to handle gracefully — for the runner's resumption
-                        // testcase the server MUST accept or the test fails
-                        // anyway, so this gate isn't load-bearing.
-                        throw QuicCodecException(
-                            "server rejected PSK; full-handshake fallback not implemented",
+                        // We offered PSK but the server picked full-
+                        // handshake. RFC 8446 §4.2.11 lets the server
+                        // do this freely (rate limit, ticket aged out,
+                        // policy mismatch). Recovering in-place
+                        // requires:
+                        //   1. Discarding the early secret derived
+                        //      from the resumption PSK.
+                        //   2. Rebuilding the transcript without the
+                        //      `pre_shared_key` extension or its
+                        //      binder bytes.
+                        //   3. Replaying the post-ClientHello derivation
+                        //      with the new transcript.
+                        // Each step is touchy enough that a
+                        // subtly-wrong transcript hash would land us on
+                        // a successful handshake with wrong keys, which
+                        // is harder to debug than a hard failure. We
+                        // surface a typed [PskRejectedException] so
+                        // application-layer reconnect logic can drop
+                        // the cached resumption state and retry from
+                        // scratch — that path is correct by
+                        // construction (fresh ClientHello, no PSK
+                        // history). Production callers wrapping this
+                        // client SHOULD catch [PskRejectedException]
+                        // and retry without [resumption].
+                        throw PskRejectedException(
+                            "server rejected PSK; application must retry without resumption",
                        )
                    }
                    val r = QuicReader(pskExt.data)
@@ -588,6 +605,22 @@ class TlsClient(
    }
 }

+/**
+ * Thrown when the server returned a ServerHello without a `pre_shared_key`
+ * extension despite the client offering one in [TlsClient.resumption].
+ * The handshake cannot proceed in-place because the early secret + transcript
+ * were already shaped around the PSK; application-layer reconnect logic
+ * should catch this, drop the cached [TlsResumptionState], and retry the
+ * handshake from scratch. See the call site in
+ * [TlsClient.handleHandshakeMessage] for the full rationale.
+ *
+ * Distinct subclass of [QuicCodecException] so callers can selectively
+ * catch the recoverable case without swallowing genuine protocol errors.
+ */
+class PskRejectedException(
+    message: String,
+) : QuicCodecException(message)
+
 /** Callback interface so the QUIC layer can react to TLS-derived secrets. */
 interface TlsSecretsListener {
    fun onHandshakeKeysReady(
@@ -91,13 +91,13 @@ class RetransmitIntegrationTest {
                client.application.largestAckedPn = futurePn
                client.application.largestAckedSentTimeMs = 1L
                // 4. Run loss detection — msuPn is < futurePn - 3 ⇒ lost.
-                val lost =
+                val result =
                    client.lossDetection.detectAndRemoveLost(
                        sentPackets = client.application.sentPackets,
                        largestAckedPn = futurePn,
                        nowMs = 2L,
                    )
-                val lostMsuPacket = lost.firstOrNull { it.packetNumber == msuPn }
+                val lostMsuPacket = result.lost.firstOrNull { it.packetNumber == msuPn }
                assertNotNull(lostMsuPacket, "msuPn=$msuPn must be declared lost (largestAckedPn=$futurePn, threshold=3)")
                // 5. Dispatch lost tokens — pendingMaxStreamsUni gets set.
                client.onTokensLost(lostMsuPacket.tokens)
@@ -113,10 +113,10 @@ class QuicLossDetectionTest {
        // (Pn 10 is the largest-acked itself — not in flight; it was just removed by drain.
        //  We model "after-drain" by removing 10 from sent before calling detectAndRemoveLost.)
        sent.remove(10L)
-        val lost = ld.detectAndRemoveLost(sent, largestAckedPn = 10L, nowMs = 1L)
+        val result = ld.detectAndRemoveLost(sent, largestAckedPn = 10L, nowMs = 1L)
        // sentAt=0, nowMs=1, lossDelayMs=374 → time threshold cutoff at -373: nothing lost by time threshold
        // (sentAt 0 is NOT <= -373). But packet threshold removes pns 0..6.
-        assertEquals(setOf(0L, 1L, 2L, 3L, 4L, 5L, 6L), lost.map { it.packetNumber }.toSet())
+        assertEquals(setOf(0L, 1L, 2L, 3L, 4L, 5L, 6L), result.lost.map { it.packetNumber }.toSet())
        assertEquals(setOf(7L, 8L, 9L), sent.keys, "pns 7..9 remain in flight; pn 10 was drained earlier")
    }

@@ -132,8 +132,8 @@ class QuicLossDetectionTest {
        sent[1L] = sentPacket(1L, sentAt = 195L) // recent: 195 + 11 = 206, now=200 ⇒ not lost
        // pn 2 is the largest-acked, drained earlier.

-        val lost = ld.detectAndRemoveLost(sent, largestAckedPn = 2L, nowMs = 200L)
-        assertEquals(listOf(0L), lost.map { it.packetNumber })
+        val result = ld.detectAndRemoveLost(sent, largestAckedPn = 2L, nowMs = 200L)
+        assertEquals(listOf(0L), result.lost.map { it.packetNumber })
        assertEquals(setOf(1L), sent.keys)
    }

@@ -145,8 +145,8 @@ class QuicLossDetectionTest {
        val ld = QuicLossDetection()
        val sent = mutableMapOf<Long, SentPacket>()
        sent[5L] = sentPacket(5L, sentAt = 0L)
-        val lost = ld.detectAndRemoveLost(sent, largestAckedPn = 5L, nowMs = 1L)
-        assertTrue(lost.isEmpty())
+        val result = ld.detectAndRemoveLost(sent, largestAckedPn = 5L, nowMs = 1L)
+        assertTrue(result.lost.isEmpty())
        assertNotNull(sent[5L])
    }

@@ -154,8 +154,9 @@ class QuicLossDetectionTest {
    fun emptyMap_returnsEmpty() {
        val ld = QuicLossDetection()
        val sent = mutableMapOf<Long, SentPacket>()
-        val lost = ld.detectAndRemoveLost(sent, largestAckedPn = 100L, nowMs = 1L)
-        assertTrue(lost.isEmpty())
+        val result = ld.detectAndRemoveLost(sent, largestAckedPn = 100L, nowMs = 1L)
+        assertTrue(result.lost.isEmpty())
+        assertEquals(null, result.nextLossTimeMs)
    }

    @Test
@@ -175,11 +176,16 @@ class QuicLossDetectionTest {
                    ),
            )
        // PN 0 is below largestAckedPn=10 - threshold(3) = 7, so lost by packet threshold.
-        val lost = ld.detectAndRemoveLost(sent, largestAckedPn = 10L, nowMs = 1L)
-        assertEquals(1, lost.size)
-        assertEquals(2, lost.single().tokens.size)
-        assertEquals(RecoveryToken.MaxStreamsUni(150L), lost.single().tokens[0])
-        assertEquals(RecoveryToken.MaxData(5_000L), lost.single().tokens[1])
+        val result = ld.detectAndRemoveLost(sent, largestAckedPn = 10L, nowMs = 1L)
+        assertEquals(1, result.lost.size)
+        assertEquals(
+            2,
+            result.lost
+                .single()
+                .tokens.size,
+        )
+        assertEquals(RecoveryToken.MaxStreamsUni(150L), result.lost.single().tokens[0])
+        assertEquals(RecoveryToken.MaxData(5_000L), result.lost.single().tokens[1])
        assertNull(sent[0L])
    }

@@ -192,8 +198,8 @@ class QuicLossDetectionTest {
        ld.onRttSample(largestAckedSentTimeMs = 0L, ackDelayMs = 0L, nowMs = 10L)
        val sent = mutableMapOf<Long, SentPacket>()
        sent[0L] = sentPacket(0L, sentAt = 0L) // old AND below threshold
-        val lost = ld.detectAndRemoveLost(sent, largestAckedPn = 10L, nowMs = 200L)
-        assertEquals(1, lost.size)
+        val result = ld.detectAndRemoveLost(sent, largestAckedPn = 10L, nowMs = 200L)
+        assertEquals(1, result.lost.size)
        assertTrue(sent.isEmpty())
    }
 }