feat: add option to configure provider whitelist (#192)

p2p/cmd/main.go

@@ -8,6 +8,7 @@ import (
 	"strings"
 	"time"
 
+	"github.com/ethereum/go-ethereum/common"
 	contracts "github.com/primev/mev-commit/contracts-abi/config"
 	mevcommit "github.com/primev/mev-commit/p2p"
 	"github.com/primev/mev-commit/p2p/pkg/node"
@@ -176,27 +177,51 @@ var (
 		Usage:   "address of the bidder registry contract",
 		EnvVars: []string{"MEV_COMMIT_BIDDER_REGISTRY_ADDR"},
 		Value:   contracts.TestnetContracts.BidderRegistry,
+		Action: func(ctx *cli.Context, s string) error {
+			if !common.IsHexAddress(s) {
+				return fmt.Errorf("invalid bidder registry address: %s", s)
+			}
+			return nil
+		},
 	})
 
 	optionProviderRegistryAddr = altsrc.NewStringFlag(&cli.StringFlag{
 		Name:    "provider-registry-contract",
 		Usage:   "address of the provider registry contract",
 		EnvVars: []string{"MEV_COMMIT_PROVIDER_REGISTRY_ADDR"},
 		Value:   contracts.TestnetContracts.ProviderRegistry,
+		Action: func(ctx *cli.Context, s string) error {
+			if !common.IsHexAddress(s) {
+				return fmt.Errorf("invalid provider registry address: %s", s)
+			}
+			return nil
+		},
 	})
 
 	optionPreconfStoreAddr = altsrc.NewStringFlag(&cli.StringFlag{
 		Name:    "preconf-contract",
 		Usage:   "address of the preconfirmation commitment store contract",
 		EnvVars: []string{"MEV_COMMIT_PRECONF_ADDR"},
 		Value:   contracts.TestnetContracts.PreconfCommitmentStore,
+		Action: func(ctx *cli.Context, s string) error {
+			if !common.IsHexAddress(s) {
+				return fmt.Errorf("invalid preconfirmation commitment store address: %s", s)
+			}
+			return nil
+		},
 	})
 
 	optionBlockTrackerAddr = altsrc.NewStringFlag(&cli.StringFlag{
 		Name:    "block-tracker-contract",
 		Usage:   "address of the block tracker contract",
 		EnvVars: []string{"MEV_COMMIT_BLOCK_TRACKER_ADDR"},
 		Value:   contracts.TestnetContracts.BlockTracker,
+		Action: func(ctx *cli.Context, s string) error {
+			if !common.IsHexAddress(s) {
+				return fmt.Errorf("invalid block tracker address: %s", s)
+			}
+			return nil
+		},
 	})
 
 	optionSettlementRPCEndpoint = altsrc.NewStringFlag(&cli.StringFlag{
@@ -236,6 +261,20 @@ var (
 		Usage:   "Path to the server TLS private key",
 		EnvVars: []string{"MEV_COMMIT_SERVER_TLS_PRIVATE_KEY"},
 	})
+
+	optionProviderWhitelist = altsrc.NewStringSliceFlag(&cli.StringSliceFlag{
+		Name:    "provider-whitelist",
+		Usage:   "list of provider addresses to whitelist for bids",
+		EnvVars: []string{"MEV_COMMIT_PROVIDER_WHITELIST"},
+		Action: func(ctx *cli.Context, vals []string) error {
+			for i, v := range vals {
+				if !common.IsHexAddress(v) {
+					return fmt.Errorf("invalid provider address at index %d: %s", i, v)
+				}
+			}
+			return nil
+		},
+	})
 )
 
 func main() {
@@ -266,6 +305,7 @@ func main() {
 		optionNATPort,
 		optionServerTLSCert,
 		optionServerTLSPrivateKey,
+		optionProviderWhitelist,
 	}
 
 	app := &cli.App{
@@ -332,6 +372,11 @@ func launchNodeWithConfig(c *cli.Context) error {
 		return fmt.Errorf("both -%s and -%s must be provided to enable TLS", optionServerTLSCert.Name, optionServerTLSPrivateKey.Name)
 	}
 
+	whitelist := make([]common.Address, 0, len(c.StringSlice(optionProviderWhitelist.Name)))
+	for _, addr := range c.StringSlice(optionProviderWhitelist.Name) {
+		whitelist = append(whitelist, common.HexToAddress(addr))
+	}
+
 	nd, err := node.NewNode(&node.Options{
 		KeySigner:                keysigner,
 		Secret:                   c.String(optionSecret.Name),
@@ -351,6 +396,7 @@ func launchNodeWithConfig(c *cli.Context) error {
 		NatAddr:                  natAddr,
 		TLSCertificateFile:       crtFile,
 		TLSPrivateKeyFile:        keyFile,
+		ProviderWhitelist:        whitelist,
 	})
 	if err != nil {
 		return fmt.Errorf("failed starting node: %w", err)

p2p/pkg/keyexchange/keyexchange.go

@@ -7,9 +7,11 @@ import (
 	"errors"
 	"fmt"
 	"log/slog"
+	"slices"
 	"sync"
 	"time"
 
+	"github.com/ethereum/go-ethereum/common"
 	"github.com/ethereum/go-ethereum/crypto/ecies"
 	keyexchangepb "github.com/primev/mev-commit/p2p/gen/go/keyexchange/v1"
 	"github.com/primev/mev-commit/p2p/pkg/crypto"
@@ -28,16 +30,18 @@ func New(
 	store Store,
 	logger *slog.Logger,
 	signer signer.Signer,
+	providerWhitelist []common.Address,
 ) *KeyExchange {
 	return &KeyExchange{
-		topo:      topo,
-		streamer:  streamer,
-		keySigner: keySigner,
-		aesKey:    aesKey,
-		address:   keySigner.GetAddress(),
-		store:     store,
-		logger:    logger,
-		signer:    signer,
+		topo:              topo,
+		streamer:          streamer,
+		keySigner:         keySigner,
+		aesKey:            aesKey,
+		address:           keySigner.GetAddress(),
+		store:             store,
+		logger:            logger,
+		signer:            signer,
+		providerWhitelist: providerWhitelist,
 	}
 }
 
@@ -74,6 +78,16 @@ func (ke *KeyExchange) SendTimestampMessage() error {
 
 func (ke *KeyExchange) getProviders() ([]p2p.Peer, error) {
 	providers := ke.topo.GetPeers(topology.Query{Type: p2p.PeerTypeProvider})
+	if len(ke.providerWhitelist) > 0 {
+		for i := 0; i < len(providers); i++ {
+			if !slices.ContainsFunc(ke.providerWhitelist, func(e common.Address) bool {
+				return providers[i].EthAddress.Cmp(e) == 0
+			}) {
+				ke.logger.Debug("ignoring provider for key exchange", "provider", providers[i].EthAddress)
+				providers = append(providers[:i], providers[i+1:]...)
+			}
+		}
+	}
 	if len(providers) == 0 {
 		return nil, ErrNoProvidersAvailable
 	}

p2p/pkg/keyexchange/keyexchange_test.go

@@ -5,13 +5,15 @@ import (
 	"crypto/ecdh"
 	"crypto/elliptic"
 	"crypto/rand"
+	"errors"
 	"io"
 	"os"
 	"testing"
 	"time"
 
 	"log/slog"
 
+	"github.com/ethereum/go-ethereum/common"
 	"github.com/ethereum/go-ethereum/crypto"
 	"github.com/ethereum/go-ethereum/crypto/ecies"
 	p2pcrypto "github.com/primev/mev-commit/p2p/pkg/crypto"
@@ -107,8 +109,8 @@ func TestKeyExchange_SendAndHandleTimestampMessage(t *testing.T) {
 		t.Fatal(err)
 	}
 
-	ke1 := keyexchange.New(topo1, svc1, ks, aesKey, bidderStore, logger, signer)
-	ke2 := keyexchange.New(topo2, svc2, ks, nil, providerStore, logger, signer)
+	ke1 := keyexchange.New(topo1, svc1, ks, aesKey, bidderStore, logger, signer, nil)
+	ke2 := keyexchange.New(topo2, svc2, ks, nil, providerStore, logger, signer, nil)
 	if err != nil {
 		t.Fatalf("keyexchange new failed: %v", err)
 	}
@@ -141,3 +143,88 @@ func TestKeyExchange_SendAndHandleTimestampMessage(t *testing.T) {
 		time.Sleep(100 * time.Millisecond)
 	}
 }
+
+func TestKeyExchange_Whitelist(t *testing.T) {
+	t.Parallel()
+
+	privKey, err := crypto.GenerateKey()
+	if err != nil {
+		t.Fatal(err)
+	}
+	address := crypto.PubkeyToAddress(privKey.PublicKey)
+	ks := mockkeysigner.NewMockKeySigner(privKey, address)
+
+	bidderPeer := p2p.Peer{
+		EthAddress: ks.GetAddress(),
+		Type:       p2p.PeerTypeBidder,
+	}
+
+	bidderStore := store.NewStore()
+	providerStore := store.NewStore()
+
+	encryptionPrivateKey, err := ecies.GenerateKey(rand.Reader, elliptic.P256(), nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	err = providerStore.SetECIESPrivateKey(encryptionPrivateKey)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	nikePrivateKey, err := ecdh.P256().GenerateKey(rand.Reader)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	err = providerStore.SetNikePrivateKey(nikePrivateKey)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	providerPeer := p2p.Peer{
+		EthAddress: ks.GetAddress(),
+		Type:       p2p.PeerTypeProvider,
+		Keys:       &p2p.Keys{PKEPublicKey: &encryptionPrivateKey.PublicKey, NIKEPublicKey: nikePrivateKey.PublicKey()},
+	}
+	topo1 := &testTopology{peers: []p2p.Peer{providerPeer}}
+	topo2 := &testTopology{peers: []p2p.Peer{bidderPeer}}
+
+	logger := newTestLogger(t, os.Stdout)
+
+	signer := signer.New()
+	svc1 := p2ptest.New(
+		&bidderPeer,
+	)
+
+	svc2 := p2ptest.New(
+		&providerPeer,
+	)
+
+	aesKey, err := p2pcrypto.GenerateAESKey()
+	if err != nil {
+		t.Fatal(err)
+	}
+	err = bidderStore.SetAESKey(ks.GetAddress(), aesKey)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	randomWhitelistedPeerKey, err := crypto.GenerateKey()
+	if err != nil {
+		t.Fatal(err)
+	}
+	randomWhitelistedPeerAddress := crypto.PubkeyToAddress(randomWhitelistedPeerKey.PublicKey)
+
+	ke1 := keyexchange.New(topo1, svc1, ks, aesKey, bidderStore, logger, signer, []common.Address{randomWhitelistedPeerAddress})
+	ke2 := keyexchange.New(topo2, svc2, ks, nil, providerStore, logger, signer, nil)
+	if err != nil {
+		t.Fatalf("keyexchange new failed: %v", err)
+	}
+	svc1.SetPeerHandler(bidderPeer, ke2.Streams()[0])
+
+	err = ke1.SendTimestampMessage()
+	if !errors.Is(err, keyexchange.ErrNoProvidersAvailable) {
+		t.Fatalf("SendTimestampMessage should have failed")
+	}
+}

p2p/pkg/keyexchange/models.go

@@ -29,14 +29,15 @@ var (
 
 // KeyExchange manages the key exchange process.
 type KeyExchange struct {
-	keySigner keysigner.KeySigner
-	address   common.Address
-	aesKey    []byte
-	topo      Topology
-	streamer  p2p.Streamer
-	signer    signer.Signer
-	store     Store
-	logger    *slog.Logger
+	keySigner         keysigner.KeySigner
+	address           common.Address
+	aesKey            []byte
+	topo              Topology
+	streamer          p2p.Streamer
+	signer            signer.Signer
+	store             Store
+	providerWhitelist []common.Address
+	logger            *slog.Logger
 }
 
 // Topology interface to get peers.

p2p/pkg/node/node.go

@@ -78,6 +78,7 @@ type Options struct {
 	NatAddr                  string
 	TLSCertificateFile       string
 	TLSPrivateKeyFile        string
+	ProviderWhitelist        []common.Address
 }
 
 type Node struct {
@@ -399,6 +400,7 @@ func NewNode(opts *Options) (*Node, error) {
 				store,
 				opts.Logger.With("component", "keyexchange_protocol"),
 				signer.New(),
+				nil,
 			)
 			p2pSvc.AddStreamHandlers(keyexchange.Streams()...)
 			srv.RegisterMetricsCollectors(preconfProto.Metrics()...)
@@ -456,6 +458,7 @@ func NewNode(opts *Options) (*Node, error) {
 				store,
 				opts.Logger.With("component", "keyexchange_protocol"),
 				signer.New(),
+				opts.ProviderWhitelist,
 			)
 			topo.SubscribePeer(func(p p2p.Peer) {
 				if p.Type == p2p.PeerTypeProvider {