From 951477af928cfb1a48584cd4d78f944c3bcb2d08 Mon Sep 17 00:00:00 2001
From: Marten Seemann <martenseemann@gmail.com>
Date: Thu, 15 Aug 2019 12:38:25 +0700
Subject: [PATCH] queue 0-RTT packets for retransmission after receiving a
 Retry

---
 internal/ackhandler/sent_packet_handler.go    |  8 +++++
 .../ackhandler/sent_packet_handler_test.go    | 30 +++++++++++++++++++
 2 files changed, 38 insertions(+)

diff --git a/internal/ackhandler/sent_packet_handler.go b/internal/ackhandler/sent_packet_handler.go
index 730d135e1..febd89f63 100644
--- a/internal/ackhandler/sent_packet_handler.go
+++ b/internal/ackhandler/sent_packet_handler.go
@@ -598,7 +598,15 @@ func (h *sentPacketHandler) ResetForRetry() error {
 		h.queueFramesForRetransmission(p)
 		return true, nil
 	})
+	// All application data packets sent at this point are 0-RTT packets.
+	// In the case of a Retry, we can assume that the server dropped all of them.
+	h.appDataPackets.history.Iterate(func(p *Packet) (bool, error) {
+		h.queueFramesForRetransmission(p)
+		return true, nil
+	})
+
 	h.initialPackets = newPacketNumberSpace(h.initialPackets.pns.Pop())
+	h.appDataPackets = newPacketNumberSpace(h.appDataPackets.pns.Pop())
 	h.setLossDetectionTimer()
 	return nil
 }
diff --git a/internal/ackhandler/sent_packet_handler_test.go b/internal/ackhandler/sent_packet_handler_test.go
index 2ef5dc027..1294da135 100644
--- a/internal/ackhandler/sent_packet_handler_test.go
+++ b/internal/ackhandler/sent_packet_handler_test.go
@@ -875,5 +875,35 @@ var _ = Describe("SentPacketHandler", func() {
 			Expect(handler.GetLossDetectionTimeout()).To(BeZero())
 			Expect(handler.SendMode()).To(Equal(SendAny))
 		})
+
+		It("queues outstanding frames for retransmission and cancels alarms", func() {
+			var lostInitial, lost0RTT bool
+			handler.SentPacket(&Packet{
+				PacketNumber:    13,
+				EncryptionLevel: protocol.EncryptionInitial,
+				Frames: []Frame{
+					{Frame: &wire.CryptoFrame{Data: []byte("foobar")}, OnLost: func(wire.Frame) { lostInitial = true }},
+				},
+				Length: 100,
+			})
+			pn := handler.PopPacketNumber(protocol.Encryption0RTT)
+			handler.SentPacket(&Packet{
+				PacketNumber:    pn,
+				EncryptionLevel: protocol.Encryption0RTT,
+				Frames: []Frame{
+					{Frame: &wire.StreamFrame{Data: []byte("foobar")}, OnLost: func(wire.Frame) { lost0RTT = true }},
+				},
+				Length: 999,
+			})
+			Expect(handler.bytesInFlight).ToNot(BeZero())
+			// now receive a Retry
+			Expect(handler.ResetForRetry()).To(Succeed())
+			Expect(handler.bytesInFlight).To(BeZero())
+			Expect(lostInitial).To(BeTrue())
+			Expect(lost0RTT).To(BeTrue())
+
+			// make sure we keep increasing the packet number for 0-RTT packets
+			Expect(handler.PopPacketNumber(protocol.Encryption0RTT)).To(BeNumerically(">", pn))
+		})
 	})
 })