swarm/network: fixed dynamic test, repeat tests

swarm/network/discovery_test.go

@@ -18,18 +18,18 @@ package network
 
 import (
 	"bytes"
-	"context"
 	"crypto/ecdsa"
 	"crypto/rand"
 	"fmt"
 	"net"
+	"strings"
 	"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"
 	"github.com/ethereum/go-ethereum/swarm/pot"
 )
@@ -71,86 +71,102 @@ func TestDiscovery(t *testing.T) {
 	}
 }
 
-// TestSubpeersMsg tests that the correct set of peers is suggested
+// TestSubpeersMsg runs testSubpeersMsg multiple times
-// to another peer with a given depth and kademlia
 func TestSubpeersMsg(t *testing.T) {
+	repetitions := 10
+	for r := 0; r < repetitions; r++ {
+		t.Run(fmt.Sprintf("Test Run number %d", r), testSubpeersMsg)
+	}
+}
+
+// testSubpeersMsg tests that the correct set of peers is suggested
+// to another peer with a given depth and kademlia
+func testSubpeersMsg(t *testing.T) {
 
 	// construct ProtocolTester and hive
 	params := NewHiveParams()
-	// setup
+	// setup the hive
 	prvkey, err := crypto.GenerateKey()
 	if err != nil {
 		t.Fatal(err)
 	}
 	addr := PrivateKeyToBzzKey(prvkey)
 	to := NewKademlia(addr, NewKadParams())
-	hive := NewHive(params, to, nil) // hive
+	// create the hive
+	hive := NewHive(params, to, nil)
 
-	numOfTestNodes := 12
+	// we will use a set of preconstructed addresses
+	var bzzAddrs []*BzzAddr
+	// define a number of nodes for the test
+	nodeCount := 12
+	// for every of these nodes, add to the hive's connections...
+	for i := 0; i < nodeCount; i++ {
+		// ...create a BzzAddr and connect it in the hive (`hive.On(p)`)
+		a, err := registerBzzAddr(hive)
+		if err != nil {
+			t.Fatal(err)
+		}
+		// also store the actual bzzAddr into the slice
+		bzzAddrs = append(bzzAddrs, a)
+	}
 
-	s, err := newBzzBaseTester(t, numOfTestNodes, prvkey, DiscoverySpec, hive.Run)
+	// create a channel. We will wait later...
+	waitC := make(chan struct{})
+
+	// create a special bzzBaseTester in which we can associate `enode.ID` to the `bzzAddr` we created above
+	s, err := newPreconnectedBzzBaseTester(t, waitC, bzzAddrs, nodeCount, prvkey, DiscoverySpec, hive.Run)
 	if err != nil {
 		t.Fatal(err)
 	}
 
-	// the control node is the only one from the ProtocolTester
-	control := s.Nodes[numOfTestNodes-1]
-
-	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
-	defer cancel()
-
 	// start the hive
 	hive.Start(s.Server)
 	defer hive.Stop()
 
-	// we need to wait until the control node is actually connected to our hive
+	// ...so we wait here until all connections have been established
-WAIT_PIVOT:
+	<-waitC
-	for {
+
-		select {
+	// choose a control node
-		case <-ctx.Done():
+	control := s.Nodes[0]
-			t.Fatal("Timed out waiting for the control node to connect")
-		case <-time.After(100 * time.Millisecond):
-			if len(hive.peers) == len(s.BzzAddrs) {
-				break WAIT_PIVOT
-			}
-		}
-	}
 
 	// get BzzAddr of the control
-	controlBzz := s.BzzAddrs[control.ID()].Over()
+	controlBzz := hive.peers[control.ID()].Over()
-
+	// build a control kademlia for the control node from the address pool
+	// we use this so we can identify the actual `controlDepth` of the control node
 	controlKad := NewKademlia(controlBzz, NewKadParams())
-	for _, p := range s.BzzAddrs {
+	for _, p := range bzzAddrs {
 		if !bytes.Equal(p.Over(), controlBzz) {
 			controlKad.On(NewPeer(&BzzPeer{nil, p, time.Now(), false}, controlKad))
 		}
 	}
+	// to be fully correct, even the pivot's hive should be added
+	controlKad.On(NewPeer(&BzzPeer{nil, &BzzAddr{OAddr: hive.BaseAddr(), UAddr: hive.BaseAddr()}, time.Now(), false}, controlKad))
+	// now we can evaluate the depth of the control node, which we need...
 	controlDepth := controlKad.NeighbourhoodDepth()
 
-	// now we need to identify which peers are expected
+	// ...to identify which peers are expected to be advertized
-	// iterate the hive's connection and only add peers below testDepth
+	// iterate the hive's connection and only add peers below controlDepth
 	var expectedPeers []*BzzAddr
 	hive.EachConn(controlBzz, 255, func(p *Peer, po int) bool {
 		if po < controlDepth {
 			return false
 		}
-		// don't add the control node itself to expectedPeers;
+		expectedPeers = append(expectedPeers, p.BzzAddr)
-		// the control node was not added as
-		if !bytes.Equal(p.BzzAddr.Over(), controlBzz) {
-			expectedPeers = append(expectedPeers, p.BzzAddr)
-		}
 		return true
 	})
 
+	// this is the hive's depth, which will be sent first to the control node initiating the test exchanges
 	hiveDepth := hive.NeighbourhoodDepth()
 
-	// the test exchange is as follows:
+	// if the controlDepth is 0, nothing will happen, so in this case artificially set it to 2
-	// 1. Trigger a subPeersMsg from control to our hive
-	// 2. Hive will respond with peersMsg with the set of expected peers
 	if controlDepth == 0 {
-		controlDepth = 1
+		controlDepth = 2
 	}
 
+	// the test exchange is as follows:
+	// 1. Wait for a `subPeersMsg` advertizing the hive's depth from our hive to our control node
+	// 2. Trigger a `suPeersMsg` from the control node advertizing its own depth
+	// 3. Hive will respond with peersMsg with the set of expected peers
 	err = s.TestExchanges(p2ptest.Exchange{
 		Label: "incoming subPeersMsg",
 		Expects: []p2ptest.Expect{
@@ -195,27 +211,31 @@ WAIT_PIVOT:
 	if err != nil {
 		t.Fatal(err)
 	}
+
+	close(waitC)
 }
 
 // add the BzzAddr to the hive
-func registerBzzAddr(po int, hive *Hive, on bool) {
+func registerBzzAddr(hive *Hive) (*BzzAddr, error) {
-	a := pot.RandomAddressAt(pot.NewAddressFromBytes(hive.BaseAddr()), po)
+	addr := pot.RandomAddress()
-	bzzAddr := &BzzAddr{OAddr: a.Bytes(), UAddr: a.Bytes()}
+	pKey, err := ecdsa.GenerateKey(crypto.S256(), rand.Reader)
-	if on {
+	if err != nil {
-		// actually connect
+		return nil, err
-		peer := newDiscPeer(bzzAddr, a.String(), hive)
-		hive.On(peer)
-	} else {
-		// only add to address book
-		hive.Register(bzzAddr)
 	}
+	pubKey := pKey.PublicKey
+	nod := enode.NewV4(&pubKey, net.IPv4(127, 0, 0, 1), 0, 0)
+	bzzAddr := &BzzAddr{OAddr: addr.Bytes(), UAddr: []byte(nod.String())}
+	// actually connect
+	peer := newDiscPeer(bzzAddr, nod.ID(), hive)
+	hive.On(peer)
+	return bzzAddr, nil
 }
 
 // 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 {
+func newDiscPeer(bzzAddr *BzzAddr, id enode.ID, hive *Hive) *Peer {
-	p2pPeer := p2p.NewPeer(adapters.RandomNodeConfig().ID, name, nil)
+	p2pPeer := p2p.NewPeer(id, id.String(), nil)
 	return NewPeer(&BzzPeer{
 		Peer:      protocols.NewPeer(p2pPeer, &dummyMsgRW{}, DiscoverySpec),
 		BzzAddr:   bzzAddr,

swarm/network/protocol_test.go

@@ -21,6 +21,7 @@ import (
 	"flag"
 	"fmt"
 	"os"
+	"sync"
 	"testing"
 	"time"
 
@@ -70,9 +71,9 @@ func HandshakeMsgExchange(lhs, rhs *HandshakeMsg, id enode.ID) []p2ptest.Exchang
 	}
 }
 
-func newBzzBaseTester(t *testing.T, n int, prvkey *ecdsa.PrivateKey, spec *protocols.Spec, run func(*BzzPeer) error) (*bzzTester, error) {
+// get some protocol entities
+func getProtocolAttrs(run func(*BzzPeer) error) (map[string]chan bool, func(p *BzzPeer) error) {
 	cs := make(map[string]chan bool)
-	bzzAddrs := make(map[enode.ID]*BzzAddr)
 
 	srv := func(p *BzzPeer) error {
 		defer func() {
@@ -83,13 +84,11 @@ func newBzzBaseTester(t *testing.T, n int, prvkey *ecdsa.PrivateKey, spec *proto
 		return run(p)
 	}
 
-	protocol := func(p *p2p.Peer, rw p2p.MsgReadWriter) error {
+	return cs, srv
-		bzzAddr := NewAddr(p.Node())
+}
-		bzzAddrs[p.Node().ID()] = bzzAddr
-		return srv(&BzzPeer{Peer: protocols.NewPeer(p, rw, spec), BzzAddr: bzzAddr})
-	}
 
-	s := p2ptest.NewProtocolTester(prvkey, n, protocol)
+// get a BzzAddr for the hive
+func getBzzAddr(prvkey *ecdsa.PrivateKey) (*BzzAddr, error) {
 	var record enr.Record
 	bzzKey := PrivateKeyToBzzKey(prvkey)
 	record.Set(NewENRAddrEntry(bzzKey))
@@ -101,7 +100,65 @@ func newBzzBaseTester(t *testing.T, n int, prvkey *ecdsa.PrivateKey, spec *proto
 	if err != nil {
 		return nil, fmt.Errorf("unable to create enode: %v", err)
 	}
-	addr := getENRBzzAddr(nod)
+	return getENRBzzAddr(nod), nil
+
+}
+
+// this is a custom version of the bzzBaseTester, which we can use to
+// associate preconstructed BzzAddr addresses to the `enode.ID` created in the `ProtocolTester`
+// this allows control over the address space in tests needing the protocol agnostic `ProtocolTester`
+func newPreconnectedBzzBaseTester(t *testing.T, waitC chan struct{}, preConns []*BzzAddr, n int, prvkey *ecdsa.PrivateKey, spec *protocols.Spec, run func(*BzzPeer) error) (*bzzTester, error) {
+	var lock sync.Mutex
+	cs, srv := getProtocolAttrs(run)
+
+	// in this custom version of the protocol, we can associate the preconstructed BzzAddr to `enode.ID`s
+	protocol := func(p *p2p.Peer, rw p2p.MsgReadWriter) error {
+		lock.Lock()
+		// take one address from the "stack"
+		n := len(preConns) - 1
+		// associate it to bzzAddr, which is used further down to create the BzzPeer
+		bzzAddr := preConns[n]
+		// pop the address from the "stack"
+		preConns = preConns[:n]
+		// if there are no more addresses left, inform the caller
+		if len(preConns) == 0 {
+			waitC <- struct{}{}
+		}
+		lock.Unlock()
+		// run the protocol, create the protocol node
+		return srv(&BzzPeer{Peer: protocols.NewPeer(p, rw, spec), BzzAddr: bzzAddr})
+	}
+
+	s := p2ptest.NewProtocolTester(prvkey, n, protocol)
+
+	for _, node := range s.Nodes {
+		log.Warn("node", "node", node)
+		cs[node.ID().String()] = make(chan bool)
+	}
+
+	return &bzzTester{
+		// this version of the bzzTester assumes that the hive has been built outside of this function already
+		addr:           nil, // ...so no need to create a new address
+		ProtocolTester: s,
+		cs:             cs,
+	}, nil
+}
+
+// standard bzzTester
+func newBzzBaseTester(t *testing.T, n int, prvkey *ecdsa.PrivateKey, spec *protocols.Spec, run func(*BzzPeer) error) (*bzzTester, error) {
+
+	cs, srv := getProtocolAttrs(run)
+
+	protocol := func(p *p2p.Peer, rw p2p.MsgReadWriter) error {
+		return srv(&BzzPeer{Peer: protocols.NewPeer(p, rw, spec), BzzAddr: NewAddr(p.Node())})
+	}
+
+	s := p2ptest.NewProtocolTester(prvkey, n, protocol)
+
+	addr, err := getBzzAddr(prvkey)
+	if err != nil {
+		t.Fatal(err)
+	}
 
 	for _, node := range s.Nodes {
 		log.Warn("node", "node", node)
@@ -112,16 +169,14 @@ func newBzzBaseTester(t *testing.T, n int, prvkey *ecdsa.PrivateKey, spec *proto
 		addr:           addr,
 		ProtocolTester: s,
 		cs:             cs,
-		BzzAddrs:       bzzAddrs,
 	}, nil
 }
 
 type bzzTester struct {
 	*p2ptest.ProtocolTester
-	BzzAddrs map[enode.ID]*BzzAddr
+	addr *BzzAddr
-	addr     *BzzAddr
+	cs   map[string]chan bool
-	cs       map[string]chan bool
+	bzz  *Bzz
-	bzz      *Bzz
 }
 
 func newBzz(addr *BzzAddr, lightNode bool) *Bzz {