feat: cggmp21 presign pool with auto cleanup

Author: vietddude

INSTALLATION.md

@@ -168,6 +168,8 @@ Update `config.yaml`:
 event_initiator_pubkey: "09be5d070816aadaa1b6638cad33e819a8aed7101626f6bf1e0b427412c3408a"
 ```
 
+> 💡 **Note**: If you plan to use the presign pool worker (see [Presign Pool Worker](#presign-pool-worker) section), you'll need the `event_initiator.key` file (or `event_initiator.key.age` if encrypted) to be available. The private key file is generated alongside the identity file.
+
 ---
 
 ## Configure Node Identities
@@ -274,6 +276,49 @@ mpcium start -n node2
 
 ---
 
+## Presign Pool Worker
+
+The presign pool worker automatically maintains a pool of presignatures for hot wallets, ensuring they're ready for immediate use.
+
+### Setup
+
+To enable the presign pool worker on a node:
+
+1. **Copy the event initiator private key** to the node directory:
+   
+   If you generated the initiator with encryption:
+   ```bash
+   # Decrypt the key first
+   age --decrypt -o event_initiator.key event_initiator.key.age
+   ```
+   
+   Then copy it to the node directory:
+   ```bash
+   cp event_initiator.key node0/
+   ```
+
+   If you generated the initiator without encryption:
+   ```bash
+   cp event_initiator.key node0/
+   ```
+
+2. **Start the node with the `--presign-pool-worker` flag**:
+   
+   ```bash
+   cd node0
+   mpcium start -n node0 --presign-pool-worker
+   ```
+
+> ⚠️ **Important**: Only one node in the cluster should run the presign pool worker. The node running this worker must have the `event_initiator.key` file in its working directory.
+
+### How It Works
+
+- The worker monitors hot wallet activity and automatically generates presignatures when needed
+- It maintains a pool of presignatures between `MinPoolSize` (default: 5) and `MaxPoolSize` (default: 20)
+- The worker subscribes to hot wallet events and proactively refills the presignature pool
+
+---
+
 ## Production Deployment (High Security)
 
 1. Use production-grade **NATS** and **Consul** clusters.

pkg/eventconsumer/event_consumer.go

@@ -663,7 +663,7 @@ func (ec *eventConsumer) consumePresignEvent() error {
 		}
 
 		ctx := context.Background()
-		success, err := session.Presign(ctx, msg.WalletID)
+		success, err := session.Presign(ctx, msg.TxID)
 		if err != nil {
 			ec.handlePresignSessionError(msg.WalletID,
 				err, "Presign operation failed",
@@ -708,7 +708,7 @@ func (ec *eventConsumer) handlePresignSessionSuccess(walletID string, txID strin
 	}
 
 	err = ec.presignResultQueue.Enqueue(event.PresignResultTopic, presignResultBytes, &messaging.EnqueueOptions{
-		IdempotententKey: composePresignIdempotentKey(walletID, natMsg),
+		IdempotententKey: composePresignIdempotentKey(txID, natMsg),
 	})
 	if err != nil {
 		logger.Error("Failed to enqueue presign result event", err,
@@ -717,7 +717,7 @@ func (ec *eventConsumer) handlePresignSessionSuccess(walletID string, txID strin
 		)
 	}
 	// Presign events don't use reply inboxes, so no need to send reply
-	logger.Info("[COMPLETED PRESIGN] Presign completed successfully", "walletID", walletID)
+	logger.Info("[COMPLETED PRESIGN] Presign completed successfully", "walletID", walletID, "txID", txID)
 }
 
 // handlePresignSessionError handles errors that occur during presign operations

pkg/mpc/node.go

@@ -159,7 +159,7 @@ func (p *Node) CreateTaurusSession(
 	switch protocol {
 	case types.ProtocolCGGMP21:
 		tr := taurus.NewNATSTransport(walletID, selfPartyID, act, taurus.CGGMP21, p.pubSub, p.direct, p.identityStore)
-		session = taurus.NewCGGMP21Session(walletID, selfPartyID, allPartyIDs, threshold, nil, tr, p.kvstore, p.keyinfoStore)
+		session = taurus.NewCGGMP21Session(walletID, selfPartyID, allPartyIDs, threshold, p.presignInfoStore, tr, p.kvstore, p.keyinfoStore)
 	case types.ProtocolTaproot:
 		tr := taurus.NewNATSTransport(walletID, selfPartyID, act, taurus.FROSTTaproot, p.pubSub, p.direct, p.identityStore)
 		session = taurus.NewTaprootSession(walletID, selfPartyID, allPartyIDs, threshold, tr, p.kvstore, p.keyinfoStore)

pkg/mpc/taurus/cggmp21.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	cryptoEcdsa "crypto/ecdsa"
 	"crypto/sha256"
+	"encoding/binary"
 	"errors"
 	"fmt"
 	"math/big"
@@ -145,40 +146,26 @@ func (p *CGGMP21Session) Sign(ctx context.Context, msg *big.Int) ([]byte, error)
 	if p.savedData == nil {
 		return nil, errors.New("no key loaded")
 	}
+	logger.Info("Starting CGGMP21 sign", "walletID", p.sessionID)
 
-	logger.Info("starting CGGMP21 sign", "walletID", p.sessionID)
 	msgHash := msg.Bytes()
-
 	var (
-		result any
-		err    error
+		sigResult any
+		err       error
 	)
 
-	// Try deterministic presign selection across nodes
+	// Try presign path if store available
 	if p.presignInfoStore != nil {
-		presig, txID := p.selectAndLoadPresign(msgHash)
-		if presig != nil && txID != "" {
-			logger.Info("using presign for signing", "walletID", p.sessionID, "txID", txID)
-			result, err = p.run(ctx, cmp.PresignOnline(p.savedData, presig, msgHash, p.workerPool))
-			if err != nil {
-				return nil, fmt.Errorf("presign online failed: %w", err)
-			}
-			// Mark used (best-effort)
-			_ = p.markPresignUsed(txID)
-		} else {
-			result, err = p.run(ctx, cmp.Sign(p.savedData, p.peerIDs, msgHash, p.workerPool))
-			if err != nil {
-				return nil, fmt.Errorf("full sign failed: %w", err)
-			}
-		}
+		sigResult, err = p.signWithPresign(ctx, msgHash)
 	} else {
-		result, err = p.run(ctx, cmp.Sign(p.savedData, p.peerIDs, msgHash, p.workerPool))
-		if err != nil {
-			return nil, fmt.Errorf("full sign failed: %w", err)
-		}
+		sigResult, err = p.signFull(ctx, msgHash)
+	}
+	if err != nil {
+		return nil, err
 	}
 
-	sig, ok := result.(*ecdsa.Signature)
+	// Cast and verify
+	sig, ok := sigResult.(*ecdsa.Signature)
 	if !ok {
 		return nil, errors.New("unexpected result type")
 	}
@@ -274,46 +261,93 @@ func (p *CGGMP21Session) composePresignKey(sid, txID string) string {
 	return fmt.Sprintf("cggmp21:%s:%s", sid, txID)
 }
 
-// selectAndLoadPresign deterministically chooses a presign txID and loads its PreSignature from KV.
-// Selection: sort by CreatedAt asc, TxID asc; pick index = hash(msgHash) mod len(list).
-func (p *CGGMP21Session) selectAndLoadPresign(msgHash []byte) (*ecdsa.PreSignature, string) {
+// signWithPresign tries to select and use an existing presign for signing.
+// If no valid presign exists, falls back to full sign.
+func (p *CGGMP21Session) signWithPresign(ctx context.Context, msgHash []byte) (any, error) {
+	presig, txID := p.selectAndLoadPresign()
+	if presig == nil || txID == "" {
+		logger.Debug("No presign found, fallback to full sign", "walletID", p.sessionID)
+		return p.signFull(ctx, msgHash)
+	}
+
+	logger.Info("Using presign for signing", "walletID", p.sessionID, "txID", txID)
+	result, err := p.run(ctx, cmp.PresignOnline(p.savedData, presig, msgHash, p.workerPool))
+	if err != nil {
+		return nil, fmt.Errorf("presign online failed: %w", err)
+	}
+
+	// Mark and cleanup in background (best effort)
+	go func() {
+		if err := p.markPresignUsed(txID); err != nil {
+			logger.Warn("mark presign used failed", "walletID", p.sessionID, "txID", txID, "err", err)
+		}
+		if err := p.deletePresign(txID); err != nil {
+			logger.Warn("delete presign failed", "walletID", p.sessionID, "txID", txID, "err", err)
+		}
+	}()
+
+	return result, nil
+}
+
+// signFull executes a full CGGMP21 signing round.
+func (p *CGGMP21Session) signFull(ctx context.Context, msgHash []byte) (any, error) {
+	logger.Info("Executing full CGGMP21 signing", "walletID", p.sessionID)
+	result, err := p.run(ctx, cmp.Sign(p.savedData, p.peerIDs, msgHash, p.workerPool))
+	if err != nil {
+		return nil, fmt.Errorf("full sign failed: %w", err)
+	}
+	return result, nil
+}
+
+func (p *CGGMP21Session) selectAndLoadPresign() (*ecdsa.PreSignature, string) {
 	infos, err := p.presignInfoStore.ListPendingPresigns(p.sessionID)
 	if err != nil || len(infos) == 0 {
 		return nil, ""
 	}
-	// filter by active + protocol/keytype
-	filtered := make([]*presigninfo.PresignInfo, 0, len(infos))
+
+	// Filter usable presigns
+	var filtered []*presigninfo.PresignInfo
 	for _, inf := range infos {
-		if inf.Status == presigninfo.PresignStatusActive && inf.Protocol == types.ProtocolCGGMP21 && inf.KeyType == types.KeyTypeSecp256k1 {
+		if inf.Status == presigninfo.PresignStatusActive &&
+			inf.Protocol == types.ProtocolCGGMP21 &&
+			inf.KeyType == types.KeyTypeSecp256k1 {
 			filtered = append(filtered, inf)
 		}
 	}
 	if len(filtered) == 0 {
 		return nil, ""
 	}
-	// sort
+
+	// Sort + pick deterministically
 	sort.Slice(filtered, func(i, j int) bool {
 		if filtered[i].CreatedAt.Equal(filtered[j].CreatedAt) {
 			return filtered[i].TxID < filtered[j].TxID
 		}
 		return filtered[i].CreatedAt.Before(filtered[j].CreatedAt)
 	})
-	// pick index via hash
-	h := sha256.Sum256(msgHash)
-	idx := int(h[0]) % len(filtered)
+	h := sha256.Sum256([]byte(p.sessionID))
+	idx := int(binary.BigEndian.Uint64(h[:8]) % uint64(len(filtered)))
 	chosen := filtered[idx]
-	// load material
-	bytes, err := p.kvstore.Get(p.composePresignKey(p.sessionID, chosen.TxID))
-	if err != nil || len(bytes) == 0 {
+
+	// Load presign from KV
+	key := p.composePresignKey(p.sessionID, chosen.TxID)
+	data, err := p.kvstore.Get(key)
+	if err != nil || len(data) == 0 {
+		logger.Warn("presign missing", "walletID", p.sessionID, "txID", chosen.TxID, "err", err)
 		return nil, ""
 	}
-	presig := new(ecdsa.PreSignature)
-	if err := cbor.Unmarshal(bytes, presig); err != nil {
+
+	presig := ecdsa.EmptyPreSignature(curve.Secp256k1{})
+	if err := cbor.Unmarshal(data, presig); err != nil {
+		logger.Warn("unmarshal presign failed", "walletID", p.sessionID, "txID", chosen.TxID, "err", err)
 		return nil, ""
 	}
 	if err := presig.Validate(); err != nil {
+		logger.Warn("presign invalid", "walletID", p.sessionID, "txID", chosen.TxID, "err", err)
 		return nil, ""
 	}
+
+	logger.Debug("Presign chosen", "walletID", p.sessionID, "txID", chosen.TxID, "idx", idx)
 	return presig, chosen.TxID
 }
 
@@ -327,3 +361,7 @@ func (p *CGGMP21Session) markPresignUsed(txID string) error {
 	info.UsedAt = &now
 	return p.presignInfoStore.Save(p.sessionID, info)
 }
+
+func (p *CGGMP21Session) deletePresign(txID string) error {
+	return p.kvstore.Delete(p.composePresignKey(p.sessionID, txID))
+}