feat: optional multi-authorizer checks for key generation, signing, and resharing

Author: sivo4kin

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,
@@ -197,21 +193,16 @@ func runNode(ctx context.Context, c *cli.Command) error {
 
 	timeoutConsumer.Run()
 	defer timeoutConsumer.Close()
-	keygenConsumer := eventconsumer.NewKeygenConsumer(natsConn, keygenBroker, pubsub, peerRegistry)
-	signingConsumer := eventconsumer.NewSigningConsumer(natsConn, signingBroker, 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
 	if err := peerRegistry.Ready(); err != nil {
 		logger.Error("Failed to mark peer registry as ready", err)
 	}
 	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")

config.prod.yaml.template

@@ -17,3 +17,19 @@ backup_dir: backups
 max_concurrent_keygen: 2
 max_concurrent_signing: 10
 session_warm_up_delay_ms: 100
+
+# Authorization (optional)
+authorization:
+  enabled: false
+  default_threshold: 0
+  operation_policies:
+    keygen:
+      required_authorizers: 0
+      authorizer_ids: []
+    signing:
+      required_authorizers: 0
+      authorizer_ids: []
+    reshare:
+      required_authorizers: 0
+      authorizer_ids: []
+  authorizers: {}

config.yaml.template

@@ -14,3 +14,19 @@ backup_dir: backups
 max_concurrent_keygen: 2
 max_concurrent_signing: 10
 session_warm_up_delay_ms: 100
+
+# Authorization (optional)
+authorization:
+  enabled: false
+  default_threshold: 0
+  operation_policies:
+    keygen:
+      required_authorizers: 0
+      authorizer_ids: []
+    signing:
+      required_authorizers: 0
+      authorizer_ids: []
+    reshare:
+      required_authorizers: 0
+      authorizer_ids: []
+  authorizers: {}

pkg/event/types.go

@@ -91,6 +91,8 @@ const (
 	// Context and cancellation errors
 	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/event_consumer.go

@@ -166,6 +166,13 @@ func (ec *eventConsumer) handleKeyGenEvent(natMsg *nats.Msg) {
 		return
 	}
 
+	// Optional multi-authorizer check for keygen
+	if err := ec.identityStore.AuthorizeInitiatorMessage("keygen", &msg); err != nil {
+		logger.Error("Authorization failed for keygen", err)
+		ec.handleKeygenSessionError(msg.WalletID, err, "Authorization failed for keygen", natMsg)
+		return
+	}
+
 	walletID := msg.WalletID
 	ecdsaSession, err := ec.node.CreateKeyGenSession(mpc.SessionTypeECDSA, walletID, ec.mpcThreshold, ec.genKeyResultQueue)
 	if err != nil {
@@ -353,6 +360,12 @@ func (ec *eventConsumer) handleSigningEvent(natMsg *nats.Msg) {
 		return
 	}
 
+	// Optional multi-authorizer check for signing
+	if err := ec.identityStore.AuthorizeInitiatorMessage("signing", &msg); err != nil {
+		logger.Error("Authorization failed for signing", err)
+		return
+	}
+
 	logger.Info(
 		"Received signing event",
 		"waleltID",
@@ -589,6 +602,13 @@ func (ec *eventConsumer) consumeReshareEvent() error {
 			return
 		}
 
+		// Optional multi-authorizer check for reshare
+		if err := ec.identityStore.AuthorizeInitiatorMessage("reshare", &msg); err != nil {
+			logger.Error("Authorization failed for reshare", err)
+			ec.handleReshareSessionError(msg.WalletID, msg.KeyType, msg.NewThreshold, err, "Authorization failed for reshare", natMsg)
+			return
+		}
+
 		walletID := msg.WalletID
 		keyType := msg.KeyType
 

pkg/eventconsumer/events.go

@@ -17,20 +17,30 @@ type InitiatorMessage interface {
 	Sig() []byte
 	// InitiatorID returns the ID whose public key we have to look up.
 	InitiatorID() string
+	// AuthorizerSigs returns the optional list of authorizer signatures.
+	AuthorizerSigs() []AuthorizerSignature
+}
+
+// AuthorizerSignature represents an approval signature from an external authorizer.
+type AuthorizerSignature struct {
+	AuthorizerID string `json:"authorizer_id"`
+	Signature    []byte `json:"signature"`
 }
 
 type GenerateKeyMessage struct {
-	WalletID  string `json:"wallet_id"`
-	Signature []byte `json:"signature"`
+	WalletID             string                `json:"wallet_id"`
+	Signature            []byte                `json:"signature"`
+	AuthorizerSignatures []AuthorizerSignature `json:"authorizer_signatures,omitempty"`
 }
 
 type SignTxMessage struct {
-	KeyType             KeyType `json:"key_type"`
-	WalletID            string  `json:"wallet_id"`
-	NetworkInternalCode string  `json:"network_internal_code"`
-	TxID                string  `json:"tx_id"`
-	Tx                  []byte  `json:"tx"`
-	Signature           []byte  `json:"signature"`
+	KeyType              KeyType               `json:"key_type"`
+	WalletID             string                `json:"wallet_id"`
+	NetworkInternalCode  string                `json:"network_internal_code"`
+	TxID                 string                `json:"tx_id"`
+	Tx                   []byte                `json:"tx"`
+	Signature            []byte                `json:"signature"`
+	AuthorizerSignatures []AuthorizerSignature `json:"authorizer_signatures,omitempty"`
 }
 
 func (m *SignTxMessage) Raw() ([]byte, error) {
@@ -59,6 +69,10 @@ func (m *SignTxMessage) InitiatorID() string {
 	return m.TxID
 }
 
+func (m *SignTxMessage) AuthorizerSigs() []AuthorizerSignature {
+	return m.AuthorizerSignatures
+}
+
 func (m *GenerateKeyMessage) Raw() ([]byte, error) {
 	return []byte(m.WalletID), nil
 }
@@ -70,3 +84,7 @@ func (m *GenerateKeyMessage) Sig() []byte {
 func (m *GenerateKeyMessage) InitiatorID() string {
 	return m.WalletID
 }
+
+func (m *GenerateKeyMessage) AuthorizerSigs() []AuthorizerSignature {
+	return m.AuthorizerSignatures
+}

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,
 	}
 }
 
@@ -60,6 +71,9 @@ func (sc *keygenConsumer) waitForAllPeersReadyToGenKey(ctx context.Context) erro
 	for {
 		select {
 		case <-ctx.Done():
+			if ctx.Err() == context.Canceled {
+				return nil
+			}
 			return ctx.Err()
 		case <-ticker.C:
 			allPeersReady := sc.peerRegistry.ArePeersReady()
@@ -80,6 +94,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)
 	}
 
@@ -104,9 +121,22 @@ 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.Ack()
+		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)
+		_ = msg.Ack()
 		return
 	}
 
@@ -161,6 +191,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 {