Use own pointer to packetConn

libp2p · Jul 12, 2023 · 8560e4e · 8560e4e
1 parent 40011e7
commit 8560e4e
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 6 deletions.
diff --git a/p2p/transport/quicreuse/connmgr.go b/p2p/transport/quicreuse/connmgr.go
@@ -164,7 +164,7 @@ func (c *ConnManager) transportForListen(network string, laddr *net.UDPAddr) (re
 	if err != nil {
 		return nil, err
 	}
-	return &singleOwnerTransport{tr: quic.Transport{Conn: conn, StatelessResetKey: &c.srk}}, nil
+	return &singleOwnerTransport{tr: quic.Transport{Conn: conn, StatelessResetKey: &c.srk}, packetConn: conn}, nil
 }
 
 func (c *ConnManager) DialQUIC(ctx context.Context, raddr ma.Multiaddr, tlsConf *tls.Config, allowWindowIncrease func(conn quic.Connection, delta uint64) bool) (quic.Connection, error) {
@@ -221,7 +221,7 @@ func (c *ConnManager) TransportForDial(network string, raddr *net.UDPAddr) (refC
 	if err != nil {
 		return nil, err
 	}
-	return &singleOwnerTransport{tr: quic.Transport{Conn: conn, StatelessResetKey: &c.srk}}, nil
+	return &singleOwnerTransport{tr: quic.Transport{Conn: conn, StatelessResetKey: &c.srk}, packetConn: conn}, nil
 }
 
 func (c *ConnManager) Protocols() []int {

diff --git a/p2p/transport/quicreuse/reuse.go b/p2p/transport/quicreuse/reuse.go
@@ -24,6 +24,9 @@ type refCountedQuicTransport interface {
 
 type singleOwnerTransport struct {
 	tr quic.Transport
+
+	// Used to write packets directly around QUIC.
+	packetConn net.PacketConn
 }
 
 func (c *singleOwnerTransport) IncreaseCount() {}
@@ -41,7 +44,7 @@ func (c *singleOwnerTransport) LocalAddr() net.Addr {
 
 func (c *singleOwnerTransport) WriteTo(b []byte, addr net.Addr) (int, error) {
 	// Safe because we called quic.OptimizeConn ourselves.
-	return c.tr.Conn.WriteTo(b, addr)
+	return c.packetConn.WriteTo(b, addr)
 }
 
 func (c *singleOwnerTransport) Close() error {
@@ -57,6 +60,9 @@ var (
 type refcountedTransport struct {
 	tr quic.Transport
 
+	// Used to write packets directly around QUIC.
+	packetConn net.PacketConn
+
 	mutex       sync.Mutex
 	refCount    int
 	unusedSince time.Time
@@ -79,7 +85,7 @@ func (c *refcountedTransport) Close() error {
 
 func (c *refcountedTransport) WriteTo(b []byte, addr net.Addr) (int, error) {
 	// Safe because we called quic.OptimizeConn ourselves.
-	return c.tr.Conn.WriteTo(b, addr)
+	return c.packetConn.WriteTo(b, addr)
 }
 
 func (c *refcountedTransport) LocalAddr() net.Addr {
@@ -265,7 +271,7 @@ func (r *reuse) transportForDialLocked(network string, source *net.IP) (*refcoun
 		Conn:              conn,
 		StatelessResetKey: r.statelessResetKey,
 		Tracer:            r.metricsTracer,
-	}}
+	}, packetConn: conn}
 	r.globalDialers[conn.LocalAddr().(*net.UDPAddr).Port] = rconn
 	return rconn, nil
 }
@@ -313,7 +319,7 @@ func (r *reuse) TransportForListen(network string, laddr *net.UDPAddr) (*refcoun
 		Conn:              conn,
 		StatelessResetKey: r.statelessResetKey,
 		Tracer:            r.metricsTracer,
-	}}
+	}, packetConn: conn}
 
 	tr.IncreaseCount()