From fd550f5408a1a3278b20bd5118af7e3bdc93890b Mon Sep 17 00:00:00 2001
From: Chen Kai <281165273grape@gmail.com>
Date: Sun, 2 Feb 2025 12:58:28 +0800
Subject: [PATCH] feat:change the discv5 challenge behavior to send same node
 only one time.

Signed-off-by: Chen Kai <281165273grape@gmail.com>
---
 p2p/discover/v5wire/encoding.go      | 29 +++++++++++++++-----------
 p2p/discover/v5wire/encoding_test.go | 31 ++++++++++++++++++++++++++++
 2 files changed, 48 insertions(+), 12 deletions(-)

diff --git a/p2p/discover/v5wire/encoding.go b/p2p/discover/v5wire/encoding.go
index 904a3ddec6..a6cc278bba 100644
--- a/p2p/discover/v5wire/encoding.go
+++ b/p2p/discover/v5wire/encoding.go
@@ -104,18 +104,19 @@ var DefaultProtocolID = [6]byte{'d', 'i', 's', 'c', 'v', '5'}
 
 // Errors.
 var (
-	errTooShort            = errors.New("packet too short")
-	errInvalidHeader       = errors.New("invalid packet header")
-	errInvalidFlag         = errors.New("invalid flag value in header")
-	errMinVersion          = errors.New("version of packet header below minimum")
-	errMsgTooShort         = errors.New("message/handshake packet below minimum size")
-	errAuthSize            = errors.New("declared auth size is beyond packet length")
-	errUnexpectedHandshake = errors.New("unexpected auth response, not in handshake")
-	errInvalidAuthKey      = errors.New("invalid ephemeral pubkey")
-	errNoRecord            = errors.New("expected ENR in handshake but none sent")
-	errInvalidNonceSig     = errors.New("invalid ID nonce signature")
-	errMessageTooShort     = errors.New("message contains no data")
-	errMessageDecrypt      = errors.New("cannot decrypt message")
+	errTooShort             = errors.New("packet too short")
+	errInvalidHeader        = errors.New("invalid packet header")
+	errInvalidFlag          = errors.New("invalid flag value in header")
+	errMinVersion           = errors.New("version of packet header below minimum")
+	errMsgTooShort          = errors.New("message/handshake packet below minimum size")
+	errAuthSize             = errors.New("declared auth size is beyond packet length")
+	errUnexpectedHandshake  = errors.New("unexpected auth response, not in handshake")
+	errInvalidAuthKey       = errors.New("invalid ephemeral pubkey")
+	errNoRecord             = errors.New("expected ENR in handshake but none sent")
+	errInvalidNonceSig      = errors.New("invalid ID nonce signature")
+	errMessageTooShort      = errors.New("message contains no data")
+	errMessageDecrypt       = errors.New("cannot decrypt message")
+	errChallengeAlreadySent = errors.New("challenge already sent to this node")
 )
 
 // Public errors.
@@ -217,6 +218,10 @@ func (c *Codec) Encode(id enode.ID, addr string, packet Packet, challenge *Whoar
 
 	// Store sent WHOAREYOU challenges.
 	if challenge, ok := packet.(*Whoareyou); ok {
+		// check if we already sent a challenge to this node
+		if c.sc.getHandshake(id, addr) != nil {
+			return nil, Nonce{}, errChallengeAlreadySent
+		}
 		challenge.ChallengeData = bytesCopy(&c.buf)
 		c.sc.storeSentHandshake(id, addr, challenge)
 	} else if msgData == nil {
diff --git a/p2p/discover/v5wire/encoding_test.go b/p2p/discover/v5wire/encoding_test.go
index df97e40e89..bbcab7e23f 100644
--- a/p2p/discover/v5wire/encoding_test.go
+++ b/p2p/discover/v5wire/encoding_test.go
@@ -155,6 +155,37 @@ func TestHandshake_norecord(t *testing.T) {
 	net.nodeA.expectDecode(t, NodesMsg, nodes)
 }
 
+func TestEncodeWhoareyouAlreadySent(t *testing.T) {
+	t.Parallel()
+	net := newHandshakeTest()
+	defer net.close()
+
+	// A -> B   RANDOM PACKET
+	packet, _ := net.nodeA.encode(t, net.nodeB, &Findnode{})
+	resp := net.nodeB.expectDecode(t, UnknownPacket, packet)
+
+	// A <- B   First WHOAREYOU
+	challenge1 := &Whoareyou{
+		Nonce:     resp.(*Unknown).Nonce,
+		IDNonce:   testIDnonce,
+		RecordSeq: 0,
+	}
+	whoareyou1, _ := net.nodeB.encode(t, net.nodeA, challenge1)
+	net.nodeA.expectDecode(t, WhoareyouPacket, whoareyou1)
+
+	// A <- B   Second WHOAREYOU (should fail)
+	challenge2 := &Whoareyou{
+		Nonce:     resp.(*Unknown).Nonce,
+		IDNonce:   testIDnonce,
+		RecordSeq: 0,
+	}
+
+	_, _, err := net.nodeB.c.Encode(net.nodeA.id(), net.nodeA.addr(), challenge2, nil)
+	if !errors.Is(err, errChallengeAlreadySent) {
+		t.Fatalf("second challenge should fail with errChallengeAlreadySent, got %v", err)
+	}
+}
+
 // In this test, A tries to send FINDNODE with existing secrets but B doesn't know
 // anything about A.
 func TestHandshake_rekey(t *testing.T) {