WIP: fix some unit tests

cmd/tss/mock_tss_server.go

@@ -3,12 +3,12 @@ package main
 import (
 	"errors"
 
+	"github.com/libp2p/go-libp2p/core/peer"
 	"gitlab.com/thorchain/tss/go-tss/blame"
 	"gitlab.com/thorchain/tss/go-tss/common"
 	"gitlab.com/thorchain/tss/go-tss/conversion"
 	"gitlab.com/thorchain/tss/go-tss/keygen"
 	"gitlab.com/thorchain/tss/go-tss/keysign"
-	"gitlab.com/thorchain/tss/go-tss/tss"
 )
 
 type MockTssServer struct {
@@ -31,8 +31,8 @@ func (mts *MockTssServer) GetLocalPeerID() string {
 	return conversion.GetRandomPeerID().String()
 }
 
-func (mts *MockTssServer) GetKnownPeers() []tss.PeerInfo {
-	return []tss.PeerInfo{}
+func (mts *MockTssServer) GetKnownPeers() []peer.AddrInfo {
+	return []peer.AddrInfo{}
 }
 
 func (mts *MockTssServer) Keygen(req keygen.Request) (keygen.Response, error) {

p2p/communication.go

@@ -59,6 +59,7 @@ type Communication struct {
 	externalAddr     maddr.Multiaddr
 	streamMgr        *StreamMgr
 	whitelistedPeers []peer.ID
+	discovery        *PeerDiscovery
 }
 
 // NewCommunication create a new instance of Communication
@@ -350,6 +351,7 @@ func (c *Communication) startChannel(privKeyBytes []byte) error {
 		bootstrapPeerAddrInfos = append(bootstrapPeerAddrInfos, *peerInfo)
 	}
 	discovery := NewPeerDiscovery(c.host, bootstrapPeerAddrInfos)
+	c.discovery = discovery
 	discovery.Start(context.Background())
 
 	return nil
@@ -420,6 +422,9 @@ func (c *Communication) Start(priKeyBytes []byte) error {
 
 // Stop communication
 func (c *Communication) Stop() error {
+	if c.discovery != nil {
+		c.discovery.Stop()
+	}
 	// we need to stop the handler and the p2p services firstly, then terminate the our communication threads
 	if err := c.host.Close(); err != nil {
 		c.logger.Err(err).Msg("fail to close host network")

p2p/communication_test.go

@@ -7,6 +7,7 @@ import (
 
 	"github.com/libp2p/go-libp2p/core/peer"
 	maddr "github.com/multiformats/go-multiaddr"
+	"github.com/rs/zerolog"
 	"gitlab.com/thorchain/tss/go-tss/messages"
 
 	"github.com/libp2p/go-libp2p/core/crypto"
@@ -39,6 +40,9 @@ func checkExist(a []maddr.Multiaddr, b string) bool {
 }
 
 func (CommunicationTestSuite) TestEstablishP2pCommunication(c *C) {
+	// Set the log level to Info
+	zerolog.SetGlobalLevel(zerolog.InfoLevel)
+
 	bootstrapPeerID, err := peer.Decode("16Uiu2HAm4TmEzUqy3q3Dv7HvdoSboHk5sFj2FH3npiN5vDbJC6gh")
 	c.Assert(err, IsNil)
 	sk1, _, err := crypto.GenerateSecp256k1Key(rand.Reader)
@@ -97,7 +101,7 @@ func (CommunicationTestSuite) TestEstablishP2pCommunication(c *C) {
 	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{})
+	comm5, err := NewCommunication([]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")

p2p/discovery.go

@@ -25,6 +25,7 @@ type PeerDiscovery struct {
 	bootstrapPeers []peer.AddrInfo
 	mu             sync.RWMutex
 	logger         zerolog.Logger
+	closeChan      chan struct{}
 }
 
 func NewPeerDiscovery(h host.Host, bootstrapPeers []peer.AddrInfo) *PeerDiscovery {
@@ -33,6 +34,7 @@ func NewPeerDiscovery(h host.Host, bootstrapPeers []peer.AddrInfo) *PeerDiscover
 		knownPeers:     make(map[peer.ID]peer.AddrInfo),
 		bootstrapPeers: bootstrapPeers,
 		logger:         log.With().Str("module", "peer-discovery").Logger(),
+		closeChan:      make(chan struct{}),
 	}
 
 	// Set up discovery protocol handler
@@ -57,6 +59,10 @@ func (pd *PeerDiscovery) Start(ctx context.Context) {
 	go pd.startGossip(ctx)
 }
 
+func (pd *PeerDiscovery) Stop() {
+	close(pd.closeChan)
+}
+
 // addPeer adds a peer to known peers
 func (pd *PeerDiscovery) addPeer(pinfo peer.AddrInfo) {
 	pd.mu.Lock()
@@ -113,10 +119,15 @@ func (pd *PeerDiscovery) startGossip(ctx context.Context) {
 
 	for {
 		select {
+		case _, ok := <-pd.closeChan:
+			if !ok {
+				pd.logger.Info().Msg("Peer discovery stopped")
+				return
+			}
+			pd.logger.Warn().Msgf("Should not receive from closed channel!")
 		case <-ctx.Done():
 			return
 		case <-ticker.C:
-
 			pd.gossipPeers(ctx)
 		}
 	}
@@ -156,7 +167,7 @@ func (pd *PeerDiscovery) gossipPeers(ctx context.Context) {
 			pd.logger.Error().Err(err).Msgf("Failed to read from stream")
 			continue
 		}
-		pd.logger.Info().Msgf("Received peer data: %s", string(buf))
+		pd.logger.Debug().Msgf("Received peer data: %s", string(buf))
 
 		// Parse received peer info and add to known peers
 		var recvPeers []peer.AddrInfo

p2p/discovery_test.go

@@ -7,54 +7,67 @@ import (
 	"time"
 
 	"github.com/libp2p/go-libp2p/core/crypto"
+	"github.com/libp2p/go-libp2p/core/peer"
 	maddr "github.com/multiformats/go-multiaddr"
 	"github.com/stretchr/testify/assert"
 )
 
 func TestDiscovery(t *testing.T) {
 	GossipInterval = 1 * time.Second
 	bootstrapPeer := "/ip4/127.0.0.1/tcp/2220/p2p/16Uiu2HAm4TmEzUqy3q3Dv7HvdoSboHk5sFj2FH3npiN5vDbJC6gh"
+	bootstrapPeerID, err := peer.Decode("16Uiu2HAm4TmEzUqy3q3Dv7HvdoSboHk5sFj2FH3npiN5vDbJC6gh")
+
 	bootstrapPrivKey := "6LABmWB4iXqkqOJ9H0YFEA2CSSx6bA7XAKGyI/TDtas="
 	externalIP := "127.0.0.1"
+	var whitelistedPeers []peer.ID
+
 	//fakeExternalMultiAddr := "/ip4/127.0.0.1/tcp/2220"
 	validMultiAddr, err := maddr.NewMultiaddr(bootstrapPeer)
 	assert.NoError(t, err)
-
 	privKey, err := base64.StdEncoding.DecodeString(bootstrapPrivKey)
 	assert.NoError(t, err)
-	comm, err := NewCommunication(nil, 2220, externalIP)
-	assert.NoError(t, err)
-	assert.NoError(t, comm.Start(privKey))
-	defer comm.Stop()
 
 	sk1, _, err := crypto.GenerateSecp256k1Key(rand.Reader)
 	sk1raw, _ := sk1.Raw()
 	assert.NoError(t, err)
-	comm2, err := NewCommunication([]maddr.Multiaddr{validMultiAddr}, 2221, externalIP)
+
+	sk2, _, err := crypto.GenerateSecp256k1Key(rand.Reader)
+	assert.NoError(t, err)
+	sk2raw, _ := sk2.Raw()
+	sk3, _, err := crypto.GenerateSecp256k1Key(rand.Reader)
+	assert.NoError(t, err)
+	sk3raw, _ := sk3.Raw()
+	id1, err := peer.IDFromPrivateKey(sk1)
+	id2, err := peer.IDFromPrivateKey(sk2)
+	id3, err := peer.IDFromPrivateKey(sk3)
+	whitelistedPeers = append(whitelistedPeers, id1, id2, id3, bootstrapPeerID)
+
+	comm, err := NewCommunication(nil, 2220, externalIP, whitelistedPeers)
+	assert.NoError(t, err)
+	assert.NoError(t, comm.Start(privKey))
+	defer comm.Stop()
+
+	comm2, err := NewCommunication([]maddr.Multiaddr{validMultiAddr}, 2221, externalIP, whitelistedPeers)
 	assert.NoError(t, err)
 	err = comm2.Start(sk1raw)
 	assert.NoError(t, err)
 	defer comm2.Stop()
 
 	// we connect to an invalid peer and see
-	sk2, _, err := crypto.GenerateSecp256k1Key(rand.Reader)
-	assert.NoError(t, err)
-	sk2raw, _ := sk2.Raw()
+
 	//id, err := peer.IDFromPrivateKey(sk2)
 	//assert.NoError(t, err)
 	//invalidAddr := "/ip4/127.0.0.1/tcp/2220/p2p/" + id.String()
 	//invalidMultiAddr, err := maddr.NewMultiaddr(invalidAddr)
 	assert.NoError(t, err)
-	comm3, err := NewCommunication([]maddr.Multiaddr{validMultiAddr}, 2222, externalIP)
+	comm3, err := NewCommunication([]maddr.Multiaddr{validMultiAddr}, 2222, externalIP, whitelistedPeers)
 	assert.NoError(t, err)
 	err = comm3.Start(sk2raw)
 	defer comm3.Stop()
 
 	// we connect to one invalid and one valid address
-	sk3, _, err := crypto.GenerateSecp256k1Key(rand.Reader)
-	assert.NoError(t, err)
-	sk3raw, _ := sk3.Raw()
-	comm4, err := NewCommunication([]maddr.Multiaddr{validMultiAddr}, 2223, externalIP)
+
+	comm4, err := NewCommunication([]maddr.Multiaddr{validMultiAddr}, 2223, externalIP, whitelistedPeers)
 	assert.NoError(t, err)
 	err = comm4.Start(sk3raw)
 	assert.NoError(t, err)