feat(quicreuse): Support custom listenUDP methods

libp2p · Jan 21, 2025 · 3a951e1 · 3a951e1
1 parent 6afce8c
commit 3a951e1
Show file tree

Hide file tree

Showing 3 changed files with 50 additions and 7 deletions.
diff --git a/p2p/transport/quicreuse/connmgr.go b/p2p/transport/quicreuse/connmgr.go
@@ -37,6 +37,8 @@ type ConnManager struct {
 	reuseUDP6       *reuse
 	enableReuseport bool
 
+	customListenUDP listenUDP
+
 	enableMetrics bool
 	registerer    prometheus.Registerer
 
@@ -77,7 +79,9 @@ func NewConnManager(statelessResetKey quic.StatelessResetKey, tokenKey quic.Toke
 	cm.serverConfig = serverConfig
 	if cm.enableReuseport {
 		cm.reuseUDP4 = newReuse(&statelessResetKey, &tokenKey)
+		cm.reuseUDP4.customListenUDP = cm.customListenUDP
 		cm.reuseUDP6 = newReuse(&statelessResetKey, &tokenKey)
+		cm.reuseUDP6.customListenUDP = cm.customListenUDP
 	}
 	return cm, nil
 }
@@ -238,7 +242,13 @@ func (c *ConnManager) transportForListen(association any, network string, laddr
 		return tr, nil
 	}
 
-	conn, err := net.ListenUDP(network, laddr)
+	var conn net.PacketConn
+	var err error
+	if c.customListenUDP != nil {
+		conn, err = c.customListenUDP(network, laddr)
+	} else {
+		conn, err = net.ListenUDP(network, laddr)
+	}
 	if err != nil {
 		return nil, err
 	}
@@ -320,7 +330,14 @@ func (c *ConnManager) TransportWithAssociationForDial(association any, network s
 	case "udp6":
 		laddr = &net.UDPAddr{IP: net.IPv6zero, Port: 0}
 	}
-	conn, err := net.ListenUDP(network, laddr)
+	var conn net.PacketConn
+	var err error
+	if c.customListenUDP != nil {
+		conn, err = c.customListenUDP(network, laddr)
+	} else {
+		conn, err = net.ListenUDP(network, laddr)
+	}
+
 	if err != nil {
 		return nil, err
 	}

diff --git a/p2p/transport/quicreuse/options.go b/p2p/transport/quicreuse/options.go
@@ -1,9 +1,22 @@
 package quicreuse
 
-import "github.com/prometheus/client_golang/prometheus"
+import (
+	"net"
+
+	"github.com/prometheus/client_golang/prometheus"
+)
 
 type Option func(*ConnManager) error
 
+type listenUDP func(network string, laddr *net.UDPAddr) (net.PacketConn, error)
+
+func CustomListenUDP(f listenUDP) Option {
+	return func(m *ConnManager) error {
+		m.customListenUDP = f
+		return nil
+	}
+}
+
 func DisableReuseport() Option {
 	return func(m *ConnManager) error {
 		m.enableReuseport = false

diff --git a/p2p/transport/quicreuse/reuse.go b/p2p/transport/quicreuse/reuse.go
@@ -165,8 +165,9 @@ func (c *refcountedTransport) ShouldGarbageCollect(now time.Time) bool {
 type reuse struct {
 	mutex sync.Mutex
 
-	closeChan  chan struct{}
-	gcStopChan chan struct{}
+	closeChan       chan struct{}
+	gcStopChan      chan struct{}
+	customListenUDP listenUDP
 
 	routes  routing.Router
 	unicast map[string] /* IP.String() */ map[int] /* port */ *refcountedTransport
@@ -323,7 +324,13 @@ func (r *reuse) transportForDialLocked(association any, network string, source *
 	case "udp6":
 		addr = &net.UDPAddr{IP: net.IPv6zero, Port: 0}
 	}
-	conn, err := net.ListenUDP(network, addr)
+	var conn net.PacketConn
+	var err error
+	if r.customListenUDP != nil {
+		conn, err = r.customListenUDP(network, addr)
+	} else {
+		conn, err = net.ListenUDP(network, addr)
+	}
 	if err != nil {
 		return nil, err
 	}
@@ -389,7 +396,13 @@ func (r *reuse) TransportForListen(network string, laddr *net.UDPAddr) (*refcoun
 		}
 	}
 
-	conn, err := net.ListenUDP(network, laddr)
+	var conn net.PacketConn
+	var err error
+	if r.customListenUDP != nil {
+		conn, err = r.customListenUDP(network, laddr)
+	} else {
+		conn, err = net.ListenUDP(network, laddr)
+	}
 	if err != nil {
 		return nil, err
 	}