From b6634fe1242c38169062224cc10e3be489ecf7ba Mon Sep 17 00:00:00 2001
From: Marten Seemann <martenseemann@gmail.com>
Date: Mon, 29 Mar 2021 14:53:35 +0700
Subject: [PATCH] reset the loss detection timer when the client's address is
 validated

---
 internal/ackhandler/sent_packet_handler.go    | 24 ++++++++++++-------
 .../ackhandler/sent_packet_handler_test.go    | 15 ++++++++++++
 2 files changed, 30 insertions(+), 9 deletions(-)

diff --git a/internal/ackhandler/sent_packet_handler.go b/internal/ackhandler/sent_packet_handler.go
index 2c31fd89..5f25659e 100644
--- a/internal/ackhandler/sent_packet_handler.go
+++ b/internal/ackhandler/sent_packet_handler.go
@@ -201,9 +201,10 @@ func (h *sentPacketHandler) ReceivedBytes(n protocol.ByteCount) {
 	}
 }
 
-func (h *sentPacketHandler) ReceivedPacket(encLevel protocol.EncryptionLevel) {
-	if h.perspective == protocol.PerspectiveServer && encLevel == protocol.EncryptionHandshake {
+func (h *sentPacketHandler) ReceivedPacket(l protocol.EncryptionLevel) {
+	if h.perspective == protocol.PerspectiveServer && l == protocol.EncryptionHandshake && !h.peerAddressValidated {
 		h.peerAddressValidated = true
+		h.setLossDetectionTimer()
 	}
 }
 
@@ -485,7 +486,8 @@ func (h *sentPacketHandler) hasOutstandingPackets() bool {
 
 func (h *sentPacketHandler) setLossDetectionTimer() {
 	oldAlarm := h.alarm // only needed in case tracing is enabled
-	if lossTime, encLevel := h.getLossTimeAndSpace(); !lossTime.IsZero() {
+	lossTime, encLevel := h.getLossTimeAndSpace()
+	if !lossTime.IsZero() {
 		// Early retransmit timer or time loss detection.
 		h.alarm = lossTime
 		if h.tracer != nil && h.alarm != oldAlarm {
@@ -497,9 +499,11 @@ func (h *sentPacketHandler) setLossDetectionTimer() {
 	// Cancel the alarm if amplification limited.
 	if h.isAmplificationLimited() {
 		h.alarm = time.Time{}
-		h.logger.Debugf("Canceling loss detection timer. Amplification limited.")
-		if h.tracer != nil && !oldAlarm.IsZero() {
-			h.tracer.LossTimerCanceled()
+		if !oldAlarm.IsZero() {
+			h.logger.Debugf("Canceling loss detection timer. Amplification limited.")
+			if h.tracer != nil {
+				h.tracer.LossTimerCanceled()
+			}
 		}
 		return
 	}
@@ -507,9 +511,11 @@ func (h *sentPacketHandler) setLossDetectionTimer() {
 	// Cancel the alarm if no packets are outstanding
 	if !h.hasOutstandingPackets() && h.peerCompletedAddressValidation {
 		h.alarm = time.Time{}
-		h.logger.Debugf("Canceling loss detection timer. No packets in flight.")
-		if h.tracer != nil && !oldAlarm.IsZero() {
-			h.tracer.LossTimerCanceled()
+		if !oldAlarm.IsZero() {
+			h.logger.Debugf("Canceling loss detection timer. No packets in flight.")
+			if h.tracer != nil {
+				h.tracer.LossTimerCanceled()
+			}
 		}
 		return
 	}
diff --git a/internal/ackhandler/sent_packet_handler_test.go b/internal/ackhandler/sent_packet_handler_test.go
index afc51194..5fef0dc1 100644
--- a/internal/ackhandler/sent_packet_handler_test.go
+++ b/internal/ackhandler/sent_packet_handler_test.go
@@ -843,6 +843,21 @@ var _ = Describe("SentPacketHandler", func() {
 			handler.ReceivedBytes(1)
 			Expect(handler.GetLossDetectionTimeout()).ToNot(BeZero())
 		})
+
+		It("resets the loss detection timer when the client's address is validated", func() {
+			handler.ReceivedBytes(300)
+			handler.SentPacket(&Packet{
+				PacketNumber:    1,
+				Length:          900,
+				EncryptionLevel: protocol.EncryptionHandshake,
+				Frames:          []Frame{{Frame: &wire.PingFrame{}}},
+				SendTime:        time.Now(),
+			})
+			// Amplification limited. We don't need to set a timer now.
+			Expect(handler.GetLossDetectionTimeout()).To(BeZero())
+			handler.ReceivedPacket(protocol.EncryptionHandshake)
+			Expect(handler.GetLossDetectionTimeout()).ToNot(BeZero())
+		})
 	})
 
 	Context("amplification limit, for the client", func() {