quicreuse: pass the token generator key to the reuse

libp2p · Oct 22, 2023 · bf96b7f · bf96b7f
1 parent dd036ae
commit bf96b7f
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 11 deletions.
diff --git a/p2p/transport/quicreuse/connmgr.go b/p2p/transport/quicreuse/connmgr.go
@@ -61,8 +61,8 @@ func NewConnManager(statelessResetKey quic.StatelessResetKey, tokenKey quic.Toke
 	cm.clientConfig = quicConf
 	cm.serverConfig = serverConfig
 	if cm.enableReuseport {
-		cm.reuseUDP4 = newReuse(&statelessResetKey)
-		cm.reuseUDP6 = newReuse(&statelessResetKey)
+		cm.reuseUDP4 = newReuse(&statelessResetKey, &tokenKey)
+		cm.reuseUDP6 = newReuse(&statelessResetKey, &tokenKey)
 	}
 	return cm, nil
 }

diff --git a/p2p/transport/quicreuse/reuse.go b/p2p/transport/quicreuse/reuse.go
@@ -123,16 +123,18 @@ type reuse struct {
 	globalDialers map[int]*refcountedTransport
 
 	statelessResetKey *quic.StatelessResetKey
+	tokenGeneratorKey *quic.TokenGeneratorKey
 }
 
-func newReuse(srk *quic.StatelessResetKey) *reuse {
+func newReuse(srk *quic.StatelessResetKey, tokenKey *quic.TokenGeneratorKey) *reuse {
 	r := &reuse{
 		unicast:           make(map[string]map[int]*refcountedTransport),
 		globalListeners:   make(map[int]*refcountedTransport),
 		globalDialers:     make(map[int]*refcountedTransport),
 		closeChan:         make(chan struct{}),
 		gcStopChan:        make(chan struct{}),
 		statelessResetKey: srk,
+		tokenGeneratorKey: tokenKey,
 	}
 	go r.gc()
 	return r
@@ -268,6 +270,7 @@ func (r *reuse) transportForDialLocked(network string, source *net.IP) (*refcoun
 	tr := &refcountedTransport{Transport: quic.Transport{
 		Conn:              conn,
 		StatelessResetKey: r.statelessResetKey,
+		TokenGeneratorKey: r.tokenGeneratorKey,
 	}, packetConn: conn}
 	r.globalDialers[conn.LocalAddr().(*net.UDPAddr).Port] = tr
 	return tr, nil

diff --git a/p2p/transport/quicreuse/reuse_test.go b/p2p/transport/quicreuse/reuse_test.go
@@ -61,7 +61,7 @@ func cleanup(t *testing.T, reuse *reuse) {
 }
 
 func TestReuseListenOnAllIPv4(t *testing.T) {
-	reuse := newReuse(nil)
+	reuse := newReuse(nil, nil)
 	require.Eventually(t, isGarbageCollectorRunning, 500*time.Millisecond, 50*time.Millisecond, "expected garbage collector to be running")
 	cleanup(t, reuse)
 
@@ -73,7 +73,7 @@ func TestReuseListenOnAllIPv4(t *testing.T) {
 }
 
 func TestReuseListenOnAllIPv6(t *testing.T) {
-	reuse := newReuse(nil)
+	reuse := newReuse(nil, nil)
 	require.Eventually(t, isGarbageCollectorRunning, 500*time.Millisecond, 50*time.Millisecond, "expected garbage collector to be running")
 	cleanup(t, reuse)
 
@@ -86,7 +86,7 @@ func TestReuseListenOnAllIPv6(t *testing.T) {
 }
 
 func TestReuseCreateNewGlobalConnOnDial(t *testing.T) {
-	reuse := newReuse(nil)
+	reuse := newReuse(nil, nil)
 	cleanup(t, reuse)
 
 	addr, err := net.ResolveUDPAddr("udp4", "1.1.1.1:1234")
@@ -100,7 +100,7 @@ func TestReuseCreateNewGlobalConnOnDial(t *testing.T) {
 }
 
 func TestReuseConnectionWhenDialing(t *testing.T) {
-	reuse := newReuse(nil)
+	reuse := newReuse(nil, nil)
 	cleanup(t, reuse)
 
 	addr, err := net.ResolveUDPAddr("udp4", "0.0.0.0:0")
@@ -117,7 +117,7 @@ func TestReuseConnectionWhenDialing(t *testing.T) {
 }
 
 func TestReuseConnectionWhenListening(t *testing.T) {
-	reuse := newReuse(nil)
+	reuse := newReuse(nil, nil)
 	cleanup(t, reuse)
 
 	raddr, err := net.ResolveUDPAddr("udp4", "1.1.1.1:1234")
@@ -132,7 +132,7 @@ func TestReuseConnectionWhenListening(t *testing.T) {
 }
 
 func TestReuseConnectionWhenDialBeforeListen(t *testing.T) {
-	reuse := newReuse(nil)
+	reuse := newReuse(nil, nil)
 	cleanup(t, reuse)
 
 	// dial any address
@@ -166,7 +166,7 @@ func TestReuseListenOnSpecificInterface(t *testing.T) {
 	if platformHasRoutingTables() {
 		t.Skip("this test only works on platforms that support routing tables")
 	}
-	reuse := newReuse(nil)
+	reuse := newReuse(nil, nil)
 	cleanup(t, reuse)
 
 	router, err := netroute.New()
@@ -203,7 +203,7 @@ func TestReuseGarbageCollect(t *testing.T) {
 		maxUnusedDuration = 10 * maxUnusedDuration
 	}
 
-	reuse := newReuse(nil)
+	reuse := newReuse(nil, nil)
 	cleanup(t, reuse)
 
 	numGlobals := func() int {