From 7a63faf734d7150e764c0c08de7f83ceaf64e272 Mon Sep 17 00:00:00 2001
From: Nishant Das <nish1993@hotmail.com>
Date: Thu, 16 Apr 2020 21:58:37 +0800
Subject: [PATCH] p2p/discover: add helper methods to UDPv5 (#20918)

This adds two new methods to UDPv5, AllNodes and LocalNode.

AllNodes returns all the nodes stored in the local table; this is
useful for the purposes of metrics collection and also debugging any
potential issues with other discovery v5 implementations.

LocalNode returns the local node object. The reason for exposing this
is so that users can modify and set/delete new key-value entries in
the local record.
---
 p2p/discover/v5_udp.go      | 20 ++++++++++++++++++++
 p2p/discover/v5_udp_test.go | 22 ++++++++++++++++++++++
 2 files changed, 42 insertions(+)

diff --git a/p2p/discover/v5_udp.go b/p2p/discover/v5_udp.go
index e667be169..d53375b48 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 15ea0402c..7d3915e2d 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 {