Refactor ecdh service to become resilient to node disconnect and rejoin

Author: anhthii

cmd/mpcium/main.go

@@ -129,7 +129,6 @@ func runNode(ctx context.Context, c *cli.Command) error {
 	if err != nil {
 		logger.Fatal("Failed to connect to NATS", err)
 	}
-	defer natsConn.Close()
 
 	pubsub := messaging.NewNATSPubSub(natsConn)
 	keygenBroker, err := messaging.NewJetStreamBroker(ctx, natsConn, event.KeygenBrokerStream, []string{
@@ -162,7 +161,7 @@ func runNode(ctx context.Context, c *cli.Command) error {
 	logger.Info("Node is running", "ID", nodeID, "name", nodeName)
 
 	peerNodeIDs := GetPeerIDs(peers)
-	peerRegistry := mpc.NewRegistry(nodeID, peerNodeIDs, consulClient.KV(), directMessaging)
+	peerRegistry := mpc.NewRegistry(nodeID, peerNodeIDs, consulClient.KV(), directMessaging, pubsub, identityStore)
 
 	mpcNode := mpc.NewNode(
 		nodeID,
@@ -176,9 +175,6 @@ func runNode(ctx context.Context, c *cli.Command) error {
 	)
 	defer mpcNode.Close()
 
-	// ECDH session for DH key exchange
-	ecdhSession := mpcNode.GetECDHSession()
-
 	eventConsumer := eventconsumer.NewEventConsumer(
 		mpcNode,
 		pubsub,
@@ -206,12 +202,7 @@ func runNode(ctx context.Context, c *cli.Command) error {
 	}
 	logger.Info("[READY] Node is ready", "nodeID", nodeID)
 
-	logger.Info("Waiting for ECDH key exchange to complete...", "nodeID", nodeID)
-	if err := ecdhSession.WaitForExchangeComplete(); err != nil {
-		logger.Fatal("ECDH exchange failed", err)
-	}
-
-	logger.Info("ECDH key exchange completed successfully, starting consumers...", "nodeID", nodeID)
+	logger.Info("Starting consumers", "nodeID", nodeID)
 	appContext, cancel := context.WithCancel(context.Background())
 	//Setup signal handling to cancel context on termination signals.
 	go func() {
@@ -221,6 +212,11 @@ func runNode(ctx context.Context, c *cli.Command) error {
 		logger.Warn("Shutdown signal received, canceling context...")
 		cancel()
 
+		// Resign from peer registry first (before closing NATS)
+		if err := peerRegistry.Resign(); err != nil {
+			logger.Error("Failed to resign from peer registry", err)
+		}
+
 		// Gracefully close consumers
 		if err := keygenConsumer.Close(); err != nil {
 			logger.Error("Failed to close keygen consumer", err)
@@ -229,10 +225,6 @@ func runNode(ctx context.Context, c *cli.Command) error {
 			logger.Error("Failed to close signing consumer", err)
 		}
 
-		if err := ecdhSession.Close(); err != nil {
-			logger.Error("Failed to close ECDH session", err)
-		}
-
 		err := natsConn.Drain()
 		if err != nil {
 			logger.Error("Failed to drain NATS connection", err)
@@ -264,21 +256,6 @@ func runNode(ctx context.Context, c *cli.Command) error {
 		logger.Info("Signing consumer finished successfully")
 	}()
 
-	go func() {
-		for {
-			select {
-			case <-appContext.Done():
-				return
-			case err := <-ecdhSession.ErrChan():
-				if err != nil {
-					logger.Error("ECDH session error", err)
-					errChan <- fmt.Errorf("ecdh session error: %w", err)
-					return
-				}
-			}
-		}
-	}()
-
 	go func() {
 		wg.Wait()
 		logger.Info("All consumers have finished")

pkg/eventconsumer/keygen_consumer.go

@@ -83,6 +83,9 @@ func (sc *keygenConsumer) waitForAllPeersReadyToGenKey(ctx context.Context) erro
 func (sc *keygenConsumer) Run(ctx context.Context) error {
 	// Wait for sufficient peers before starting to consume messages
 	if err := sc.waitForAllPeersReadyToGenKey(ctx); err != nil {
+		if err == context.Canceled {
+			return nil
+		}
 		return fmt.Errorf("failed to wait for sufficient peers: %w", err)
 	}
 

pkg/eventconsumer/sign_consumer.go

@@ -93,6 +93,9 @@ func (sc *signingConsumer) waitForSufficientPeers(ctx context.Context) error {
 func (sc *signingConsumer) Run(ctx context.Context) error {
 	// Wait for sufficient peers before starting to consume messages
 	if err := sc.waitForSufficientPeers(ctx); err != nil {
+		if err == context.Canceled {
+			return nil
+		}
 		return fmt.Errorf("failed to wait for sufficient peers: %w", err)
 	}
 

pkg/identity/identity.go

@@ -44,6 +44,8 @@ type Store interface {
 
 	SetSymmetricKey(peerID string, key []byte)
 	GetSymmetricKey(peerID string) ([]byte, error)
+	RemoveSymmetricKey(peerID string)
+	GetSymetricKeyCount() int
 	CheckSymmetricKeyComplete(desired int) bool
 
 	EncryptMessage(plaintext []byte, peerID string) ([]byte, error)
@@ -238,18 +240,22 @@ func (s *fileStore) GetSymmetricKey(peerID string) ([]byte, error) {
 	return nil, fmt.Errorf("SymmetricKey key not found for node ID: %s", peerID)
 }
 
-func (s *fileStore) CheckSymmetricKeyComplete(desired int) bool {
+func (s *fileStore) RemoveSymmetricKey(peerID string) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	delete(s.symmetricKeys, peerID)
+}
+
+func (s *fileStore) GetSymetricKeyCount() int {
 	s.mu.RLock()
 	defer s.mu.RUnlock()
+	return len(s.symmetricKeys)
+}
 
-	completeCount := 0
-	for _, value := range s.symmetricKeys {
-		if len(value) > 0 {
-			completeCount++
-		}
-	}
-
-	return completeCount == desired
+func (s *fileStore) CheckSymmetricKeyComplete(desired int) bool {
+	s.mu.RLock()
+	defer s.mu.RUnlock()
+	return len(s.symmetricKeys) == desired
 }
 
 // GetPublicKey retrieves a node's public key by its ID

pkg/mpc/key_exchange_session.go

@@ -17,8 +17,6 @@ import (
 
 	"encoding/json"
 
-	"sync"
-
 	"github.com/nats-io/nats.go"
 )
 
@@ -30,24 +28,21 @@ const (
 type ECDHSession interface {
 	ListenKeyExchange() error
 	BroadcastPublicKey() error
-	WaitForExchangeComplete() error
-	ResetLocalKeys()
+	RemovePeer(peerID string)
+	GetReadyPeersCount() int
 	ErrChan() <-chan error
 	Close() error
 }
 
 type ecdhSession struct {
-	nodeID           string
-	peerIDs          []string
-	pubSub           messaging.PubSub
-	ecdhSub          messaging.Subscription
-	identityStore    identity.Store
-	privateKey       *ecdh.PrivateKey
-	publicKey        *ecdh.PublicKey
-	exchangeComplete chan struct{}
-	errCh            chan error
-	exchangeDone     bool
-	mu               sync.RWMutex
+	nodeID        string
+	peerIDs       []string
+	pubSub        messaging.PubSub
+	ecdhSub       messaging.Subscription
+	identityStore identity.Store
+	privateKey    *ecdh.PrivateKey
+	publicKey     *ecdh.PublicKey
+	errCh         chan error
 }
 
 func NewECDHSession(
@@ -57,20 +52,24 @@ func NewECDHSession(
 	identityStore identity.Store,
 ) *ecdhSession {
 	return &ecdhSession{
-		nodeID:           nodeID,
-		peerIDs:          peerIDs,
-		pubSub:           pubSub,
-		identityStore:    identityStore,
-		exchangeComplete: make(chan struct{}, 1),
-		errCh:            make(chan error, 1),
+		nodeID:        nodeID,
+		peerIDs:       peerIDs,
+		pubSub:        pubSub,
+		identityStore: identityStore,
+		errCh:         make(chan error, 1),
 	}
 }
 
-func (e *ecdhSession) ResetLocalKeys() {
-	// Set a specific key to an empty []byte
-	for _, peerID := range e.peerIDs {
-		e.identityStore.SetSymmetricKey(peerID, []byte{})
-	}
+func (e *ecdhSession) RemovePeer(peerID string) {
+	e.identityStore.RemoveSymmetricKey(peerID)
+}
+
+func (e *ecdhSession) GetReadyPeersCount() int {
+	return e.identityStore.GetSymetricKeyCount()
+}
+
+func (e *ecdhSession) ErrChan() <-chan error {
+	return e.errCh
 }
 
 func (e *ecdhSession) ListenKeyExchange() error {
@@ -114,21 +113,7 @@ func (e *ecdhSession) ListenKeyExchange() error {
 		// Derive symmetric key using HKDF
 		symmetricKey := e.deriveSymmetricKey(sharedSecret, ecdhMsg.From)
 		e.identityStore.SetSymmetricKey(ecdhMsg.From, symmetricKey)
-
-		requiredKeyCount := len(e.peerIDs) - 1
-		logger.Info("ECDH progress", "peer", ecdhMsg.From, "required", requiredKeyCount)
-
-		if e.identityStore.CheckSymmetricKeyComplete(requiredKeyCount) {
-			logger.Info("Completed ECDH!", "symmetric key counts of peers", requiredKeyCount)
-			logger.Info("ALL PEERS ARE READY! Starting to accept MPC requests")
-
-			e.mu.Lock()
-			if !e.exchangeDone {
-				e.exchangeDone = true
-				e.exchangeComplete <- struct{}{}
-			}
-			e.mu.Unlock()
-		}
+		logger.Debug("ECDH progress", "peer", ecdhMsg.From, "current", e.identityStore.GetSymetricKeyCount())
 	})
 
 	e.ecdhSub = sub
@@ -138,10 +123,6 @@ func (e *ecdhSession) ListenKeyExchange() error {
 	return nil
 }
 
-func (s *ecdhSession) ErrChan() <-chan error {
-	return s.errCh
-}
-
 func (s *ecdhSession) Close() error {
 	err := s.ecdhSub.Unsubscribe()
 	if err != nil {
@@ -173,25 +154,6 @@ func (e *ecdhSession) BroadcastPublicKey() error {
 	return nil
 }
 
-func (e *ecdhSession) WaitForExchangeComplete() error {
-	e.mu.RLock()
-	if e.exchangeDone {
-		e.mu.RUnlock()
-		return nil
-	}
-	e.mu.RUnlock()
-	timeout := time.After(ECDHExchangeTimeout) // 2 minutes timeout
-
-	select {
-	case <-e.exchangeComplete:
-		return nil
-	case err := <-e.errCh:
-		return err
-	case <-timeout:
-		return fmt.Errorf("ECDH exchange timeout!")
-	}
-}
-
 func deriveConsistentInfo(a, b string) []byte {
 	if a < b {
 		return []byte(a + b)

pkg/mpc/node.go

@@ -39,7 +39,6 @@ type Node struct {
 	identityStore  identity.Store
 
 	peerRegistry PeerRegistry
-	ecdhSession  ECDHSession
 }
 
 func NewNode(
@@ -55,11 +54,6 @@ func NewNode(
 	start := time.Now()
 	elapsed := time.Since(start)
 	logger.Info("Starting new node, preparams is generated successfully!", "elapsed", elapsed.Milliseconds())
-	// Each node initiates the DH key exchange listener at the beginning and invoke message sending when all peers are ready
-	dhSession := NewECDHSession(nodeID, peerIDs, pubSub, identityStore)
-	if err := dhSession.ListenKeyExchange(); err != nil {
-		logger.Fatal("Failed to listen to DH key exchange", err)
-	}
 
 	node := &Node{
 		nodeID:        nodeID,
@@ -70,22 +64,11 @@ func NewNode(
 		keyinfoStore:  keyinfoStore,
 		peerRegistry:  peerRegistry,
 		identityStore: identityStore,
-		ecdhSession:   dhSession,
 	}
 	node.ecdsaPreParams = node.generatePreParams()
 
-	// we define two types of tasks, initTask and resetTask
-	ecdhTasks := func(isInit bool) {
-		if isInit {
-			if err := dhSession.BroadcastPublicKey(); err != nil {
-				logger.Fatal("DH key broadcast failed", err)
-			}
-		} else {
-			dhSession.ResetLocalKeys()
-		}
-	}
-
-	go peerRegistry.WatchPeersReady(ecdhTasks)
+	// Start watching peers - ECDH is now handled by the registry
+	go peerRegistry.WatchPeersReady()
 	return node
 }
 
@@ -430,9 +413,6 @@ func (p *Node) Close() {
 	}
 }
 
-func (p *Node) GetECDHSession() ECDHSession {
-	return p.ecdhSession
-}
 
 func (p *Node) generatePreParams() []*keygen.LocalPreParams {
 	start := time.Now()