From b928ad7e4fc66ab4fc0ce9eb8e323b1781bc1ef4 Mon Sep 17 00:00:00 2001
From: Marten Seemann <martenseemann@gmail.com>
Date: Thu, 18 Jun 2020 09:55:24 +0700
Subject: [PATCH] fix flaky INVALID_TOKEN integration test

---
 integrationtests/self/handshake_test.go | 52 ++++++++++++++-----------
 1 file changed, 29 insertions(+), 23 deletions(-)

diff --git a/integrationtests/self/handshake_test.go b/integrationtests/self/handshake_test.go
index aaa7702d0..e88a424d1 100644
--- a/integrationtests/self/handshake_test.go
+++ b/integrationtests/self/handshake_test.go
@@ -461,31 +461,37 @@ var _ = Describe("Handshake tests", func() {
 
 			Eventually(done).Should(BeClosed())
 		})
-	})
 
-	It("rejects invalid Retry token with the INVALID_TOKEN error", func() {
-		tokenChan := make(chan *quic.Token, 10)
-		serverConfig.AcceptToken = func(addr net.Addr, token *quic.Token) bool {
-			tokenChan <- token
-			return false
-		}
+		It("rejects invalid Retry token with the INVALID_TOKEN error", func() {
+			tokenChan := make(chan *quic.Token, 10)
+			serverConfig.AcceptToken = func(addr net.Addr, token *quic.Token) bool {
+				tokenChan <- token
+				return false
+			}
 
-		server, err := quic.ListenAddr("localhost:0", getTLSConfig(), serverConfig)
-		Expect(err).ToNot(HaveOccurred())
-		defer server.Close()
+			server, err := quic.ListenAddr("localhost:0", getTLSConfig(), serverConfig)
+			Expect(err).ToNot(HaveOccurred())
+			defer server.Close()
 
-		_, err = quic.DialAddr(
-			fmt.Sprintf("localhost:%d", server.Addr().(*net.UDPAddr).Port),
-			getTLSClientConfig(),
-			nil,
-		)
-		Expect(err).To(HaveOccurred())
-		Expect(err.Error()).To(ContainSubstring("INVALID_TOKEN"))
-		Expect(tokenChan).To(HaveLen(2))
-		token := <-tokenChan
-		Expect(token).To(BeNil())
-		token = <-tokenChan
-		Expect(token).ToNot(BeNil())
-		Expect(token.IsRetryToken).To(BeTrue())
+			_, err = quic.DialAddr(
+				fmt.Sprintf("localhost:%d", server.Addr().(*net.UDPAddr).Port),
+				getTLSClientConfig(),
+				nil,
+			)
+			Expect(err).To(HaveOccurred())
+			Expect(err.Error()).To(ContainSubstring("INVALID_TOKEN"))
+			// Receiving a Retry might lead the client to measure a very small RTT.
+			// Then, it sometimes would retransmit the ClientHello before receiving the ServerHello.
+			Expect(len(tokenChan)).To(BeNumerically(">=", 2))
+			token := <-tokenChan
+			Expect(token).To(BeNil())
+			token = <-tokenChan
+			Expect(token).ToNot(BeNil())
+			// If the ClientHello was retransmitted, make sure that it contained the same Retry token.
+			for i := 2; i < len(tokenChan); i++ {
+				Expect(<-tokenChan).To(Equal(token))
+			}
+			Expect(token.IsRetryToken).To(BeTrue())
+		})
 	})
 })