fix rejoining ecdh bug and add ecdh cache key cleaning

Author: nann-cheng

pkg/identity/identity.go

@@ -239,7 +239,17 @@ func (s *fileStore) GetSymmetricKey(peerID string) ([]byte, error) {
 }
 
 func (s *fileStore) CheckSymmetricKeyComplete(desired int) bool {
-	return len(s.symmetricKeys) == desired
+	s.mu.RLock()
+	defer s.mu.RUnlock()
+
+	completeCount := 0
+	for _, value := range s.symmetricKeys {
+		if len(value) > 0 {
+			completeCount++
+		}
+	}
+
+	return completeCount == desired
 }
 
 // GetPublicKey retrieves a node's public key by its ID

pkg/mpc/key_exchange_session.go

@@ -28,9 +28,10 @@ const (
 )
 
 type ECDHSession interface {
-	StartKeyExchange() error
+	ListenKeyExchange() error
 	BroadcastPublicKey() error
 	WaitForExchangeComplete() error
+	ResetLocalKeys()
 	ErrChan() <-chan error
 	Close() error
 }
@@ -65,7 +66,14 @@ func NewECDHSession(
 	}
 }
 
-func (e *ecdhSession) StartKeyExchange() error {
+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) ListenKeyExchange() error {
 	// Generate an ephemeral ECDH key pair
 	privateKey, err := ecdh.X25519().GenerateKey(rand.Reader)
 	if err != nil {
@@ -85,7 +93,7 @@ func (e *ecdhSession) StartKeyExchange() error {
 		if ecdhMsg.From == e.nodeID {
 			return
 		}
-		logger.Info("Received ECDH message from", "node", ecdhMsg.From)
+
 		//TODO: consider how to avoid replay attack
 		if err := e.identityStore.VerifySignature(&ecdhMsg); err != nil {
 			e.errCh <- err
@@ -115,10 +123,11 @@ func (e *ecdhSession) StartKeyExchange() error {
 			logger.Info("ALL PEERS ARE READY! Starting to accept MPC requests")
 
 			e.mu.Lock()
-			e.exchangeDone = true
+			if !e.exchangeDone {
+				e.exchangeDone = true
+				e.exchangeComplete <- struct{}{}
+			}
 			e.mu.Unlock()
-
-			e.exchangeComplete <- struct{}{}
 		}
 	})
 

pkg/mpc/node.go

@@ -57,8 +57,8 @@ func NewNode(
 	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.StartKeyExchange(); err != nil {
-		logger.Fatal("Failed to start DH key exchange", err)
+	if err := dhSession.ListenKeyExchange(); err != nil {
+		logger.Fatal("Failed to listen to DH key exchange", err)
 	}
 
 	node := &Node{
@@ -74,13 +74,18 @@ func NewNode(
 	}
 	node.ecdsaPreParams = node.generatePreParams()
 
-	ecdhTask := func() {
-		if err := dhSession.BroadcastPublicKey(); err != nil {
-			logger.Fatal("DH key broadcast failed", err)
+	// 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(ecdhTask)
+	go peerRegistry.WatchPeersReady(ecdhTasks)
 	return node
 }
 
@@ -429,10 +434,6 @@ func (p *Node) GetECDHSession() ECDHSession {
 	return p.ecdhSession
 }
 
-func (p *Node) GetDHSession() ECDHSession {
-	return p.ecdhSession
-}
-
 func (p *Node) generatePreParams() []*keygen.LocalPreParams {
 	start := time.Now()
 	// Try to load from kvstore

pkg/mpc/registry.go

@@ -21,7 +21,7 @@ const (
 type PeerRegistry interface {
 	Ready() error
 	ArePeersReady() bool
-	WatchPeersReady(callback func())
+	WatchPeersReady(callback func(bool))
 	// Resign is called by the node when it is going to shutdown
 	Resign() error
 	GetReadyPeersCount() int64
@@ -71,7 +71,7 @@ func (r *registry) readyKey(nodeID string) string {
 	return fmt.Sprintf("ready/%s", nodeID)
 }
 
-func (r *registry) registerReadyPairs(peerIDs []string, callback func()) {
+func (r *registry) registerReadyPairs(peerIDs []string, callback func(isInit bool)) {
 	for _, peerID := range peerIDs {
 		ready, exist := r.readyMap[peerID]
 		if !exist {
@@ -89,9 +89,10 @@ func (r *registry) registerReadyPairs(peerIDs []string, callback func()) {
 		r.mu.Lock()
 		r.ready = true
 		r.mu.Unlock()
-		time.AfterFunc(5*time.Second, callback)
+		time.AfterFunc(5*time.Second, func() {
+			callback(true)
+		})
 	}
-
 }
 
 // Ready is called by the node when it complete generate preparams and starting to accept
@@ -122,9 +123,10 @@ func (r *registry) composeHealthCheckTopic(nodeID string) string {
 	return fmt.Sprintf("healthcheck:%s", nodeID)
 }
 
-func (r *registry) WatchPeersReady(callback func()) {
+func (r *registry) WatchPeersReady(callback func(isInit bool)) {
+	go r.checkPeersHealth()
+
 	ticker := time.NewTicker(ReadinessCheckPeriod)
-	go r.checkPeersHeath()
 	// first tick is executed immediately
 	for ; true; <-ticker.C {
 		pairs, _, err := r.consulKV.List("ready/", nil)
@@ -151,6 +153,7 @@ func (r *registry) WatchPeersReady(callback func()) {
 					logger.Warn("Peer disconnected!", "peerID", peerID)
 					r.readyMap[peerID] = false
 					atomic.AddInt64(&r.readyCount, -1)
+					callback(false)
 				}
 
 			}
@@ -161,7 +164,7 @@ func (r *registry) WatchPeersReady(callback func()) {
 
 }
 
-func (r *registry) checkPeersHeath() {
+func (r *registry) checkPeersHealth() {
 	for {
 		time.Sleep(5 * time.Second)
 		if !r.ArePeersReady() {