rwadurian/backend/mpc-system/tests/integration/mpc_full_flow_test.go

package integration_test

import (
	"crypto/ecdsa"
	"crypto/sha256"
	"encoding/hex"
	"fmt"
	"testing"

	"github.com/rwadurian/mpc-system/pkg/tss"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

// TestFull2of3MPCFlow tests the complete 2-of-3 MPC flow:
// 1. Key generation with 3 parties
// 2. Signing with 2 parties (threshold)
// 3. Signature verification
func TestFull2of3MPCFlow(t *testing.T) {
	fmt.Println("========================================")
	fmt.Println("  MPC 2-of-3 Full Flow Integration Test")
	fmt.Println("========================================")

	// ============================================
	// Step 1: Key Generation (2-of-3)
	// ============================================
	fmt.Println("\n[Step 1] Running 2-of-3 Distributed Key Generation...")
	fmt.Println("  - Threshold (t): 1 (meaning t+1=2 signers required)")
	fmt.Println("  - Total Parties (n): 3")

	// In tss-lib, threshold=1 means 2 signers are required (t+1)
	threshold := 1
	totalParties := 3

	keygenResults, err := tss.RunLocalKeygen(threshold, totalParties)
	require.NoError(t, err, "Keygen should succeed")
	require.Len(t, keygenResults, 3, "Should have 3 key shares")

	// Extract the shared public key
	publicKey := keygenResults[0].PublicKey
	require.NotNil(t, publicKey, "Public key should not be nil")

	fmt.Printf("  [OK] Key generation completed!\n")
	fmt.Printf("      Public Key X: %s...\n", publicKey.X.Text(16)[:32])
	fmt.Printf("      Public Key Y: %s...\n", publicKey.Y.Text(16)[:32])

	// Verify all parties have the same public key
	for i, result := range keygenResults {
		assert.Equal(t, publicKey.X, result.PublicKey.X, "Party %d should have same X", i)
		assert.Equal(t, publicKey.Y, result.PublicKey.Y, "Party %d should have same Y", i)
	}
	fmt.Println("      All parties have consistent public key")

	// ============================================
	// Step 2: Signing with 2 Parties (Threshold)
	// ============================================
	fmt.Println("\n[Step 2] Running Threshold Signing (2-of-3)...")

	// Create a message to sign
	message := []byte("Hello MPC World! This is a test transaction.")
	messageHash := sha256.Sum256(message)

	fmt.Printf("      Message: \"%s\"\n", string(message))
	fmt.Printf("      Message Hash: %s\n", hex.EncodeToString(messageHash[:]))

	// Test all 3 combinations of 2 parties
	combinations := []struct {
		name    string
		parties []*tss.LocalKeygenResult
	}{
		{"Party 0 + Party 1", []*tss.LocalKeygenResult{keygenResults[0], keygenResults[1]}},
		{"Party 0 + Party 2", []*tss.LocalKeygenResult{keygenResults[0], keygenResults[2]}},
		{"Party 1 + Party 2", []*tss.LocalKeygenResult{keygenResults[1], keygenResults[2]}},
	}

	for i, combo := range combinations {
		fmt.Printf("\n  [Signing %d] %s\n", i+1, combo.name)

		signResult, err := tss.RunLocalSigning(threshold, combo.parties, messageHash[:])
		require.NoError(t, err, "Signing with %s should succeed", combo.name)

		// Verify signature components
		require.NotNil(t, signResult.R, "R should not be nil")
		require.NotNil(t, signResult.S, "S should not be nil")
		require.Len(t, signResult.Signature, 64, "Signature should be 64 bytes")

		fmt.Printf("      R: %s...\n", signResult.R.Text(16)[:32])
		fmt.Printf("      S: %s...\n", signResult.S.Text(16)[:32])
		fmt.Printf("      Recovery ID: %d\n", signResult.RecoveryID)

		// ============================================
		// Step 3: Verify Signature
		// ============================================
		valid := ecdsa.Verify(publicKey, messageHash[:], signResult.R, signResult.S)
		require.True(t, valid, "Signature verification should pass for %s", combo.name)

		fmt.Printf("      [OK] Signature verified successfully!\n")
	}

	// ============================================
	// Step 4: Test Different Messages
	// ============================================
	fmt.Println("\n[Step 3] Testing with Different Messages...")

	messages := []string{
		"Transaction: Send 1.5 ETH to 0x1234...",
		"Contract call: approve(spender, amount)",
		"NFT transfer: tokenId=42",
	}

	signers := []*tss.LocalKeygenResult{keygenResults[0], keygenResults[1]}

	for _, msg := range messages {
		msgHash := sha256.Sum256([]byte(msg))
		signResult, err := tss.RunLocalSigning(threshold, signers, msgHash[:])
		require.NoError(t, err)

		valid := ecdsa.Verify(publicKey, msgHash[:], signResult.R, signResult.S)
		require.True(t, valid)

		fmt.Printf("      [OK] Message: \"%s...\"\n", msg[:min(30, len(msg))])
	}

	// ============================================
	// Summary
	// ============================================
	fmt.Println("\n========================================")
	fmt.Println("  Test Summary")
	fmt.Println("========================================")
	fmt.Println("  [OK] 2-of-3 Key Generation: PASSED")
	fmt.Println("  [OK] Threshold Signing (3 combinations): PASSED")
	fmt.Println("  [OK] Signature Verification: PASSED")
	fmt.Println("  [OK] Multi-message Signing: PASSED")
	fmt.Println("========================================")
	fmt.Println("  All MPC operations completed successfully!")
	fmt.Println("========================================")
}

// TestSecurityProperties tests security properties of the MPC system
func TestSecurityProperties(t *testing.T) {
	fmt.Println("\n========================================")
	fmt.Println("  Security Properties Test")
	fmt.Println("========================================")

	threshold := 1
	totalParties := 3

	// Generate keys
	keygenResults, err := tss.RunLocalKeygen(threshold, totalParties)
	require.NoError(t, err)

	publicKey := keygenResults[0].PublicKey
	message := []byte("Security test message")
	messageHash := sha256.Sum256(message)

	// Test 1: Single party cannot sign
	fmt.Println("\n[Test 1] Verifying single party cannot sign alone...")
	// Note: With threshold=1, minimum 2 parties are required
	// Attempting to sign with 1 party should fail
	singleParty := []*tss.LocalKeygenResult{keygenResults[0]}
	_, err = tss.RunLocalSigning(threshold, singleParty, messageHash[:])
	// This should fail because we need at least t+1=2 parties
	if err != nil {
		fmt.Println("      [OK] Single party signing correctly rejected")
	} else {
		t.Error("Single party should not be able to sign")
	}

	// Test 2: Different key shares produce same public key
	fmt.Println("\n[Test 2] Verifying key share consistency...")
	for i := 0; i < totalParties; i++ {
		assert.Equal(t, publicKey.X.Cmp(keygenResults[i].PublicKey.X), 0)
		assert.Equal(t, publicKey.Y.Cmp(keygenResults[i].PublicKey.Y), 0)
	}
	fmt.Println("      [OK] All parties have consistent public key")

	// Test 3: Signatures from different party combinations verify with same public key
	fmt.Println("\n[Test 3] Verifying signature consistency across party combinations...")
	combo1 := []*tss.LocalKeygenResult{keygenResults[0], keygenResults[1]}
	combo2 := []*tss.LocalKeygenResult{keygenResults[1], keygenResults[2]}

	sig1, err := tss.RunLocalSigning(threshold, combo1, messageHash[:])
	require.NoError(t, err)

	sig2, err := tss.RunLocalSigning(threshold, combo2, messageHash[:])
	require.NoError(t, err)

	// Both signatures should verify with the same public key
	valid1 := ecdsa.Verify(publicKey, messageHash[:], sig1.R, sig1.S)
	valid2 := ecdsa.Verify(publicKey, messageHash[:], sig2.R, sig2.S)

	assert.True(t, valid1, "Signature from combo1 should verify")
	assert.True(t, valid2, "Signature from combo2 should verify")
	fmt.Println("      [OK] All party combinations produce valid signatures")

	fmt.Println("\n========================================")
	fmt.Println("  Security tests passed!")
	fmt.Println("========================================")
}

func min(a, b int) int {
	if a < b {
		return a
	}
	return b
}