p2p/discover: fix tests

p2p/discover/table.go

@@ -583,23 +583,6 @@ func (tab *Table) addThroughPing(n *node) {
 	tab.add(n)
 }
 
-// stuff adds nodes the table to the end of their corresponding bucket
-// if the bucket is not full. The caller must not hold tab.mutex.
-func (tab *Table) stuff(nodes []*node) {
-	tab.mutex.Lock()
-	defer tab.mutex.Unlock()
-
-	for _, n := range nodes {
-		if n.ID() == tab.self().ID() {
-			continue // don't add self
-		}
-		b := tab.bucket(n.ID())
-		if len(b.entries) < bucketSize {
-			tab.bumpOrAdd(b, n)
-		}
-	}
-}
-
 // delete removes an entry from the node table. It is used to evacuate dead nodes.
 func (tab *Table) delete(node *node) {
 	tab.mutex.Lock()

p2p/discover/table_test.go

@@ -175,7 +175,7 @@ func TestTable_closest(t *testing.T) {
 		tab, db := newTestTable(transport)
 		defer tab.Close()
 		defer db.Close()
-		tab.stuff(test.All)
+		fillTable(tab, test.All)
 
 		// check that closest(Target, N) returns nodes
 		result := tab.closest(test.Target, test.N).entries
@@ -240,7 +240,7 @@ func TestTable_ReadRandomNodesGetAll(t *testing.T) {
 
 		for i := 0; i < len(buf); i++ {
 			ld := cfg.Rand.Intn(len(tab.buckets))
-			tab.stuff([]*node{nodeAtDistance(tab.self().ID(), ld, intIP(ld))})
+			fillTable(tab, []*node{nodeAtDistance(tab.self().ID(), ld, intIP(ld))})
 		}
 		gotN := tab.ReadRandomNodes(buf)
 		if gotN != tab.len() {
@@ -274,8 +274,9 @@ func (*closeTest) Generate(rand *rand.Rand, size int) reflect.Value {
 	for _, id := range gen([]enode.ID{}, rand).([]enode.ID) {
 		r := new(enr.Record)
 		r.Set(enr.IP(genIP(rand)))
-		n := enode.SignNull(r, id)
+		n := wrapNode(enode.SignNull(r, id))
-		t.All = append(t.All, wrapNode(n))
+		n.livenessChecks = 1
+		t.All = append(t.All, n)
 	}
 	return reflect.ValueOf(t)
 }
@@ -292,7 +293,8 @@ func TestTable_Lookup(t *testing.T) {
 	// seed table with initial node (otherwise lookup will terminate immediately)
 	seedKey, _ := decodePubkey(lookupTestnet.dists[256][0])
 	seed := wrapNode(enode.NewV4(seedKey, net.IP{127, 0, 0, 1}, 0, 256))
-	tab.stuff([]*node{seed})
+	seed.livenessChecks = 1
+	fillTable(tab, []*node{seed})
 
 	results := tab.lookup(lookupTestnet.target, true)
 	t.Logf("results:")

p2p/discover/table_util_test.go

@@ -83,6 +83,23 @@ func fillBucket(tab *Table, n *node) (last *node) {
 	return b.entries[bucketSize-1]
 }
 
+// fillTable adds nodes the table to the end of their corresponding bucket
+// if the bucket is not full. The caller must not hold tab.mutex.
+func fillTable(tab *Table, nodes []*node) {
+	tab.mutex.Lock()
+	defer tab.mutex.Unlock()
+
+	for _, n := range nodes {
+		if n.ID() == tab.self().ID() {
+			continue // don't add self
+		}
+		b := tab.bucket(n.ID())
+		if len(b.entries) < bucketSize {
+			tab.bumpOrAdd(b, n)
+		}
+	}
+}
+
 type pingRecorder struct {
 	mu           sync.Mutex
 	dead, pinged map[enode.ID]bool
@@ -109,10 +126,6 @@ func (t *pingRecorder) findnode(toid enode.ID, toaddr *net.UDPAddr, target encPu
 	return nil, nil
 }
 
-func (t *pingRecorder) waitping(from enode.ID) error {
-	return nil // remote always pings
-}
-
 func (t *pingRecorder) ping(toid enode.ID, toaddr *net.UDPAddr) error {
 	t.mu.Lock()
 	defer t.mu.Unlock()

p2p/discover/udp_test.go

@@ -19,6 +19,7 @@ package discover
 import (
 	"bytes"
 	"crypto/ecdsa"
+	crand "crypto/rand"
 	"encoding/binary"
 	"encoding/hex"
 	"errors"
@@ -81,12 +82,17 @@ func newUDPTest(t *testing.T) *udpTest {
 
 // handles a packet as if it had been sent to the transport.
 func (test *udpTest) packetIn(wantError error, ptype byte, data packet) error {
-	enc, _, err := encodePacket(test.remotekey, ptype, data)
+	return test.packetInFrom(wantError, test.remotekey, test.remoteaddr, ptype, data)
+}
+
+// handles a packet as if it had been sent to the transport by the key/endpoint.
+func (test *udpTest) packetInFrom(wantError error, key *ecdsa.PrivateKey, addr *net.UDPAddr, ptype byte, data packet) error {
+	enc, _, err := encodePacket(key, ptype, data)
 	if err != nil {
 		return test.errorf("packet (%d) encode error: %v", ptype, err)
 	}
 	test.sent = append(test.sent, enc)
-	if err = test.udp.handlePacket(test.remoteaddr, enc); err != wantError {
+	if err = test.udp.handlePacket(addr, enc); err != wantError {
 		return test.errorf("error mismatch: got %q, want %q", err, wantError)
 	}
 	return nil
@@ -94,19 +100,19 @@ func (test *udpTest) packetIn(wantError error, ptype byte, data packet) error {
 
 // waits for a packet to be sent by the transport.
 // validate should have type func(*udpTest, X) error, where X is a packet type.
-func (test *udpTest) waitPacketOut(validate interface{}) ([]byte, error) {
+func (test *udpTest) waitPacketOut(validate interface{}) (*net.UDPAddr, []byte, error) {
 	dgram := test.pipe.waitPacketOut()
-	p, _, hash, err := decodePacket(dgram)
+	p, _, hash, err := decodePacket(dgram.data)
 	if err != nil {
-		return hash, test.errorf("sent packet decode error: %v", err)
+		return &dgram.to, hash, test.errorf("sent packet decode error: %v", err)
 	}
 	fn := reflect.ValueOf(validate)
 	exptype := fn.Type().In(0)
 	if reflect.TypeOf(p) != exptype {
-		return hash, test.errorf("sent packet type mismatch, got: %v, want: %v", reflect.TypeOf(p), exptype)
+		return &dgram.to, hash, test.errorf("sent packet type mismatch, got: %v, want: %v", reflect.TypeOf(p), exptype)
 	}
 	fn.Call([]reflect.Value{reflect.ValueOf(p)})
-	return hash, nil
+	return &dgram.to, hash, nil
 }
 
 func (test *udpTest) errorf(format string, args ...interface{}) error {
@@ -179,7 +185,7 @@ func TestUDP_responseTimeouts(t *testing.T) {
 			p.errc = nilErr
 			test.udp.addReplyMatcher <- p
 			time.AfterFunc(randomDuration(60*time.Millisecond), func() {
-				if !test.udp.handleReply(p.from, p.ptype, nil) {
+				if !test.udp.handleReply(p.from, p.ip, p.ptype, nil) {
 					t.Logf("not matched: %v", p)
 				}
 			})
@@ -242,12 +248,20 @@ func TestUDP_findnode(t *testing.T) {
 	// distribution shouldn't matter much, although we need to
 	// take care not to overflow any bucket.
 	nodes := &nodesByDistance{target: testTarget.id()}
-	for i := 0; i < bucketSize; i++ {
+	live := make(map[enode.ID]bool)
+	numCandidates := 2 * bucketSize
+	for i := 0; i < numCandidates; i++ {
 		key := newkey()
-		n := wrapNode(enode.NewV4(&key.PublicKey, net.IP{10, 13, 0, 1}, 0, i))
+		ip := net.IP{10, 13, 0, byte(i)}
-		nodes.push(n, bucketSize)
+		n := wrapNode(enode.NewV4(&key.PublicKey, ip, 0, 2000))
+		// Ensure half of table content isn't verified live yet.
+		if i > numCandidates/2 {
+			n.livenessChecks = 1
+			live[n.ID()] = true
+		}
+		nodes.push(n, numCandidates)
 	}
-	test.table.stuff(nodes.entries)
+	fillTable(test.table, nodes.entries)
 
 	// ensure there's a bond with the test node,
 	// findnode won't be accepted otherwise.
@@ -255,17 +269,19 @@ func TestUDP_findnode(t *testing.T) {
 	test.table.db.UpdateLastPongReceived(remoteID, test.remoteaddr.IP, time.Now())
 
 	// check that closest neighbors are returned.
-	test.packetIn(nil, findnodePacket, &findnode{Target: testTarget, Expiration: futureExp})
 	expected := test.table.closest(testTarget.id(), bucketSize)
-
+	test.packetIn(nil, findnodePacket, &findnode{Target: testTarget, Expiration: futureExp})
 	waitNeighbors := func(want []*node) {
 		test.waitPacketOut(func(p *neighbors) {
 			if len(p.Nodes) != len(want) {
 				t.Errorf("wrong number of results: got %d, want %d", len(p.Nodes), bucketSize)
 			}
-			for i := range p.Nodes {
+			for i, n := range p.Nodes {
-				if p.Nodes[i].ID.id() != want[i].ID() {
+				if n.ID.id() != want[i].ID() {
-					t.Errorf("result mismatch at %d:\n  got:  %v\n  want: %v", i, p.Nodes[i], expected.entries[i])
+					t.Errorf("result mismatch at %d:\n  got:  %v\n  want: %v", i, n, expected.entries[i])
+				}
+				if !live[n.ID.id()] {
+					t.Errorf("result includes dead node %v", n.ID.id())
 				}
 			}
 		})
@@ -333,8 +349,29 @@ func TestUDP_pingMatch(t *testing.T) {
 	test := newUDPTest(t)
 	defer test.table.Close()
 
+	randToken := make([]byte, 32)
+	crand.Read(randToken)
+
 	test.packetIn(nil, pingPacket, &ping{From: testRemote, To: testLocalAnnounced, Version: 4, Expiration: futureExp})
-	test.packetIn(errUnsolicitedReply, pongPacket, &pong{To: testLocalAnnounced, Expiration: futureExp})
+	test.waitPacketOut(func(*pong) error { return nil })
+	test.waitPacketOut(func(*ping) error { return nil })
+	test.packetIn(errUnsolicitedReply, pongPacket, &pong{ReplyTok: randToken, To: testLocalAnnounced, Expiration: futureExp})
+}
+
+func TestUDP_pingMatchIP(t *testing.T) {
+	test := newUDPTest(t)
+	defer test.table.Close()
+
+	test.packetIn(nil, pingPacket, &ping{From: testRemote, To: testLocalAnnounced, Version: 4, Expiration: futureExp})
+	test.waitPacketOut(func(*pong) error { return nil })
+
+	_, hash, _ := test.waitPacketOut(func(*ping) error { return nil })
+	wrongAddr := &net.UDPAddr{IP: net.IP{33, 44, 1, 2}, Port: 30000}
+	test.packetInFrom(errUnsolicitedReply, test.remotekey, wrongAddr, pongPacket, &pong{
+		ReplyTok:   hash,
+		To:         testLocalAnnounced,
+		Expiration: futureExp,
+	})
 }
 
 func TestUDP_successfulPing(t *testing.T) {
@@ -364,7 +401,7 @@ func TestUDP_successfulPing(t *testing.T) {
 	})
 
 	// remote is unknown, the table pings back.
-	hash, _ := test.waitPacketOut(func(p *ping) error {
+	_, hash, _ := test.waitPacketOut(func(p *ping) error {
 		if !reflect.DeepEqual(p.From, test.udp.ourEndpoint()) {
 			t.Errorf("got ping.From %#v, want %#v", p.From, test.udp.ourEndpoint())
 		}
@@ -518,7 +555,12 @@ type dgramPipe struct {
 	cond    *sync.Cond
 	closing chan struct{}
 	closed  bool
-	queue   [][]byte
+	queue   []dgram
+}
+
+type dgram struct {
+	to   net.UDPAddr
+	data []byte
 }
 
 func newpipe() *dgramPipe {
@@ -539,7 +581,7 @@ func (c *dgramPipe) WriteToUDP(b []byte, to *net.UDPAddr) (n int, err error) {
 	if c.closed {
 		return 0, errors.New("closed")
 	}
-	c.queue = append(c.queue, msg)
+	c.queue = append(c.queue, dgram{*to, b})
 	c.cond.Signal()
 	return len(b), nil
 }
@@ -564,7 +606,7 @@ func (c *dgramPipe) LocalAddr() net.Addr {
 	return &net.UDPAddr{IP: testLocal.IP, Port: int(testLocal.UDP)}
 }
 
-func (c *dgramPipe) waitPacketOut() []byte {
+func (c *dgramPipe) waitPacketOut() dgram {
 	c.mu.Lock()
 	defer c.mu.Unlock()
 	for len(c.queue) == 0 {