Merge pull request #2042 from lucas-clemente/token-cache

implement a store for address validation tokes

Changelog.md

@@ -3,6 +3,7 @@
 ## v0.13.0 (unreleased)
 
 - Add an `EarlyListener` that allows sending of 0.5-RTT data.
+- Add a `TokenStore` to store address validation tokens.
 
 ## v0.12.0 (2019-08-05)
 

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() {

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())
+		})
+	})
 })

interface.go

@@ -27,6 +27,23 @@ type Token struct {
 	SentTime     time.Time
 }
 
+// A ClientToken is a token received by the client.
+// It can be used to skip address validation on future connection attempts.
+type ClientToken struct {
+	data []byte
+}
+
+type TokenStore interface {
+	// Pop searches for a ClientToken associated with the given key.
+	// Since tokens are not supposed to be reused, it must remove the token from the cache.
+	// It returns nil when no token is found.
+	Pop(key string) (token *ClientToken)
+
+	// Put adds a token to the cache with the given key. It might get called
+	// multiple times in a connection.
+	Put(key string, token *ClientToken)
+}
+
 // An ErrorCode is an application-defined error code.
 // Valid values range between 0 and MAX_UINT62.
 type ErrorCode = protocol.ApplicationErrorCode
@@ -214,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

@@ -118,6 +118,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
@@ -334,6 +335,16 @@ var newClientSession = func(
 		s.perspective,
 		s.version,
 	)
+	if len(tlsConf.ServerName) > 0 {
+		s.tokenStoreKey = tlsConf.ServerName
+	} else {
+		s.tokenStoreKey = conn.RemoteAddr().String()
+	}
+	if s.config.TokenStore != nil {
+		if token := s.config.TokenStore.Pop(s.tokenStoreKey); token != nil {
+			s.packer.SetToken(token.data)
+		}
+	}
 	return s, s.postSetup()
 }
 
@@ -775,6 +786,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
@@ -901,6 +913,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"
@@ -311,6 +312,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())
@@ -1581,6 +1589,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 {
@@ -1593,8 +1603,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(
@@ -1602,9 +1619,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,
@@ -1653,10 +1670,27 @@ 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
+			mockTokenStore.EXPECT().Pop(gomock.Any())
+			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,

token_store.go

@@ -0,0 +1,117 @@
+package quic
+
+import (
+	"container/list"
+	"sync"
+
+	"github.com/lucas-clemente/quic-go/internal/utils"
+)
+
+type singleOriginTokenStore struct {
+	tokens []*ClientToken
+	len    int
+	p      int
+}
+
+func newSingleOriginTokenStore(size int) *singleOriginTokenStore {
+	return &singleOriginTokenStore{tokens: make([]*ClientToken, size)}
+}
+
+func (s *singleOriginTokenStore) Add(token *ClientToken) {
+	s.tokens[s.p] = token
+	s.p = s.index(s.p + 1)
+	s.len = utils.Min(s.len+1, len(s.tokens))
+}
+
+func (s *singleOriginTokenStore) Pop() *ClientToken {
+	s.p = s.index(s.p - 1)
+	token := s.tokens[s.p]
+	s.tokens[s.p] = nil
+	s.len = utils.Max(s.len-1, 0)
+	return token
+}
+
+func (s *singleOriginTokenStore) Len() int {
+	return s.len
+}
+
+func (s *singleOriginTokenStore) index(i int) int {
+	mod := len(s.tokens)
+	return (i + mod) % mod
+}
+
+type lruTokenStoreEntry struct {
+	key   string
+	cache *singleOriginTokenStore
+}
+
+type lruTokenStore struct {
+	mutex sync.Mutex
+
+	m                map[string]*list.Element
+	q                *list.List
+	capacity         int
+	singleOriginSize int
+}
+
+var _ TokenStore = &lruTokenStore{}
+
+// NewLRUTokenStore creates a new LRU cache for tokens received by the client.
+// maxOrigins specifies how many origins this cache is saving tokens for.
+// tokensPerOrigin specifies the maximum number of tokens per origin.
+func NewLRUTokenStore(maxOrigins, tokensPerOrigin int) TokenStore {
+	return &lruTokenStore{
+		m:                make(map[string]*list.Element),
+		q:                list.New(),
+		capacity:         maxOrigins,
+		singleOriginSize: tokensPerOrigin,
+	}
+}
+
+func (s *lruTokenStore) Put(key string, token *ClientToken) {
+	s.mutex.Lock()
+	defer s.mutex.Unlock()
+
+	if el, ok := s.m[key]; ok {
+		entry := el.Value.(*lruTokenStoreEntry)
+		entry.cache.Add(token)
+		s.q.MoveToFront(el)
+		return
+	}
+
+	if s.q.Len() < s.capacity {
+		entry := &lruTokenStoreEntry{
+			key:   key,
+			cache: newSingleOriginTokenStore(s.singleOriginSize),
+		}
+		entry.cache.Add(token)
+		s.m[key] = s.q.PushFront(entry)
+		return
+	}
+
+	elem := s.q.Back()
+	entry := elem.Value.(*lruTokenStoreEntry)
+	delete(s.m, entry.key)
+	entry.key = key
+	entry.cache = newSingleOriginTokenStore(s.singleOriginSize)
+	entry.cache.Add(token)
+	s.q.MoveToFront(elem)
+	s.m[key] = elem
+}
+
+func (s *lruTokenStore) Pop(key string) *ClientToken {
+	s.mutex.Lock()
+	defer s.mutex.Unlock()
+
+	var token *ClientToken
+	if el, ok := s.m[key]; ok {
+		s.q.MoveToFront(el)
+		cache := el.Value.(*lruTokenStoreEntry).cache
+		token = cache.Pop()
+		if cache.Len() == 0 {
+			s.q.Remove(el)
+			delete(s.m, key)
+		}
+	}
+	return token
+}

token_store_test.go

@@ -0,0 +1,108 @@
+package quic
+
+import (
+	"fmt"
+
+	. "github.com/onsi/ginkgo"
+	. "github.com/onsi/gomega"
+)
+
+var _ = Describe("Token Cache", func() {
+	var s TokenStore
+
+	BeforeEach(func() {
+		s = NewLRUTokenStore(3, 4)
+	})
+
+	mockToken := func(num int) *ClientToken {
+		return &ClientToken{data: []byte(fmt.Sprintf("%d", num))}
+	}
+
+	Context("for a single origin", func() {
+		const origin = "localhost"
+
+		It("adds and gets tokens", func() {
+			s.Put(origin, mockToken(1))
+			s.Put(origin, mockToken(2))
+			Expect(s.Pop(origin)).To(Equal(mockToken(2)))
+			Expect(s.Pop(origin)).To(Equal(mockToken(1)))
+			Expect(s.Pop(origin)).To(BeNil())
+		})
+
+		It("overwrites old tokens", func() {
+			s.Put(origin, mockToken(1))
+			s.Put(origin, mockToken(2))
+			s.Put(origin, mockToken(3))
+			s.Put(origin, mockToken(4))
+			s.Put(origin, mockToken(5))
+			Expect(s.Pop(origin)).To(Equal(mockToken(5)))
+			Expect(s.Pop(origin)).To(Equal(mockToken(4)))
+			Expect(s.Pop(origin)).To(Equal(mockToken(3)))
+			Expect(s.Pop(origin)).To(Equal(mockToken(2)))
+			Expect(s.Pop(origin)).To(BeNil())
+		})
+
+		It("continues after getting a token", func() {
+			s.Put(origin, mockToken(1))
+			s.Put(origin, mockToken(2))
+			s.Put(origin, mockToken(3))
+			Expect(s.Pop(origin)).To(Equal(mockToken(3)))
+			s.Put(origin, mockToken(4))
+			s.Put(origin, mockToken(5))
+			Expect(s.Pop(origin)).To(Equal(mockToken(5)))
+			Expect(s.Pop(origin)).To(Equal(mockToken(4)))
+			Expect(s.Pop(origin)).To(Equal(mockToken(2)))
+			Expect(s.Pop(origin)).To(Equal(mockToken(1)))
+			Expect(s.Pop(origin)).To(BeNil())
+		})
+	})
+
+	Context("for multiple origins", func() {
+		It("adds and gets tokens", func() {
+			s.Put("host1", mockToken(1))
+			s.Put("host2", mockToken(2))
+			Expect(s.Pop("host1")).To(Equal(mockToken(1)))
+			Expect(s.Pop("host1")).To(BeNil())
+			Expect(s.Pop("host2")).To(Equal(mockToken(2)))
+			Expect(s.Pop("host2")).To(BeNil())
+		})
+
+		It("evicts old entries", func() {
+			s.Put("host1", mockToken(1))
+			s.Put("host2", mockToken(2))
+			s.Put("host3", mockToken(3))
+			s.Put("host4", mockToken(4))
+			Expect(s.Pop("host1")).To(BeNil())
+			Expect(s.Pop("host2")).To(Equal(mockToken(2)))
+			Expect(s.Pop("host3")).To(Equal(mockToken(3)))
+			Expect(s.Pop("host4")).To(Equal(mockToken(4)))
+		})
+
+		It("moves old entries to the front, when new tokens are added", func() {
+			s.Put("host1", mockToken(1))
+			s.Put("host2", mockToken(2))
+			s.Put("host3", mockToken(3))
+			s.Put("host1", mockToken(11))
+			// make sure one is evicted
+			s.Put("host4", mockToken(4))
+			Expect(s.Pop("host2")).To(BeNil())
+			Expect(s.Pop("host1")).To(Equal(mockToken(11)))
+			Expect(s.Pop("host1")).To(Equal(mockToken(1)))
+			Expect(s.Pop("host3")).To(Equal(mockToken(3)))
+			Expect(s.Pop("host4")).To(Equal(mockToken(4)))
+		})
+
+		It("deletes hosts that are empty", func() {
+			s.Put("host1", mockToken(1))
+			s.Put("host2", mockToken(2))
+			s.Put("host3", mockToken(3))
+			Expect(s.Pop("host2")).To(Equal(mockToken(2)))
+			Expect(s.Pop("host2")).To(BeNil())
+			// host2 is now empty and should have been deleted, making space for host4
+			s.Put("host4", mockToken(4))
+			Expect(s.Pop("host1")).To(Equal(mockToken(1)))
+			Expect(s.Pop("host3")).To(Equal(mockToken(3)))
+			Expect(s.Pop("host4")).To(Equal(mockToken(4)))
+		})
+	})
+})