swarm/pss: forwarding function refactoring

swarm/pss/forwarding_test.go

@@ -0,0 +1,356 @@
+package pss
+
+import (
+	"fmt"
+	"math/rand"
+	"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/swarm/network"
+	"github.com/ethereum/go-ethereum/swarm/pot"
+	whisper "github.com/ethereum/go-ethereum/whisper/whisperv5"
+)
+
+type testCase struct {
+	name      string
+	recipient []byte
+	peers     []pot.Address
+	expected  []int
+	exclusive bool
+	nFails    int
+	success   bool
+	errors    string
+}
+
+var testCases []testCase
+
+// the purpose of this test is to see that pss.forward() function correctly
+// selects the peers for message forwarding, depending on the message address
+// and kademlia constellation.
+func TestForwardBasic(t *testing.T) {
+	baseAddrBytes := make([]byte, 32)
+	for i := 0; i < len(baseAddrBytes); i++ {
+		baseAddrBytes[i] = 0xFF
+	}
+	var c testCase
+	base := pot.NewAddressFromBytes(baseAddrBytes)
+	var peerAddresses []pot.Address
+	const depth = 10
+	for i := 0; i <= depth; i++ {
+		// add two peers for each proximity order
+		a := pot.RandomAddressAt(base, i)
+		peerAddresses = append(peerAddresses, a)
+		a = pot.RandomAddressAt(base, i)
+		peerAddresses = append(peerAddresses, a)
+	}
+
+	// skip one level, add one peer at one level deeper.
+	// as a result, we will have an edge case of three peers in nearest neighbours' bin.
+	peerAddresses = append(peerAddresses, pot.RandomAddressAt(base, depth+2))
+
+	kad := network.NewKademlia(base[:], network.NewKadParams())
+	ps := createPss(t, kad)
+	addPeers(kad, peerAddresses)
+
+	const firstNearest = depth * 2 // shallowest peer in the nearest neighbours' bin
+	nearestNeighbours := []int{firstNearest, firstNearest + 1, firstNearest + 2}
+	var all []int // indices of all the peers
+	for i := 0; i < len(peerAddresses); i++ {
+		all = append(all, i)
+	}
+
+	for i := 0; i < len(peerAddresses); i++ {
+		// send msg directly to the known peers (recipient address == peer address)
+		c = testCase{
+			name:      fmt.Sprintf("Send direct to known, id: [%d]", i),
+			recipient: peerAddresses[i][:],
+			peers:     peerAddresses,
+			expected:  []int{i},
+			exclusive: false,
+		}
+		testCases = append(testCases, c)
+	}
+
+	for i := 0; i < firstNearest; i++ {
+		// send random messages with proximity orders, corresponding to PO of each bin,
+		// with one peer being closer to the recipient address
+		a := pot.RandomAddressAt(peerAddresses[i], 64)
+		c = testCase{
+			name:      fmt.Sprintf("Send random to each PO, id: [%d]", i),
+			recipient: a[:],
+			peers:     peerAddresses,
+			expected:  []int{i},
+			exclusive: false,
+		}
+		testCases = append(testCases, c)
+	}
+
+	for i := 0; i < firstNearest; i++ {
+		// send random messages with proximity orders, corresponding to PO of each bin,
+		// with random proximity relative to the recipient address
+		po := i / 2
+		a := pot.RandomAddressAt(base, po)
+		c = testCase{
+			name:      fmt.Sprintf("Send direct to known, id: [%d]", i),
+			recipient: a[:],
+			peers:     peerAddresses,
+			expected:  []int{po * 2, po*2 + 1},
+			exclusive: true,
+		}
+		testCases = append(testCases, c)
+	}
+
+	for i := firstNearest; i < len(peerAddresses); i++ {
+		// recipient address falls into the nearest neighbours' bin
+		a := pot.RandomAddressAt(base, i)
+		c = testCase{
+			name:      fmt.Sprintf("recipient address falls into the nearest neighbours' bin, id: [%d]", i),
+			recipient: a[:],
+			peers:     peerAddresses,
+			expected:  nearestNeighbours,
+			exclusive: false,
+		}
+		testCases = append(testCases, c)
+	}
+
+	// send msg with proximity order much deeper than the deepest nearest neighbour
+	a2 := pot.RandomAddressAt(base, 77)
+	c = testCase{
+		name:      "proximity order much deeper than the deepest nearest neighbour",
+		recipient: a2[:],
+		peers:     peerAddresses,
+		expected:  nearestNeighbours,
+		exclusive: false,
+	}
+	testCases = append(testCases, c)
+
+	// test with partial addresses
+	const part = 12
+
+	for i := 0; i < firstNearest; i++ {
+		// send messages with partial address falling into different proximity orders
+		po := i / 2
+		if i%8 != 0 {
+			c = testCase{
+				name:      fmt.Sprintf("partial address falling into different proximity orders, id: [%d]", i),
+				recipient: peerAddresses[i][:i],
+				peers:     peerAddresses,
+				expected:  []int{po * 2, po*2 + 1},
+				exclusive: true,
+			}
+			testCases = append(testCases, c)
+		}
+		c = testCase{
+			name:      fmt.Sprintf("extended partial address falling into different proximity orders, id: [%d]", i),
+			recipient: peerAddresses[i][:part],
+			peers:     peerAddresses,
+			expected:  []int{po * 2, po*2 + 1},
+			exclusive: true,
+		}
+		testCases = append(testCases, c)
+	}
+
+	for i := firstNearest; i < len(peerAddresses); i++ {
+		// partial address falls into the nearest neighbours' bin
+		c = testCase{
+			name:      fmt.Sprintf("partial address falls into the nearest neighbours' bin, id: [%d]", i),
+			recipient: peerAddresses[i][:part],
+			peers:     peerAddresses,
+			expected:  nearestNeighbours,
+			exclusive: false,
+		}
+		testCases = append(testCases, c)
+	}
+
+	// partial address with proximity order deeper than any of the nearest neighbour
+	a3 := pot.RandomAddressAt(base, part)
+	c = testCase{
+		name:      "partial address with proximity order deeper than any of the nearest neighbour",
+		recipient: a3[:part],
+		peers:     peerAddresses,
+		expected:  nearestNeighbours,
+		exclusive: false,
+	}
+	testCases = append(testCases, c)
+
+	// special cases where partial address matches a large group of peers
+
+	// zero bytes of address is given, msg should be delivered to all the peers
+	c = testCase{
+		name:      "zero bytes of address is given",
+		recipient: []byte{},
+		peers:     peerAddresses,
+		expected:  all,
+		exclusive: false,
+	}
+	testCases = append(testCases, c)
+
+	// luminous radius of 8 bits, proximity order 8
+	indexAtPo8 := 16
+	c = testCase{
+		name:      "luminous radius of 8 bits",
+		recipient: []byte{0xFF},
+		peers:     peerAddresses,
+		expected:  all[indexAtPo8:],
+		exclusive: false,
+	}
+	testCases = append(testCases, c)
+
+	// luminous radius of 256 bits, proximity order 8
+	a4 := pot.Address{}
+	a4[0] = 0xFF
+	c = testCase{
+		name:      "luminous radius of 256 bits",
+		recipient: a4[:],
+		peers:     peerAddresses,
+		expected:  []int{indexAtPo8, indexAtPo8 + 1},
+		exclusive: true,
+	}
+	testCases = append(testCases, c)
+
+	// check correct behaviour in case send fails
+	for i := 2; i < firstNearest-3; i += 2 {
+		po := i / 2
+		// send random messages with proximity orders, corresponding to PO of each bin,
+		// with different numbers of failed attempts.
+		// msg should be received by only one of the deeper peers.
+		a := pot.RandomAddressAt(base, po)
+		c = testCase{
+			name:      fmt.Sprintf("Send direct to known, id: [%d]", i),
+			recipient: a[:],
+			peers:     peerAddresses,
+			expected:  all[i+1:],
+			exclusive: true,
+			nFails:    rand.Int()%3 + 2,
+		}
+		testCases = append(testCases, c)
+	}
+
+	for _, c := range testCases {
+		testForwardMsg(t, ps, &c)
+	}
+}
+
+// this function tests the forwarding of a single message. the recipient address is passed as param,
+// along with addresses of all peers, and indices of those peers which are expected to receive the message.
+func testForwardMsg(t *testing.T, ps *Pss, c *testCase) {
+	recipientAddr := c.recipient
+	peers := c.peers
+	expected := c.expected
+	exclusive := c.exclusive
+	nFails := c.nFails
+	tries := 0 // number of previous failed tries
+
+	resultMap := make(map[pot.Address]int)
+
+	defer func() { sendFunc = sendMsg }()
+	sendFunc = func(_ *Pss, sp *network.Peer, _ *PssMsg) bool {
+		if tries < nFails {
+			tries++
+			return false
+		}
+		a := pot.NewAddressFromBytes(sp.Address())
+		resultMap[a]++
+		return true
+	}
+
+	msg := newTestMsg(recipientAddr)
+	ps.forward(msg)
+
+	// check test results
+	var fail bool
+	precision := len(recipientAddr)
+	if precision > 4 {
+		precision = 4
+	}
+	s := fmt.Sprintf("test [%s]\nmsg address: %x..., radius: %d", c.name, recipientAddr[:precision], 8*len(recipientAddr))
+
+	// false negatives (expected message didn't reach peer)
+	if exclusive {
+		var cnt int
+		for _, i := range expected {
+			a := peers[i]
+			cnt += resultMap[a]
+			resultMap[a] = 0
+		}
+		if cnt != 1 {
+			s += fmt.Sprintf("\n%d messages received by %d peers with indices: [%v]", cnt, len(expected), expected)
+			fail = true
+		}
+	} else {
+		for _, i := range expected {
+			a := peers[i]
+			received := resultMap[a]
+			if received != 1 {
+				s += fmt.Sprintf("\npeer number %d [%x...] received %d messages", i, a[:4], received)
+				fail = true
+			}
+			resultMap[a] = 0
+		}
+	}
+
+	// false positives (unexpected message reached peer)
+	for k, v := range resultMap {
+		if v != 0 {
+			// find the index of the false positive peer
+			var j int
+			for j = 0; j < len(peers); j++ {
+				if peers[j] == k {
+					break
+				}
+			}
+			s += fmt.Sprintf("\npeer number %d [%x...] received %d messages", j, k[:4], v)
+			fail = true
+		}
+	}
+
+	if fail {
+		t.Fatal(s)
+	}
+}
+
+func addPeers(kad *network.Kademlia, addresses []pot.Address) {
+	for _, a := range addresses {
+		p := newTestDiscoveryPeer(a, kad)
+		kad.On(p)
+	}
+}
+
+func createPss(t *testing.T, kad *network.Kademlia) *Pss {
+	privKey, err := crypto.GenerateKey()
+	pssp := NewPssParams().WithPrivateKey(privKey)
+	ps, err := NewPss(kad, pssp)
+	if err != nil {
+		t.Fatal(err.Error())
+	}
+	return ps
+}
+
+func newTestDiscoveryPeer(addr pot.Address, kad *network.Kademlia) *network.Peer {
+	rw := &p2p.MsgPipeRW{}
+	p := p2p.NewPeer(enode.ID{}, "test", []p2p.Cap{})
+	pp := protocols.NewPeer(p, rw, &protocols.Spec{})
+	bp := &network.BzzPeer{
+		Peer: pp,
+		BzzAddr: &network.BzzAddr{
+			OAddr: addr.Bytes(),
+			UAddr: []byte(fmt.Sprintf("%x", addr[:])),
+		},
+	}
+	return network.NewPeer(bp, kad)
+}
+
+func newTestMsg(addr []byte) *PssMsg {
+	msg := newPssMsg(&msgParams{})
+	msg.To = addr[:]
+	msg.Expire = uint32(time.Now().Add(time.Second * 60).Unix())
+	msg.Payload = &whisper.Envelope{
+		Topic: [4]byte{},
+		Data:  []byte("i have nothing to hide"),
+	}
+	return msg
+}