From c5e94b67dbf83d1b57b0d2fdb758215069a0aab6 Mon Sep 17 00:00:00 2001
From: thinkAfCod <q315xia@163.com>
Date: Wed, 2 Apr 2025 16:57:46 +0800
Subject: [PATCH] p2p/discover: cache raw bytes of whoareyou packet

---
 p2p/discover/v5_udp_test.go     | 47 ++++++++++++++++++++++++---------
 p2p/discover/v5wire/encoding.go | 18 ++++++++++---
 p2p/discover/v5wire/msg.go      |  3 +++
 3 files changed, 52 insertions(+), 16 deletions(-)

diff --git a/p2p/discover/v5_udp_test.go b/p2p/discover/v5_udp_test.go
index 3026dff538..3d96b7434b 100644
--- a/p2p/discover/v5_udp_test.go
+++ b/p2p/discover/v5_udp_test.go
@@ -181,29 +181,38 @@ func TestUDPv5_handshakeRepeatChallenge(t *testing.T) {
 	nonce1 := v5wire.Nonce{1}
 	nonce2 := v5wire.Nonce{2}
 	nonce3 := v5wire.Nonce{3}
-	check := func(p *v5wire.Whoareyou, wantNonce v5wire.Nonce) {
+	var authTag1 v5wire.Nonce
+	//var record *v5wire.Whoareyou
+	check := func(p *v5wire.Whoareyou, actualAuthTag, wantNonce v5wire.Nonce) {
 		t.Helper()
 		if p.Nonce != wantNonce {
 			t.Error("wrong nonce in WHOAREYOU:", p.Nonce, wantNonce)
 		}
+		if authTag1 == (v5wire.Nonce{}) {
+			authTag1 = actualAuthTag
+		} else {
+			if authTag1 != authTag1 {
+				t.Error("wrong auth tag in WHOAREYOU header:", authTag1, wantNonce)
+			}
+		}
 	}
 
 	// Unknown packet from unknown node.
 	test.packetIn(&v5wire.Unknown{Nonce: nonce1})
-	test.waitPacketOut(func(p *v5wire.Whoareyou, addr netip.AddrPort, _ v5wire.Nonce) {
-		check(p, nonce1)
+	test.waitPacketOut(func(p *v5wire.Whoareyou, addr netip.AddrPort, authTag v5wire.Nonce) {
+		check(p, authTag, nonce1)
 	})
 
 	// Second unknown packet. Here we expect the response to reference the
 	// first unknown packet.
 	test.packetIn(&v5wire.Unknown{Nonce: nonce2})
-	test.waitPacketOut(func(p *v5wire.Whoareyou, addr netip.AddrPort, _ v5wire.Nonce) {
-		check(p, nonce1)
+	test.waitPacketOut(func(p *v5wire.Whoareyou, addr netip.AddrPort, authTag v5wire.Nonce) {
+		check(p, authTag, nonce1)
 	})
 	// Third unknown packet. This should still return the first nonce.
 	test.packetIn(&v5wire.Unknown{Nonce: nonce3})
-	test.waitPacketOut(func(p *v5wire.Whoareyou, addr netip.AddrPort, _ v5wire.Nonce) {
-		check(p, nonce1)
+	test.waitPacketOut(func(p *v5wire.Whoareyou, addr netip.AddrPort, authTag v5wire.Nonce) {
+		check(p, authTag, nonce1)
 	})
 }
 
@@ -766,20 +775,32 @@ type testCodecFrame struct {
 }
 
 func (c *testCodec) Encode(toID enode.ID, addr string, p v5wire.Packet, _ *v5wire.Whoareyou) ([]byte, v5wire.Nonce, error) {
-	c.ctr++
-	var authTag v5wire.Nonce
-	binary.BigEndian.PutUint64(authTag[:], c.ctr)
-
-	if w, ok := p.(*v5wire.Whoareyou); ok {
+	if _, ok := p.(*v5wire.Whoareyou); ok {
 		// Store recently sent Whoareyou challenges.
 		if c.sentChallenges == nil {
 			c.sentChallenges = make(map[enode.ID]*v5wire.Whoareyou)
 		}
-		c.sentChallenges[toID] = w
+		if sentWhoareyou := c.sentChallenges[toID]; sentWhoareyou != nil {
+			// If we've already sent a challenge to this node, don't send another one.
+			return sentWhoareyou.Encoded, sentWhoareyou.Nonce, nil
+		}
 	}
 
+	c.ctr++
+	var authTag v5wire.Nonce
+	binary.BigEndian.PutUint64(authTag[:], c.ctr)
+
 	penc, _ := rlp.EncodeToBytes(p)
 	frame, err := rlp.EncodeToBytes(testCodecFrame{c.id, authTag, p.Kind(), penc})
+	if err != nil {
+		return frame, authTag, err
+	}
+
+	if w, ok := p.(*v5wire.Whoareyou); ok {
+		c.sentChallenges[toID] = w
+		w.Nonce = authTag
+		w.Encoded = frame
+	}
 	return frame, authTag, err
 }
 
diff --git a/p2p/discover/v5wire/encoding.go b/p2p/discover/v5wire/encoding.go
index e50b7cd16d..15b35a1e31 100644
--- a/p2p/discover/v5wire/encoding.go
+++ b/p2p/discover/v5wire/encoding.go
@@ -189,6 +189,10 @@ func (c *Codec) Encode(id enode.ID, addr string, packet Packet, challenge *Whoar
 	)
 	switch {
 	case packet.Kind() == WhoareyouPacket:
+		// just send the WHOAREYOU packet raw again, rather than the re-encoded challenge data
+		if sentWhoareyou := c.sc.getHandshake(id, addr); sentWhoareyou != nil {
+			return sentWhoareyou.Encoded, sentWhoareyou.Nonce, nil
+		}
 		head, err = c.encodeWhoareyou(id, packet.(*Whoareyou))
 	case challenge != nil:
 		// We have an unanswered challenge, send handshake.
@@ -215,19 +219,27 @@ func (c *Codec) Encode(id enode.ID, addr string, packet Packet, challenge *Whoar
 	// Encode header data.
 	c.writeHeaders(&head)
 
+	var enc []byte
 	// Store sent WHOAREYOU challenges.
 	if challenge, ok := packet.(*Whoareyou); ok {
 		challenge.ChallengeData = bytesCopy(&c.buf)
+		enc, err = c.EncodeRaw(id, head, msgData)
+		if err != nil {
+			return nil, Nonce{}, err
+		}
+		challenge.Encoded = enc
 		c.sc.storeSentHandshake(id, addr, challenge)
-	} else if msgData == nil {
+		return enc, head.Nonce, err
+	}
+
+	if msgData == nil {
 		headerData := c.buf.Bytes()
 		msgData, err = c.encryptMessage(session, packet, &head, headerData)
 		if err != nil {
 			return nil, Nonce{}, err
 		}
 	}
-
-	enc, err := c.EncodeRaw(id, head, msgData)
+	enc, err = c.EncodeRaw(id, head, msgData)
 	return enc, head.Nonce, err
 }
 
diff --git a/p2p/discover/v5wire/msg.go b/p2p/discover/v5wire/msg.go
index 401db2f6c5..089fd4ebdc 100644
--- a/p2p/discover/v5wire/msg.go
+++ b/p2p/discover/v5wire/msg.go
@@ -73,6 +73,9 @@ type (
 		Node *enode.Node
 
 		sent mclock.AbsTime // for handshake GC.
+
+		// Encoded is packet raw data for sending out, but should not be include in the RLP encoding.
+		Encoded []byte `rlp:"-"`
 	}
 
 	// PING is sent during liveness checks.