p2p: add accessors for discovery instances

This commit is contained in:
Felix Lange 2024-05-23 12:07:16 +02:00
parent 16b4386f1e
commit 4386e8a662

View file

@ -190,8 +190,8 @@ type Server struct {
nodedb *enode.DB nodedb *enode.DB
localnode *enode.LocalNode localnode *enode.LocalNode
ntab *discover.UDPv4 discv4 *discover.UDPv4
DiscV5 *discover.UDPv5 discv5 *discover.UDPv5
discmix *enode.FairMix discmix *enode.FairMix
dialsched *dialScheduler dialsched *dialScheduler
@ -400,6 +400,16 @@ func (srv *Server) Self() *enode.Node {
return ln.Node() return ln.Node()
} }
// DiscoveryV4 returns the discovery v4 instance, if configured.
func (srv *Server) DiscoveryV4() *discover.UDPv4 {
return srv.discv4
}
// DiscoveryV4 returns the discovery v5 instance, if configured.
func (srv *Server) DiscoveryV5() *discover.UDPv5 {
return srv.discv5
}
// Stop terminates the server and all active peer connections. // Stop terminates the server and all active peer connections.
// It blocks until all active connections have been closed. // It blocks until all active connections have been closed.
func (srv *Server) Stop() { func (srv *Server) Stop() {
@ -547,13 +557,13 @@ func (srv *Server) setupDiscovery() error {
) )
// If both versions of discovery are running, setup a shared // If both versions of discovery are running, setup a shared
// connection, so v5 can read unhandled messages from v4. // connection, so v5 can read unhandled messages from v4.
if srv.DiscoveryV4 && srv.DiscoveryV5 { if srv.Config.DiscoveryV4 && srv.Config.DiscoveryV5 {
unhandled = make(chan discover.ReadPacket, 100) unhandled = make(chan discover.ReadPacket, 100)
sconn = &sharedUDPConn{conn, unhandled} sconn = &sharedUDPConn{conn, unhandled}
} }
// Start discovery services. // Start discovery services.
if srv.DiscoveryV4 { if srv.Config.DiscoveryV4 {
cfg := discover.Config{ cfg := discover.Config{
PrivateKey: srv.PrivateKey, PrivateKey: srv.PrivateKey,
NetRestrict: srv.NetRestrict, NetRestrict: srv.NetRestrict,
@ -565,17 +575,17 @@ func (srv *Server) setupDiscovery() error {
if err != nil { if err != nil {
return err return err
} }
srv.ntab = ntab srv.discv4 = ntab
srv.discmix.AddSource(ntab.RandomNodes()) srv.discmix.AddSource(ntab.RandomNodes())
} }
if srv.DiscoveryV5 { if srv.Config.DiscoveryV5 {
cfg := discover.Config{ cfg := discover.Config{
PrivateKey: srv.PrivateKey, PrivateKey: srv.PrivateKey,
NetRestrict: srv.NetRestrict, NetRestrict: srv.NetRestrict,
Bootnodes: srv.BootstrapNodesV5, Bootnodes: srv.BootstrapNodesV5,
Log: srv.log, Log: srv.log,
} }
srv.DiscV5, err = discover.ListenV5(sconn, srv.localnode, cfg) srv.discv5, err = discover.ListenV5(sconn, srv.localnode, cfg)
if err != nil { if err != nil {
return err return err
} }
@ -602,8 +612,8 @@ func (srv *Server) setupDialScheduler() {
dialer: srv.Dialer, dialer: srv.Dialer,
clock: srv.clock, clock: srv.clock,
} }
if srv.ntab != nil { if srv.discv4 != nil {
config.resolver = srv.ntab config.resolver = srv.discv4
} }
if config.dialer == nil { if config.dialer == nil {
config.dialer = tcpDialer{&net.Dialer{Timeout: defaultDialTimeout}} config.dialer = tcpDialer{&net.Dialer{Timeout: defaultDialTimeout}}
@ -799,11 +809,11 @@ running:
srv.log.Trace("P2P networking is spinning down") srv.log.Trace("P2P networking is spinning down")
// Terminate discovery. If there is a running lookup it will terminate soon. // Terminate discovery. If there is a running lookup it will terminate soon.
if srv.ntab != nil { if srv.discv4 != nil {
srv.ntab.Close() srv.discv4.Close()
} }
if srv.DiscV5 != nil { if srv.discv5 != nil {
srv.DiscV5.Close() srv.discv5.Close()
} }
// Disconnect all peers. // Disconnect all peers.
for _, p := range peers { for _, p := range peers {