queue 0-RTT packets for retransmission when 0-RTT is rejected

internal/ackhandler/received_packet_handler.go

@@ -93,6 +93,9 @@ func (h *receivedPacketHandler) DropPackets(encLevel protocol.EncryptionLevel) {
 		h.initialPackets = nil
 	case protocol.EncryptionHandshake:
 		h.handshakePackets = nil
+	case protocol.Encryption0RTT:
+		// Nothing to do here.
+		// If we are rejecting 0-RTT, no 0-RTT packets will have been decrypted.
 	default:
 		panic(fmt.Sprintf("Cannot drop keys for encryption level %s", encLevel))
 	}

internal/ackhandler/received_packet_handler_test.go

@@ -91,4 +91,8 @@ var _ = Describe("Received Packet Handler", func() {
 		Expect(handler.GetAckFrame(protocol.EncryptionHandshake)).To(BeNil())
 		Expect(handler.GetAckFrame(protocol.Encryption1RTT)).ToNot(BeNil())
 	})
+
+	It("does nothing when droping 0-RTT packets", func() {
+		handler.DropPackets(protocol.Encryption0RTT)
+	})
 })

internal/ackhandler/sent_packet_handler.go

@@ -102,19 +102,34 @@ func NewSentPacketHandler(
 
 func (h *sentPacketHandler) DropPackets(encLevel protocol.EncryptionLevel) {
 	// remove outstanding packets from bytes_in_flight
-	pnSpace := h.getPacketNumberSpace(encLevel)
-	pnSpace.history.Iterate(func(p *Packet) (bool, error) {
-		if p.includedInBytesInFlight {
-			h.bytesInFlight -= p.Length
-		}
-		return true, nil
-	})
+	if encLevel == protocol.EncryptionInitial || encLevel == protocol.EncryptionHandshake {
+		pnSpace := h.getPacketNumberSpace(encLevel)
+		pnSpace.history.Iterate(func(p *Packet) (bool, error) {
+			if p.includedInBytesInFlight {
+				h.bytesInFlight -= p.Length
+			}
+			return true, nil
+		})
+	}
 	// drop the packet history
 	switch encLevel {
 	case protocol.EncryptionInitial:
 		h.initialPackets = nil
 	case protocol.EncryptionHandshake:
 		h.handshakePackets = nil
+	case protocol.Encryption0RTT:
+		// TODO(#2067): invalidate sent data
+		h.appDataPackets.history.Iterate(func(p *Packet) (bool, error) {
+			if p.EncryptionLevel != protocol.Encryption0RTT {
+				return false, nil
+			}
+			h.queueFramesForRetransmission(p)
+			if p.includedInBytesInFlight {
+				h.bytesInFlight -= p.Length
+			}
+			h.appDataPackets.history.Remove(p.PacketNumber)
+			return true, nil
+		})
 	default:
 		panic(fmt.Sprintf("Cannot drop keys for encryption level %s", encLevel))
 	}

internal/ackhandler/sent_packet_handler_test.go

@@ -838,6 +838,23 @@ var _ = Describe("SentPacketHandler", func() {
 			Expect(handler.bytesInFlight).To(Equal(protocol.ByteCount(10)))
 			Expect(handler.handshakePackets).To(BeNil())
 		})
+
+		// TODO(#2067): invalidate 0-RTT data when 0-RTT is rejected
+		It("retransmits 0-RTT packets when 0-RTT keys are dropped", func() {
+			for i := protocol.PacketNumber(0); i < 6; i++ {
+				handler.SentPacket(ackElicitingPacket(&Packet{
+					PacketNumber:    i,
+					EncryptionLevel: protocol.Encryption0RTT,
+				}))
+			}
+			for i := protocol.PacketNumber(0); i < 6; i++ {
+				handler.SentPacket(ackElicitingPacket(&Packet{PacketNumber: i}))
+			}
+			Expect(handler.bytesInFlight).To(Equal(protocol.ByteCount(12)))
+			handler.DropPackets(protocol.Encryption0RTT)
+			Expect(lostPackets).To(Equal([]protocol.PacketNumber{0, 1, 2, 3, 4, 5}))
+			Expect(handler.bytesInFlight).To(Equal(protocol.ByteCount(6)))
+		})
 	})
 
 	Context("peeking and popping packet number", func() {