feat (WIP): implement presign pool worker and enhance presign handling with new data structures

Author: vietddude

cmd/mpcium/main.go

@@ -11,6 +11,7 @@ import (
 	"syscall"
 	"time"
 
+	"github.com/fystack/mpcium/pkg/client"
 	"github.com/fystack/mpcium/pkg/config"
 	"github.com/fystack/mpcium/pkg/constant"
 	"github.com/fystack/mpcium/pkg/event"
@@ -22,7 +23,10 @@ import (
 	"github.com/fystack/mpcium/pkg/logger"
 	"github.com/fystack/mpcium/pkg/messaging"
 	"github.com/fystack/mpcium/pkg/mpc"
+	"github.com/fystack/mpcium/pkg/presign"
+	"github.com/fystack/mpcium/pkg/presigninfo"
 	"github.com/fystack/mpcium/pkg/security"
+	"github.com/fystack/mpcium/pkg/types"
 	"github.com/hashicorp/consul/api"
 	"github.com/nats-io/nats.go"
 	"github.com/spf13/viper"
@@ -77,6 +81,11 @@ func main() {
 						Aliases: []string{"k"},
 						Usage:   "Path to file containing password for decrypting .age encrypted node private key",
 					},
+					&cli.BoolFlag{
+						Name:  "presign-pool-worker",
+						Usage: "Enable presign pool worker",
+						Value: false,
+					},
 					&cli.BoolFlag{
 						Name:  "debug",
 						Usage: "Enable debug logging",
@@ -109,6 +118,7 @@ func runNode(ctx context.Context, c *cli.Command) error {
 	usePrompts := c.Bool("prompt-credentials")
 	passwordFile := c.String("password-file")
 	agePasswordFile := c.String("identity-password-file")
+	presignPoolWorker := c.Bool("presign-pool-worker")
 	debug := c.Bool("debug")
 
 	viper.SetDefault("backup_enabled", true)
@@ -193,6 +203,7 @@ func runNode(ctx context.Context, c *cli.Command) error {
 
 	peerNodeIDs := GetPeerIDs(peers)
 	peerRegistry := mpc.NewRegistry(nodeID, peerNodeIDs, consulClient.KV(), directMessaging, pubsub, identityStore)
+	presignInfoStore := presigninfo.NewStore(consulClient.KV())
 
 	mpcNode := mpc.NewNode(
 		nodeID,
@@ -201,6 +212,7 @@ func runNode(ctx context.Context, c *cli.Command) error {
 		directMessaging,
 		badgerKV,
 		keyinfoStore,
+		presignInfoStore,
 		peerRegistry,
 		identityStore,
 	)
@@ -293,14 +305,44 @@ func runNode(ctx context.Context, c *cli.Command) error {
 		logger.Info("All consumers have finished")
 		close(errChan)
 	}()
+
+	// Start presign pool worker before entering the blocking error loop
+	if presignPoolWorker {
+		presignPoolCtx, presignPoolCancel := context.WithCancel(appContext)
+		defer presignPoolCancel()
+		localSigner, err := client.NewLocalSigner(types.EventInitiatorKeyTypeEd25519, client.LocalSignerOptions{
+			KeyPath: "./event_initiator.key",
+		})
+		if err != nil {
+			logger.Fatal("Failed to create local signer", err)
+		}
+		mpcClient := client.NewMPCClient(client.Options{
+			NatsConn: natsConn,
+			Signer:   localSigner,
+		})
+		presignPool := presign.NewPresignPool(nil, mpcClient, presignInfoStore)
+
+		_, err = pubsub.Subscribe(eventconsumer.MPCHotWalletEvent, func(nm *nats.Msg) {
+			walletID := string(nm.Data)
+			if walletID != "" {
+				presignPool.TouchHot(walletID)
+			}
+		})
+		if err != nil {
+			logger.Warn("Failed to subscribe to hot wallet events", "err", err)
+		}
+
+		presignPool.Start(presignPoolCtx)
+		defer presignPool.Stop()
+	}
+
 	for err := range errChan {
 		if err != nil {
 			logger.Error("Consumer error received", err)
 			cancel()
 			return err
 		}
 	}
-
 	return nil
 }
 

examples/generate/main.go

@@ -134,7 +134,11 @@ func main() {
 	for _, walletID := range walletIDs {
 		wg.Add(1) // Add to WaitGroup BEFORE attempting to create wallet
 
-		if err := mpcClient.CreateWallet(walletID); err != nil {
+		if err := mpcClient.CreateWallet(&types.GenerateKeyMessage{
+			WalletID:      walletID,
+			ECDSAProtocol: types.ProtocolCGGMP21,
+			EdDSAProtocol: types.ProtocolGG18,
+		}); err != nil {
 			logger.Error("CreateWallet failed", err)
 			walletStartTimes.Delete(walletID)
 			// Mark this wallet as processed to prevent callback from processing it

examples/sign/main.go

@@ -71,8 +71,8 @@ func main() {
 
 	txMsg := &types.SignTxMessage{
 		KeyType:             types.KeyTypeSecp256k1,
-		Protocol:            types.ProtocolFROST,
-		WalletID:            "6d553e80-a1dc-4894-9eaf-b81e3fe0c94a", // Use the generated wallet ID
+		Protocol:            types.ProtocolCGGMP21,
+		WalletID:            "7ae6ae1c-7663-4dc4-b982-33fb0a3602c3", // Use the generated wallet ID
 		NetworkInternalCode: "solana-devnet",
 		TxID:                txID,
 		Tx:                  dummyTx,

pkg/eventconsumer/event_consumer.go

@@ -26,6 +26,9 @@ const (
 	MPCReshareEvent  = "mpc:reshare"
 	MPCPresignEvent  = "mpc:presign"
 
+	// Internal event to notify presign pool of a hot wallet
+	MPCHotWalletEvent = "mpc:wallet_hot"
+
 	DefaultConcurrentKeygen   = 2
 	DefaultConcurrentSigning  = 20
 	DefaultSessionWarmUpDelay = 200
@@ -66,6 +69,12 @@ type eventConsumer struct {
 	cleanupInterval time.Duration // How often to run cleanup
 	sessionTimeout  time.Duration // How long before a session is considered stale
 	cleanupStopChan chan struct{} // Signal to stop cleanup goroutine
+
+	// Track recent signing activity to detect hot wallets
+	hotMu        sync.Mutex
+	recentSigns  map[string][]time.Time // key: walletID|keyType|protocol → timestamps within window
+	hotWindow    time.Duration          // window for counting signs (e.g., 5 minutes)
+	hotThreshold int                    // signs needed to mark as hot
 }
 
 func NewEventConsumer(
@@ -120,6 +129,9 @@ func NewEventConsumer(
 		keygenMsgBuffer:      make(chan *nats.Msg, 100),
 		signingMsgBuffer:     make(chan *nats.Msg, 200), // Larger buffer for signing
 		sessionWarmUpDelayMs: sessionWarmUpDelayMs,
+		recentSigns:          make(map[string][]time.Time),
+		hotWindow:            5 * time.Minute,
+		hotThreshold:         2,
 	}
 
 	go ec.startKeyGenEventWorker()
@@ -393,6 +405,9 @@ func (ec *eventConsumer) handleSigningEvent(natMsg *nats.Msg) {
 		ec.node.ID(),
 	)
 
+	// Track activity to detect hot wallets
+	ec.trackAndMaybeNotifyHot(msg)
+
 	// Check for duplicate session and track if new
 	if ec.checkDuplicateSession(msg.WalletID, msg.TxID) {
 		duplicateErr := fmt.Errorf(
@@ -658,7 +673,7 @@ func (ec *eventConsumer) consumePresignEvent() error {
 		}
 
 		if success {
-			ec.handlePresignSessionSuccess(msg.WalletID, natMsg)
+			ec.handlePresignSessionSuccess(msg.WalletID, msg.TxID, natMsg)
 		} else {
 			ec.handlePresignSessionError(msg.WalletID,
 				fmt.Errorf("presign operation returned false"),
@@ -676,10 +691,11 @@ func (ec *eventConsumer) consumePresignEvent() error {
 }
 
 // handlePresignSessionSuccess handles successful presign operations
-func (ec *eventConsumer) handlePresignSessionSuccess(walletID string, natMsg *nats.Msg) {
+func (ec *eventConsumer) handlePresignSessionSuccess(walletID string, txID string, natMsg *nats.Msg) {
 	presignResult := event.PresignResultEvent{
 		ResultType: event.ResultTypeSuccess,
 		WalletID:   walletID,
+		TxID:       txID,
 		Status:     "success",
 	}
 
@@ -863,3 +879,38 @@ func composeSigningIdempotentKey(txID string, natMsg *nats.Msg) string {
 func composeReshareIdempotentKey(sessionID string, natMsg *nats.Msg) string {
 	return composeIdempotentKey(sessionID, natMsg, mpc.TypeReshareWalletResultFmt)
 }
+
+// trackAndMaybeNotifyHot records a signing event and publishes a hot wallet event
+// if at least hotThreshold signs occur within hotWindow for the same
+// (walletID, keyType, protocol) tuple.
+func (ec *eventConsumer) trackAndMaybeNotifyHot(msg types.SignTxMessage) {
+	if msg.Protocol != types.ProtocolCGGMP21 {
+		return
+	}
+	key := fmt.Sprintf("%s:%s:%s", msg.WalletID, string(msg.KeyType), string(msg.Protocol))
+	now := time.Now()
+
+	ec.hotMu.Lock()
+	// prune old entries
+	list := ec.recentSigns[key]
+	pruned := list[:0]
+	cutoff := now.Add(-ec.hotWindow)
+	for _, t := range list {
+		if t.After(cutoff) {
+			pruned = append(pruned, t)
+		}
+	}
+
+	ec.recentSigns[key] = append([]time.Time(nil), pruned...)
+	currentCount := len(ec.recentSigns[key])
+
+	// If this push reaches the threshold, publish hot wallet once
+	shouldPublish := currentCount+1 == ec.hotThreshold
+	ec.recentSigns[key] = append(ec.recentSigns[key], now)
+	ec.hotMu.Unlock()
+
+	if shouldPublish {
+		_ = ec.pubsub.Publish(MPCHotWalletEvent, []byte(msg.WalletID))
+		logger.Info("Published hot wallet event", "walletID", msg.WalletID)
+	}
+}

pkg/mpc/node.go

@@ -15,6 +15,7 @@ import (
 	"github.com/fystack/mpcium/pkg/logger"
 	"github.com/fystack/mpcium/pkg/messaging"
 	"github.com/fystack/mpcium/pkg/mpc/taurus"
+	"github.com/fystack/mpcium/pkg/presigninfo"
 	"github.com/fystack/mpcium/pkg/types"
 	"github.com/taurusgroup/multi-party-sig/pkg/party"
 )
@@ -34,13 +35,13 @@ type Node struct {
 	nodeID  string
 	peerIDs []string
 
-	pubSub         messaging.PubSub
-	direct         messaging.DirectMessaging
-	kvstore        kvstore.KVStore
-	keyinfoStore   keyinfo.Store
-	ecdsaPreParams []*keygen.LocalPreParams
-	identityStore  identity.Store
-	presignCache   *taurus.PresignCache
+	pubSub           messaging.PubSub
+	direct           messaging.DirectMessaging
+	kvstore          kvstore.KVStore
+	keyinfoStore     keyinfo.Store
+	presignInfoStore presigninfo.Store
+	ecdsaPreParams   []*keygen.LocalPreParams
+	identityStore    identity.Store
 
 	peerRegistry PeerRegistry
 }
@@ -52,6 +53,7 @@ func NewNode(
 	direct messaging.DirectMessaging,
 	kvstore kvstore.KVStore,
 	keyinfoStore keyinfo.Store,
+	presignInfoStore presigninfo.Store,
 	peerRegistry PeerRegistry,
 	identityStore identity.Store,
 ) *Node {
@@ -60,15 +62,15 @@ func NewNode(
 	logger.Info("Starting new node, preparams is generated successfully!", "elapsed", elapsed.Milliseconds())
 
 	node := &Node{
-		nodeID:        nodeID,
-		peerIDs:       peerIDs,
-		pubSub:        pubSub,
-		direct:        direct,
-		kvstore:       kvstore,
-		keyinfoStore:  keyinfoStore,
-		peerRegistry:  peerRegistry,
-		identityStore: identityStore,
-		presignCache:  taurus.NewPresignCache(10 * time.Minute),
+		nodeID:           nodeID,
+		peerIDs:          peerIDs,
+		pubSub:           pubSub,
+		direct:           direct,
+		kvstore:          kvstore,
+		keyinfoStore:     keyinfoStore,
+		presignInfoStore: presignInfoStore,
+		peerRegistry:     peerRegistry,
+		identityStore:    identityStore,
 	}
 	node.ecdsaPreParams = node.generatePreParams()
 
@@ -157,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, p.presignCache, tr, p.kvstore, p.keyinfoStore)
+		session = taurus.NewCGGMP21Session(walletID, selfPartyID, allPartyIDs, threshold, nil, 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)