Fix ECDH key exchange race condition in distributed deployments

Author: anhthii

pkg/mpc/key_exchange_session.go

@@ -32,17 +32,19 @@ type ECDHSession interface {
 	GetReadyPeersCount() int
 	ErrChan() <-chan error
 	Close() error
+	OnKeyExchangeComplete(callback func())
 }
 
 type ecdhSession struct {
-	nodeID        string
-	peerIDs       []string
-	pubSub        messaging.PubSub
-	ecdhSub       messaging.Subscription
-	identityStore identity.Store
-	privateKey    *ecdh.PrivateKey
-	publicKey     *ecdh.PublicKey
-	errCh         chan error
+	nodeID                string
+	peerIDs               []string
+	pubSub                messaging.PubSub
+	ecdhSub               messaging.Subscription
+	identityStore         identity.Store
+	privateKey            *ecdh.PrivateKey
+	publicKey             *ecdh.PublicKey
+	errCh                 chan error
+	onKeyExchangeComplete func()
 }
 
 func NewECDHSession(
@@ -72,6 +74,10 @@ func (e *ecdhSession) ErrChan() <-chan error {
 	return e.errCh
 }
 
+func (e *ecdhSession) OnKeyExchangeComplete(callback func()) {
+	e.onKeyExchangeComplete = callback
+}
+
 func (e *ecdhSession) ListenKeyExchange() error {
 	// Generate an ephemeral ECDH key pair
 	privateKey, err := ecdh.X25519().GenerateKey(rand.Reader)
@@ -113,7 +119,15 @@ func (e *ecdhSession) ListenKeyExchange() error {
 		// Derive symmetric key using HKDF
 		symmetricKey := e.deriveSymmetricKey(sharedSecret, ecdhMsg.From)
 		e.identityStore.SetSymmetricKey(ecdhMsg.From, symmetricKey)
-		logger.Debug("ECDH progress", "peer", ecdhMsg.From, "current", e.identityStore.GetSymetricKeyCount())
+
+		currentKeyCount := e.identityStore.GetSymetricKeyCount()
+		logger.Debug("ECDH progress", "peer", ecdhMsg.From, "current", currentKeyCount, "expected", len(e.peerIDs))
+
+		// Check if ECDH exchange is complete and notify callback
+		if currentKeyCount == len(e.peerIDs) && e.onKeyExchangeComplete != nil {
+			logger.Info("ECDH key exchange completed successfully", "totalKeys", currentKeyCount)
+			e.onKeyExchangeComplete()
+		}
 	})
 
 	e.ecdhSub = sub

pkg/mpc/registry.go

@@ -85,6 +85,11 @@ func NewRegistry(
 		mpcThreshold:  mpcThreshold,
 	}
 
+	// Set up callback to check ready state when ECDH completes
+	ecdhSession.OnKeyExchangeComplete(func() {
+		reg.checkAndUpdateReadyState()
+	})
+
 	go reg.consumeECDHErrors()
 
 	return reg
@@ -126,11 +131,28 @@ func (r *registry) registerReadyPairs(peerIDs []string) {
 		r.readyMap[peerID] = true
 	}
 
-	if len(peerIDs) == len(r.peerNodeIDs) && !r.ready {
+	// Check if we should update ready state
+	r.checkAndUpdateReadyState()
+}
+
+// checkAndUpdateReadyState checks if all conditions are met to mark the registry as ready
+func (r *registry) checkAndUpdateReadyState() {
+	// Count ready peers in readyMap
+	readyPeersCount := 0
+	for _, isReady := range r.readyMap {
+		if isReady {
+			readyPeersCount++
+		}
+	}
+
+	// Only mark as ready when both conditions are met:
+	// 1. All peers are registered in Consul
+	// 2. ECDH key exchange is complete
+	if readyPeersCount == len(r.peerNodeIDs) && r.isECDHReady() && !r.ready {
 		r.mu.Lock()
 		r.ready = true
 		r.mu.Unlock()
-		logger.Info("All peers are ready including ECDH exchange completion")
+		logger.Info("[READY] All peers are ready including ECDH exchange completion")
 	}
 }