simplify generation of stateless reset tokens (#4858)

Author: marten-seemann

client.go

@@ -22,9 +22,10 @@ type client struct {
 	tlsConf *tls.Config
 	config  *Config
 
-	connIDGenerator ConnectionIDGenerator
-	srcConnID       protocol.ConnectionID
-	destConnID      protocol.ConnectionID
+	connIDGenerator   ConnectionIDGenerator
+	statelessResetter *statelessResetter
+	srcConnID         protocol.ConnectionID
+	destConnID        protocol.ConnectionID
 
 	initialPacketNumber  protocol.PacketNumber
 	hasNegotiatedVersion bool
@@ -137,13 +138,14 @@ func dial(
 	ctx context.Context,
 	conn sendConn,
 	connIDGenerator ConnectionIDGenerator,
+	statelessResetter *statelessResetter,
 	packetHandlers packetHandlerManager,
 	tlsConf *tls.Config,
 	config *Config,
 	onClose func(),
 	use0RTT bool,
 ) (quicConn, error) {
-	c, err := newClient(conn, connIDGenerator, config, tlsConf, onClose, use0RTT)
+	c, err := newClient(conn, connIDGenerator, statelessResetter, config, tlsConf, onClose, use0RTT)
 	if err != nil {
 		return nil, err
 	}
@@ -162,7 +164,15 @@ func dial(
 	return c.conn, nil
 }
 
-func newClient(sendConn sendConn, connIDGenerator ConnectionIDGenerator, config *Config, tlsConf *tls.Config, onClose func(), use0RTT bool) (*client, error) {
+func newClient(
+	sendConn sendConn,
+	connIDGenerator ConnectionIDGenerator,
+	statelessResetter *statelessResetter,
+	config *Config,
+	tlsConf *tls.Config,
+	onClose func(),
+	use0RTT bool,
+) (*client, error) {
 	srcConnID, err := connIDGenerator.GenerateConnectionID()
 	if err != nil {
 		return nil, err
@@ -172,17 +182,18 @@ func newClient(sendConn sendConn, connIDGenerator ConnectionIDGenerator, config
 		return nil, err
 	}
 	c := &client{
-		connIDGenerator: connIDGenerator,
-		srcConnID:       srcConnID,
-		destConnID:      destConnID,
-		sendConn:        sendConn,
-		use0RTT:         use0RTT,
-		onClose:         onClose,
-		tlsConf:         tlsConf,
-		config:          config,
-		version:         config.Versions[0],
-		handshakeChan:   make(chan struct{}),
-		logger:          utils.DefaultLogger.WithPrefix("client"),
+		connIDGenerator:   connIDGenerator,
+		statelessResetter: statelessResetter,
+		srcConnID:         srcConnID,
+		destConnID:        destConnID,
+		sendConn:          sendConn,
+		use0RTT:           use0RTT,
+		onClose:           onClose,
+		tlsConf:           tlsConf,
+		config:            config,
+		version:           config.Versions[0],
+		handshakeChan:     make(chan struct{}),
+		logger:            utils.DefaultLogger.WithPrefix("client"),
 	}
 	return c, nil
 }
@@ -197,6 +208,7 @@ func (c *client) dial(ctx context.Context) error {
 		c.destConnID,
 		c.srcConnID,
 		c.connIDGenerator,
+		c.statelessResetter,
 		c.config,
 		c.tlsConf,
 		c.initialPacketNumber,

client_test.go

@@ -33,6 +33,7 @@ var _ = Describe("Client", func() {
 			destConnID protocol.ConnectionID,
 			srcConnID protocol.ConnectionID,
 			connIDGenerator ConnectionIDGenerator,
+			statelessResetToken *statelessResetter,
 			conf *Config,
 			tlsConf *tls.Config,
 			initialPacketNumber protocol.PacketNumber,
@@ -107,6 +108,7 @@ var _ = Describe("Client", func() {
 				_ protocol.ConnectionID,
 				_ protocol.ConnectionID,
 				_ ConnectionIDGenerator,
+				_ *statelessResetter,
 				_ *Config,
 				_ *tls.Config,
 				_ protocol.PacketNumber,
@@ -124,7 +126,15 @@ var _ = Describe("Client", func() {
 				conn.EXPECT().HandshakeComplete().Return(c)
 				return conn
 			}
-			cl, err := newClient(packetConn, &protocol.DefaultConnectionIDGenerator{}, populateConfig(config), tlsConf, nil, false)
+			cl, err := newClient(
+				packetConn,
+				&protocol.DefaultConnectionIDGenerator{},
+				newStatelessResetter(nil),
+				populateConfig(config),
+				tlsConf,
+				nil,
+				false,
+			)
 			Expect(err).ToNot(HaveOccurred())
 			cl.packetHandlers = manager
 			Expect(cl).ToNot(BeNil())
@@ -144,6 +154,7 @@ var _ = Describe("Client", func() {
 				_ protocol.ConnectionID,
 				_ protocol.ConnectionID,
 				_ ConnectionIDGenerator,
+				_ *statelessResetter,
 				_ *Config,
 				_ *tls.Config,
 				_ protocol.PacketNumber,
@@ -161,7 +172,15 @@ var _ = Describe("Client", func() {
 				return conn
 			}
 
-			cl, err := newClient(packetConn, &protocol.DefaultConnectionIDGenerator{}, populateConfig(config), tlsConf, nil, true)
+			cl, err := newClient(
+				packetConn,
+				&protocol.DefaultConnectionIDGenerator{},
+				newStatelessResetter(nil),
+				populateConfig(config),
+				tlsConf,
+				nil,
+				true,
+			)
 			Expect(err).ToNot(HaveOccurred())
 			cl.packetHandlers = manager
 			Expect(cl).ToNot(BeNil())
@@ -181,6 +200,7 @@ var _ = Describe("Client", func() {
 				_ protocol.ConnectionID,
 				_ protocol.ConnectionID,
 				_ ConnectionIDGenerator,
+				_ *statelessResetter,
 				_ *Config,
 				_ *tls.Config,
 				_ protocol.PacketNumber,
@@ -197,7 +217,13 @@ var _ = Describe("Client", func() {
 				return conn
 			}
 			var closed bool
-			cl, err := newClient(packetConn, &protocol.DefaultConnectionIDGenerator{}, populateConfig(config), tlsConf, func() { closed = true }, true)
+			cl, err := newClient(
+				packetConn,
+				&protocol.DefaultConnectionIDGenerator{},
+				newStatelessResetter(nil),
+				populateConfig(config), tlsConf, func() { closed = true },
+				true,
+			)
 			Expect(err).ToNot(HaveOccurred())
 			cl.packetHandlers = manager
 			Expect(cl).ToNot(BeNil())
@@ -266,6 +292,7 @@ var _ = Describe("Client", func() {
 				_ protocol.ConnectionID,
 				_ protocol.ConnectionID,
 				_ ConnectionIDGenerator,
+				_ *statelessResetter,
 				configP *Config,
 				_ *tls.Config,
 				_ protocol.PacketNumber,
@@ -309,6 +336,7 @@ var _ = Describe("Client", func() {
 				_ protocol.ConnectionID,
 				connID protocol.ConnectionID,
 				_ ConnectionIDGenerator,
+				_ *statelessResetter,
 				configP *Config,
 				_ *tls.Config,
 				pn protocol.PacketNumber,

conn_id_generator.go

@@ -15,34 +15,34 @@ type connIDGenerator struct {
 	activeSrcConnIDs        map[uint64]protocol.ConnectionID
 	initialClientDestConnID *protocol.ConnectionID // nil for the client
 
-	addConnectionID        func(protocol.ConnectionID)
-	getStatelessResetToken func(protocol.ConnectionID) protocol.StatelessResetToken
-	removeConnectionID     func(protocol.ConnectionID)
-	retireConnectionID     func(protocol.ConnectionID)
-	replaceWithClosed      func([]protocol.ConnectionID, []byte)
-	queueControlFrame      func(wire.Frame)
+	addConnectionID    func(protocol.ConnectionID)
+	statelessResetter  *statelessResetter
+	removeConnectionID func(protocol.ConnectionID)
+	retireConnectionID func(protocol.ConnectionID)
+	replaceWithClosed  func([]protocol.ConnectionID, []byte)
+	queueControlFrame  func(wire.Frame)
 }
 
 func newConnIDGenerator(
 	initialConnectionID protocol.ConnectionID,
 	initialClientDestConnID *protocol.ConnectionID, // nil for the client
 	addConnectionID func(protocol.ConnectionID),
-	getStatelessResetToken func(protocol.ConnectionID) protocol.StatelessResetToken,
+	statelessResetter *statelessResetter,
 	removeConnectionID func(protocol.ConnectionID),
 	retireConnectionID func(protocol.ConnectionID),
 	replaceWithClosed func([]protocol.ConnectionID, []byte),
 	queueControlFrame func(wire.Frame),
 	generator ConnectionIDGenerator,
 ) *connIDGenerator {
 	m := &connIDGenerator{
-		generator:              generator,
-		activeSrcConnIDs:       make(map[uint64]protocol.ConnectionID),
-		addConnectionID:        addConnectionID,
-		getStatelessResetToken: getStatelessResetToken,
-		removeConnectionID:     removeConnectionID,
-		retireConnectionID:     retireConnectionID,
-		replaceWithClosed:      replaceWithClosed,
-		queueControlFrame:      queueControlFrame,
+		generator:          generator,
+		activeSrcConnIDs:   make(map[uint64]protocol.ConnectionID),
+		addConnectionID:    addConnectionID,
+		statelessResetter:  statelessResetter,
+		removeConnectionID: removeConnectionID,
+		retireConnectionID: retireConnectionID,
+		replaceWithClosed:  replaceWithClosed,
+		queueControlFrame:  queueControlFrame,
 	}
 	m.activeSrcConnIDs[0] = initialConnectionID
 	m.initialClientDestConnID = initialClientDestConnID
@@ -104,7 +104,7 @@ func (m *connIDGenerator) issueNewConnID() error {
 	m.queueControlFrame(&wire.NewConnectionIDFrame{
 		SequenceNumber:      m.highestSeq + 1,
 		ConnectionID:        connID,
-		StatelessResetToken: m.getStatelessResetToken(connID),
+		StatelessResetToken: m.statelessResetter.GetStatelessResetToken(connID),
 	})
 	m.highestSeq++
 	return nil

conn_id_generator_test.go

@@ -19,14 +19,11 @@ var _ = Describe("Connection ID Generator", func() {
 		replacedWithClosed []protocol.ConnectionID
 		queuedFrames       []wire.Frame
 		g                  *connIDGenerator
+		statelessResetter  *statelessResetter
 	)
 	initialConnID := protocol.ParseConnectionID([]byte{1, 2, 3, 4, 5, 6, 7})
 	initialClientDestConnID := protocol.ParseConnectionID([]byte{0xa, 0xb, 0xc, 0xd, 0xe})
-
-	connIDToToken := func(c protocol.ConnectionID) protocol.StatelessResetToken {
-		b := c.Bytes()[0]
-		return protocol.StatelessResetToken{b, b, b, b, b, b, b, b, b, b, b, b, b, b, b, b}
-	}
+	statelessResetter = newStatelessResetter(nil)
 
 	BeforeEach(func() {
 		addedConnIDs = nil
@@ -38,7 +35,7 @@ var _ = Describe("Connection ID Generator", func() {
 			initialConnID,
 			&initialClientDestConnID,
 			func(c protocol.ConnectionID) { addedConnIDs = append(addedConnIDs, c) },
-			connIDToToken,
+			statelessResetter,
 			func(c protocol.ConnectionID) { removedConnIDs = append(removedConnIDs, c) },
 			func(c protocol.ConnectionID) { retiredConnIDs = append(retiredConnIDs, c) },
 			func(cs []protocol.ConnectionID, _ []byte) { replacedWithClosed = append(replacedWithClosed, cs...) },
@@ -61,7 +58,7 @@ var _ = Describe("Connection ID Generator", func() {
 			nf := f.(*wire.NewConnectionIDFrame)
 			Expect(nf.SequenceNumber).To(BeEquivalentTo(i + 1))
 			Expect(nf.ConnectionID.Len()).To(Equal(7))
-			Expect(nf.StatelessResetToken).To(Equal(connIDToToken(nf.ConnectionID)))
+			Expect(nf.StatelessResetToken).To(Equal(statelessResetter.GetStatelessResetToken(nf.ConnectionID)))
 		}
 	})
 

connection.go

@@ -85,7 +85,6 @@ func (p *receivedPacket) Clone() *receivedPacket {
 
 type connRunner interface {
 	Add(protocol.ConnectionID, packetHandler) bool
-	GetStatelessResetToken(protocol.ConnectionID) protocol.StatelessResetToken
 	Retire(protocol.ConnectionID)
 	Remove(protocol.ConnectionID)
 	ReplaceWithClosed([]protocol.ConnectionID, []byte)
@@ -225,7 +224,7 @@ var newConnection = func(
 	destConnID protocol.ConnectionID,
 	srcConnID protocol.ConnectionID,
 	connIDGenerator ConnectionIDGenerator,
-	statelessResetToken protocol.StatelessResetToken,
+	statelessResetter *statelessResetter,
 	conf *Config,
 	tlsConf *tls.Config,
 	tokenGenerator *handshake.TokenGenerator,
@@ -263,7 +262,7 @@ var newConnection = func(
 		srcConnID,
 		&clientDestConnID,
 		func(connID protocol.ConnectionID) { runner.Add(connID, s) },
-		runner.GetStatelessResetToken,
+		statelessResetter,
 		runner.Remove,
 		runner.Retire,
 		runner.ReplaceWithClosed,
@@ -282,6 +281,7 @@ var newConnection = func(
 		s.logger,
 	)
 	s.maxPayloadSizeEstimate.Store(uint32(estimateMaxPayloadSize(protocol.ByteCount(s.config.InitialPacketSize))))
+	statelessResetToken := statelessResetter.GetStatelessResetToken(srcConnID)
 	params := &wire.TransportParameters{
 		InitialMaxStreamDataBidiLocal:   protocol.ByteCount(s.config.InitialStreamReceiveWindow),
 		InitialMaxStreamDataBidiRemote:  protocol.ByteCount(s.config.InitialStreamReceiveWindow),
@@ -340,6 +340,7 @@ var newClientConnection = func(
 	destConnID protocol.ConnectionID,
 	srcConnID protocol.ConnectionID,
 	connIDGenerator ConnectionIDGenerator,
+	statelessResetter *statelessResetter,
 	conf *Config,
 	tlsConf *tls.Config,
 	initialPacketNumber protocol.PacketNumber,
@@ -372,7 +373,7 @@ var newClientConnection = func(
 		srcConnID,
 		nil,
 		func(connID protocol.ConnectionID) { runner.Add(connID, s) },
-		runner.GetStatelessResetToken,
+		statelessResetter,
 		runner.Remove,
 		runner.Retire,
 		runner.ReplaceWithClosed,

connection_test.go

@@ -125,7 +125,7 @@ func newServerTestConnection(
 		protocol.ConnectionID{},
 		srcConnID,
 		&protocol.DefaultConnectionIDGenerator{},
-		protocol.StatelessResetToken{},
+		newStatelessResetter(nil),
 		populateConfig(config),
 		&tls.Config{},
 		handshake.NewTokenGenerator(handshake.TokenProtectorKey{}),
@@ -180,6 +180,7 @@ func newClientTestConnection(
 		destConnID,
 		srcConnID,
 		&protocol.DefaultConnectionIDGenerator{},
+		newStatelessResetter(nil),
 		populateConfig(config),
 		&tls.Config{ServerName: "quic-go.net"},
 		0,