add an integration test for the token store

integrationtests/self/handshake_test.go

@@ -19,6 +19,32 @@ type versioner interface {
 	GetVersion() protocol.VersionNumber
 }
 
+type tokenStore struct {
+	store quic.TokenStore
+	gets  chan<- string
+	puts  chan<- string
+}
+
+var _ quic.TokenStore = &tokenStore{}
+
+func newTokenStore(gets, puts chan<- string) quic.TokenStore {
+	return &tokenStore{
+		store: quic.NewLRUTokenStore(10, 4),
+		gets:  gets,
+		puts:  puts,
+	}
+}
+
+func (c *tokenStore) Put(key string, token *quic.ClientToken) {
+	c.puts <- key
+	c.store.Put(key, token)
+}
+
+func (c *tokenStore) Pop(key string) *quic.ClientToken {
+	c.gets <- key
+	return c.store.Pop(key)
+}
+
 var _ = Describe("Handshake tests", func() {
 	var (
 		server        quic.Listener
@@ -328,4 +354,62 @@ var _ = Describe("Handshake tests", func() {
 			Expect(server.Close()).To(Succeed())
 		})
 	})
+
+	Context("using tokens", func() {
+		It("uses tokens provided in NEW_TOKEN frames", func() {
+			tokenChan := make(chan *quic.Token, 100)
+			serverConfig.AcceptToken = func(addr net.Addr, token *quic.Token) bool {
+				if token != nil && !token.IsRetryToken {
+					tokenChan <- token
+				}
+				return true
+			}
+
+			server, err := quic.ListenAddr("localhost:0", getTLSConfig(), serverConfig)
+			Expect(err).ToNot(HaveOccurred())
+
+			// dial the first session and receive the token
+			go func() {
+				defer GinkgoRecover()
+				_, err := server.Accept(context.Background())
+				Expect(err).ToNot(HaveOccurred())
+			}()
+
+			gets := make(chan string, 100)
+			puts := make(chan string, 100)
+			tokenStore := newTokenStore(gets, puts)
+			quicConf := &quic.Config{TokenStore: tokenStore}
+			sess, err := quic.DialAddr(
+				fmt.Sprintf("localhost:%d", server.Addr().(*net.UDPAddr).Port),
+				getTLSClientConfig(),
+				quicConf,
+			)
+			Expect(err).ToNot(HaveOccurred())
+			Expect(gets).To(Receive())
+			Eventually(puts).Should(Receive())
+			Expect(tokenChan).ToNot(Receive())
+			// received a token. Close this session.
+			Expect(sess.Close()).To(Succeed())
+
+			// dial the second session and verify that the token was used
+			done := make(chan struct{})
+			go func() {
+				defer GinkgoRecover()
+				defer close(done)
+				_, err := server.Accept(context.Background())
+				Expect(err).ToNot(HaveOccurred())
+			}()
+			sess, err = quic.DialAddr(
+				fmt.Sprintf("localhost:%d", server.Addr().(*net.UDPAddr).Port),
+				getTLSClientConfig(),
+				quicConf,
+			)
+			Expect(err).ToNot(HaveOccurred())
+			defer sess.Close()
+			Expect(gets).To(Receive())
+			Expect(tokenChan).To(Receive())
+
+			Eventually(done).Should(BeClosed())
+		})
+	})
 })