quic: initiate key updates

For golang/go#58547

Change-Id: If27c0745fc49cb9e8cb9906733ce2f453926b893
Reviewed-on: https://go-review.googlesource.com/c/net/+/529595
LUCI-TryBot-Result: Go LUCI <[email protected]>
Reviewed-by: Jonathan Amsterdam <[email protected]>

Author: neild

internal/quic/conn.go

@@ -100,6 +100,7 @@ func newConn(now time.Time, side connSide, initialConnID []byte, peerAddr netip.
 	// The smallest allowed maximum QUIC datagram size is 1200 bytes.
 	// TODO: PMTU discovery.
 	const maxDatagramSize = 1200
+	c.keysAppData.init()
 	c.loss.init(c.side, maxDatagramSize, now)
 	c.streamsInit()
 	c.lifetimeInit()

internal/quic/conn_test.go

@@ -242,6 +242,7 @@ func newTestConn(t *testing.T, side connSide, opts ...any) *testConn {
 	}
 	tc.conn = conn
 
+	conn.keysAppData.updateAfter = maxPacketNumber // disable key updates
 	tc.keysInitial.r = conn.keysInitial.w
 	tc.keysInitial.w = conn.keysInitial.r
 
@@ -687,6 +688,7 @@ func (tc *testConn) encodeTestPacket(p *testPacket, pad int) []byte {
 					tc.wkeyAppData.pkt[p.keyNumber],
 				},
 			},
+			updateAfter: maxPacketNumber,
 		}
 		if p.keyPhaseBit {
 			k.phase |= keyPhaseBit

internal/quic/key_update_test.go

@@ -161,3 +161,74 @@ func TestKeyUpdateRejectPacketFromPriorPhase(t *testing.T) {
 			},
 		})
 }
+
+func TestKeyUpdateLocallyInitiated(t *testing.T) {
+	const updateAfter = 4 // initiate key update after 1-RTT packet 4
+	tc := newTestConn(t, serverSide)
+	tc.conn.keysAppData.updateAfter = updateAfter
+	tc.handshake()
+
+	for {
+		tc.writeFrames(packetType1RTT, debugFramePing{})
+		tc.advanceToTimer()
+		tc.wantFrameType("conn ACKs last packet",
+			packetType1RTT, debugFrameAck{})
+		if tc.lastPacket.num > updateAfter {
+			break
+		}
+		if got, want := tc.lastPacket.keyNumber, 0; got != want {
+			t.Errorf("before key update, conn sent packet with key %v, want %v", got, want)
+		}
+		if tc.lastPacket.keyPhaseBit {
+			t.Errorf("before key update, keyPhaseBit is set, want unset")
+		}
+	}
+	if got, want := tc.lastPacket.keyNumber, 1; got != want {
+		t.Errorf("after key update, conn sent packet with key %v, want %v", got, want)
+	}
+	if !tc.lastPacket.keyPhaseBit {
+		t.Errorf("after key update, keyPhaseBit is unset, want set")
+	}
+	tc.wantFrame("first packet after a key update is always ack-eliciting",
+		packetType1RTT, debugFramePing{})
+	tc.wantIdle("no more frames")
+
+	// Peer sends another packet using the prior phase keys.
+	tc.writeFrames(packetType1RTT, debugFramePing{})
+	tc.advanceToTimer()
+	tc.wantFrameType("conn ACKs packet in prior phase",
+		packetType1RTT, debugFrameAck{})
+	tc.wantIdle("packet is not ack-eliciting")
+	if got, want := tc.lastPacket.keyNumber, 1; got != want {
+		t.Errorf("after key update, conn sent packet with key %v, want %v", got, want)
+	}
+
+	// Peer updates to the next phase.
+	tc.sendKeyNumber = 1
+	tc.sendKeyPhaseBit = true
+	tc.writeAckForAll()
+	tc.writeFrames(packetType1RTT, debugFramePing{})
+	tc.advanceToTimer()
+	tc.wantFrameType("conn ACKs packet in current phase",
+		packetType1RTT, debugFrameAck{})
+	tc.wantIdle("packet is not ack-eliciting")
+	if got, want := tc.lastPacket.keyNumber, 1; got != want {
+		t.Errorf("after key update, conn sent packet with key %v, want %v", got, want)
+	}
+
+	// Peer initiates its own update.
+	tc.sendKeyNumber = 2
+	tc.sendKeyPhaseBit = false
+	tc.writeFrames(packetType1RTT, debugFramePing{})
+	tc.advanceToTimer()
+	tc.wantFrameType("conn ACKs packet in current phase",
+		packetType1RTT, debugFrameAck{})
+	tc.wantFrame("first packet after a key update is always ack-eliciting",
+		packetType1RTT, debugFramePing{})
+	if got, want := tc.lastPacket.keyNumber, 2; got != want {
+		t.Errorf("after peer key update, conn sent packet with key %v, want %v", got, want)
+	}
+	if tc.lastPacket.keyPhaseBit {
+		t.Errorf("after peer key update, keyPhaseBit is unset, want set")
+	}
+}

internal/quic/packet_codec_test.go

@@ -153,6 +153,9 @@ func TestRoundtripEncodeShortPacket(t *testing.T) {
 	aes128Keys.w = aes128Keys.r
 	aes256Keys.w = aes256Keys.r
 	chachaKeys.w = chachaKeys.r
+	aes128Keys.updateAfter = maxPacketNumber
+	aes256Keys.updateAfter = maxPacketNumber
+	chachaKeys.updateAfter = maxPacketNumber
 	connID := make([]byte, connIDLen)
 	for i := range connID {
 		connID[i] = byte(i)

internal/quic/packet_protection.go

@@ -340,9 +340,19 @@ type updatingKeyPair struct {
 	updating    bool
 	minSent     packetNumber // min packet number sent since entering the updating state
 	minReceived packetNumber // min packet number received in the next phase
+	updateAfter packetNumber // packet number after which to initiate key update
 	r, w        updatingKeys
 }
 
+func (k *updatingKeyPair) init() {
+	// 1-RTT packets until the first key update.
+	//
+	// We perform the first key update early in the connection so a peer
+	// which does not support key updates will fail rapidly,
+	// rather than after the connection has been long established.
+	k.updateAfter = 1000
+}
+
 func (k *updatingKeyPair) canRead() bool {
 	return k.r.hdr.hp != nil
 }
@@ -371,8 +381,6 @@ func (k *updatingKeyPair) needAckEliciting() bool {
 // protect applies packet protection to a packet.
 // Parameters and returns are as for fixedKeyPair.protect.
 func (k *updatingKeyPair) protect(hdr, pay []byte, pnumOff int, pnum packetNumber) []byte {
-	// TODO: Initiate key updates as required to avoid the AEAD usage limit.
-	// https://www.rfc-editor.org/rfc/rfc9001#section-6.6
 	var pkt []byte
 	if k.updating {
 		hdr[0] |= k.phase ^ keyPhaseBit
@@ -381,6 +389,21 @@ func (k *updatingKeyPair) protect(hdr, pay []byte, pnumOff int, pnum packetNumbe
 	} else {
 		hdr[0] |= k.phase
 		pkt = k.w.pkt[0].protect(hdr, pay, pnum)
+		if pnum >= k.updateAfter {
+			// Initiate a key update, starting with the next packet we send.
+			//
+			// We do this after protecting the current packet
+			// to allow Conn.appendFrames to ensure that the first packet sent
+			// in the new phase is ack-eliciting.
+			k.updating = true
+			k.minSent = maxPacketNumber
+			k.minReceived = maxPacketNumber
+			// The lowest confidentiality limit for a supported AEAD is 2^23 packets.
+			// https://www.rfc-editor.org/rfc/rfc9001#section-6.6-5
+			//
+			// Schedule our next update for half that.
+			k.updateAfter += (1 << 22)
+		}
 	}
 	k.w.hdr.protect(pkt, pnumOff)
 	return pkt