fix test case

This commit is contained in:
thinkAfCod 2025-05-14 17:28:58 +08:00 committed by Chen Kai
parent c603f255c8
commit aa62418193
2 changed files with 36 additions and 8 deletions

View file

@ -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) {
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.
if wp, ok := p.(*v5wire.Whoareyou); ok && len(wp.Encoded) > 0 {
return wp.Encoded, wp.Nonce, nil
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.

View file

@ -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()