Return error when cluster is not ready when sending generate key event

Author: anhthii

cmd/mpcium/main.go

@@ -193,7 +193,7 @@ func runNode(ctx context.Context, c *cli.Command) error {
 
 	timeoutConsumer.Run()
 	defer timeoutConsumer.Close()
-	keygenConsumer := eventconsumer.NewKeygenConsumer(natsConn, keygenBroker, pubsub, peerRegistry)
+	keygenConsumer := eventconsumer.NewKeygenConsumer(natsConn, keygenBroker, pubsub, peerRegistry, genKeyResultQueue)
 	signingConsumer := eventconsumer.NewSigningConsumer(natsConn, signingBroker, pubsub, peerRegistry, singingResultQueue)
 
 	// Make the node ready before starting the signing consumer

pkg/event/types.go

@@ -92,6 +92,7 @@ const (
 	ErrorCodeContextCancelled ErrorCode = "ERROR_CONTEXT_CANCELLED"
 	ErrorCodeOperationAborted ErrorCode = "ERROR_OPERATION_ABORTED"
 	ErrorCodeNotMajority      ErrorCode = "ERROR_NOT_MAJORITY"
+	ErrorCodeClusterNotReady  ErrorCode = "ERROR_CLUSTER_NOT_READY"
 )
 
 // GetErrorCodeFromError attempts to categorize a generic error into a specific error code

pkg/eventconsumer/keygen_consumer.go

@@ -2,13 +2,16 @@ package eventconsumer
 
 import (
 	"context"
+	"encoding/json"
+	"errors"
 	"fmt"
 	"time"
 
 	"github.com/fystack/mpcium/pkg/event"
 	"github.com/fystack/mpcium/pkg/logger"
 	"github.com/fystack/mpcium/pkg/messaging"
 	"github.com/fystack/mpcium/pkg/mpc"
+	"github.com/fystack/mpcium/pkg/types"
 	"github.com/google/uuid"
 	"github.com/nats-io/nats.go"
 	"github.com/nats-io/nats.go/jetstream"
@@ -31,22 +34,30 @@ type KeygenConsumer interface {
 
 // keygenConsumer implements KeygenConsumer.
 type keygenConsumer struct {
-	natsConn     *nats.Conn
-	pubsub       messaging.PubSub
-	jsBroker     messaging.MessageBroker
-	peerRegistry mpc.PeerRegistry
+	natsConn          *nats.Conn
+	pubsub            messaging.PubSub
+	jsBroker          messaging.MessageBroker
+	peerRegistry      mpc.PeerRegistry
+	keygenResultQueue messaging.MessageQueue
 
 	// jsSub holds the JetStream subscription, so it can be cleaned up during Close().
 	jsSub messaging.MessageSubscription
 }
 
 // NewKeygenConsumer returns a new instance of KeygenConsumer.
-func NewKeygenConsumer(natsConn *nats.Conn, jsBroker messaging.MessageBroker, pubsub messaging.PubSub, peerRegistry mpc.PeerRegistry) KeygenConsumer {
+func NewKeygenConsumer(
+	natsConn *nats.Conn,
+	jsBroker messaging.MessageBroker,
+	pubsub messaging.PubSub,
+	peerRegistry mpc.PeerRegistry,
+	keygenResultQueue messaging.MessageQueue,
+) KeygenConsumer {
 	return &keygenConsumer{
-		natsConn:     natsConn,
-		pubsub:       pubsub,
-		jsBroker:     jsBroker,
-		peerRegistry: peerRegistry,
+		natsConn:          natsConn,
+		pubsub:            pubsub,
+		jsBroker:          jsBroker,
+		peerRegistry:      peerRegistry,
+		keygenResultQueue: keygenResultQueue,
 	}
 }
 
@@ -110,9 +121,21 @@ func (sc *keygenConsumer) Run(ctx context.Context) error {
 }
 
 func (sc *keygenConsumer) handleKeygenEvent(msg jetstream.Msg) {
+	raw := msg.Data()
+	var keygenMsg types.GenerateKeyMessage
+	sessionID := msg.Headers().Get("SessionID")
+
+	err := json.Unmarshal(raw, &keygenMsg)
+	if err != nil {
+		logger.Error("SigningConsumer: Failed to unmarshal keygen message", err)
+		sc.handleKeygenError(keygenMsg, event.ErrorCodeUnmarshalFailure, err, sessionID)
+		_ = msg.Nak()
+		return
+	}
 
 	if !sc.peerRegistry.ArePeersReady() {
-		logger.Warn("KeygenConsumer: Not all peers are ready to sign, skipping message processing")
+		logger.Warn("KeygenConsumer: Not all peers are ready to gen key, skipping message processing")
+		sc.handleKeygenError(keygenMsg, event.ErrorCodeClusterNotReady, errors.New("not all peers are ready"), sessionID)
 		return
 	}
 
@@ -167,6 +190,33 @@ func (sc *keygenConsumer) handleKeygenEvent(msg jetstream.Msg) {
 	_ = msg.Nak()
 }
 
+func (sc *keygenConsumer) handleKeygenError(keygenMsg types.GenerateKeyMessage, errorCode event.ErrorCode, err error, sessionID string) {
+	keygenResult := event.KeygenResultEvent{
+		ResultType:  event.ResultTypeError,
+		ErrorCode:   string(errorCode),
+		WalletID:    keygenMsg.WalletID,
+		ErrorReason: err.Error(),
+	}
+
+	keygenResultBytes, err := json.Marshal(keygenResult)
+	if err != nil {
+		logger.Error("Failed to marshal keygen result event", err,
+			"walletID", keygenResult.WalletID,
+		)
+		return
+	}
+
+	topic := fmt.Sprintf(mpc.TypeGenerateWalletResultFmt, keygenResult.WalletID)
+	err = sc.keygenResultQueue.Enqueue(topic, keygenResultBytes, &messaging.EnqueueOptions{
+		IdempotententKey: buildIdempotentKey(keygenMsg.WalletID, sessionID, mpc.TypeGenerateWalletResultFmt),
+	})
+	if err != nil {
+		logger.Error("Failed to enqueue keygen result event", err,
+			"walletID", keygenMsg.WalletID,
+		)
+	}
+}
+
 // Close unsubscribes from the JetStream subject and cleans up resources.
 func (sc *keygenConsumer) Close() error {
 	if sc.jsSub != nil {

pkg/eventconsumer/sign_consumer.go

@@ -230,7 +230,7 @@ func (sc *signingConsumer) handleSigningError(signMsg types.SignTxMessage, error
 	}
 
 	err = sc.signingResultQueue.Enqueue(event.SigningResultCompleteTopic, signingResultBytes, &messaging.EnqueueOptions{
-		IdempotententKey: buildSigningIdempotentKey(signMsg.TxID, sessionID, mpc.TypeSigningResultFmt),
+		IdempotententKey: buildIdempotentKey(signMsg.TxID, sessionID, mpc.TypeSigningResultFmt),
 	})
 	if err != nil {
 		logger.Error("Failed to enqueue signing result event", err,
@@ -252,7 +252,7 @@ func (sc *signingConsumer) Close() error {
 	return nil
 }
 
-func buildSigningIdempotentKey(baseID string, sessionID string, formatTemplate string) string {
+func buildIdempotentKey(baseID string, sessionID string, formatTemplate string) string {
 	var uniqueKey string
 	if sessionID != "" {
 		uniqueKey = fmt.Sprintf("%s:%s", baseID, sessionID)

pkg/mpc/registry.go

@@ -68,8 +68,8 @@ func NewRegistry(
 ) *registry {
 	ecdhSession := NewECDHSession(nodeID, peerNodeIDs, pubSub, identityStore)
 	mpcThreshold := viper.GetInt("mpc_threshold")
-	if mpcThreshold <= 2 {
-		logger.Fatal("mpc_threshold must be greater than 2", nil)
+	if mpcThreshold < 1 {
+		logger.Fatal("mpc_threshold must be greater than 0", nil)
 	}
 
 	return &registry{