mpc/pkg/protocol/cggmp21/adapter.go

package cggmp21

import (
	"context"
	"crypto/ecdsa"
	"crypto/elliptic"
	"encoding/json"
	"errors"
	"fmt"
	"math/big"
	"sync"

	log "github.com/luxfi/log"
	mpsEcdsa "github.com/luxfi/threshold/pkg/ecdsa"
	"github.com/luxfi/threshold/pkg/math/curve"
	"github.com/luxfi/threshold/pkg/party"
	"github.com/luxfi/threshold/pkg/pool"
	mpsProtocol "github.com/luxfi/threshold/pkg/protocol"
	"github.com/luxfi/threshold/protocols/cmp"
	"github.com/luxfi/threshold/protocols/cmp/config"

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

// CGGMP21Protocol implements the Protocol interface using CGGMP21
type CGGMP21Protocol struct {
	pool   *pool.Pool
	logger log.Logger
}

// NewCGGMP21Protocol creates a new CGGMP21 protocol adapter
func NewCGGMP21Protocol() *CGGMP21Protocol {
	return &CGGMP21Protocol{
		pool:   pool.NewPool(0), // Use max threads
		logger: log.NewTestLogger(log.InfoLevel),
	}
}

// Close cleans up resources
func (p *CGGMP21Protocol) Close() {
	if p.pool != nil {
		p.pool.TearDown()
	}
}

// Name returns the protocol name
func (p *CGGMP21Protocol) Name() string {
	return "CGGMP21"
}

// KeyGen starts a distributed key generation
func (p *CGGMP21Protocol) KeyGen(selfID string, partyIDs []string, threshold int) (protocol.Party, error) {
	// Convert string IDs to party.ID
	ids := make([]party.ID, len(partyIDs))
	for i, id := range partyIDs {
		ids[i] = party.ID(id)
	}

	// Create the keygen protocol
	startFunc := cmp.Keygen(curve.Secp256k1{}, party.ID(selfID), ids, threshold, p.pool)

	// Create handler with proper context, logging, and session ID
	ctx := context.Background()
	sessionID := []byte(fmt.Sprintf("cggmp21-keygen-%s", selfID))
	handler, err := mpsProtocol.NewHandler(
		ctx,
		p.logger,
		nil, // No prometheus registry
		startFunc,
		sessionID,
		mpsProtocol.DefaultConfig(),
	)
	if err != nil {
		return nil, fmt.Errorf("failed to create keygen handler: %w", err)
	}

	return &partyAdapter{
		handler: handler,
		selfID:  selfID,
	}, nil
}

// Refresh refreshes shares from an existing config
func (p *CGGMP21Protocol) Refresh(cfg protocol.KeyGenConfig) (protocol.Party, error) {
	// Convert to CGGMP21 config
	cmpConfig, err := toCMPConfig(cfg)
	if err != nil {
		return nil, err
	}

	// Create refresh protocol
	startFunc := cmp.Refresh(cmpConfig, p.pool)

	// Create handler with proper context, logging, and session ID
	ctx := context.Background()
	sessionID := []byte(fmt.Sprintf("cggmp21-refresh-%s", cfg.GetPartyID()))
	handler, err := mpsProtocol.NewHandler(
		ctx,
		p.logger,
		nil, // No prometheus registry
		startFunc,
		sessionID,
		mpsProtocol.DefaultConfig(),
	)
	if err != nil {
		return nil, fmt.Errorf("failed to create refresh handler: %w", err)
	}

	return &partyAdapter{
		handler: handler,
		selfID:  cfg.GetPartyID(),
	}, nil
}

// Sign starts a signing protocol
func (p *CGGMP21Protocol) Sign(cfg protocol.KeyGenConfig, signers []string, messageHash []byte) (protocol.Party, error) {
	// Convert to CGGMP21 config
	cmpConfig, err := toCMPConfig(cfg)
	if err != nil {
		return nil, err
	}

	// Convert signer IDs
	signerIDs := make([]party.ID, len(signers))
	for i, id := range signers {
		signerIDs[i] = party.ID(id)
	}

	// Create sign protocol
	startFunc := cmp.Sign(cmpConfig, signerIDs, messageHash, p.pool)

	// Create handler with proper context, logging, and session ID
	ctx := context.Background()
	sessionID := []byte(fmt.Sprintf("cggmp21-sign-%s-%x", cfg.GetPartyID(), messageHash[:8]))
	handler, err := mpsProtocol.NewHandler(
		ctx,
		p.logger,
		nil, // No prometheus registry
		startFunc,
		sessionID,
		mpsProtocol.DefaultConfig(),
	)
	if err != nil {
		return nil, fmt.Errorf("failed to create sign handler: %w", err)
	}

	return &partyAdapter{
		handler: handler,
		selfID:  cfg.GetPartyID(),
	}, nil
}

// PreSign starts a presigning protocol
func (p *CGGMP21Protocol) PreSign(cfg protocol.KeyGenConfig, signers []string) (protocol.Party, error) {
	// Convert to CGGMP21 config
	cmpConfig, err := toCMPConfig(cfg)
	if err != nil {
		return nil, err
	}

	// Convert signer IDs
	signerIDs := make([]party.ID, len(signers))
	for i, id := range signers {
		signerIDs[i] = party.ID(id)
	}

	// Create presign protocol
	startFunc := cmp.Presign(cmpConfig, signerIDs, p.pool)

	// Create handler with proper context, logging, and session ID
	ctx := context.Background()
	sessionID := []byte(fmt.Sprintf("cggmp21-presign-%s", cfg.GetPartyID()))
	handler, err := mpsProtocol.NewHandler(
		ctx,
		p.logger,
		nil, // No prometheus registry
		startFunc,
		sessionID,
		mpsProtocol.DefaultConfig(),
	)
	if err != nil {
		return nil, fmt.Errorf("failed to create presign handler: %w", err)
	}

	return &partyAdapter{
		handler: handler,
		selfID:  cfg.GetPartyID(),
	}, nil
}

// PreSignOnline completes a signature with a presignature
func (p *CGGMP21Protocol) PreSignOnline(cfg protocol.KeyGenConfig, preSig protocol.PreSignature, messageHash []byte) (protocol.Party, error) {
	// Convert to CGGMP21 types
	cmpConfig, err := toCMPConfig(cfg)
	if err != nil {
		return nil, err
	}

	cmpPreSig, ok := preSig.(*preSignatureAdapter)
	if !ok {
		return nil, errors.New("invalid presignature type")
	}

	// Create presign online protocol
	startFunc := cmp.PresignOnline(cmpConfig, cmpPreSig.preSig, messageHash, p.pool)

	// Create handler with proper context, logging, and session ID
	ctx := context.Background()
	sessionID := []byte(fmt.Sprintf("cggmp21-presign-online-%s-%x", cfg.GetPartyID(), messageHash[:8]))
	handler, err := mpsProtocol.NewHandler(
		ctx,
		p.logger,
		nil, // No prometheus registry
		startFunc,
		sessionID,
		mpsProtocol.DefaultConfig(),
	)
	if err != nil {
		return nil, fmt.Errorf("failed to create presign online handler: %w", err)
	}

	return &partyAdapter{
		handler: handler,
		selfID:  cfg.GetPartyID(),
	}, nil
}

// partyAdapter adapts mpsProtocol.Handler to protocol.Party
type partyAdapter struct {
	handler *mpsProtocol.Handler
	selfID  string
	mu      sync.Mutex
	done    bool
	result  interface{}
	err     error
}

func (p *partyAdapter) Update(msg protocol.Message) error {
	// Convert to MPS message format
	// If broadcast, To is nil. Otherwise, it's the first recipient
	var to party.ID
	if !msg.IsBroadcast() && len(msg.GetTo()) > 0 {
		to = party.ID(msg.GetTo()[0])
	}

	mpsMsg := &mpsProtocol.Message{
		From:      party.ID(msg.GetFrom()),
		To:        to,
		Broadcast: msg.IsBroadcast(),
		Data:      msg.GetData(),
	}

	// Check if handler can accept the message
	if !p.handler.CanAccept(mpsMsg) {
		return errors.New("message rejected by handler")
	}

	// Update handler with message
	// Note: MultiHandler doesn't have Update method, we need to send via Accept
	p.handler.Accept(mpsMsg)
	return nil
}

func (p *partyAdapter) Messages() <-chan protocol.Message {
	ch := make(chan protocol.Message)

	go func() {
		defer close(ch)

		for {
			select {
			case msg, ok := <-p.handler.Listen():
				if !ok {
					// Protocol finished
					p.mu.Lock()
					p.done = true
					p.result, p.err = p.handler.Result()
					p.mu.Unlock()
					return
				}

				// Convert and send message
				var toList []string
				if !msg.Broadcast && msg.To != "" {
					toList = []string{string(msg.To)}
				}
				ch <- &messageAdapter{
					from:      string(msg.From),
					to:        toList,
					data:      msg.Data,
					broadcast: msg.Broadcast,
				}
			}
		}
	}()

	return ch
}

func (p *partyAdapter) Errors() <-chan error {
	// CGGMP21 doesn't have a separate error channel
	// Errors are returned in Result()
	ch := make(chan error)
	close(ch)
	return ch
}

func (p *partyAdapter) Done() bool {
	p.mu.Lock()
	defer p.mu.Unlock()
	return p.done
}

func (p *partyAdapter) Result() (interface{}, error) {
	p.mu.Lock()
	defer p.mu.Unlock()

	if !p.done {
		return nil, errors.New("protocol not finished")
	}

	if p.err != nil {
		return nil, p.err
	}

	// Convert result to appropriate type
	switch r := p.result.(type) {
	case *config.Config:
		return &configAdapter{config: r}, nil
	case *mpsEcdsa.Signature:
		return &signatureAdapter{sig: r}, nil
	case *mpsEcdsa.PreSignature:
		return &preSignatureAdapter{preSig: r}, nil
	default:
		return p.result, nil
	}
}

// messageAdapter implements protocol.Message
type messageAdapter struct {
	from      string
	to        []string
	data      []byte
	broadcast bool
}

func (m *messageAdapter) GetFrom() string   { return m.from }
func (m *messageAdapter) GetTo() []string   { return m.to }
func (m *messageAdapter) GetData() []byte   { return m.data }
func (m *messageAdapter) IsBroadcast() bool { return m.broadcast }

// configAdapter implements protocol.KeyGenConfig
type configAdapter struct {
	config *config.Config
}

func (c *configAdapter) GetPartyID() string {
	return string(c.config.ID)
}

func (c *configAdapter) GetThreshold() int {
	return c.config.Threshold
}

func (c *configAdapter) GetPublicKey() *ecdsa.PublicKey {
	point := c.config.PublicPoint()
	if point == nil || point.IsIdentity() {
		return nil
	}
	return curvePointToECDSA(point)
}

func (c *configAdapter) GetShare() *big.Int {
	// Get ECDSA scalar share and convert to big.Int
	if c.config.ECDSA != nil {
		bytes, _ := c.config.ECDSA.MarshalBinary()
		return new(big.Int).SetBytes(bytes)
	}
	return nil
}

func (c *configAdapter) GetSharePublicKey() *ecdsa.PublicKey {
	public, ok := c.config.Public[c.config.ID]
	if !ok || public.ECDSA == nil || public.ECDSA.IsIdentity() {
		return nil
	}
	return curvePointToECDSA(public.ECDSA)
}

func (c *configAdapter) GetPartyIDs() []string {
	ids := c.config.PartyIDs()
	result := make([]string, len(ids))
	for i, id := range ids {
		result[i] = string(id)
	}
	return result
}

func (c *configAdapter) Serialize() ([]byte, error) {
	// Use CBOR-based MarshalBinary for proper serialization of crypto types.
	// json.Marshal corrupts curve.Scalar and curve.Point fields.
	return c.config.MarshalBinary()
}

// signatureAdapter implements protocol.Signature
type signatureAdapter struct {
	sig *mpsEcdsa.Signature
}

func (s *signatureAdapter) GetR() *big.Int {
	// Convert curve.Point R to big.Int
	if s.sig.R != nil && s.sig.R.XScalar() != nil {
		bytes, _ := s.sig.R.XScalar().MarshalBinary()
		return new(big.Int).SetBytes(bytes)
	}
	return nil
}

func (s *signatureAdapter) GetS() *big.Int {
	// Convert curve.Scalar S to big.Int
	if s.sig.S != nil {
		bytes, _ := s.sig.S.MarshalBinary()
		return new(big.Int).SetBytes(bytes)
	}
	return nil
}

// Verify checks the threshold ECDSA signature against pubKey over message.
//
// message is the 32-byte digest that was signed, not a preimage — CGGMP21
// signs a digest supplied by the caller, so hashing here would verify a
// different message than the one the ring signed.
func (s *signatureAdapter) Verify(pubKey *ecdsa.PublicKey, message []byte) bool {
	if s == nil || s.sig == nil {
		return false
	}
	point := protocol.CurvePoint(pubKey)
	if point == nil {
		return false
	}
	return s.sig.Verify(point, message)
}

func (s *signatureAdapter) Serialize() ([]byte, error) {
	return json.Marshal(s.sig)
}

// preSignatureAdapter implements protocol.PreSignature
type preSignatureAdapter struct {
	preSig *mpsEcdsa.PreSignature
}

func (p *preSignatureAdapter) GetID() string {
	// Convert RID (byte slice) to hex string
	return fmt.Sprintf("%x", p.preSig.ID)
}

func (p *preSignatureAdapter) Validate() error {
	return p.preSig.Validate()
}

// Helper functions

func convertToPartyIDs(ids []string) []party.ID {
	if ids == nil {
		return nil
	}
	result := make([]party.ID, len(ids))
	for i, id := range ids {
		result[i] = party.ID(id)
	}
	return result
}

func convertFromPartyIDs(ids []party.ID) []string {
	if ids == nil {
		return nil
	}
	result := make([]string, len(ids))
	for i, id := range ids {
		result[i] = string(id)
	}
	return result
}

func toCMPConfig(cfg protocol.KeyGenConfig) (*config.Config, error) {
	// Try to cast directly first -- zero-copy path.
	if adapter, ok := cfg.(*configAdapter); ok {
		return adapter.config, nil
	}

	// Fallback: serialize through the interface's Serialize() method
	// (which uses CBOR via MarshalBinary) and deserialize into a fresh Config.
	data, err := cfg.Serialize()
	if err != nil {
		return nil, fmt.Errorf("failed to serialize config for conversion: %w", err)
	}

	// EmptyConfig initialises a Config with the correct curve so that
	// UnmarshalBinary can decode curve-specific scalars and points.
	out := config.EmptyConfig(curve.Secp256k1{})
	if err := out.UnmarshalBinary(data); err != nil {
		return nil, fmt.Errorf("failed to deserialize CMP config: %w", err)
	}

	return out, nil
}

// curvePointToECDSA converts a curve.Point (secp256k1) to a standard crypto/ecdsa.PublicKey
// by marshalling to compressed form and recovering both X and Y coordinates.
func curvePointToECDSA(point curve.Point) *ecdsa.PublicKey {
	compressed, err := point.MarshalBinary()
	if err != nil || len(compressed) != 33 {
		return nil
	}

	x := new(big.Int).SetBytes(compressed[1:33])

	// secp256k1: y^2 = x^3 + 7 mod p
	p, _ := new(big.Int).SetString("FFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFEFFFFFC2F", 16)
	x3 := new(big.Int).Mul(x, x)
	x3.Mod(x3, p)
	x3.Mul(x3, x)
	x3.Mod(x3, p)
	y2 := new(big.Int).Add(x3, big.NewInt(7))
	y2.Mod(y2, p)

	// p mod 4 == 3, so sqrt = y2^((p+1)/4) mod p
	exp := new(big.Int).Add(p, big.NewInt(1))
	exp.Rsh(exp, 2)
	y := new(big.Int).Exp(y2, exp, p)

	// Verify the square root
	check := new(big.Int).Mul(y, y)
	check.Mod(check, p)
	if check.Cmp(y2) != 0 {
		return nil
	}

	// Select correct Y parity: 0x02 = even, 0x03 = odd
	isOdd := compressed[0] == 0x03
	if y.Bit(0) == 1 != isOdd {
		y.Sub(p, y)
	}

	return &ecdsa.PublicKey{
		Curve: ellipticSecp256k1(),
		X:     x,
		Y:     y,
	}
}

// secp256k1Params provides a secp256k1 elliptic.Curve for use with crypto/ecdsa.PublicKey.
var secp256k1CurveParams *elliptic.CurveParams

func init() {
	secp256k1CurveParams = &elliptic.CurveParams{
		Name:    "secp256k1",
		BitSize: 256,
	}
	secp256k1CurveParams.P, _ = new(big.Int).SetString("FFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFEFFFFFC2F", 16)
	secp256k1CurveParams.N, _ = new(big.Int).SetString("FFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFEBAAEDCE6AF48A03BBFD25E8CD0364141", 16)
	secp256k1CurveParams.B = big.NewInt(7)
	secp256k1CurveParams.Gx, _ = new(big.Int).SetString("79BE667EF9DCBBAC55A06295CE870B07029BFCDB2DCE28D959F2815B16F81798", 16)
	secp256k1CurveParams.Gy, _ = new(big.Int).SetString("483ADA7726A3C4655DA4FBFC0E1108A8FD17B448A68554199C47D08FFB10D4B8", 16)
}

// ellipticSecp256k1 returns an elliptic.Curve representing secp256k1.
func ellipticSecp256k1() elliptic.Curve {
	return secp256k1CurveParams
}