kms/pkg/keys/manager.go

package keys

import (
	"context"
	"fmt"
	"log"
	"time"

	"github.com/luxfi/kms/pkg/mpc"
)

// Store is the interface the manager uses to persist key metadata.
type Store interface {
	Put(ks *ValidatorKeySet) error
	Update(ks *ValidatorKeySet) error
	Get(validatorID string) (*ValidatorKeySet, error)
	List() []*ValidatorKeySet
	Delete(validatorID string) error
}

// Signer handles MPC signing operations (M-Chain / MPC daemon).
// FROST (Ed25519) and CGGMP21 (secp256k1) threshold signatures.
type Signer interface {
	Keygen(ctx context.Context, vaultID string, req mpc.KeygenRequest) (*mpc.KeygenResult, error)
	Sign(ctx context.Context, req mpc.SignRequest) (*mpc.SignResult, error)
	Reshare(ctx context.Context, walletID string, req mpc.ReshareRequest) error
	GetWallet(ctx context.Context, walletID string) (*mpc.Wallet, error)
	Status(ctx context.Context) (*mpc.ClusterStatus, error)
}

// MPCBackend is what a manager needs from the daemon.
//
// It named an Encryptor too — Encrypt and Decrypt over "threshold FHE
// (T-Chain)". No node ever served either: one op had no handler at all and the
// other reached the ring's reveal handler with the wrong request. Nothing
// called them, and the interface advertising them is what made the real
// opening operation look like it was already there.
type MPCBackend interface {
	Signer
}

// Manager orchestrates validator key lifecycle.
// K-Chain concern: key registry, policy, metadata.
type Manager struct {
	signer  Signer
	store   Store
	vaultID string
}

// NewManager creates a key manager.
func NewManager(backend MPCBackend, store Store, vaultID string) *Manager {
	return &Manager{signer: backend, store: store, vaultID: vaultID}
}

// GenerateValidatorKeys creates a new validator key set via MPC DKG.
// It generates two MPC wallets: one for BLS (secp256k1/CGGMP21) and one for
// Corona (ed25519/FROST), then stores the mapping.
func (m *Manager) GenerateValidatorKeys(ctx context.Context, req GenerateRequest) (*ValidatorKeySet, error) {
	if req.ValidatorID == "" {
		return nil, fmt.Errorf("keys: validator_id is required")
	}
	if req.Threshold < 2 {
		return nil, fmt.Errorf("keys: threshold must be >= 2")
	}
	if req.Parties < req.Threshold {
		return nil, fmt.Errorf("keys: parties must be >= threshold")
	}

	// Check for duplicate.
	if _, err := m.store.Get(req.ValidatorID); err == nil {
		return nil, fmt.Errorf("keys: validator %s already exists", req.ValidatorID)
	}

	// Generate BLS key (secp256k1 via CGGMP21 protocol).
	blsResult, err := m.signer.Keygen(ctx, m.vaultID, mpc.KeygenRequest{
		Name:     fmt.Sprintf("validator-%s-bls", req.ValidatorID),
		KeyType:  "secp256k1",
		Protocol: "cggmp21",
	})
	if err != nil {
		return nil, fmt.Errorf("keys: bls keygen failed: %w", err)
	}

	// Generate Corona key (ed25519 via FROST protocol).
	coronaResult, err := m.signer.Keygen(ctx, m.vaultID, mpc.KeygenRequest{
		Name:     fmt.Sprintf("validator-%s-corona", req.ValidatorID),
		KeyType:  "ed25519",
		Protocol: "frost",
	})
	if err != nil {
		// BLS keygen succeeded but Corona failed. DKG cannot be rolled back —
		// the BLS wallet is now orphaned in the MPC cluster. Log for manual cleanup.
		log.Printf("keys: CRITICAL: corona keygen failed after BLS keygen succeeded; orphaned BLS wallet_id=%s for validator=%s — manual cleanup required: %v",
			blsResult.WalletID, req.ValidatorID, err)
		return nil, fmt.Errorf("keys: corona keygen failed (orphaned bls wallet %s): %w", blsResult.WalletID, err)
	}

	// Verify we got the threshold we asked for. Keygen takes no per-request
	// threshold — the MPC ring's own --threshold governs — so requesting
	// 3-of-5 and receiving something weaker is silent unless we check.
	// Fail closed: an explicit error beats recording an unverified key.
	if err := verifyThreshold("bls", req, blsResult.Threshold, blsResult.Participants); err != nil {
		return nil, err
	}
	if err := verifyThreshold("corona", req, coronaResult.Threshold, coronaResult.Participants); err != nil {
		return nil, err
	}

	blsPub := ""
	if blsResult.ECDSAPubkey != nil {
		blsPub = *blsResult.ECDSAPubkey
	}
	coronaPub := ""
	if coronaResult.EDDSAPubkey != nil {
		coronaPub = *coronaResult.EDDSAPubkey
	}

	now := time.Now().UTC()
	ks := &ValidatorKeySet{
		ValidatorID:     req.ValidatorID,
		BLSWalletID:     blsResult.WalletID,
		CoronaWalletID:  coronaResult.WalletID,
		BLSPublicKey:    blsPub,
		CoronaPublicKey: coronaPub,
		Threshold:       blsResult.Threshold,
		Parties:         len(blsResult.Participants),
		Status:          "active",
		CreatedAt:       now,
		UpdatedAt:       now,
	}

	if err := m.store.Put(ks); err != nil {
		return nil, fmt.Errorf("keys: store put: %w", err)
	}

	return ks, nil
}

// verifyThreshold checks that a freshly generated key actually has the t-of-n
// the caller asked for, and refuses the key otherwise.
//
// Why this must exist: mpc.KeygenRequest has no threshold field. The MPC ring's
// own --threshold decides the CGGMP21 polynomial degree, so GenerateValidatorKeys
// validating req.Threshold proves nothing about the key it gets back — the
// request is advisory at best. Before mpcd reported these values it emitted
// neither field, so every key set was stored as Threshold:0 Parties:0 while its
// caller had asked for 3-of-5 and passed validation.
//
// Fail closed in both directions: a mismatch is an error, and so is a ring that
// declines to state what it did. Silently trusting an unstated threshold is the
// exact failure this guards.
func verifyThreshold(keyName string, req GenerateRequest, gotThreshold int, gotParticipants []string) error {
	if gotThreshold == 0 && len(gotParticipants) == 0 {
		return fmt.Errorf(
			"keys: %s keygen did not report its threshold or party set, so the requested %d-of-%d cannot be verified; "+
				"upgrade the MPC ring to mpcd >= v1.17.15 (older builds omit both fields and this silently recorded 0-of-0)",
			keyName, req.Threshold, req.Parties)
	}
	if gotThreshold != req.Threshold {
		return fmt.Errorf("keys: %s keygen produced a %d-of-%d key but %d-of-%d was requested; refusing to record it",
			keyName, gotThreshold, len(gotParticipants), req.Threshold, req.Parties)
	}
	if len(gotParticipants) != req.Parties {
		return fmt.Errorf("keys: %s keygen produced a %d-of-%d key but %d-of-%d was requested; refusing to record it",
			keyName, gotThreshold, len(gotParticipants), req.Threshold, req.Parties)
	}
	return nil
}

// SignWithBLS signs a message using the validator's BLS key via MPC threshold signing.
func (m *Manager) SignWithBLS(ctx context.Context, validatorID string, message []byte) (*SignResponse, error) {
	ks, err := m.store.Get(validatorID)
	if err != nil {
		return nil, fmt.Errorf("keys: validator %s: %w", validatorID, err)
	}

	result, err := m.signer.Sign(ctx, mpc.SignRequest{
		VaultID:  m.vaultID,
		WalletID: ks.BLSWalletID,
		KeyType:  "secp256k1",
		Payload:  message,
	})
	if err != nil {
		return nil, fmt.Errorf("keys: bls sign: %w", err)
	}

	return &SignResponse{
		Signature: result.Signature,
		R:         result.R,
		S:         result.S,
		V:         result.V,
	}, nil
}

// SignWithCorona signs a message using the validator's Corona key via MPC threshold signing.
func (m *Manager) SignWithCorona(ctx context.Context, validatorID string, message []byte) (*SignResponse, error) {
	ks, err := m.store.Get(validatorID)
	if err != nil {
		return nil, fmt.Errorf("keys: validator %s: %w", validatorID, err)
	}

	result, err := m.signer.Sign(ctx, mpc.SignRequest{
		VaultID:  m.vaultID,
		WalletID: ks.CoronaWalletID,
		KeyType:  "ed25519",
		Payload:  message,
	})
	if err != nil {
		return nil, fmt.Errorf("keys: corona sign: %w", err)
	}

	return &SignResponse{
		Signature: result.Signature,
		R:         result.R,
		S:         result.S,
		V:         result.V,
	}, nil
}

// Rotate reshares a validator's keys with new threshold or participants.
func (m *Manager) Rotate(ctx context.Context, validatorID string, req RotateRequest) (*ValidatorKeySet, error) {
	ks, err := m.store.Get(validatorID)
	if err != nil {
		return nil, fmt.Errorf("keys: validator %s: %w", validatorID, err)
	}

	reshareReq := mpc.ReshareRequest{
		NewThreshold:    req.NewThreshold,
		NewParticipants: req.NewParticipants,
	}

	// Build rollback params from current state.
	rollbackReq := mpc.ReshareRequest{
		NewThreshold:    ks.Threshold,
		NewParticipants: nil, // same participant set
	}

	// Reshare BLS key.
	if err := m.signer.Reshare(ctx, ks.BLSWalletID, reshareReq); err != nil {
		return nil, fmt.Errorf("keys: bls reshare: %w", err)
	}

	// Reshare Corona key.
	if err := m.signer.Reshare(ctx, ks.CoronaWalletID, reshareReq); err != nil {
		// BLS was reshared but Corona failed — keys are now inconsistent.
		// Attempt to roll BLS back to previous threshold/participants.
		log.Printf("keys: WARNING: corona reshare failed after BLS reshare succeeded for validator=%s, attempting BLS rollback: %v",
			validatorID, err)
		if rbErr := m.signer.Reshare(ctx, ks.BLSWalletID, rollbackReq); rbErr != nil {
			log.Printf("keys: CRITICAL: BLS rollback also failed for validator=%s — keys are in inconsistent state, manual intervention required: %v",
				validatorID, rbErr)
			return nil, fmt.Errorf("keys: corona reshare failed AND bls rollback failed (inconsistent state): corona=%w, bls_rollback=%v", err, rbErr)
		}
		log.Printf("keys: BLS rollback succeeded for validator=%s after corona reshare failure", validatorID)
		return nil, fmt.Errorf("keys: corona reshare failed (bls rolled back): %w", err)
	}

	if req.NewThreshold > 0 {
		ks.Threshold = req.NewThreshold
	}
	if len(req.NewParticipants) > 0 {
		ks.Parties = len(req.NewParticipants)
	}
	ks.UpdatedAt = time.Now().UTC()

	if err := m.store.Update(ks); err != nil {
		return nil, fmt.Errorf("keys: store update: %w", err)
	}

	return ks, nil
}

// Get retrieves a validator key set.
func (m *Manager) Get(validatorID string) (*ValidatorKeySet, error) {
	return m.store.Get(validatorID)
}

// List returns all validator key sets.
func (m *Manager) List() []*ValidatorKeySet {
	return m.store.List()
}