bzzeth: protocol runloop now quits on peer disconnect

Author: jmozah

bzzeth/bzzeth.go

@@ -18,6 +18,7 @@ package bzzeth
 
 import (
 	"context"
+	"errors"
 
 	"github.com/ethereum/go-ethereum/node"
 	"github.com/ethereum/go-ethereum/p2p"
@@ -26,6 +27,10 @@ import (
 	"github.com/ethersphere/swarm/p2p/protocols"
 )
 
+var (
+	errRcvdMsgFromSwarmNode = errors.New("received message from Swarm node")
+)
+
 // BzzEth implements node.Service
 var _ node.Service = &BzzEth{}
 
@@ -61,14 +66,14 @@ func (b *BzzEth) Run(p *p2p.Peer, rw p2p.MsgReadWriter) error {
 	bp.serveHeaders = handshake.(*Handshake).ServeHeaders
 	log.Debug("handshake", "hs", handshake, "peer", bp)
 
+	b.peers.add(bp)
+	defer b.peers.remove(bp)
+
 	// This protocol is all about interaction between an Eth node and a Swarm Node.
 	// If another swarm node tries to connect then the protocol goes into idle
 	if isSwarmNodeFunc(bp) {
-		<-b.quit
-		return nil
+		return peer.Run(b.handleMsgFromSwarmNode(bp))
 	}
-	b.peers.add(bp)
-	defer b.peers.remove(bp)
 
 	return peer.Run(b.handleMsg(bp))
 }
@@ -85,6 +90,15 @@ func (b *BzzEth) handleMsg(p *Peer) func(context.Context, interface{}) error {
 	}
 }
 
+// handleMsgFromSwarmNode is used in the case if this node is connected to a Swarm node
+// If any message is received in this case, the peer needs to be dropped
+func (b *BzzEth) handleMsgFromSwarmNode(p *Peer) func(context.Context, interface{}) error {
+	return func(ctx context.Context, msg interface{}) error {
+		p.logger.Warn("bzzeth.handleMsgFromSwarmNode")
+		return errRcvdMsgFromSwarmNode
+	}
+}
+
 // Protocols returns the p2p protocol
 func (b *BzzEth) Protocols() []p2p.Protocol {
 	return []p2p.Protocol{

bzzeth/bzzeth_test.go

@@ -72,6 +72,23 @@ func handshakeExchange(tester *p2ptest.ProtocolTester, peerID enode.ID, serveHea
 		})
 }
 
+// This message is exchanged between two Swarm nodes to check if the connection drops
+func dummyHandshakeMessage(tester *p2ptest.ProtocolTester, peerID enode.ID) error {
+	return tester.TestExchanges(
+		p2ptest.Exchange{
+			Label: "Handshake",
+			Triggers: []p2ptest.Trigger{
+				{
+					Code: 0,
+					Msg: Handshake{
+						ServeHeaders: true,
+					},
+					Peer: peerID,
+				},
+			},
+		})
+}
+
 // tests handshake between eth node and swarm node
 // on successful handshake the protocol does not go idle
 // peer added to the pool and serves headers is registered
@@ -89,18 +106,10 @@ func TestBzzEthHandshake(t *testing.T) {
 	}
 
 	// after successful handshake, expect peer added to peer pool
-	var p *Peer
-	for i := 0; i < 10; i++ {
-		p = b.peers.get(node.ID())
-		if p != nil {
-			break
-		}
-		time.Sleep(100 * time.Millisecond)
-	}
+	p := getPeerAfterConnection(node.ID(), b)
 	if p == nil {
 		t.Fatal("bzzeth peer not added")
 	}
-
 	if !p.serveHeaders {
 		t.Fatal("bzzeth peer serveHeaders not set")
 	}
@@ -114,35 +123,103 @@ func TestBzzEthHandshake(t *testing.T) {
 
 // TestBzzBzzHandshake tests that a handshake between two Swarm nodes
 func TestBzzBzzHandshake(t *testing.T) {
+	// redefine isSwarmNodeFunc to force recognise remote peer as swarm node
+	defer func(f func(*Peer) bool) {
+		isSwarmNodeFunc = f
+	}(isSwarmNodeFunc)
+	isSwarmNodeFunc = func(_ *Peer) bool { return true }
+
 	tester, b, teardown, err := newBzzEthTester()
 	if err != nil {
 		t.Fatal(err)
 	}
 	defer teardown()
 
+	node := tester.Nodes[0]
+	err = handshakeExchange(tester, node.ID(), false, true)
+	if err != nil {
+		t.Fatalf("expected no error, got %v", err)
+	}
+
+	// after successful handshake, expect peer added to peer pool
+	p := getPeerAfterConnection(node.ID(), b)
+	if p == nil {
+		t.Fatal("bzzeth peer not added")
+	}
+
+	// after closing the protocol, expect disconnect
+	close(b.quit)
+	err = tester.TestDisconnected(&p2ptest.Disconnect{Peer: node.ID(), Error: errors.New("?")})
+	if err == nil || err.Error() != "timed out waiting for peers to disconnect" {
+		t.Fatal(err)
+	}
+}
+
+// TestBzzBzzHandshakeWithMessage tests that a handshake between two Swarm nodes and message exchange
+// disconnects the peer
+func TestBzzBzzHandshakeWithMessage(t *testing.T) {
 	// redefine isSwarmNodeFunc to force recognise remote peer as swarm node
 	defer func(f func(*Peer) bool) {
 		isSwarmNodeFunc = f
 	}(isSwarmNodeFunc)
 	isSwarmNodeFunc = func(_ *Peer) bool { return true }
 
+	tester, b, teardown, err := newBzzEthTester()
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer teardown()
+
 	node := tester.Nodes[0]
 	err = handshakeExchange(tester, node.ID(), false, true)
 	if err != nil {
 		t.Fatalf("expected no error, got %v", err)
 	}
 
-	// after handshake expect protocol to hang, peer not added to pool
-	p := b.peers.get(node.ID())
-	if p != nil {
-		t.Fatal("bzzeth swarm peer incorrectly added")
+	// after successful handshake, expect peer added to peer pool
+	var p *Peer
+	for i := 0; i < 10; i++ {
+		p = b.peers.get(node.ID())
+		if p != nil {
+			break
+		}
+		time.Sleep(100 * time.Millisecond)
+	}
+	if p == nil {
+		t.Fatal("bzzeth peer not added")
 	}
 
-	// after closing the ptotocall, expect disconnect
-	close(b.quit)
-	err = tester.TestDisconnected(&p2ptest.Disconnect{Peer: node.ID(), Error: errors.New("protocol returned")})
+	// Send a dummy handshake message, wait for sometime and check if peer is dropped
+	err = dummyHandshakeMessage(tester, node.ID())
 	if err != nil {
 		t.Fatal(err)
 	}
+	// after successful handshake, expect peer added to peer pool
+	p1 := isPeerDisconnected(node.ID(), b)
+	if p1 != nil {
+		t.Fatal("bzzeth peer still connected")
+	}
+}
+
+func getPeerAfterConnection(id enode.ID, b *BzzEth) (p *Peer) {
+	for i := 0; i < 10; i++ {
+		p = b.peers.get(id)
+		if p != nil {
+			break
+		}
+		time.Sleep(100 * time.Millisecond)
+	}
+	return
+}
 
+func isPeerDisconnected(id enode.ID, b *BzzEth) (p *Peer) {
+	var p1 *Peer
+	for i := 0; i < 10; i++ {
+		p1 = b.peers.get(id)
+		if p1 == nil {
+			break
+		}
+		time.Sleep(100 * time.Millisecond)
+	}
+	return
 }