diff --git a/p2p/transport/quicreuse/connmgr.go b/p2p/transport/quicreuse/connmgr.go index def7c5da56..7a7c1434c2 100644 --- 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 index e87574df76..48dd5ee940 100644 --- 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 index e329ea49dd..49e5aa2e88 100644 --- 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 }