swarm/pss: bugfix (msg should only be sent to one nearest peer)

swarm/pss/forwarding_test.go

@@ -2,6 +2,7 @@ package pss
 
 import (
 	"fmt"
+	"math/rand"
 	"testing"
 	"time"
 
@@ -14,132 +15,286 @@ import (
 	whisper "github.com/ethereum/go-ethereum/whisper/whisperv5"
 )
 
-var testResMap map[pot.Address]int
+type testCase struct {
-
+	name      string
-// this function substitutes the real send function, since
+	recipient []byte
-// we only want to test the peer selection functionality
+	peers     []pot.Address
-func dummySendMsg(_ *Pss, sp *network.Peer, _ *PssMsg) bool {
+	expected  []int
-	a := pot.NewAddressFromBytes(sp.Address())
+	exclusive bool
-	testResMap[a]++
+	nFails    int
-	return true
+	success   bool
+	errors    string
 }
 
-// setDummySendMsg replaces sendMessage function for testing purposes
+var testCases []testCase
-func setDummySendMsg() {
-	sendMessage = dummySendMsg
-}
-
-// resetSendMsgProduction resets sendMessage function to production version
-func resetSendMsgProduction() {
-	sendMessage = sendMessageProd
-}
 
 // 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) {
-	setDummySendMsg()
-	defer resetSendMsgProduction()
-
 	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
-	var a pot.Address
 	const depth = 10
 	for i := 0; i <= depth; i++ {
-		// add one peer for each proximity order
+		// 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)
 	}
 
-	// add one peer to the "depth" level, then skip one level, add one peer at one level below.
+	// 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))
 	peerAddresses = append(peerAddresses, pot.RandomAddressAt(base, depth+2))
 
 	kad := network.NewKademlia(base[:], network.NewKadParams())
 	ps := createPss(t, kad)
 	addPeers(kad, peerAddresses)
 
-	const firstNearest = depth // shallowest peer in the nearest neighbours' bin
+	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)
-		testForwardMsg(100+i, t, ps, peerAddresses[i][:], peerAddresses, []int{i})
+		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
+		// send random messages with proximity orders, corresponding to PO of each bin,
-		a = pot.RandomAddressAt(base, i)
+		// with one peer being closer to the recipient address
-		testForwardMsg(200+i, t, ps, a[:], peerAddresses, []int{i})
+		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)
+		a := pot.RandomAddressAt(base, i)
-		testForwardMsg(300+i, t, ps, a[:], peerAddresses, nearestNeighbours)
+		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
-	a = pot.RandomAddressAt(base, 77)
+	a2 := pot.RandomAddressAt(base, 77)
-	testForwardMsg(400, t, ps, a[:], peerAddresses, nearestNeighbours)
+	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 {
-			testForwardMsg(500+i, t, ps, peerAddresses[i][:i], peerAddresses, []int{i})
+			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)
 		}
-		testForwardMsg(550+i, t, ps, peerAddresses[i][:part], peerAddresses, []int{i})
+		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
-		testForwardMsg(600+i, t, ps, peerAddresses[i][:part], peerAddresses, nearestNeighbours)
+		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
-	a = pot.RandomAddressAt(base, part)
+	a3 := pot.RandomAddressAt(base, part)
-	testForwardMsg(700, t, ps, a[:part], peerAddresses, nearestNeighbours)
+	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
-	all := []int{0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12}
-	testForwardMsg(800, t, ps, []byte{}, peerAddresses, all)
 
-	// luminous radius of one byte (8 bits)
+	// zero bytes of address is given, msg should be delivered to all the peers
-	testForwardMsg(900, t, ps, baseAddrBytes[:1], peerAddresses, all[8:])
+	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(testID int, t *testing.T, ps *Pss, recipientAddr []byte, peers []pot.Address, expected []int) {
+func testForwardMsg(t *testing.T, ps *Pss, c *testCase) {
-	testResMap = make(map[pot.Address]int)
+	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 = sendMessageProd }()
+	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
-	s := fmt.Sprintf("test id: %d, msg address: %x..., radius: %d", testID, recipientAddr[:len(recipientAddr)%4], 8*len(recipientAddr))
+	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)
-	for _, i := range expected {
+	if exclusive {
-		a := peers[i]
+		var cnt int
-		received := testResMap[a]
+		for _, i := range expected {
-		if received != 1 {
+			a := peers[i]
-			s += fmt.Sprintf("\npeer number %d [%x...] received %d messages", i, a[:4], received)
+			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
 		}
-		testResMap[a] = 0
+	} 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 testResMap {
+	for k, v := range resultMap {
 		if v != 0 {
 			// find the index of the false positive peer
 			var j int

swarm/pss/pss.go

@@ -886,14 +886,9 @@ func (p *Pss) send(to []byte, topic Topic, msg []byte, asymmetric bool, key []by
 	return nil
 }
 
-// sendMessage is a helper function that tries to send a message and returns true on success
+// sendFunc is a helper function that tries to send a message and returns true on success.
-// It is set in the init function for usage in production, and optionally overridden in tests
+// It is set here for usage in production, and optionally overridden in tests.
-// for data validation.
+var sendFunc func(p *Pss, sp *network.Peer, msg *PssMsg) bool = sendMessageProd
-var sendMessage func(p *Pss, sp *network.Peer, msg *PssMsg) bool
-
-func init() {
-	sendMessage = sendMessageProd
-}
 
 // tries to send a message, returns true if successful
 func sendMessageProd(p *Pss, sp *network.Peer, msg *PssMsg) bool {
@@ -954,19 +949,25 @@ func (p *Pss) forward(msg *PssMsg) error {
 		broadcastThreshold = luminosityRadius
 	}
 
+	var onlySendOnce bool // indicates if the message should only be sent to one peer with closest address
+
 	// if measured from the recipient address as opposed to the base address (see Kademlia.EachConn
 	// call below), then peers that fall in the same proximity bin as recipient address will appear
 	// [at least] one bit closer, but only if these additional bits are given in the recipient address.
 	if broadcastThreshold < luminosityRadius && broadcastThreshold < neighbourhoodDepth {
 		broadcastThreshold++
+		onlySendOnce = true
 	}
 
 	p.Kademlia.EachConn(to, addressLength*8, func(sp *network.Peer, po int, _ bool) bool {
 		if po < broadcastThreshold && sent > 0 {
 			return false // stop iterating
 		}
-		if sendMessage(p, sp, msg) {
+		if sendFunc(p, sp, msg) {
 			sent++
+			if onlySendOnce {
+				return false
+			}
 			if po == addressLength*8 {
 				// stop iterating if successfully sent to the exact recipient (perfect match of full address)
 				return false