threshold/protocols/frost/keygen_shares_test.go

package frost_test

import (
	"testing"

	"github.com/cronokirby/saferith"
	"github.com/luxfi/threshold/internal/test"
	"github.com/luxfi/threshold/pkg/math/curve"
	"github.com/luxfi/threshold/pkg/math/polynomial"
	"github.com/luxfi/threshold/pkg/party"
	"github.com/luxfi/threshold/pkg/protocol"
	"github.com/luxfi/threshold/protocols/frost"
	"github.com/stretchr/testify/require"
)

// TestKeygenSharesSimple tests with a simple manual keygen simulation
func TestKeygenSharesSimple(t *testing.T) {
	group := curve.Secp256k1{}
	n := 3
	threshold := 2
	partyIDs := test.PartyIDs(n) // [a, b, c]

	t.Logf("Simulating FROST keygen manually for parties: %v", partyIDs)

	// Each party creates a polynomial f_i(x) of degree threshold-1
	secrets := make([]curve.Scalar, n)
	polynomials := make([]*polynomial.Polynomial, n)

	for i := 0; i < n; i++ {
		// Each party chooses a random secret (constant term)
		secrets[i] = group.NewScalar()
		secrets[i].SetNat(new(saferith.Nat).SetUint64(uint64(100 + i))) // Deterministic for testing
		polynomials[i] = polynomial.NewPolynomial(group, threshold-1, secrets[i])
		t.Logf("Party %s: secret = %v", partyIDs[i], secrets[i])
	}

	// Each party evaluates their polynomial at every party ID and sends shares
	shares := make(map[party.ID]curve.Scalar) // shares[i] = sum of f_j(i) for all j

	for i, id := range partyIDs {
		share := group.NewScalar()
		for j := 0; j < n; j++ {
			// Party j evaluates their polynomial at party i's ID
			eval := polynomials[j].Evaluate(id.Scalar(group))
			share.Add(eval)
		}
		shares[id] = share
		t.Logf("Party %s: combined share = %v", id, share)

		// Verify: share * G should equal the verification share
		verificationShare := share.ActOnBase()
		t.Logf("Party %s: verification share = s_i * G", id)
		_ = verificationShare
		_ = i
	}

	// The group secret is the sum of all individual secrets
	groupSecret := group.NewScalar()
	for _, secret := range secrets {
		groupSecret.Add(secret)
	}
	groupPublicKey := groupSecret.ActOnBase()
	t.Logf("Group secret (sum of secrets) = %v", groupSecret)
	t.Logf("Group public key = secret * G")

	// Test: Any threshold subset should reconstruct the group secret
	subset := partyIDs[:threshold] // [a, b]
	t.Logf("\nTesting reconstruction with subset: %v", subset)

	lambdas := polynomial.Lagrange(group, subset)
	reconstructedSecret := group.NewScalar()
	for _, id := range subset {
		reconstructedSecret.Add(lambdas[id].Mul(shares[id]))
	}

	require.True(t, reconstructedSecret.Equal(groupSecret),
		"Reconstructed secret should match group secret")
	t.Logf("✓ Successfully reconstructed group secret from threshold subset")

	// The reconstructed public key should also match
	reconstructedPK := reconstructedSecret.ActOnBase()
	require.True(t, reconstructedPK.Equal(groupPublicKey),
		"Reconstructed public key should match")
	t.Logf("✓ Reconstructed public key matches")

	// Now test with actual FROST keygen
	t.Logf("\n=== Running actual FROST keygen ===")

	keygenResults, err := test.RunProtocol(t, partyIDs, []byte("keygen-shares-simple"), func(id party.ID) protocol.StartFunc {
		return frost.Keygen(group, id, partyIDs, threshold)
	})
	require.NoError(t, err, "keygen failed")

	// Extract configs
	configs := make(map[party.ID]*frost.Config, n)
	for id, result := range keygenResults {
		cfg := result.(*frost.Config)
		configs[id] = cfg
	}

	// Get the public key from keygen
	keygenPK := configs[partyIDs[0]].PublicKey

	// Test reconstruction with keygen shares
	keygenLambdas := polynomial.Lagrange(group, subset)
	keygenReconstructed := group.NewPoint()

	for _, id := range subset {
		yShare := configs[partyIDs[0]].VerificationShares.Points[id]
		keygenReconstructed = keygenReconstructed.Add(keygenLambdas[id].Act(yShare))
	}

	// In FROST, shares are from summed polynomials, so we need all n parties
	// for reconstruction. This is expected behavior, not a bug.
	if !keygenReconstructed.Equal(keygenPK) {
		t.Logf("FROST keygen: Threshold reconstruction doesn't work (expected - FROST needs all n parties)")

		// Debug: Check individual verification shares
		for _, id := range partyIDs {
			privateShare := configs[id].PrivateShare
			expectedYShare := privateShare.ActOnBase()
			actualYShare := configs[id].VerificationShares.Points[id]

			if !expectedYShare.Equal(actualYShare) {
				t.Logf("Party %s: verification share mismatch!", id)
				t.Logf("  Expected (s_i * G): %v", expectedYShare)
				t.Logf("  Actual from keygen: %v", actualYShare)
			} else {
				t.Logf("Party %s: verification share matches private share ✓", id)
			}
		}

		// Debug: Check if all parties have the same view
		t.Logf("\nChecking if all parties have consistent verification shares:")
		for _, id1 := range partyIDs {
			for _, id2 := range partyIDs {
				share1 := configs[id1].VerificationShares.Points[id2]
				share2 := configs[partyIDs[0]].VerificationShares.Points[id2]
				if !share1.Equal(share2) {
					t.Logf("INCONSISTENT: Party %s has different share for %s", id1, id2)
				}
			}
		}

		// Debug: Manual reconstruction with private shares
		t.Logf("\nManual reconstruction with private shares:")
		manualReconstruct := group.NewScalar()
		for _, id := range subset {
			manualReconstruct.Add(keygenLambdas[id].Mul(configs[id].PrivateShare))
		}
		manualPK := manualReconstruct.ActOnBase()
		if manualPK.Equal(keygenPK) {
			t.Logf("✓ Manual reconstruction with private shares works!")
		} else {
			t.Logf("✗ Manual reconstruction with private shares also fails (expected)")
		}

		// Test with all n parties - this should work
		t.Logf("\nReconstruction with ALL parties:")
		allLambdas := polynomial.Lagrange(group, partyIDs)
		allReconstruct := group.NewPoint()
		for _, id := range partyIDs {
			yShare := configs[partyIDs[0]].VerificationShares.Points[id]
			allReconstruct = allReconstruct.Add(allLambdas[id].Act(yShare))
		}
		if allReconstruct.Equal(keygenPK) {
			t.Logf("✓ Reconstruction with all n parties works (as expected for FROST)")
		}
	} else {
		t.Logf("✓ FROST keygen: Successfully reconstructed public key from threshold subset")
	}
}