From 617e19f6e8bc7ce4e0e7a59187550e1f02d0a3b5 Mon Sep 17 00:00:00 2001
From: thinkAfCod <q315xia@163.com>
Date: Mon, 12 May 2025 19:10:17 +0800
Subject: [PATCH] read notify

---
 p2p/discover/v5_udp.go          |  8 +++++---
 p2p/discover/v5wire/encoding.go | 35 +++++++++++++++++++++------------
 p2p/discover/v5wire/msg.go      | 16 +++++++--------
 3 files changed, 35 insertions(+), 24 deletions(-)

diff --git a/p2p/discover/v5_udp.go b/p2p/discover/v5_udp.go
index 677143fd60..23801b923a 100644
--- a/p2p/discover/v5_udp.go
+++ b/p2p/discover/v5_udp.go
@@ -620,9 +620,6 @@ func (t *UDPv5) dispatch() {
 
 		case p := <-t.packetInCh:
 			t.handlePacket(p.Data, p.Addr)
-			// Arm next read.
-			t.readNextCh <- struct{}{}
-
 		case <-t.closeCtx.Done():
 			close(t.readNextCh)
 			for id, queue := range t.callQueue {
@@ -811,6 +808,11 @@ func (t *UDPv5) handlePacket(rawpacket []byte, fromAddr netip.AddrPort) error {
 	addr := fromAddr.String()
 	t.log.Trace("<< "+addr, "rawPacket", hexutil.Encode(rawpacket))
 	fromID, fromNode, packet, err := t.codec.Decode(rawpacket, addr)
+	select {
+	case t.readNextCh <- struct{}{}:
+	case <-t.closeCtx.Done():
+	}
+
 	if err != nil {
 		if t.unhandled != nil && v5wire.IsInvalidHeader(err) {
 			// The packet seems unrelated to discv5, send it to the next protocol.
diff --git a/p2p/discover/v5wire/encoding.go b/p2p/discover/v5wire/encoding.go
index ec5ef8a261..e1067378ab 100644
--- a/p2p/discover/v5wire/encoding.go
+++ b/p2p/discover/v5wire/encoding.go
@@ -191,8 +191,15 @@ func (c *Codec) Encode(id enode.ID, addr string, packet Packet, challenge *Whoar
 	case packet.Kind() == WhoareyouPacket:
 		// just send the WHOAREYOU packet raw again, rather than the re-encoded challenge data
 		w := packet.(*Whoareyou)
-		if len(w.Encoded) > 0 {
-			return w.Encoded, w.Nonce, nil
+		if len(w.ChallengeData) > 0 {
+			// This WHOAREYOU packet was encoded before, so it's a resend.
+			// The unmasked packet content is stored in w.ChallengeData.
+			// Just apply the masking again to finish encoding.
+			c.buf.Reset()
+			c.buf.Write(w.ChallengeData)
+			copy(head.IV[:], w.ChallengeData)
+			enc := applyMasking(id, head.IV, c.buf.Bytes())
+			return enc, w.Nonce, nil
 		}
 		head, err = c.encodeWhoareyou(id, packet.(*Whoareyou))
 	case challenge != nil:
@@ -227,7 +234,6 @@ func (c *Codec) Encode(id enode.ID, addr string, packet Packet, challenge *Whoar
 		if err != nil {
 			return nil, Nonce{}, err
 		}
-		challenge.Encoded = bytes.Clone(enc)
 		c.sc.storeSentHandshake(id, addr, challenge)
 		return enc, head.Nonce, err
 	}
@@ -245,13 +251,9 @@ func (c *Codec) Encode(id enode.ID, addr string, packet Packet, challenge *Whoar
 
 // EncodeRaw encodes a packet with the given header.
 func (c *Codec) EncodeRaw(id enode.ID, head Header, msgdata []byte) ([]byte, error) {
+	// header
 	c.writeHeaders(&head)
-
-	// Apply masking.
-	masked := c.buf.Bytes()[sizeofMaskingIV:]
-	mask := head.mask(id)
-	mask.XORKeyStream(masked[:], masked[:])
-
+	applyMasking(id, head.IV, c.buf.Bytes())
 	// Write message data.
 	c.buf.Write(msgdata)
 	return c.buf.Bytes(), nil
@@ -463,7 +465,7 @@ func (c *Codec) Decode(inputData []byte, addr string) (src enode.ID, n *enode.No
 	// Unmask the static header.
 	var head Header
 	copy(head.IV[:], input[:sizeofMaskingIV])
-	mask := head.mask(c.localnode.ID())
+	mask := createMask(c.localnode.ID(), head.IV)
 	staticHeader := input[sizeofMaskingIV:sizeofStaticPacketData]
 	mask.XORKeyStream(staticHeader, staticHeader)
 
@@ -678,13 +680,20 @@ func (h *StaticHeader) checkValid(packetLen int, protocolID [6]byte) error {
 	return nil
 }
 
-// mask returns a cipher for 'masking' / 'unmasking' packet headers.
-func (h *Header) mask(destID enode.ID) cipher.Stream {
+// createMask returns a cipher for 'masking' / 'unmasking' packet headers.
+func createMask(destID enode.ID, iv [16]byte) cipher.Stream {
 	block, err := aes.NewCipher(destID[:16])
 	if err != nil {
 		panic("can't create cipher")
 	}
-	return cipher.NewCTR(block, h.IV[:])
+	return cipher.NewCTR(block, iv[:])
+}
+
+func applyMasking(destID enode.ID, iv [16]byte, packet []byte) []byte {
+	masked := packet[sizeofMaskingIV:]
+	mask := createMask(destID, iv)
+	mask.XORKeyStream(masked[:], masked[:])
+	return packet
 }
 
 func bytesCopy(r *bytes.Buffer) []byte {
diff --git a/p2p/discover/v5wire/msg.go b/p2p/discover/v5wire/msg.go
index 089fd4ebdc..eb4123bba1 100644
--- a/p2p/discover/v5wire/msg.go
+++ b/p2p/discover/v5wire/msg.go
@@ -63,19 +63,19 @@ type (
 
 	// WHOAREYOU contains the handshake challenge.
 	Whoareyou struct {
-		ChallengeData []byte   // Encoded challenge
-		Nonce         Nonce    // Nonce of request packet
-		IDNonce       [16]byte // Identity proof data
-		RecordSeq     uint64   // ENR sequence number of recipient
+		Nonce     Nonce    // Nonce of request packet
+		IDNonce   [16]byte // Identity proof data
+		RecordSeq uint64   // ENR sequence number of recipient
 
 		// Node is the locally known node record of recipient.
 		// This must be set by the caller of Encode.
-		Node *enode.Node
+		Node *enode.Node `rlp:"-"`
+		// ChallengeData stores the unmasked encoding of the whole packet. This is the
+		// input data for verification. It is assigned by both Encode and Decode
+		// operations.
+		ChallengeData []byte `rlp:"-"`
 
 		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.