swarm/network: fixed 0 depth bug and test

swarm/network/discovery.go

@@ -160,12 +160,11 @@ func (d *Peer) handleSubPeersMsg(msg *subPeersMsg) error {
 	if !d.sentPeers {
 		d.setDepth(msg.Depth)
 		var peers []*BzzAddr
-		d.kad.EachConn(d.Over(), 255, func(p *Peer, po int) bool {
+		d.kad.EachConn(d.kad.BaseAddr(), 255, func(p *Peer, po int) bool {
 			if uint8(po) < msg.Depth {
-				return false
-			}
-			if !d.seen(p.BzzAddr) {
-				peers = append(peers, p.BzzAddr)
+				if !d.seen(p.BzzAddr) {
+					peers = append(peers, p.BzzAddr)
+				}
 			}
 			return true
 		})

swarm/network/discovery_test.go

@@ -17,12 +17,12 @@
 package network
 
 import (
-	"net"
+	"bytes"
+	"context"
 	"testing"
+	"time"
 
-	"github.com/ethereum/go-ethereum/crypto"
 	"github.com/ethereum/go-ethereum/p2p"
-	"github.com/ethereum/go-ethereum/p2p/enode"
 	"github.com/ethereum/go-ethereum/p2p/protocols"
 	"github.com/ethereum/go-ethereum/p2p/simulations/adapters"
 	p2ptest "github.com/ethereum/go-ethereum/p2p/testing"
@@ -66,50 +66,73 @@ func TestDiscovery(t *testing.T) {
 	}
 }
 
+// TestSubpeersMsg tests that the correct set of peers is suggested
+// to another peer with a given depth and kademlia
 func TestSubpeersMsg(t *testing.T) {
 
+	// This is the defined depth
 	testDepth := 2
 
+	// construct ProtocolTester and hive
 	params := NewHiveParams()
-	key, err := crypto.GenerateKey()
-	if err != nil {
-		panic("unable to generate key")
-	}
-	node := enode.NewV4(&key.PublicKey, net.IP{127, 0, 0, 1}, 30303, 30303)
-	addr := NewAddr(node)
-	kad := NewKademlia(addr.OAddr, NewKadParams())
-	hive := NewHive(params, kad, nil) // hive
-	s := newBzzBaseTester(t, 1, addr, DiscoverySpec, hive.Run)
-	pivot := s.Nodes[0]
+	s, hive := newHiveTester(t, params, 1, nil)
 
+	// register some addresses in specific bins (must coincide with testDepth)
 	registerBzzAddr(0, hive, true)  // bin 0
+	registerBzzAddr(0, hive, true)  // bin 0
+	registerBzzAddr(1, hive, true)  // bin 1
 	registerBzzAddr(1, hive, true)  // bin 1
 	registerBzzAddr(3, hive, true)  // bin 3
 	registerBzzAddr(4, hive, true)  // bin 4
 	registerBzzAddr(3, hive, false) // add a known but not connected peer
 	registerBzzAddr(1, hive, false) // add a known but not connected peer
 
+	// start the hive
+	hive.Start(s.Server)
+	defer hive.Stop()
+
+	// the pivot node is the only one from the ProtocolTester
+	pivot := s.Nodes[0]
+
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer cancel()
+
+	// we need to wait until the pivot node is actually connected to our hive
+WAIT_PIVOT:
+	for {
+		select {
+		case <-ctx.Done():
+			t.Fatal("Timed out waiting for the pivot node to connect")
+		case <-time.After(100 * time.Millisecond):
+			if _, ok := hive.peers[pivot.ID()]; ok {
+				break WAIT_PIVOT
+			}
+		}
+	}
+
+	// get BzzAddr of the pivot
+	pivotAddress := hive.peers[pivot.ID()]
+	pivotBzz := pivotAddress.BzzAddr.Over()
+
+	// now we need to identify which peers are expected
+	// iterate the hive's connection and only add peers below testDepth
 	var expectedPeers []*BzzAddr
 	hive.EachConn(hive.BaseAddr(), 256, func(p *Peer, po int) bool {
 		if po < testDepth {
-			expectedPeers = append(expectedPeers, p.BzzAddr)
+			// don't add the pivot node itself to expectedPeers;
+			// the pivot node was not added as
+			if !bytes.Equal(p.BzzAddr.Over(), b) {
+				expectedPeers = append(expectedPeers, p.BzzAddr)
+			}
 		}
 		return true
 	})
 
-	// start the hive and wait for the connection
-	hive.Start(s.Server)
-	defer hive.Stop()
-
-	err = s.TestExchanges(p2ptest.Exchange{
+	// the test exchange is as follows:
+	// 1. Trigger a subPeersMsg from pivot to our hive
+	// 2. Hive will respond with peersMsg with the set of expected peers
+	err := s.TestExchanges(p2ptest.Exchange{
 		Label: "incoming subPeersMsg",
-		Expects: []p2ptest.Expect{
-			{
-				Code: 0,
-				Msg:  &peersMsg{Peers: expectedPeers},
-				Peer: pivot.ID(),
-			},
-		},
 		Triggers: []p2ptest.Trigger{
 			{
 				Code: 1,
@@ -117,6 +140,18 @@ func TestSubpeersMsg(t *testing.T) {
 				Peer: pivot.ID(),
 			},
 		},
+		Expects: []p2ptest.Expect{
+			{
+				Code: 1,
+				Msg:  &subPeersMsg{Depth: uint8(testDepth)},
+				Peer: pivot.ID(),
+			},
+			{
+				Code: 0,
+				Msg:  &peersMsg{Peers: expectedPeers},
+				Peer: pivot.ID(),
+			},
+		},
 	})
 
 	if err != nil {
@@ -125,17 +160,23 @@ func TestSubpeersMsg(t *testing.T) {
 	return
 }
 
+// add the BzzAddr to the hive
 func registerBzzAddr(po int, hive *Hive, on bool) {
 	a := pot.RandomAddressAt(pot.NewAddressFromBytes(hive.BaseAddr()), po)
 	bzzAddr := &BzzAddr{OAddr: a.Bytes(), UAddr: a.Bytes()}
 	if on {
+		// actually connect
 		peer := newDiscPeer(bzzAddr, a.String(), hive)
 		hive.On(peer)
 	} else {
+		// only add to address book
 		hive.Register(bzzAddr)
 	}
 }
 
+// as we are not creating a real node via the protocol,
+// we need to create the discovery peer objects for the additional kademlia
+// nodes manually
 func newDiscPeer(bzzAddr *BzzAddr, name string, hive *Hive) *Peer {
 	p2pPeer := p2p.NewPeer(adapters.RandomNodeConfig().Node().ID(), name, nil)
 	peer := NewPeer(&BzzPeer{
@@ -146,6 +187,9 @@ func newDiscPeer(bzzAddr *BzzAddr, name string, hive *Hive) *Peer {
 	return peer
 }
 
+// we also need this dummy object otherwise at hive.Stop(),
+// which will call `Drop` on all nodes, we will have null pointer errors,
+// as the underlying `p2p.Peer` objects were not created
 type dummyMsgRW struct{}
 
 func (d *dummyMsgRW) ReadMsg() (p2p.Msg, error) {