communication tests

p2p/communication_test.go

@@ -3,14 +3,14 @@ package p2p
 import (
 	"crypto/rand"
 	"encoding/base64"
+	"fmt"
 
 	"github.com/libp2p/go-libp2p/core/peer"
 	maddr "github.com/multiformats/go-multiaddr"
+	"gitlab.com/thorchain/tss/go-tss/messages"
 
 	"github.com/libp2p/go-libp2p/core/crypto"
 	. "gopkg.in/check.v1"
-
-	"gitlab.com/thorchain/tss/go-tss/messages"
 )
 
 type CommunicationTestSuite struct{}
@@ -39,44 +39,51 @@ func checkExist(a []maddr.Multiaddr, b string) bool {
 }
 
 func (CommunicationTestSuite) TestEstablishP2pCommunication(c *C) {
-	bootstrapPeer := "/ip4/127.0.0.1/tcp/2220/p2p/16Uiu2HAm4TmEzUqy3q3Dv7HvdoSboHk5sFj2FH3npiN5vDbJC6gh"
+	bootstrapPeerID, err := peer.Decode("16Uiu2HAm4TmEzUqy3q3Dv7HvdoSboHk5sFj2FH3npiN5vDbJC6gh")
+	c.Assert(err, IsNil)
+	sk1, _, err := crypto.GenerateSecp256k1Key(rand.Reader)
+	c.Assert(err, IsNil)
+	sk1raw, _ := sk1.Raw()
+	id1, err := peer.IDFromPrivateKey(sk1)
+	c.Assert(err, IsNil)
+	sk2, _, err := crypto.GenerateSecp256k1Key(rand.Reader)
+	c.Assert(err, IsNil)
+	id2, err := peer.IDFromPrivateKey(sk2)
+	c.Assert(err, IsNil)
+
+	bootstrapPeer := fmt.Sprintf("/ip4/127.0.0.1/tcp/2220/p2p/%s", bootstrapPeerID.String())
+	whitelistedPeers := []peer.ID{bootstrapPeerID, id1, id2}
 	bootstrapPrivKey := "6LABmWB4iXqkqOJ9H0YFEA2CSSx6bA7XAKGyI/TDtas="
 	fakeExternalIP := "11.22.33.44"
 	fakeExternalMultiAddr := "/ip4/11.22.33.44/tcp/2220"
 	validMultiAddr, err := maddr.NewMultiaddr(bootstrapPeer)
 	c.Assert(err, IsNil)
 	privKey, err := base64.StdEncoding.DecodeString(bootstrapPrivKey)
 	c.Assert(err, IsNil)
-	comm, err := NewCommunication("commTest", nil, 2220, fakeExternalIP, []peer.ID{})
+	comm, err := NewCommunication("commTest", nil, 2220, fakeExternalIP, whitelistedPeers)
 	c.Assert(err, IsNil)
 	c.Assert(comm.Start(privKey), IsNil)
 
 	defer comm.Stop()
-	sk1, _, err := crypto.GenerateSecp256k1Key(rand.Reader)
-	sk1raw, _ := sk1.Raw()
 	c.Assert(err, IsNil)
-	comm2, err := NewCommunication("commTest", []maddr.Multiaddr{validMultiAddr}, 2221, "", []peer.ID{})
+	comm2, err := NewCommunication("commTest", []maddr.Multiaddr{validMultiAddr}, 2221, "", whitelistedPeers)
 	c.Assert(err, IsNil)
 	err = comm2.Start(sk1raw)
 	c.Assert(err, IsNil)
 	defer comm2.Stop()
 
 	// we connect to an invalid peer and see
-	sk2, _, err := crypto.GenerateSecp256k1Key(rand.Reader)
-	c.Assert(err, IsNil)
-	id, err := peer.IDFromPrivateKey(sk2)
-	c.Assert(err, IsNil)
-	invalidAddr := "/ip4/127.0.0.1/tcp/2220/p2p/" + id.String()
+	invalidAddr := "/ip4/127.0.0.1/tcp/2220/p2p/" + id2.String()
 	invalidMultiAddr, err := maddr.NewMultiaddr(invalidAddr)
 	c.Assert(err, IsNil)
-	comm3, err := NewCommunication("commTest", []maddr.Multiaddr{invalidMultiAddr}, 2222, "", []peer.ID{})
+	comm3, err := NewCommunication("commTest", []maddr.Multiaddr{invalidMultiAddr}, 2222, "", whitelistedPeers)
 	c.Assert(err, IsNil)
 	err = comm3.Start(sk1raw)
 	c.Assert(err, ErrorMatches, "fail to connect to bootstrap peer: fail to connect to any peer")
 	defer comm3.Stop()
 
 	// we connect to one invalid and one valid address
-	comm4, err := NewCommunication("commTest", []maddr.Multiaddr{invalidMultiAddr, validMultiAddr}, 2223, "", []peer.ID{})
+	comm4, err := NewCommunication("commTest", []maddr.Multiaddr{invalidMultiAddr, validMultiAddr}, 2223, "", whitelistedPeers)
 	c.Assert(err, IsNil)
 	err = comm4.Start(sk1raw)
 	c.Assert(err, IsNil)
@@ -88,4 +95,11 @@ func (CommunicationTestSuite) TestEstablishP2pCommunication(c *C) {
 	c.Assert(checkExist(ps.Addrs(comm.host.ID()), fakeExternalMultiAddr), Equals, true)
 	ps = comm4.host.Peerstore()
 	c.Assert(checkExist(ps.Addrs(comm.host.ID()), fakeExternalMultiAddr), Equals, true)
+
+	// same as above, just without whitelisted peers
+	comm5, err := NewCommunication("commTest", []maddr.Multiaddr{invalidMultiAddr, validMultiAddr}, 2224, "", []peer.ID{})
+	c.Assert(err, IsNil)
+	err = comm5.Start(sk1raw)
+	c.Assert(err, ErrorMatches, "fail to connect to bootstrap peer: fail to connect to any peer")
+	defer comm5.Stop()
 }