Add CMP reshare handling in event consumer

Author: vietddude

pkg/eventconsumer/event_consumer.go

@@ -712,6 +712,11 @@ func (ec *eventConsumer) consumeReshareEvent() error {
 			ec.handleReshareSessionError(walletID, keyType, msg.NewThreshold, err, "Failed to get session type", natMsg)
 			return
 		}
+		// Handle CMP reshare separately
+		if keyType == types.KeyTypeTaurusCmp {
+			ec.handleCMPReshare(msg, natMsg)
+			return
+		}
 
 		createSession := func(isNewPeer bool) (mpc.ReshareSession, error) {
 			return ec.node.CreateReshareSession(
@@ -851,6 +856,72 @@ func (ec *eventConsumer) consumeReshareEvent() error {
 	return err
 }
 
+// NOTE: In CMP reshare, it just refresh the keyshare of each node but keep the same public key and threshold.
+// Therefore, we don't need to create new party sessions for CMP reshare.
+func (ec *eventConsumer) handleCMPReshare(msg types.ResharingMessage, natMsg *nats.Msg) {
+	logger.Info("Starting CMP reshare", "walletID", msg.WalletID, "sessionID", msg.SessionID)
+
+	// Create CMP session for reshare
+	taurusSession, err := ec.node.CreateCMPSession(msg.WalletID, msg.NewThreshold, taurus.ActReshare)
+	if err != nil {
+		logger.Error("Failed to create Taurus CMP reshare session", err, "walletID", msg.WalletID)
+		ec.handleReshareSessionError(msg.WalletID, msg.KeyType, msg.NewThreshold, err, "Failed to create Taurus CMP reshare session", natMsg)
+		return
+	}
+
+	// Load the existing key for reshare
+	if err := taurusSession.LoadKey(msg.WalletID); err != nil {
+		logger.Error("Failed to load key for CMP reshare", err, "walletID", msg.WalletID)
+		ec.handleReshareSessionError(msg.WalletID, msg.KeyType, msg.NewThreshold, err, "Failed to load key for CMP reshare", natMsg)
+		return
+	}
+
+	// Create context for reshare
+	ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second) // Longer timeout for reshare
+	defer cancel()
+
+	// Perform CMP reshare
+	keyData, err := taurusSession.Reshare(ctx)
+	if err != nil {
+		logger.Error("CMP reshare failed", err, "walletID", msg.WalletID, "sessionID", msg.SessionID)
+		ec.handleReshareSessionError(msg.WalletID, msg.KeyType, msg.NewThreshold, err, "CMP reshare failed", natMsg)
+		return
+	}
+
+	// Create reshare result event
+	reshareResult := event.ResharingResultEvent{
+		ResultType:   event.ResultTypeSuccess,
+		WalletID:     msg.WalletID,
+		NewThreshold: keyData.Threshold,
+		KeyType:      msg.KeyType,
+		PubKey:       keyData.PubKeyBytes,
+	}
+
+	// Marshal and enqueue the result
+	reshareResultBytes, err := json.Marshal(reshareResult)
+	if err != nil {
+		logger.Error("Failed to marshal CMP reshare result event", err, "walletID", msg.WalletID, "sessionID", msg.SessionID)
+		ec.handleReshareSessionError(msg.WalletID, msg.KeyType, msg.NewThreshold, err, "Failed to marshal CMP reshare result", natMsg)
+		return
+	}
+
+	// Enqueue the reshare result
+	key := fmt.Sprintf(mpc.TypeReshareWalletResultFmt, msg.SessionID)
+	err = ec.reshareResultQueue.Enqueue(key, reshareResultBytes, &messaging.EnqueueOptions{
+		IdempotententKey: composeReshareIdempotentKey(msg.SessionID, natMsg),
+	})
+	if err != nil {
+		logger.Error("Failed to enqueue CMP reshare result event", err, "walletID", msg.WalletID, "sessionID", msg.SessionID)
+		ec.handleReshareSessionError(msg.WalletID, msg.KeyType, msg.NewThreshold, err, "Failed to enqueue CMP reshare result", natMsg)
+		return
+	}
+
+	// Remove this line - don't send reply for reshare messages
+	// ec.sendReplyToRemoveMsg(natMsg)
+
+	logger.Info("[COMPLETED CMP RESHARE] CMP reshare completed successfully", "walletID", msg.WalletID, "sessionID", msg.SessionID)
+}
+
 // handleReshareSessionError handles errors that occur during reshare operations
 func (ec *eventConsumer) handleReshareSessionError(
 	walletID string,

pkg/mpc/node.go

@@ -154,9 +154,10 @@ func (p *Node) CreateCMPSession(
 	adapter := taurus.NewTaurusNetworkAdapter(walletID, selfPartyID, tr, allPartyIDs)
 	pl := pool.NewPool(0)
 	session := taurus.NewCmpParty(walletID, selfPartyID, allPartyIDs, threshold, pl, adapter, p.keyinfoStore, p.kvstore)
-	if act == taurus.ActSign {
+	if act == taurus.ActSign || act == taurus.ActReshare {
 		session.LoadKey(walletID)
 	}
+
 	return session, nil
 }
 

pkg/mpc/taurus/cmp.go

@@ -154,21 +154,48 @@ func (p *CmpParty) Sign(ctx context.Context, msg *big.Int) ([]byte, error) {
 	return sig.SigEthereum()
 }
 
-func (p *CmpParty) Reshare(ctx context.Context) (types.KeyData, error) {
+func (p *CmpParty) Reshare(ctx context.Context) (res types.ReshareData, err error) {
 	if p.savedData == nil {
-		return types.KeyData{}, errors.New("no key loaded")
+		return res, errors.New("no key loaded")
 	}
 	cfg, err := p.run(ctx, cmp.Refresh(p.savedData, p.pl))
 	if err != nil {
-		return types.KeyData{}, err
+		return res, err
 	}
 	savedData, ok := cfg.(*cmp.Config)
 	if !ok {
-		return types.KeyData{}, errors.New("unexpected result type")
+		return res, errors.New("unexpected result type")
 	}
 	p.savedData = savedData
 	packed, _ := p.savedData.MarshalBinary()
-	return types.KeyData{SID: p.sid, Type: "taurus_cmp", Payload: packed}, nil
+
+	key := p.composeKey(p.sid)
+	// Store updated key share
+	if p.kvstore != nil {
+		if err := p.kvstore.Put(key, packed); err != nil {
+			return res, fmt.Errorf("store key: %w", err)
+		}
+	}
+
+	// Extract public key coordinates
+	x, y, err := ExtractXYFromPoint(p.savedData.PublicPoint())
+	if err != nil {
+		return res, fmt.Errorf("extract pubkey: %w", err)
+	}
+
+	// Use secp256k1 curve, not P256
+	pubKey := &cryptoEcdsa.PublicKey{
+		Curve: btcec.S256(),
+		X:     x,
+		Y:     y,
+	}
+
+	pubKeyBytes, err := encoding.EncodeS256PubKey(pubKey)
+	if err != nil {
+		return res, fmt.Errorf("encode pubkey: %w", err)
+	}
+
+	return types.ReshareData{KeyData: types.KeyData{SID: p.sid, Type: "taurus_cmp", PubKeyBytes: pubKeyBytes}, Threshold: p.threshold}, nil
 }
 
 func (p *CmpParty) run(ctx context.Context, proto protocol.StartFunc) (any, error) {

pkg/types/taurus.go

@@ -38,3 +38,8 @@ type KeyData struct {
 	Payload     []byte
 	PubKeyBytes []byte
 }
+
+type ReshareData struct {
+	KeyData
+	Threshold int
+}