add the TokenStore to the quic.Config, store received tokens

client.go

@@ -253,6 +253,7 @@ func populateClientConfig(config *Config, createdPacketConn bool) *Config {
 		KeepAlive:                             config.KeepAlive,
 		StatelessResetKey:                     config.StatelessResetKey,
 		QuicTracer:                            config.QuicTracer,
+		TokenStore:                            config.TokenStore,
 	}
 }
 

client_test.go

@@ -444,6 +444,7 @@ var _ = Describe("Client", func() {
 		Context("quic.Config", func() {
 			It("setups with the right values", func() {
 				tracer := quictrace.NewTracer()
+				tokenStore := NewLRUTokenStore(10, 4)
 				config := &Config{
 					HandshakeTimeout:      1337 * time.Minute,
 					IdleTimeout:           42 * time.Hour,
@@ -452,6 +453,7 @@ var _ = Describe("Client", func() {
 					ConnectionIDLength:    13,
 					StatelessResetKey:     []byte("foobar"),
 					QuicTracer:            tracer,
+					TokenStore:            tokenStore,
 				}
 				c := populateClientConfig(config, false)
 				Expect(c.HandshakeTimeout).To(Equal(1337 * time.Minute))
@@ -461,6 +463,7 @@ var _ = Describe("Client", func() {
 				Expect(c.ConnectionIDLength).To(Equal(13))
 				Expect(c.StatelessResetKey).To(Equal([]byte("foobar")))
 				Expect(c.QuicTracer).To(Equal(tracer))
+				Expect(c.TokenStore).To(Equal(tokenStore))
 			})
 
 			It("errors when the Config contains an invalid version", func() {

interface.go

@@ -231,6 +231,11 @@ type Config struct {
 	//   * else, that it was issued within the last 24 hours.
 	// This option is only valid for the server.
 	AcceptToken func(clientAddr net.Addr, token *Token) bool
+	// The TokenStore stores tokens received from the server.
+	// Tokens are used to skip address validation on future connection attempts.
+	// The key used to store tokens is the ServerName from the tls.Config, if set
+	// otherwise the token is associated with the server's IP address.
+	TokenStore TokenStore
 	// MaxReceiveStreamFlowControlWindow is the maximum stream-level flow control window for receiving data.
 	// If this value is zero, it will default to 1 MB for the server and 6 MB for the client.
 	MaxReceiveStreamFlowControlWindow uint64

mock_token_store_test.go

@@ -0,0 +1,60 @@
+// Code generated by MockGen. DO NOT EDIT.
+// Source: github.com/lucas-clemente/quic-go (interfaces: TokenStore)
+
+// Package quic is a generated GoMock package.
+package quic
+
+import (
+	reflect "reflect"
+
+	gomock "github.com/golang/mock/gomock"
+)
+
+// MockTokenStore is a mock of TokenStore interface
+type MockTokenStore struct {
+	ctrl     *gomock.Controller
+	recorder *MockTokenStoreMockRecorder
+}
+
+// MockTokenStoreMockRecorder is the mock recorder for MockTokenStore
+type MockTokenStoreMockRecorder struct {
+	mock *MockTokenStore
+}
+
+// NewMockTokenStore creates a new mock instance
+func NewMockTokenStore(ctrl *gomock.Controller) *MockTokenStore {
+	mock := &MockTokenStore{ctrl: ctrl}
+	mock.recorder = &MockTokenStoreMockRecorder{mock}
+	return mock
+}
+
+// EXPECT returns an object that allows the caller to indicate expected use
+func (m *MockTokenStore) EXPECT() *MockTokenStoreMockRecorder {
+	return m.recorder
+}
+
+// Pop mocks base method
+func (m *MockTokenStore) Pop(arg0 string) *ClientToken {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "Pop", arg0)
+	ret0, _ := ret[0].(*ClientToken)
+	return ret0
+}
+
+// Pop indicates an expected call of Pop
+func (mr *MockTokenStoreMockRecorder) Pop(arg0 interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Pop", reflect.TypeOf((*MockTokenStore)(nil).Pop), arg0)
+}
+
+// Put mocks base method
+func (m *MockTokenStore) Put(arg0 string, arg1 *ClientToken) {
+	m.ctrl.T.Helper()
+	m.ctrl.Call(m, "Put", arg0, arg1)
+}
+
+// Put indicates an expected call of Put
+func (mr *MockTokenStoreMockRecorder) Put(arg0, arg1 interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Put", reflect.TypeOf((*MockTokenStore)(nil).Put), arg0, arg1)
+}

mockgen.go

@@ -19,3 +19,4 @@ package quic
 //go:generate sh -c "./mockgen_private.sh quic mock_unknown_packet_handler_test.go github.com/lucas-clemente/quic-go unknownPacketHandler"
 //go:generate sh -c "./mockgen_private.sh quic mock_packet_handler_manager_test.go github.com/lucas-clemente/quic-go packetHandlerManager"
 //go:generate sh -c "./mockgen_private.sh quic mock_multiplexer_test.go github.com/lucas-clemente/quic-go multiplexer"
+//go:generate sh -c "mockgen -package quic -destination mock_token_store_test.go github.com/lucas-clemente/quic-go TokenStore && sed -i '' 's/quic_go.//g' mock_token_store_test.go && goimports -w mock_token_store_test.go"

session.go

@@ -117,6 +117,7 @@ type session struct {
 	framer                framer
 	windowUpdateQueue     *windowUpdateQueue
 	connFlowController    flowcontrol.ConnectionFlowController
+	tokenStoreKey         string                    // only set for the client
 	tokenGenerator        *handshake.TokenGenerator // only set for the server
 
 	unpacker    unpacker
@@ -333,6 +334,11 @@ var newClientSession = func(
 		s.perspective,
 		s.version,
 	)
+	if len(tlsConf.ServerName) > 0 {
+		s.tokenStoreKey = tlsConf.ServerName
+	} else {
+		s.tokenStoreKey = conn.RemoteAddr().String()
+	}
 	return s, s.postSetup()
 }
 
@@ -767,6 +773,7 @@ func (s *session) handleFrame(f wire.Frame, pn protocol.PacketNumber, encLevel p
 		// since we don't send PATH_CHALLENGEs, we don't expect PATH_RESPONSEs
 		err = errors.New("unexpected PATH_RESPONSE frame")
 	case *wire.NewTokenFrame:
+		err = s.handleNewTokenFrame(frame)
 	case *wire.NewConnectionIDFrame:
 	case *wire.RetireConnectionIDFrame:
 		// since we don't send new connection IDs, we don't expect retirements
@@ -893,6 +900,16 @@ func (s *session) handlePathChallengeFrame(frame *wire.PathChallengeFrame) {
 	s.queueControlFrame(&wire.PathResponseFrame{Data: frame.Data})
 }
 
+func (s *session) handleNewTokenFrame(frame *wire.NewTokenFrame) error {
+	if s.perspective == protocol.PerspectiveServer {
+		return qerr.Error(qerr.ProtocolViolation, "Received NEW_TOKEN frame from the client.")
+	}
+	if s.config.TokenStore != nil {
+		s.config.TokenStore.Put(s.tokenStoreKey, &ClientToken{data: frame.Token})
+	}
+	return nil
+}
+
 func (s *session) handleAckFrame(frame *wire.AckFrame, pn protocol.PacketNumber, encLevel protocol.EncryptionLevel) error {
 	if err := s.sentPacketHandler.ReceivedAck(frame, pn, encLevel, s.lastPacketReceivedTime); err != nil {
 		return err

session_test.go

@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"context"
 	"crypto/rand"
+	"crypto/tls"
 	"errors"
 	"net"
 	"runtime/pprof"
@@ -300,6 +301,13 @@ var _ = Describe("Session", func() {
 			Expect(frames).To(Equal([]wire.Frame{&wire.PathResponseFrame{Data: data}}))
 		})
 
+		It("rejects NEW_TOKEN frames", func() {
+			err := sess.handleNewTokenFrame(&wire.NewTokenFrame{})
+			Expect(err).To(HaveOccurred())
+			Expect(err).To(BeAssignableToTypeOf(&qerr.QuicError{}))
+			Expect(err.(*qerr.QuicError).ErrorCode).To(Equal(qerr.ProtocolViolation))
+		})
+
 		It("handles BLOCKED frames", func() {
 			err := sess.handleFrame(&wire.DataBlockedFrame{}, 0, protocol.EncryptionUnspecified)
 			Expect(err).NotTo(HaveOccurred())
@@ -1564,6 +1572,8 @@ var _ = Describe("Client Session", func() {
 		packer        *MockPacker
 		mconn         *mockConnection
 		cryptoSetup   *mocks.MockCryptoSetup
+		tlsConf       *tls.Config
+		quicConf      *Config
 	)
 
 	getPacket := func(hdr *wire.ExtendedHeader, data []byte) *receivedPacket {
@@ -1576,8 +1586,15 @@ var _ = Describe("Client Session", func() {
 	}
 
 	BeforeEach(func() {
+		quicConf = populateClientConfig(&Config{}, true)
+	})
+
+	JustBeforeEach(func() {
 		Eventually(areSessionsRunning).Should(BeFalse())
 
+		if tlsConf == nil {
+			tlsConf = &tls.Config{}
+		}
 		mconn = newMockConnection()
 		sessionRunner = NewMockSessionRunner(mockCtrl)
 		sessP, err := newClientSession(
@@ -1585,9 +1602,9 @@ var _ = Describe("Client Session", func() {
 			sessionRunner,
 			protocol.ConnectionID{8, 7, 6, 5, 4, 3, 2, 1},
 			protocol.ConnectionID{1, 2, 3, 4, 5, 6, 7, 8},
-			populateClientConfig(&Config{}, true),
-			nil, // tls.Config
-			42,  // initial packet number
+			quicConf,
+			tlsConf,
+			42, // initial packet number
 			&handshake.TransportParameters{},
 			protocol.VersionTLS,
 			utils.DefaultLogger,
@@ -1636,10 +1653,25 @@ var _ = Describe("Client Session", func() {
 		Eventually(sess.Context().Done()).Should(BeClosed())
 	})
 
+	Context("handling tokens", func() {
+		var mockTokenStore *MockTokenStore
+
+		BeforeEach(func() {
+			mockTokenStore = NewMockTokenStore(mockCtrl)
+			tlsConf = &tls.Config{ServerName: "server"}
+			quicConf.TokenStore = mockTokenStore
+		})
+
+		It("handles NEW_TOKEN frames", func() {
+			mockTokenStore.EXPECT().Put("server", &ClientToken{data: []byte("foobar")})
+			Expect(sess.handleNewTokenFrame(&wire.NewTokenFrame{Token: []byte("foobar")})).To(Succeed())
+		})
+	})
+
 	Context("handling Retry", func() {
 		var validRetryHdr *wire.ExtendedHeader
 
-		BeforeEach(func() {
+		JustBeforeEach(func() {
 			validRetryHdr = &wire.ExtendedHeader{
 				Header: wire.Header{
 					IsLongHeader:         true,