diff --git a/p2p/discover/v5_udp.go b/p2p/discover/v5_udp.go index e667be1690..d53375b48b 100644 --- a/p2p/discover/v5_udp.go +++ b/p2p/discover/v5_udp.go @@ -216,6 +216,26 @@ func (t *UDPv5) Resolve(n *enode.Node) *enode.Node { return n } +// AllNodes returns all the nodes stored in the local table. +func (t *UDPv5) AllNodes() []*enode.Node { + t.tab.mutex.Lock() + defer t.tab.mutex.Unlock() + nodes := make([]*enode.Node, 0) + + for _, b := range &t.tab.buckets { + for _, n := range b.entries { + nodes = append(nodes, unwrapNode(n)) + } + } + return nodes +} + +// LocalNode returns the current local node running the +// protocol. +func (t *UDPv5) LocalNode() *enode.LocalNode { + return t.localNode +} + func (t *UDPv5) RandomNodes() enode.Iterator { if t.tab.len() == 0 { // All nodes were dropped, refresh. The very first query will hit this diff --git a/p2p/discover/v5_udp_test.go b/p2p/discover/v5_udp_test.go index 15ea0402c2..7d3915e2dc 100644 --- a/p2p/discover/v5_udp_test.go +++ b/p2p/discover/v5_udp_test.go @@ -445,6 +445,28 @@ func TestUDPv5_lookup(t *testing.T) { checkLookupResults(t, lookupTestnet, <-resultC) } +// This test checks the local node can be utilised to set key-values. +func TestUDPv5_LocalNode(t *testing.T) { + t.Parallel() + var cfg Config + node := startLocalhostV5(t, cfg) + defer node.Close() + localNd := node.LocalNode() + + // set value in node's local record + testVal := [4]byte{'A', 'B', 'C', 'D'} + localNd.Set(enr.WithEntry("testing", &testVal)) + + // retrieve the value from self to make sure it matches. + outputVal := [4]byte{} + if err := node.Self().Load(enr.WithEntry("testing", &outputVal)); err != nil { + t.Errorf("Could not load value from record: %v", err) + } + if testVal != outputVal { + t.Errorf("Wanted %#x to be retrieved from the record but instead got %#x", testVal, outputVal) + } +} + // udpV5Test is the framework for all tests above. // It runs the UDPv5 transport on a virtual socket and allows testing outgoing packets. type udpV5Test struct {