From dc49f5673bf9340dfe9019b5b05d8aa3ab244170 Mon Sep 17 00:00:00 2001
From: Marten Seemann <martenseemann@gmail.com>
Date: Sat, 3 Feb 2024 13:19:33 +0700
Subject: [PATCH] fix flaky 0-RTT packet loss recovery test (#4306)

---
 integrationtests/self/zero_rtt_test.go | 26 ++++++++------------------
 1 file changed, 8 insertions(+), 18 deletions(-)

diff --git a/integrationtests/self/zero_rtt_test.go b/integrationtests/self/zero_rtt_test.go
index 85ad015e2..55307679f 100644
--- a/integrationtests/self/zero_rtt_test.go
+++ b/integrationtests/self/zero_rtt_test.go
@@ -5,7 +5,6 @@ import (
 	"crypto/tls"
 	"fmt"
 	"io"
-	mrand "math/rand"
 	"net"
 	"sync"
 	"sync/atomic"
@@ -380,7 +379,7 @@ var _ = Describe("0-RTT", func() {
 	})
 
 	It("transfers 0-RTT data, when 0-RTT packets are lost", func() {
-		var num0RTTPackets, num0RTTDropped atomic.Uint32
+		var num0RTTPackets, numDropped atomic.Uint32
 
 		tlsConf := getTLSConfig()
 		clientConf := getTLSClientConfig()
@@ -399,17 +398,8 @@ var _ = Describe("0-RTT", func() {
 		defer ln.Close()
 
 		proxy, err := quicproxy.NewQuicProxy("localhost:0", &quicproxy.Opts{
-			RemoteAddr: fmt.Sprintf("localhost:%d", ln.Addr().(*net.UDPAddr).Port),
-			DelayPacket: func(_ quicproxy.Direction, data []byte) time.Duration {
-				if wire.IsLongHeaderPacket(data[0]) {
-					hdr, _, _, err := wire.ParsePacket(data)
-					Expect(err).ToNot(HaveOccurred())
-					if hdr.Type == protocol.PacketType0RTT {
-						num0RTTPackets.Add(1)
-					}
-				}
-				return rtt / 2
-			},
+			RemoteAddr:  fmt.Sprintf("localhost:%d", ln.Addr().(*net.UDPAddr).Port),
+			DelayPacket: func(_ quicproxy.Direction, data []byte) time.Duration { return rtt / 2 },
 			DropPacket: func(_ quicproxy.Direction, data []byte) bool {
 				if !wire.IsLongHeaderPacket(data[0]) {
 					return false
@@ -417,10 +407,11 @@ var _ = Describe("0-RTT", func() {
 				hdr, _, _, err := wire.ParsePacket(data)
 				Expect(err).ToNot(HaveOccurred())
 				if hdr.Type == protocol.PacketType0RTT {
+					count := num0RTTPackets.Add(1)
 					// drop 25% of the 0-RTT packets
-					drop := mrand.Intn(4) == 0
+					drop := count%4 == 0
 					if drop {
-						num0RTTDropped.Add(1)
+						numDropped.Add(1)
 					}
 					return drop
 				}
@@ -433,9 +424,8 @@ var _ = Describe("0-RTT", func() {
 		transfer0RTTData(ln, proxy.LocalPort(), protocol.DefaultConnectionIDLength, clientConf, nil, PRData)
 
 		num0RTT := num0RTTPackets.Load()
-		numDropped := num0RTTDropped.Load()
-		fmt.Fprintf(GinkgoWriter, "Sent %d 0-RTT packets. Dropped %d of those.", num0RTT, numDropped)
-		Expect(numDropped).ToNot(BeZero())
+		fmt.Fprintf(GinkgoWriter, "Sent %d 0-RTT packets. Dropped %d of those.", num0RTT, numDropped.Load())
+		Expect(numDropped.Load()).ToNot(BeZero())
 		Expect(num0RTT).ToNot(BeZero())
 		Expect(get0RTTPackets(counter.getRcvdLongHeaderPackets())).ToNot(BeEmpty())
 	})