2 files changed, 268 insertions, 0 deletions
diff --git a/swarm/network/simulation/kademlia.go b/swarm/network/simulation/kademlia.go
index a3419c03f..4b880aa0c 100644
--- a/swarm/network/simulation/kademlia.go
+++ b/swarm/network/simulation/kademlia.go
@@ -18,12 +18,14 @@ package simulation
 
 import (
 	"context"
+	"encoding/binary"
 	"encoding/hex"
 	"time"
 
 	"github.com/ethereum/go-ethereum/common"
 	"github.com/ethereum/go-ethereum/log"
 	"github.com/ethereum/go-ethereum/p2p/enode"
+	"github.com/ethereum/go-ethereum/p2p/simulations"
 	"github.com/ethereum/go-ethereum/swarm/network"
 )
 
@@ -96,3 +98,106 @@ func (s *Simulation) kademlias() (ks map[enode.ID]*network.Kademlia) {
 	}
 	return ks
 }
+
+// WaitTillSnapshotRecreated is blocking until all the connections specified
+// in the snapshot are registered in the kademlia.
+// It differs from WaitTillHealthy, which waits only until all the kademlias are
+// healthy (it might happen even before all the connections are established).
+func (s *Simulation) WaitTillSnapshotRecreated(ctx context.Context, snap simulations.Snapshot) error {
+	expected := getSnapshotConnections(snap.Conns)
+	ticker := time.NewTicker(150 * time.Millisecond)
+	defer ticker.Stop()
+
+	for {
+		select {
+		case <-ctx.Done():
+			return ctx.Err()
+		case <-ticker.C:
+			actual := s.getActualConnections()
+			if isAllDeployed(expected, actual) {
+				return nil
+			}
+		}
+	}
+}
+
+func (s *Simulation) getActualConnections() (res []uint64) {
+	kademlias := s.kademlias()
+	for base, k := range kademlias {
+		k.EachConn(base[:], 256, func(p *network.Peer, _ int) bool {
+			res = append(res, getConnectionHash(base, p.ID()))
+			return true
+		})
+	}
+
+	// only list those connections that appear twice (both peers should recognize connection as active)
+	res = removeDuplicatesAndSingletons(res)
+	return res
+}
+
+func getSnapshotConnections(conns []simulations.Conn) (res []uint64) {
+	for _, c := range conns {
+		res = append(res, getConnectionHash(c.One, c.Other))
+	}
+	return res
+}
+
+// returns an integer connection identifier (similar to 8-byte hash)
+func getConnectionHash(a, b enode.ID) uint64 {
+	var h [8]byte
+	for i := 0; i < 8; i++ {
+		h[i] = a[i] ^ b[i]
+	}
+	res := binary.LittleEndian.Uint64(h[:])
+	return res
+}
+
+// returns true if all connections in expected are listed in actual
+func isAllDeployed(expected []uint64, actual []uint64) bool {
+	if len(expected) == 0 {
+		return true
+	}
+
+	exp := make([]uint64, len(expected))
+	copy(exp, expected)
+	for _, c := range actual {
+		// remove value c from exp
+		for i := 0; i < len(exp); i++ {
+			if exp[i] == c {
+				exp = removeListElement(exp, i)
+				if len(exp) == 0 {
+					return true
+				}
+			}
+		}
+	}
+	return len(exp) == 0
+}
+
+func removeListElement(arr []uint64, i int) []uint64 {
+	last := len(arr) - 1
+	arr[i] = arr[last]
+	arr = arr[:last]
+	return arr
+}
+
+func removeDuplicatesAndSingletons(arr []uint64) []uint64 {
+	for i := 0; i < len(arr); {
+		found := false
+		for j := i + 1; j < len(arr); j++ {
+			if arr[i] == arr[j] {
+				arr = removeListElement(arr, j) // remove duplicate
+				found = true
+				break
+			}
+		}
+
+		if found {
+			i++
+		} else {
+			arr = removeListElement(arr, i) // remove singleton
+		}
+	}
+
+	return arr
+}
diff --git a/swarm/network/simulation/kademlia_test.go b/swarm/network/simulation/kademlia_test.go
index 4cfcecd8e..9cbc39da5 100644
--- a/swarm/network/simulation/kademlia_test.go
+++ b/swarm/network/simulation/kademlia_test.go
@@ -144,3 +144,166 @@ func createSimServiceMap(discovery bool) map[string]ServiceFunc {
 		},
 	}
 }
+
+// TestWaitTillSnapshotRecreated tests that we indeed have a network
+// configuration specified in the snapshot file, after we wait for it.
+//
+// First we create a first simulation
+// Run it as nodes connected in a ring
+// Wait until the network is healthy
+// Then we create a snapshot
+// With this snapshot we create a new simulation
+// Call WaitTillSnapshotRecreated() function and wait until it returns
+// Iterate the nodes and check if all the connections are successfully recreated
+func TestWaitTillSnapshotRecreated(t *testing.T) {
+	var err error
+	sim := New(createSimServiceMap(true))
+	_, err = sim.AddNodesAndConnectRing(16)
+	if err != nil {
+		t.Fatal(err)
+	}
+	ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
+	defer cancel()
+	_, err = sim.WaitTillHealthy(ctx)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	originalConnections := sim.getActualConnections()
+	snap, err := sim.Net.Snapshot()
+	sim.Close()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	controlSim := New(createSimServiceMap(false))
+	defer controlSim.Close()
+	err = controlSim.Net.Load(snap)
+	if err != nil {
+		t.Fatal(err)
+	}
+	err = controlSim.WaitTillSnapshotRecreated(ctx, *snap)
+	if err != nil {
+		t.Fatal(err)
+	}
+	controlConnections := controlSim.getActualConnections()
+
+	for _, c := range originalConnections {
+		if !exist(controlConnections, c) {
+			t.Fatal("connection was not recreated")
+		}
+	}
+}
+
+// exist returns true if val is found in arr
+func exist(arr []uint64, val uint64) bool {
+	for _, c := range arr {
+		if c == val {
+			return true
+		}
+	}
+	return false
+}
+
+func TestRemoveDuplicatesAndSingletons(t *testing.T) {
+	singletons := []uint64{
+		0x3c127c6f6cb026b0,
+		0x0f45190d72e71fc5,
+		0xb0184c02449e0bb6,
+		0xa85c7b84239c54d3,
+		0xe3b0c44298fc1c14,
+		0x9afbf4c8996fb924,
+		0x27ae41e4649b934c,
+		0xa495991b7852b855,
+	}
+
+	doubles := []uint64{
+		0x1b879f878de7fc7a,
+		0xc6791470521bdab4,
+		0xdd34b0ee39bbccc6,
+		0x4d904fbf0f31da10,
+		0x6403c2560432c8f8,
+		0x18954e33cf3ad847,
+		0x90db00e98dc7a8a6,
+		0x92886b0dfcc1809b,
+	}
+
+	var arr []uint64
+	arr = append(arr, doubles...)
+	arr = append(arr, singletons...)
+	arr = append(arr, doubles...)
+	arr = removeDuplicatesAndSingletons(arr)
+
+	for _, i := range singletons {
+		if exist(arr, i) {
+			t.Fatalf("singleton not removed: %d", i)
+		}
+	}
+
+	for _, i := range doubles {
+		if !exist(arr, i) {
+			t.Fatalf("wrong value removed: %d", i)
+		}
+	}
+
+	for j := 0; j < len(doubles); j++ {
+		v := doubles[j] + singletons[j]
+		if exist(arr, v) {
+			t.Fatalf("non-existing value found, index: %d", j)
+		}
+	}
+}
+
+func TestIsAllDeployed(t *testing.T) {
+	a := []uint64{
+		0x3c127c6f6cb026b0,
+		0x0f45190d72e71fc5,
+		0xb0184c02449e0bb6,
+		0xa85c7b84239c54d3,
+		0xe3b0c44298fc1c14,
+		0x9afbf4c8996fb924,
+		0x27ae41e4649b934c,
+		0xa495991b7852b855,
+	}
+
+	b := []uint64{
+		0x1b879f878de7fc7a,
+		0xc6791470521bdab4,
+		0xdd34b0ee39bbccc6,
+		0x4d904fbf0f31da10,
+		0x6403c2560432c8f8,
+		0x18954e33cf3ad847,
+		0x90db00e98dc7a8a6,
+		0x92886b0dfcc1809b,
+	}
+
+	var c []uint64
+	c = append(c, a...)
+	c = append(c, b...)
+
+	if !isAllDeployed(a, c) {
+		t.Fatal("isAllDeployed failed")
+	}
+
+	if !isAllDeployed(b, c) {
+		t.Fatal("isAllDeployed failed")
+	}
+
+	if isAllDeployed(c, a) {
+		t.Fatal("isAllDeployed failed: false positive")
+	}
+
+	if isAllDeployed(c, b) {
+		t.Fatal("isAllDeployed failed: false positive")
+	}
+
+	c = c[2:]
+
+	if isAllDeployed(a, c) {
+		t.Fatal("isAllDeployed failed: false positive")
+	}
+
+	if !isAllDeployed(b, c) {
+		t.Fatal("isAllDeployed failed")
+	}
+}