diff --git a/quic/src/commonMain/kotlin/com/vitorpamplona/quic/connection/QuicConnectionWriter.kt b/quic/src/commonMain/kotlin/com/vitorpamplona/quic/connection/QuicConnectionWriter.kt index f1934431f..615eb0a5e 100644 --- a/quic/src/commonMain/kotlin/com/vitorpamplona/quic/connection/QuicConnectionWriter.kt +++ b/quic/src/commonMain/kotlin/com/vitorpamplona/quic/connection/QuicConnectionWriter.kt @@ -68,18 +68,42 @@ fun drainOutbound( } // Drain destructive frame sources into local lists, ONCE. - val initialFrames = collectHandshakeLevelFrames(conn, EncryptionLevel.INITIAL, nowMillis) - val handshakeFrames = collectHandshakeLevelFrames(conn, EncryptionLevel.HANDSHAKE, nowMillis) + val initialContents = collectHandshakeLevelFrames(conn, EncryptionLevel.INITIAL, nowMillis) + val handshakeContents = collectHandshakeLevelFrames(conn, EncryptionLevel.HANDSHAKE, nowMillis) val applicationPkt = buildApplicationPacket(conn, nowMillis) val initialState = conn.initial val handshakeState = conn.handshake - val initialHasContent = initialFrames != null && initialState.sendProtection != null - val handshakeHasContent = handshakeFrames != null && handshakeState.sendProtection != null + val initialHasContent = initialContents != null && initialState.sendProtection != null + val handshakeHasContent = handshakeContents != null && handshakeState.sendProtection != null // Build natural-size first. - val initialNatural = if (initialHasContent) buildLongHeaderFromFrames(conn, EncryptionLevel.INITIAL, initialFrames, padBytes = 0) else null - val handshakeNatural = if (handshakeHasContent) buildLongHeaderFromFrames(conn, EncryptionLevel.HANDSHAKE, handshakeFrames, padBytes = 0) else null + val initialNatural = + if (initialContents != null && initialState.sendProtection != null) { + buildLongHeaderFromFrames( + conn = conn, + level = EncryptionLevel.INITIAL, + frames = initialContents.frames, + tokens = initialContents.tokens, + nowMillis = nowMillis, + padBytes = 0, + ) + } else { + null + } + val handshakeNatural = + if (handshakeContents != null && handshakeState.sendProtection != null) { + buildLongHeaderFromFrames( + conn = conn, + level = EncryptionLevel.HANDSHAKE, + frames = handshakeContents.frames, + tokens = handshakeContents.tokens, + nowMillis = nowMillis, + padBytes = 0, + ) + } else { + null + } val firstPass = listOfNotNull(initialNatural, handshakeNatural, applicationPkt) if (firstPass.isEmpty()) return null @@ -92,9 +116,19 @@ fun drainOutbound( if (natural < 1200) { val deficit = 1200 - natural // Rewind the Initial PN — we'll reissue with the same PN and the - // same captured frames plus padding. + // same captured frames plus padding. The SentPacket entry recorded + // by the natural-size build will be overwritten by the rebuild + // below since both use the same PN. initialState.pnSpace.rewindOutboundForRebuild() - val paddedInitial = buildLongHeaderFromFrames(conn, EncryptionLevel.INITIAL, initialFrames!!, padBytes = deficit) + val paddedInitial = + buildLongHeaderFromFrames( + conn = conn, + level = EncryptionLevel.INITIAL, + frames = initialContents!!.frames, // null-safe: gated by `initialNatural != null` above + tokens = initialContents.tokens, + nowMillis = nowMillis, + padBytes = deficit, + ) return concat(listOfNotNull(paddedInitial, handshakeNatural, applicationPkt)) } } @@ -176,27 +210,59 @@ private fun buildLongHeaderPacket( ) } +/** + * Frames + the matching retransmit-tokens collected for a single + * encryption-level packet build. Carried out of + * [collectHandshakeLevelFrames] so the caller can pass the tokens + * to the SentPacket retention recorded in + * [buildLongHeaderFromFrames]. + */ +private data class HandshakeLevelContents( + val frames: List, + val tokens: List, +) + /** * Drain ACK + CRYPTO frames for [level] into a fresh list. Destructive on - * the CRYPTO send buffer, but caller-controlled — call exactly once per - * outbound datagram. Returns null if there are no frames to send and the - * level has no protection installed. + * the CRYPTO send buffer (which now retains bytes until ACK per + * [com.vitorpamplona.quic.stream.SendBuffer]'s retain-until-ACK + * semantics — see commit B). Returns null if there are no frames to + * send and the level has no protection installed. + * + * Each emitted frame contributes a matching [RecoveryToken] so the + * caller can record a [com.vitorpamplona.quic.connection.recovery.SentPacket] + * keyed by packet number, enabling RFC 9002 retransmit at this + * encryption level (CRYPTO bytes are reliable per RFC 9000 §13.3). */ private fun collectHandshakeLevelFrames( conn: QuicConnection, level: EncryptionLevel, nowMillis: Long, -): List? { +): HandshakeLevelContents? { val state = conn.levelState(level) if (state.sendProtection == null) return null val frames = mutableListOf() - state.ackTracker.buildAckFrame(nowMillis, conn.config.ackDelayExponent.toInt())?.let { frames += it } + val tokens = mutableListOf() + state.ackTracker.buildAckFrame(nowMillis, conn.config.ackDelayExponent.toInt())?.let { + frames += it + tokens += RecoveryToken.Ack + } val cryptoChunk = state.cryptoSend.takeChunk(maxBytes = 1100) if (cryptoChunk != null && cryptoChunk.data.isNotEmpty()) { frames += CryptoFrame(cryptoChunk.offset, cryptoChunk.data) + // Step E: record a Crypto retransmit token at the matching + // encryption level so a lost handshake packet's CRYPTO bytes + // get re-emitted at the same offset on next drain (RFC 9000 + // §13.3 — handshake bytes are reliable). + tokens += + RecoveryToken.Crypto( + level = level, + offset = cryptoChunk.offset, + length = cryptoChunk.data.size.toLong(), + ) } if (frames.isEmpty()) return null - return frames + return HandshakeLevelContents(frames = frames, tokens = tokens) } /** @@ -210,6 +276,8 @@ private fun buildLongHeaderFromFrames( conn: QuicConnection, level: EncryptionLevel, frames: List, + tokens: List, + nowMillis: Long, padBytes: Int, ): ByteArray { val state = conn.levelState(level) @@ -223,22 +291,42 @@ private fun buildLongHeaderFromFrames( } val pn = state.pnSpace.allocateOutbound() val type = if (level == EncryptionLevel.INITIAL) LongHeaderType.INITIAL else LongHeaderType.HANDSHAKE - return LongHeaderPacket.build( - LongHeaderPlaintextPacket( - type = type, - version = QuicVersion.V1, - dcid = conn.destinationConnectionId, - scid = conn.sourceConnectionId, + val packet = + LongHeaderPacket.build( + LongHeaderPlaintextPacket( + type = type, + version = QuicVersion.V1, + dcid = conn.destinationConnectionId, + scid = conn.sourceConnectionId, + packetNumber = pn, + payload = payload, + ), + proto.aead, + proto.key, + proto.iv, + proto.hp, + proto.hpKey, + largestAckedInSpace = -1L, + ) + // Step E: retain the packet for RFC 9002 retransmit. Initial / + // Handshake packets carry CRYPTO frames; loss detection runs at + // each level and routes lost Crypto tokens back to the level's + // cryptoSend buffer via QuicConnection.onTokensLost (commit A). + // + // Initial-level rebuilds with padding (RFC 9000 §14.1) call + // pnSpace.rewindOutboundForRebuild() to reuse the same PN. The + // map insert below overwrites the prior entry, so the recorded + // SentPacket reflects the final padded packet's size — correct + // for retransmit purposes. + state.sentPackets[pn] = + SentPacket( packetNumber = pn, - payload = payload, - ), - proto.aead, - proto.key, - proto.iv, - proto.hp, - proto.hpKey, - largestAckedInSpace = -1L, - ) + sentAtMillis = nowMillis, + ackEliciting = isAckEliciting(frames), + sizeBytes = packet.size, + tokens = tokens, + ) + return packet } private fun buildApplicationPacket( diff --git a/quic/src/commonTest/kotlin/com/vitorpamplona/quic/connection/CryptoRetransmitTest.kt b/quic/src/commonTest/kotlin/com/vitorpamplona/quic/connection/CryptoRetransmitTest.kt new file mode 100644 index 000000000..e3663062e --- /dev/null +++ b/quic/src/commonTest/kotlin/com/vitorpamplona/quic/connection/CryptoRetransmitTest.kt @@ -0,0 +1,166 @@ +/* + * Copyright (c) 2025 Vitor Pamplona + * + * Permission is hereby granted, free of charge, to any person obtaining a copy of + * this software and associated documentation files (the "Software"), to deal in + * the Software without restriction, including without limitation the rights to use, + * copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the + * Software, and to permit persons to whom the Software is furnished to do so, + * subject to the following conditions: + * + * The above copyright notice and this permission notice shall be included in all + * copies or substantial portions of the Software. + * + * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS + * FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR + * COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN + * AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION + * WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. + */ +package com.vitorpamplona.quic.connection + +import com.vitorpamplona.quic.connection.recovery.RecoveryToken +import kotlinx.coroutines.runBlocking +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertNotNull +import kotlin.test.assertTrue + +/** + * Step E of the deferred-follow-ups pass: CRYPTO retransmit per + * encryption level. The handshake produces ClientHello / Finished + * CRYPTO bytes at Initial / Handshake levels; a lost packet + * carrying any of those must be retransmitted at the same level. + * + * Mirrors the structure of [StreamRetransmitTest] but exercises the + * Initial / Handshake-level paths in [QuicConnectionWriter]. + */ +class CryptoRetransmitTest { + @Test + fun handshakePacket_carriesCryptoToken_inSentPacket() = + runBlocking { + val client = newClientWithStartedHandshake() + // First drain produces the ClientHello at Initial level. + runCatching { drainOutbound(client, nowMillis = 1L) } + + val initialEntries = + client.initial.sentPackets.entries + .toList() + val withCrypto = + initialEntries.firstOrNull { entry -> + entry.value.tokens.any { it is RecoveryToken.Crypto } + } + assertNotNull( + withCrypto, + "Initial-level SentPacket must carry a Crypto token; saw ${initialEntries.map { it.value.tokens.map { t -> t::class.simpleName } }}", + ) + val cryptoToken = + withCrypto.value.tokens + .filterIsInstance() + .single() + assertEquals(EncryptionLevel.INITIAL, cryptoToken.level, "Crypto token's level must match the packet's level") + assertEquals(0L, cryptoToken.offset, "first ClientHello starts at offset 0") + assertTrue(cryptoToken.length > 0L, "ClientHello has non-zero length") + } + + @Test + fun cryptoData_lostAndRetransmittedAtSameLevel() = + runBlocking { + val client = newClientWithStartedHandshake() + runCatching { drainOutbound(client, nowMillis = 1L) } + + val firstEntry = + client.initial.sentPackets.entries + .firstOrNull { entry -> entry.value.tokens.any { it is RecoveryToken.Crypto } } + assertNotNull(firstEntry) + val firstPn = firstEntry.key + val cryptoToken = + firstEntry.value.tokens + .filterIsInstance() + .single() + + // Simulate loss via direct dispatch. + client.lock.lock() + try { + client.onTokensLost(listOf(cryptoToken)) + client.initial.sentPackets.remove(firstPn) + } finally { + client.lock.unlock() + } + + // Initial-level cryptoSend should now have re-queued bytes + // for retransmit. Verify by checking readableBytes. + assertTrue( + client.initial.cryptoSend.readableBytes > 0, + "lost CRYPTO bytes must be back in the Initial-level cryptoSend's retransmit queue", + ) + + // Next drain must produce a new Initial packet carrying + // the CRYPTO at the original offset. + runCatching { drainOutbound(client, nowMillis = 2L) } + val replayEntries = + client.initial.sentPackets.entries + .filter { it.key != firstPn } + val replay = + replayEntries.firstOrNull { entry -> + entry.value.tokens.any { it is RecoveryToken.Crypto } + } + assertNotNull(replay, "retransmit must produce a fresh Initial SentPacket carrying Crypto") + val replayToken = + replay.value.tokens + .filterIsInstance() + .single() + assertEquals(EncryptionLevel.INITIAL, replayToken.level) + assertEquals(cryptoToken.offset, replayToken.offset, "retransmit replays original offset") + assertEquals(cryptoToken.length, replayToken.length) + } + + @Test + fun cryptoAck_releasesBufferAtSameLevel() = + runBlocking { + val client = newClientWithStartedHandshake() + runCatching { drainOutbound(client, nowMillis = 1L) } + + val packet = + client.initial.sentPackets.entries + .first { it.value.tokens.any { t -> t is RecoveryToken.Crypto } } + // ACK via direct dispatch. + client.lock.lock() + try { + client.onTokensAcked(packet.value.tokens) + } finally { + client.lock.unlock() + } + // After ACK the Initial-level cryptoSend's flushedFloor should + // have advanced — we check by observing that another takeChunk + // returns null (no more bytes ready) AND finSent stays true if + // it was set (i.e. the ACK didn't reset it). + assertEquals(0, client.initial.cryptoSend.readableBytes) + } + + private fun newClientWithStartedHandshake(): QuicConnection = + runBlocking { + val client = + QuicConnection( + serverName = "example.test", + config = QuicConnectionConfig(), + tlsCertificateValidator = + com.vitorpamplona.quic.tls + .PermissiveCertificateValidator(), + ) + client.start() + // We don't drive the InProcessTlsServer here — we only need + // the client to have produced its ClientHello (which it does + // synchronously inside start()). The Initial-level cryptoSend + // is now non-empty. + assertTrue( + client.initial.cryptoSend.readableBytes > 0, + "client.start() must produce ClientHello bytes", + ) + // Need send-protection installed before drainOutbound emits + // anything. start() handles that. + assertNotNull(client.initial.sendProtection, "Initial-level send protection must be installed") + client + } +}