threshold/protocols/cmp/keygen/round4.go

package keygen

import (
	"errors"
	"fmt"

	"github.com/luxfi/threshold/internal/round"
	"github.com/luxfi/threshold/internal/types"
	"github.com/luxfi/threshold/pkg/hash"
	"github.com/luxfi/threshold/pkg/math/curve"
	"github.com/luxfi/threshold/pkg/math/polynomial"
	"github.com/luxfi/threshold/pkg/paillier"
	"github.com/luxfi/threshold/pkg/party"
	"github.com/luxfi/threshold/pkg/pedersen"
	zkfac "github.com/luxfi/threshold/pkg/zk/fac"
	zkmod "github.com/luxfi/threshold/pkg/zk/mod"
	zkprm "github.com/luxfi/threshold/pkg/zk/prm"
	"github.com/luxfi/threshold/protocols/cmp/config"
)

var _ round.Round = (*round4)(nil)

type round4 struct {
	*round3

	// RID = ⊕ⱼ RIDⱼ
	// Random ID generated by taking the XOR of all ridᵢ
	RID types.RID
	// ChainKey is a sequence of random bytes agreed upon together
	ChainKey types.RID
}

type message4 struct {
	// Share = Encᵢ(x) is the encryption of the receivers share
	Share *paillier.Ciphertext
	Fac   *zkfac.Proof
}

type broadcast4 struct {
	round.NormalBroadcastContent
	Mod *zkmod.Proof
	Prm *zkprm.Proof
}

// StoreBroadcastMessage implements round.BroadcastRound.
//
// - verify Mod, Prm proof for N.
func (r *round4) StoreBroadcastMessage(msg round.Message) error {
	from := msg.From
	body, ok := msg.Content.(*broadcast4)
	if !ok || body == nil {
		return round.ErrInvalidContent
	}

	// Load Paillier public key for from party
	paillierValue, ok := r.PaillierPublic.Load(from)
	if !ok {
		return errors.New("paillier public key not found for party")
	}
	paillierFrom, ok := paillierValue.(*paillier.PublicKey)
	if !ok || paillierFrom == nil {
		return errors.New("invalid paillier public key for party")
	}

	// Load Pedersen parameters for from party
	pedersenValue, ok := r.Pedersen.Load(from)
	if !ok {
		return errors.New("pedersen parameters not found for party")
	}
	pedersenFrom, ok := pedersenValue.(*pedersen.Parameters)
	if !ok || pedersenFrom == nil {
		return errors.New("invalid pedersen parameters for party")
	}

	// verify zkmod - use Paillier N, not Pedersen N
	// Use a fresh hash for verification to match proof generation
	hMod := hash.New()
	if !body.Mod.Verify(zkmod.Public{N: paillierFrom.N()}, hMod, r.Pool) {
		return fmt.Errorf("failed to validate mod proof from %s", from)
	}

	// verify zkprm
	hPrm := hash.New()
	if !body.Prm.Verify(zkprm.Public{Aux: pedersenFrom}, hPrm, r.Pool) {
		return errors.New("failed to validate prm proof")
	}

	return nil
}

// VerifyMessage implements round.Round.
//
// - verify validity of share ciphertext.
func (r *round4) VerifyMessage(msg round.Message) error {
	from := msg.From
	body, ok := msg.Content.(*message4)
	if !ok || body == nil {
		return round.ErrInvalidContent
	}

	// Load Paillier public key for recipient (ourselves)
	// For P2P messages, msg.To is the recipient which is ourselves when verifying
	paillierToValue, ok := r.PaillierPublic.Load(r.SelfID())
	if !ok {
		return errors.New("paillier public key not found for recipient")
	}
	paillierTo, ok := paillierToValue.(*paillier.PublicKey)
	if !ok || paillierTo == nil {
		return errors.New("invalid paillier public key for recipient")
	}

	if !paillierTo.ValidateCiphertexts(body.Share) {
		return errors.New("invalid ciphertext")
	}

	// Load Paillier public key for from party
	paillierFromValue, ok := r.PaillierPublic.Load(from)
	if !ok {
		return errors.New("paillier public key not found for sender")
	}
	paillierFrom, ok := paillierFromValue.(*paillier.PublicKey)
	if !ok || paillierFrom == nil {
		return errors.New("invalid paillier public key for sender")
	}

	// Load Pedersen parameters for recipient (ourselves)
	pedersenToValue, ok := r.Pedersen.Load(r.SelfID())
	if !ok {
		return errors.New("pedersen parameters not found for recipient")
	}
	pedersenTo, ok := pedersenToValue.(*pedersen.Parameters)
	if !ok || pedersenTo == nil {
		return errors.New("invalid pedersen parameters for recipient")
	}

	// verify zkfac
	// Use a fresh hash for verification to match proof generation
	hFac := hash.New()
	if !body.Fac.Verify(zkfac.Public{N: paillierFrom.N(), Aux: pedersenTo}, hFac) {
		return errors.New("failed to validate fac proof")
	}

	return nil
}

// StoreMessage implements round.Round.
//
// Since this message is only intended for us, we need to do the VSS verification here.
// - check that the decrypted share did not overflow.
// - check VSS condition.
// - save share.
func (r *round4) StoreMessage(msg round.Message) error {
	from, body := msg.From, msg.Content.(*message4)

	// decrypt share
	DecryptedShare, err := r.PaillierSecret.Dec(body.Share)
	if err != nil {
		return err
	}
	Share := r.Group().NewScalar().SetNat(DecryptedShare.Mod(r.Group().Order()))
	if DecryptedShare.Eq(curve.MakeInt(Share)) != 1 {
		return errors.New("decrypted share is not in correct range")
	}

	// verify share with VSS
	vssPolyValue, ok := r.VSSPolynomials.Load(from)
	if !ok {
		return errors.New("VSS polynomial not found for party")
	}
	vssPoly, ok := vssPolyValue.(*polynomial.Exponent)
	if !ok || vssPoly == nil {
		return errors.New("invalid VSS polynomial for party")
	}
	ExpectedPublicShare := vssPoly.Evaluate(r.SelfID().Scalar(r.Group())) // Fⱼ(i)
	PublicShare := Share.ActOnBase()
	// X == Fⱼ(i)
	if !PublicShare.Equal(ExpectedPublicShare) {
		return errors.New("failed to validate VSS share")
	}

	// Store using sync.Map for thread-safe concurrent access
	r.ShareReceived.Store(from, Share)
	return nil
}

// Finalize implements round.Round
//
// - sum of all received shares
// - compute group public key and individual public keys
// - recompute config SSID
// - validate Config
// - write new ssid hash to old hash state
// - create proof of knowledge of secret.
func (r *round4) Finalize(out chan<- *round.Message) (round.Session, error) {
	// Check if we have received all shares before proceeding
	// Count shares in sync.Map
	shareCount := 0
	r.ShareReceived.Range(func(_, _ interface{}) bool {
		shareCount++
		return true
	})
	if shareCount < r.N() {
		// Not ready to advance yet - return ourselves to wait for more messages
		return r, nil
	}

	// Update hash state with RID now that all verifications are complete
	// This was delayed from round3 to ensure proofs could be verified with pre-RID hash
	r.UpdateHashState(r.RID)

	// add all shares to our secret
	UpdatedSecretECDSA := r.Group().NewScalar()
	if r.PreviousSecretECDSA != nil {
		UpdatedSecretECDSA.Set(r.PreviousSecretECDSA)
	}
	for _, j := range r.PartyIDs() {
		shareValue, ok := r.ShareReceived.Load(j)
		if !ok {
			return r, errors.New("share not received from party")
		}
		share, ok := shareValue.(curve.Scalar)
		if !ok || share == nil {
			return r, errors.New("invalid share from party")
		}
		UpdatedSecretECDSA.Add(share)
	}

	// [F₁(X), …, Fₙ(X)]
	ShamirPublicPolynomials := make([]*polynomial.Exponent, 0, r.N())
	r.VSSPolynomials.Range(func(_, value interface{}) bool {
		vssPoly, _ := value.(*polynomial.Exponent)
		ShamirPublicPolynomials = append(ShamirPublicPolynomials, vssPoly)
		return true
	})

	// ShamirPublicPolynomial = F(X) = ∑Fⱼ(X)
	ShamirPublicPolynomial, err := polynomial.Sum(ShamirPublicPolynomials)
	if err != nil {
		return r, err
	}

	// compute the new public key share Xⱼ = F(j) (+X'ⱼ if doing a refresh)
	PublicData := make(map[party.ID]*config.Public, len(r.PartyIDs()))
	for _, j := range r.PartyIDs() {
		PublicECDSAShare := ShamirPublicPolynomial.Evaluate(j.Scalar(r.Group()))
		if r.PreviousPublicSharesECDSA != nil {
			PublicECDSAShare = PublicECDSAShare.Add(r.PreviousPublicSharesECDSA[j])
		}

		// Load from sync.Maps
		elGamalValue, ok := r.ElGamalPublic.Load(j)
		if !ok {
			return r, errors.New("ElGamal public key not found for party")
		}
		elGamal, ok := elGamalValue.(curve.Point)
		if !ok || elGamal == nil {
			return r, errors.New("invalid ElGamal public key for party")
		}

		paillierValue, ok := r.PaillierPublic.Load(j)
		if !ok {
			return r, errors.New("Paillier public key not found for party")
		}
		paillier, ok := paillierValue.(*paillier.PublicKey)
		if !ok || paillier == nil {
			return r, errors.New("invalid Paillier public key for party")
		}

		pedersenValue, ok := r.Pedersen.Load(j)
		if !ok {
			return r, errors.New("Pedersen parameters not found for party")
		}
		pedersen, ok := pedersenValue.(*pedersen.Parameters)
		if !ok || pedersen == nil {
			return r, errors.New("invalid Pedersen parameters for party")
		}

		PublicData[j] = &config.Public{
			ECDSA:    PublicECDSAShare,
			ElGamal:  elGamal,
			Paillier: paillier,
			Pedersen: pedersen,
		}
	}

	UpdatedConfig := &config.Config{
		Group:     r.Group(),
		ID:        r.SelfID(),
		Threshold: r.Threshold(),
		ECDSA:     UpdatedSecretECDSA,
		ElGamal:   r.ElGamalSecret,
		Paillier:  r.PaillierSecret,
		RID:       r.RID.Copy(),
		ChainKey:  r.ChainKey.Copy(),
		Public:    PublicData,
	}

	// write new ssid to hash, to bind the Schnorr proof to this new config
	// Write SSID, selfID to temporary hash
	h := r.Hash()
	_ = h.WriteAny(UpdatedConfig, r.SelfID())

	proof := r.SchnorrRand.Prove(h, PublicData[r.SelfID()].ECDSA, UpdatedSecretECDSA, nil)

	// send to all
	err = r.BroadcastMessage(out, &broadcast5{SchnorrResponse: proof})
	if err != nil {
		return r, err
	}

	r.UpdateHashState(UpdatedConfig)
	return &round5{
		round4:        r,
		UpdatedConfig: UpdatedConfig,
	}, nil
}

// RoundNumber implements round.Content.
func (message4) RoundNumber() round.Number { return 4 }

// MessageContent implements round.Round.
func (*round4) MessageContent() round.Content { return &message4{} }

// RoundNumber implements round.Content.
func (broadcast4) RoundNumber() round.Number { return 4 }

// BroadcastContent implements round.BroadcastRound.
func (*round4) BroadcastContent() round.BroadcastContent { return &broadcast4{} }

// Number implements round.Round.
func (*round4) Number() round.Number { return 4 }