p2p/discover: add new methods to udpV5

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

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 {