p2p/discover: simplify

This commit is contained in:
Felix Lange 2025-04-02 12:16:45 +02:00
parent 9b74d7912e
commit 6487e737ea

View file

@ -775,31 +775,27 @@ type testCodecFrame struct {
} }
func (c *testCodec) Encode(toID enode.ID, addr string, p v5wire.Packet, _ *v5wire.Whoareyou) ([]byte, v5wire.Nonce, error) { func (c *testCodec) Encode(toID enode.ID, addr string, p v5wire.Packet, _ *v5wire.Whoareyou) ([]byte, v5wire.Nonce, error) {
if _, ok := p.(*v5wire.Whoareyou); ok { if wp, ok := p.(*v5wire.Whoareyou); ok && len(wp.Encoded) > 0 {
// Store recently sent Whoareyou challenges. return wp.Encoded, wp.Nonce, nil
if c.sentChallenges == nil {
c.sentChallenges = make(map[enode.ID]*v5wire.Whoareyou)
}
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++ c.ctr++
var authTag v5wire.Nonce var authTag v5wire.Nonce
binary.BigEndian.PutUint64(authTag[:], c.ctr) binary.BigEndian.PutUint64(authTag[:], c.ctr)
penc, _ := rlp.EncodeToBytes(p) penc, _ := rlp.EncodeToBytes(p)
frame, err := rlp.EncodeToBytes(testCodecFrame{c.id, authTag, p.Kind(), penc}) frame, err := rlp.EncodeToBytes(testCodecFrame{c.id, authTag, p.Kind(), penc})
if err != nil { if err != nil {
return frame, authTag, err return frame, authTag, err
} }
// Store recently sent challenges.
if w, ok := p.(*v5wire.Whoareyou); ok { if w, ok := p.(*v5wire.Whoareyou); ok {
c.sentChallenges[toID] = w
w.Nonce = authTag w.Nonce = authTag
w.Encoded = frame w.Encoded = frame
if c.sentChallenges == nil {
c.sentChallenges = make(map[enode.ID]*v5wire.Whoareyou)
}
c.sentChallenges[toID] = w
} }
return frame, authTag, err return frame, authTag, err
} }