From aa62418193b8b1ef4749ad4415ebf2d282d3c6cd Mon Sep 17 00:00:00 2001
From: thinkAfCod <q315xia@163.com>
Date: Wed, 14 May 2025 17:28:58 +0800
Subject: [PATCH] fix test case

---
 p2p/discover/v5_udp_test.go          | 15 +++++++-------
 p2p/discover/v5wire/encoding_test.go | 29 ++++++++++++++++++++++++++++
 2 files changed, 36 insertions(+), 8 deletions(-)

diff --git a/p2p/discover/v5_udp_test.go b/p2p/discover/v5_udp_test.go
index 3a384aab12..e258f3c24d 100644
--- a/p2p/discover/v5_udp_test.go
+++ b/p2p/discover/v5_udp_test.go
@@ -772,10 +772,10 @@ type testCodecFrame struct {
 }
 
 func (c *testCodec) Encode(toID enode.ID, addr string, p v5wire.Packet, _ *v5wire.Whoareyou) ([]byte, v5wire.Nonce, error) {
-	// To match the behavior of v5wire.Codec, we return the cached encoding of
-	// WHOAREYOU challenges.
-	if wp, ok := p.(*v5wire.Whoareyou); ok && len(wp.Encoded) > 0 {
-		return wp.Encoded, wp.Nonce, nil
+	if wp, ok := p.(*v5wire.Whoareyou); ok && len(wp.ChallengeData) > 0 {
+		// To match the behavior of v5wire.Codec, we return the cached encoding of
+		// WHOAREYOU challenges.
+		return wp.ChallengeData, wp.Nonce, nil
 	}
 
 	c.ctr++
@@ -790,7 +790,7 @@ func (c *testCodec) Encode(toID enode.ID, addr string, p v5wire.Packet, _ *v5wir
 	// Store recently sent challenges.
 	if w, ok := p.(*v5wire.Whoareyou); ok {
 		w.Nonce = authTag
-		w.Encoded = frame
+		w.ChallengeData = frame
 		if c.sentChallenges == nil {
 			c.sentChallenges = make(map[enode.ID]*v5wire.Whoareyou)
 		}
@@ -827,6 +827,7 @@ func (c *testCodec) decodeFrame(input []byte) (frame testCodecFrame, p v5wire.Pa
 	case v5wire.WhoareyouPacket:
 		dec := new(v5wire.Whoareyou)
 		err = rlp.DecodeBytes(frame.Packet, &dec)
+		dec.ChallengeData = bytes.Clone(input)
 		p = dec
 	default:
 		p, err = v5wire.DecodeMessage(frame.Ptype, frame.Packet)
@@ -877,9 +878,7 @@ func (test *udpV5Test) packetInFrom(key *ecdsa.PrivateKey, addr netip.AddrPort,
 	if err != nil {
 		test.t.Errorf("%s encode error: %v", packet.Name(), err)
 	}
-	if test.udp.dispatchReadPacket(addr, enc) {
-		<-test.udp.readNextCh // unblock UDPv5.dispatch
-	}
+	test.udp.dispatchReadPacket(addr, enc)
 }
 
 // getNode ensures the test knows about a node at the given endpoint.
diff --git a/p2p/discover/v5wire/encoding_test.go b/p2p/discover/v5wire/encoding_test.go
index 2304d0f132..c513e94297 100644
--- a/p2p/discover/v5wire/encoding_test.go
+++ b/p2p/discover/v5wire/encoding_test.go
@@ -269,6 +269,35 @@ func TestHandshake_BadHandshakeAttack(t *testing.T) {
 	net.nodeB.expectDecodeErr(t, errUnexpectedHandshake, findnode)
 }
 
+func TestEncodeWhoareyouResend(t *testing.T) {
+	t.Parallel()
+	net := newHandshakeTest()
+	defer net.close()
+
+	// A -> B   WHOAREYOU
+	challenge := &Whoareyou{
+		Nonce:     Nonce{1, 2, 3, 4},
+		IDNonce:   testIDnonce,
+		RecordSeq: 0,
+	}
+	enc, _ := net.nodeA.encode(t, net.nodeB, challenge)
+	net.nodeB.expectDecode(t, WhoareyouPacket, enc)
+	whoareyou1 := bytes.Clone(enc)
+
+	if len(challenge.ChallengeData) == 0 {
+		t.Fatal("ChallengeData not assigned by encode")
+	}
+
+	// A -> B   WHOAREYOU
+	// Send the same challenge again. This should produce exactly
+	// the same bytes as the first send.
+	enc, _ = net.nodeA.encode(t, net.nodeB, challenge)
+	whoareyou2 := bytes.Clone(enc)
+	if !bytes.Equal(whoareyou2, whoareyou1) {
+		t.Fatal("re-encoded challenge not equal to first")
+	}
+}
+
 // This test checks some malformed packets.
 func TestDecodeErrorsV5(t *testing.T) {
 	t.Parallel()