feat: whitelist connection gater (#31)

cmd/tss/main.go

@@ -12,6 +12,7 @@ import (
 
 	"github.com/cosmos/cosmos-sdk/client/input"
 	golog "github.com/ipfs/go-log"
+	"github.com/libp2p/go-libp2p/core/peer"
 
 	"gitlab.com/thorchain/tss/go-tss/common"
 	"gitlab.com/thorchain/tss/go-tss/conversion"
@@ -63,6 +64,7 @@ func main() {
 		nil,
 		p2pConf.ExternalIP,
 		os.Getenv("PASSWORD"),
+		[]peer.ID{},
 	)
 	if nil != err {
 		log.Fatal(err)

conversion/conversion.go

@@ -212,3 +212,15 @@ func GetEDDSAPrivateKeyRawBytes(privateKey crypto2.PrivKey) ([]byte, error) {
 	copy(keyBytesArray[:], pk[:])
 	return keyBytesArray[:], nil
 }
+
+func Bech32PubkeyToPeerID(pubKey string) (peer.ID, error) {
+	bech32PubKey, err := sdk.UnmarshalPubKey(sdk.AccPK, pubKey)
+	if err != nil {
+		return "", err
+	}
+	secp256k1PubKey, err := crypto2.UnmarshalSecp256k1PublicKey(bech32PubKey.Bytes())
+	if err != nil {
+		return "", err
+	}
+	return peer.IDFromPublicKey(secp256k1PubKey)
+}

keygen/ecdsa/keygen_test.go

@@ -114,17 +114,23 @@ func (s *TssECDSAKeygenTestSuite) SetUpTest(c *C) {
 	multiAddr, err := maddr.NewMultiaddr(bootstrapPeer)
 	c.Assert(err, IsNil)
 	s.preParams = getPreparams(c)
+	whitelistedPeers := []peer.ID{}
+	for _, pk := range testPubKeys {
+		peer, err := conversion.Bech32PubkeyToPeerID(pk)
+		c.Assert(err, IsNil)
+		whitelistedPeers = append(whitelistedPeers, peer)
+	}
 	for i := 0; i < s.partyNum; i++ {
 		buf, err := base64.StdEncoding.DecodeString(testPriKeyArr[i])
 		c.Assert(err, IsNil)
 		if i == 0 {
-			comm, err := p2p.NewCommunication("asgard", nil, ports[i], "")
+			comm, err := p2p.NewCommunication("asgard", nil, ports[i], "", whitelistedPeers)
 			c.Assert(err, IsNil)
 			c.Assert(comm.Start(buf[:]), IsNil)
 			s.comms[i] = comm
 			continue
 		}
-		comm, err := p2p.NewCommunication("asgard", []maddr.Multiaddr{multiAddr}, ports[i], "")
+		comm, err := p2p.NewCommunication("asgard", []maddr.Multiaddr{multiAddr}, ports[i], "", whitelistedPeers)
 		c.Assert(err, IsNil)
 		c.Assert(comm.Start(buf[:]), IsNil)
 		s.comms[i] = comm

keygen/eddsa/keygen_test.go

@@ -100,17 +100,23 @@ func (s *EddsaKeygenTestSuite) SetUpTest(c *C) {
 	bootstrapPeer := "/ip4/127.0.0.1/tcp/19666/p2p/16Uiu2HAm4TmEzUqy3q3Dv7HvdoSboHk5sFj2FH3npiN5vDbJC6gh"
 	multiAddr, err := maddr.NewMultiaddr(bootstrapPeer)
 	c.Assert(err, IsNil)
+	whitelistedPeers := []peer.ID{}
+	for _, pk := range testPubKeys {
+		peer, err := conversion.Bech32PubkeyToPeerID(pk)
+		c.Assert(err, IsNil)
+		whitelistedPeers = append(whitelistedPeers, peer)
+	}
 	for i := 0; i < s.partyNum; i++ {
 		buf, err := base64.StdEncoding.DecodeString(testPriKeyArr[i])
 		c.Assert(err, IsNil)
 		if i == 0 {
-			comm, err := p2p.NewCommunication("asgard", nil, ports[i], "")
+			comm, err := p2p.NewCommunication("asgard", nil, ports[i], "", whitelistedPeers)
 			c.Assert(err, IsNil)
 			c.Assert(comm.Start(buf), IsNil)
 			s.comms[i] = comm
 			continue
 		}
-		comm, err := p2p.NewCommunication("asgard", []maddr.Multiaddr{multiAddr}, ports[i], "")
+		comm, err := p2p.NewCommunication("asgard", []maddr.Multiaddr{multiAddr}, ports[i], "", whitelistedPeers)
 		c.Assert(err, IsNil)
 		c.Assert(comm.Start(buf), IsNil)
 		s.comms[i] = comm

keysign/ecdsa/keysign_old_test.go

@@ -126,17 +126,23 @@ func (s *TssECDSAKeysignOldTestSuite) SetUpTest(c *C) {
 	bootstrapPeer := "/ip4/127.0.0.1/tcp/17666/p2p/16Uiu2HAm4TmEzUqy3q3Dv7HvdoSboHk5sFj2FH3npiN5vDbJC6gh"
 	multiAddr, err := maddr.NewMultiaddr(bootstrapPeer)
 	c.Assert(err, IsNil)
+	whitelistedPeers := []peer.ID{}
+	for _, pk := range testPubKeys {
+		peer, err := conversion.Bech32PubkeyToPeerID(pk)
+		c.Assert(err, IsNil)
+		whitelistedPeers = append(whitelistedPeers, peer)
+	}
 	for i := 0; i < s.partyNum; i++ {
 		buf, err := base64.StdEncoding.DecodeString(testPriKeyArr[i])
 		c.Assert(err, IsNil)
 		if i == 0 {
-			comm, err := p2p.NewCommunication("asgard", nil, ports[i], "")
+			comm, err := p2p.NewCommunication("asgard", nil, ports[i], "", whitelistedPeers)
 			c.Assert(err, IsNil)
 			c.Assert(comm.Start(buf), IsNil)
 			s.comms[i] = comm
 			continue
 		}
-		comm, err := p2p.NewCommunication("asgard", []maddr.Multiaddr{multiAddr}, ports[i], "")
+		comm, err := p2p.NewCommunication("asgard", []maddr.Multiaddr{multiAddr}, ports[i], "", whitelistedPeers)
 		c.Assert(err, IsNil)
 		c.Assert(comm.Start(buf), IsNil)
 		s.comms[i] = comm

keysign/ecdsa/keysign_test.go

@@ -138,17 +138,24 @@ func (s *TssECDSAKeysignTestSuite) SetUpTest(c *C) {
 	bootstrapPeer := "/ip4/127.0.0.1/tcp/17666/p2p/16Uiu2HAm4TmEzUqy3q3Dv7HvdoSboHk5sFj2FH3npiN5vDbJC6gh"
 	multiAddr, err := maddr.NewMultiaddr(bootstrapPeer)
 	c.Assert(err, IsNil)
+
+	whitelistedPeers := []peer.ID{}
+	for _, pk := range testPubKeys {
+		peer, err := conversion.Bech32PubkeyToPeerID(pk)
+		c.Assert(err, IsNil)
+		whitelistedPeers = append(whitelistedPeers, peer)
+	}
 	for i := 0; i < s.partyNum; i++ {
 		buf, err := base64.StdEncoding.DecodeString(testPriKeyArr[i])
 		c.Assert(err, IsNil)
 		if i == 0 {
-			comm, err := p2p.NewCommunication("asgard", nil, ports[i], "")
+			comm, err := p2p.NewCommunication("asgard", nil, ports[i], "", whitelistedPeers)
 			c.Assert(err, IsNil)
 			c.Assert(comm.Start(buf), IsNil)
 			s.comms[i] = comm
 			continue
 		}
-		comm, err := p2p.NewCommunication("asgard", []maddr.Multiaddr{multiAddr}, ports[i], "")
+		comm, err := p2p.NewCommunication("asgard", []maddr.Multiaddr{multiAddr}, ports[i], "", whitelistedPeers)
 		c.Assert(err, IsNil)
 		c.Assert(comm.Start(buf), IsNil)
 		s.comms[i] = comm

keysign/eddsa/keysign_test.go

@@ -136,17 +136,23 @@ func (s *EddsaKeysignTestSuite) SetUpTest(c *C) {
 	bootstrapPeer := "/ip4/127.0.0.1/tcp/15666/p2p/16Uiu2HAm4TmEzUqy3q3Dv7HvdoSboHk5sFj2FH3npiN5vDbJC6gh"
 	multiAddr, err := maddr.NewMultiaddr(bootstrapPeer)
 	c.Assert(err, IsNil)
+	whitelistedPeers := []peer.ID{}
+	for _, pk := range testPubKeys {
+		peer, err := conversion.Bech32PubkeyToPeerID(pk)
+		c.Assert(err, IsNil)
+		whitelistedPeers = append(whitelistedPeers, peer)
+	}
 	for i := 0; i < s.partyNum; i++ {
 		buf, err := base64.StdEncoding.DecodeString(testPriKeyArr[i])
 		c.Assert(err, IsNil)
 		if i == 0 {
-			comm, err := p2p.NewCommunication("asgard", nil, ports[i], "")
+			comm, err := p2p.NewCommunication("asgard", nil, ports[i], "", whitelistedPeers)
 			c.Assert(err, IsNil)
 			c.Assert(comm.Start(buf), IsNil)
 			s.comms[i] = comm
 			continue
 		}
-		comm, err := p2p.NewCommunication("asgard", []maddr.Multiaddr{multiAddr}, ports[i], "")
+		comm, err := p2p.NewCommunication("asgard", []maddr.Multiaddr{multiAddr}, ports[i], "", whitelistedPeers)
 		c.Assert(err, IsNil)
 		c.Assert(comm.Start(buf), IsNil)
 		s.comms[i] = comm

p2p/communication.go

@@ -62,10 +62,17 @@ type Communication struct {
 	BroadcastMsgChan chan *messages.BroadcastMsgChan
 	externalAddr     maddr.Multiaddr
 	streamMgr        *StreamMgr
+	whitelistedPeers []peer.ID
 }
 
 // NewCommunication create a new instance of Communication
-func NewCommunication(rendezvous string, bootstrapPeers []maddr.Multiaddr, port int, externalIP string) (*Communication, error) {
+func NewCommunication(
+	rendezvous string,
+	bootstrapPeers []maddr.Multiaddr,
+	port int,
+	externalIP string,
+	whitelistedPeers []peer.ID,
+) (*Communication, error) {
 	addr, err := maddr.NewMultiaddr(fmt.Sprintf("/ip4/0.0.0.0/tcp/%d", port))
 	if err != nil {
 		return nil, fmt.Errorf("fail to create listen addr: %w", err)
@@ -90,6 +97,7 @@ func NewCommunication(rendezvous string, bootstrapPeers []maddr.Multiaddr, port
 		BroadcastMsgChan: make(chan *messages.BroadcastMsgChan, 1024),
 		externalAddr:     externalAddr,
 		streamMgr:        NewStreamMgr(),
+		whitelistedPeers: whitelistedPeers,
 	}, nil
 }
 
@@ -301,6 +309,7 @@ func (c *Communication) startChannel(privKeyBytes []byte) error {
 		libp2p.AddrsFactory(addressFactory),
 		libp2p.ResourceManager(m),
 		libp2p.ConnectionManager(cmgr),
+		libp2p.ConnectionGater(NewWhitelistConnectionGater(c.whitelistedPeers, c.logger)),
 	)
 	if err != nil {
 		return fmt.Errorf("fail to create p2p host: %w", err)

p2p/communication_test.go

@@ -3,22 +3,22 @@ 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{}
 
 var _ = Suite(&CommunicationTestSuite{})
 
 func (CommunicationTestSuite) TestBasicCommunication(c *C) {
-	comm, err := NewCommunication("rendezvous", nil, 6668, "")
+	comm, err := NewCommunication("rendezvous", nil, 6668, "", []peer.ID{})
 	c.Assert(err, IsNil)
 	c.Assert(comm, NotNil)
 	comm.SetSubscribe(messages.TSSKeyGenMsg, "hello", make(chan *Message))
@@ -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)
+	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, "")
+	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, "")
+	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, "")
+	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()
 }

p2p/whitelist_connection_gater.go

@@ -0,0 +1,81 @@
+package p2p
+
+import (
+	"github.com/libp2p/go-libp2p/core/control"
+	"github.com/libp2p/go-libp2p/core/network"
+	"github.com/rs/zerolog"
+
+	"github.com/libp2p/go-libp2p/core/peer"
+	maddr "github.com/multiformats/go-multiaddr"
+)
+
+type WhitelistConnectionGater struct {
+	whitelistedPeers map[peer.ID]bool
+	logger           zerolog.Logger
+}
+
+func NewWhitelistConnectionGater(whitelistedPeers []peer.ID, logger zerolog.Logger) *WhitelistConnectionGater {
+	gater := &WhitelistConnectionGater{
+		logger:           logger,
+		whitelistedPeers: make(map[peer.ID]bool),
+	}
+
+	for _, p := range whitelistedPeers {
+		logger.Info().
+			Stringer("peer", p).
+			Msg("Adding peer to whitelist")
+		gater.whitelistedPeers[p] = true
+	}
+
+	return gater
+}
+
+func (wg *WhitelistConnectionGater) InterceptPeerDial(p peer.ID) (allow bool) {
+	return wg.peerAllowed("InterceptPeerDial", p, nil)
+}
+
+func (wg *WhitelistConnectionGater) InterceptAddrDial(p peer.ID, m maddr.Multiaddr) (allow bool) {
+	return wg.peerAllowed("InterceptAddrDial", p, &m)
+}
+
+func (wg *WhitelistConnectionGater) InterceptAccept(m network.ConnMultiaddrs) (allow bool) {
+	return true
+}
+
+func (wg *WhitelistConnectionGater) InterceptSecured(direction network.Direction, p peer.ID, m network.ConnMultiaddrs) (allow bool) {
+	remoteMultiAddr := m.RemoteMultiaddr()
+	return wg.peerAllowed("InterceptSecured", p, &remoteMultiAddr)
+}
+
+func (wg *WhitelistConnectionGater) InterceptUpgraded(network.Conn) (bool, control.DisconnectReason) {
+	// Allow connection upgrades
+	return true, 0
+}
+
+func (wg *WhitelistConnectionGater) peerAllowed(interceptor string, p peer.ID, remoteAddr *maddr.Multiaddr) bool {
+	allowed := wg.whitelistedPeers[p]
+
+	var event *zerolog.Event
+	if allowed {
+		event = wg.logger.Debug() // log allowed peers at Debug level
+	} else {
+		event = wg.logger.Info() // log denied peers at Info level
+	}
+
+	event = event.
+		Str("interceptor", interceptor).
+		Stringer("peer", p).
+		Bool("allowed", allowed)
+
+	if remoteAddr != nil {
+		event.Str("remote_address", (*remoteAddr).String())
+	}
+
+	if allowed {
+		event.Msg("Peer allowed")
+	} else {
+		event.Msg("Peer denied")
+	}
+
+	return allowed
+}