From 9d0301dd827ff1932ca451d40fdd69107444604b Mon Sep 17 00:00:00 2001
From: Felix Lange <fjl@twurst.com>
Date: Mon, 27 May 2024 18:16:24 +0200
Subject: [PATCH] p2p/discover: add more node update tests

---
 p2p/discover/table_test.go | 114 ++++++++++++++++++++++++++++---------
 1 file changed, 86 insertions(+), 28 deletions(-)

diff --git a/p2p/discover/table_test.go b/p2p/discover/table_test.go
index 0aaf7916de..0f9a20fda1 100644
--- a/p2p/discover/table_test.go
+++ b/p2p/discover/table_test.go
@@ -286,26 +286,23 @@ func TestTable_addInboundNode(t *testing.T) {
 	n2 := nodeAtDistance(tab.self().ID(), 256, net.IP{88, 77, 66, 2})
 	tab.addFoundNode(n1)
 	tab.addFoundNode(n2)
-	bucket := tab.bucket(n1.ID())
+	checkBucketContent(t, tab, []*enode.Node{n1.Node, n2.Node})
 
-	// Verify bucket content:
-	bcontent := []*node{n1, n2}
-	if !reflect.DeepEqual(unwrapNodes(bucket.entries), unwrapNodes(bcontent)) {
-		t.Fatalf("wrong bucket content: %v", bucket.entries)
-	}
-
-	// Add a changed version of n2.
+	// Add a changed version of n2. The bucket should be updated.
 	newrec := n2.Record()
 	newrec.Set(enr.IP{99, 99, 99, 99})
-	newn2 := wrapNode(enode.SignNull(newrec, n2.ID()))
-	tab.addInboundNode(newn2)
+	n2v2 := enode.SignNull(newrec, n2.ID())
+	tab.addInboundNode(wrapNode(n2v2))
+	checkBucketContent(t, tab, []*enode.Node{n1.Node, n2v2})
 
-	// Check that bucket is updated correctly.
-	newBcontent := []*node{n1, newn2}
-	if !reflect.DeepEqual(unwrapNodes(bucket.entries), unwrapNodes(newBcontent)) {
-		t.Fatalf("wrong bucket content after update: %v", bucket.entries)
-	}
-	checkIPLimitInvariant(t, tab)
+	// Try updating n2 without sequence number change. The update
+	// should not be accepted.
+	newrec = n2.Record()
+	newrec.Set(enr.IP{100, 100, 100, 100})
+	newrec.SetSeq(n2.Seq())
+	n2v3 := enode.SignNull(newrec, n2.ID())
+	tab.addInboundNode(wrapNode(n2v3))
+	checkBucketContent(t, tab, []*enode.Node{n1.Node, n2v2})
 }
 
 func TestTable_addFoundNode(t *testing.T) {
@@ -319,23 +316,84 @@ func TestTable_addFoundNode(t *testing.T) {
 	n2 := nodeAtDistance(tab.self().ID(), 256, net.IP{88, 77, 66, 2})
 	tab.addFoundNode(n1)
 	tab.addFoundNode(n2)
+	checkBucketContent(t, tab, []*enode.Node{n1.Node, n2.Node})
 
-	// Verify bucket content:
-	bcontent := []*node{n1, n2}
-	if !reflect.DeepEqual(tab.bucket(n1.ID()).entries, bcontent) {
-		t.Fatalf("wrong bucket content: %v", tab.bucket(n1.ID()).entries)
-	}
-
-	// Add a changed version of n2.
+	// Add a changed version of n2. The bucket should be updated.
 	newrec := n2.Record()
 	newrec.Set(enr.IP{99, 99, 99, 99})
-	newn2 := wrapNode(enode.SignNull(newrec, n2.ID()))
-	tab.addFoundNode(newn2)
+	n2v2 := enode.SignNull(newrec, n2.ID())
+	tab.addFoundNode(wrapNode(n2v2))
+	checkBucketContent(t, tab, []*enode.Node{n1.Node, n2v2})
 
-	// Check that bucket content is unchanged.
-	if !reflect.DeepEqual(tab.bucket(n1.ID()).entries, bcontent) {
-		t.Fatalf("wrong bucket content after update: %v", tab.bucket(n1.ID()).entries)
+	// Try updating n2 without a sequence number change.
+	// The update should not be accepted.
+	newrec = n2.Record()
+	newrec.Set(enr.IP{100, 100, 100, 100})
+	newrec.SetSeq(n2.Seq())
+	n2v3 := enode.SignNull(newrec, n2.ID())
+	tab.addFoundNode(wrapNode(n2v3))
+	checkBucketContent(t, tab, []*enode.Node{n1.Node, n2v2})
+}
+
+// This test checks that discv4 nodes can update their own endpoint via PING.
+func TestTable_addInboundNodeUpdateV4Accept(t *testing.T) {
+	tab, db := newTestTable(newPingRecorder(), Config{})
+	<-tab.initDone
+	defer db.Close()
+	defer tab.close()
+
+	// Add a v4 node.
+	key, _ := crypto.HexToECDSA("dd3757a8075e88d0f2b1431e7d3c5b1562e1c0aab9643707e8cbfcc8dae5cfe3")
+	n1 := enode.NewV4(&key.PublicKey, net.IP{88, 77, 66, 1}, 9000, 9000)
+	tab.addInboundNode(wrapNode(n1))
+	checkBucketContent(t, tab, []*enode.Node{n1})
+
+	// Add an updated version with changed IP.
+	// The update will be accepted because it has seq=0.
+	n1v2 := enode.NewV4(&key.PublicKey, net.IP{99, 99, 99, 99}, 9000, 9000)
+	tab.addInboundNode(wrapNode(n1v2))
+	checkBucketContent(t, tab, []*enode.Node{n1v2})
+}
+
+// This test checks that discv4 node entries will NOT be updated when a
+// changed record is found.
+func TestTable_addFoundNodeV4UpdateReject(t *testing.T) {
+	tab, db := newTestTable(newPingRecorder(), Config{})
+	<-tab.initDone
+	defer db.Close()
+	defer tab.close()
+
+	// Add a v4 node.
+	key, _ := crypto.HexToECDSA("dd3757a8075e88d0f2b1431e7d3c5b1562e1c0aab9643707e8cbfcc8dae5cfe3")
+	n1 := enode.NewV4(&key.PublicKey, net.IP{88, 77, 66, 1}, 9000, 9000)
+	tab.addFoundNode(wrapNode(n1))
+	checkBucketContent(t, tab, []*enode.Node{n1})
+
+	// Add an updated version with changed IP.
+	// The update won't be accepted because it isn't inbound.
+	n1v2 := enode.NewV4(&key.PublicKey, net.IP{99, 99, 99, 99}, 9000, 9000)
+	tab.addFoundNode(wrapNode(n1v2))
+	checkBucketContent(t, tab, []*enode.Node{n1})
+}
+
+func checkBucketContent(t *testing.T, tab *Table, nodes []*enode.Node) {
+	t.Helper()
+
+	b := tab.bucket(nodes[0].ID())
+	if reflect.DeepEqual(unwrapNodes(b.entries), nodes) {
+		return
 	}
+	t.Log("wrong bucket content. have nodes:")
+	for _, n := range b.entries {
+		t.Logf("  %v (seq=%v, ip=%v)", n.ID(), n.Seq(), n.IP())
+	}
+	t.Log("want nodes:")
+	for _, n := range nodes {
+		t.Logf("  %v (seq=%v, ip=%v)", n.ID(), n.Seq(), n.IP())
+	}
+	t.FailNow()
+
+	// Also check IP limits.
 	checkIPLimitInvariant(t, tab)
 }