escrow_replay_test.go

package paywall

import (
	"crypto/rand"
	"testing"
	"time"

	"github.com/opd-ai/paywall/wallet"
)

// TestSignatureReplayProtection_NonceUniqueness verifies that duplicate nonces are rejected
func TestSignatureReplayProtection_NonceUniqueness(t *testing.T) {
	store := NewMemoryStore()
	pw := &Paywall{
		Store:     store,
		HDWallets: make(map[wallet.WalletType]wallet.HDWallet),
		prices:    map[wallet.WalletType]float64{wallet.Bitcoin: 0.001},
	}

	em, err := NewEscrowManager(pw)
	if err != nil {
		t.Fatalf("NewEscrowManager() error = %v", err)
	}

	// Create a funded escrow payment with initial signature
	paymentID := "test-replay-nonce"
	nonce := make([]byte, 32)
	if _, err := rand.Read(nonce); err != nil {
		t.Fatalf("Failed to generate nonce: %v", err)
	}

	payment := &Payment{
		ID:              paymentID,
		MultisigEnabled: true,
		EscrowState:     EscrowFunded,
		Version:         1,
		Signatures: map[wallet.WalletType][]SignatureData{
			wallet.Bitcoin: {
				{
					SignerID:  "buyer-1",
					Role:      RoleBuyer,
					Signature: []byte("existing-sig"),
					PublicKey: []byte("buyer-pubkey"),
					SignedAt:  time.Now(),
					Nonce:     nonce,
					PaymentID: paymentID,
				},
			},
		},
	}

	if err := store.CreatePayment(payment); err != nil {
		t.Fatalf("Failed to create payment: %v", err)
	}

	// Attempt to validate a signature with the same nonce
	replaySig := &SignatureData{
		SignerID:  "seller-1",
		Role:      RoleSeller,
		Signature: []byte("replay-sig"),
		PublicKey: []byte("seller-pubkey"),
		SignedAt:  time.Now(),
		Nonce:     nonce, // Same nonce - should be rejected
		PaymentID: paymentID,
	}

	err = em.validateSignatureReplay(replaySig, payment)
	if err == nil {
		t.Error("validateSignatureReplay() expected error for duplicate nonce, got nil")
	}
	if err != nil && !contains(err.Error(), "nonce reuse") {
		t.Errorf("validateSignatureReplay() expected nonce reuse error, got: %v", err)
	}
}

// TestSignatureReplayProtection_PaymentIDBinding verifies that signatures are bound to specific payments
func TestSignatureReplayProtection_PaymentIDBinding(t *testing.T) {
	store := NewMemoryStore()
	pw := &Paywall{
		Store:     store,
		HDWallets: make(map[wallet.WalletType]wallet.HDWallet),
		prices:    map[wallet.WalletType]float64{wallet.Bitcoin: 0.001},
	}

	em, err := NewEscrowManager(pw)
	if err != nil {
		t.Fatalf("NewEscrowManager() error = %v", err)
	}

	paymentA := "payment-A"
	paymentB := "payment-B"

	payment := &Payment{
		ID:              paymentB,
		MultisigEnabled: true,
		EscrowState:     EscrowFunded,
		Version:         1,
	}

	// Create a signature bound to payment A
	nonce := make([]byte, 32)
	if _, err := rand.Read(nonce); err != nil {
		t.Fatalf("Failed to generate nonce: %v", err)
	}

	sigForPaymentA := &SignatureData{
		SignerID:  "buyer-1",
		Role:      RoleBuyer,
		Signature: []byte("buyer-sig"),
		PublicKey: []byte("buyer-pubkey"),
		SignedAt:  time.Now(),
		Nonce:     nonce,
		PaymentID: paymentA, // Bound to payment A
	}

	// Try to use signature on payment B
	err = em.validateSignatureReplay(sigForPaymentA, payment)
	if err == nil {
		t.Error("validateSignatureReplay() expected error for payment ID mismatch, got nil")
	}
	if err != nil && !contains(err.Error(), "payment ID mismatch") {
		t.Errorf("validateSignatureReplay() expected payment ID mismatch error, got: %v", err)
	}
}

// TestSignatureReplayProtection_CrossPaymentReplay tests that signatures from one escrow cannot be replayed on another
func TestSignatureReplayProtection_CrossPaymentReplay(t *testing.T) {
	store := NewMemoryStore()
	pw := &Paywall{
		Store:     store,
		HDWallets: make(map[wallet.WalletType]wallet.HDWallet),
		prices:    map[wallet.WalletType]float64{wallet.Bitcoin: 0.001},
	}

	em, err := NewEscrowManager(pw)
	if err != nil {
		t.Fatalf("NewEscrowManager() error = %v", err)
	}

	// Create two separate payments with same participants
	paymentA := &Payment{
		ID:              "escrow-A",
		MultisigEnabled: true,
		EscrowState:     EscrowFunded,
		Version:         1,
	}

	paymentB := &Payment{
		ID:              "escrow-B",
		MultisigEnabled: true,
		EscrowState:     EscrowFunded,
		Version:         1,
	}

	if err := store.CreatePayment(paymentA); err != nil {
		t.Fatalf("Failed to create payment A: %v", err)
	}
	if err := store.CreatePayment(paymentB); err != nil {
		t.Fatalf("Failed to create payment B: %v", err)
	}

	// Generate signature for payment A
	nonceA := make([]byte, 32)
	if _, err := rand.Read(nonceA); err != nil {
		t.Fatalf("Failed to generate nonce: %v", err)
	}

	buyerSigA := &SignatureData{
		SignerID:  "buyer-1",
		Role:      RoleBuyer,
		Signature: []byte("buyer-sig-A"),
		PublicKey: []byte("buyer-pubkey"),
		SignedAt:  time.Now(),
		Nonce:     nonceA,
		PaymentID: paymentA.ID,
	}

	// Validate signature works for payment A
	err = em.validateSignatureReplay(buyerSigA, paymentA)
	if err != nil && !contains(err.Error(), "missing nonce") {
		t.Errorf("validateSignatureReplay() unexpected error for payment A: %v", err)
	}

	// Try to replay signature on payment B (should fail due to PaymentID check)
	err = em.validateSignatureReplay(buyerSigA, paymentB)
	if err == nil {
		t.Error("validateSignatureReplay() expected error for cross-payment replay, got nil")
	}
	if err != nil && !contains(err.Error(), "payment ID mismatch") {
		t.Errorf("validateSignatureReplay() expected payment ID mismatch error, got: %v", err)
	}
}

// TestSignatureReplayProtection_SignatureDeduplication verifies that duplicate signatures from same role are rejected
func TestSignatureReplayProtection_SignatureDeduplication(t *testing.T) {
	store := NewMemoryStore()
	pw := &Paywall{
		Store:     store,
		HDWallets: make(map[wallet.WalletType]wallet.HDWallet),
		prices:    map[wallet.WalletType]float64{wallet.Bitcoin: 0.001},
	}

	em, err := NewEscrowManager(pw)
	if err != nil {
		t.Fatalf("NewEscrowManager() error = %v", err)
	}

	paymentID := "test-dedup"
	firstNonce := make([]byte, 32)
	if _, err := rand.Read(firstNonce); err != nil {
		t.Fatalf("Failed to generate nonce: %v", err)
	}

	firstSig := []byte("first-signature")
	payment := &Payment{
		ID:              paymentID,
		MultisigEnabled: true,
		EscrowState:     EscrowFunded,
		Version:         1,
		Signatures: map[wallet.WalletType][]SignatureData{
			wallet.Bitcoin: {
				{
					SignerID:  "buyer-1",
					Role:      RoleBuyer,
					Signature: firstSig,
					PublicKey: []byte("buyer-pubkey"),
					SignedAt:  time.Now(),
					Nonce:     firstNonce,
					PaymentID: paymentID,
				},
			},
		},
	}

	if err := store.CreatePayment(payment); err != nil {
		t.Fatalf("Failed to create payment: %v", err)
	}

	// Test idempotent case: same signature re-submitted (should succeed)
	idempotentSig := &SignatureData{
		SignerID:  "buyer-1",
		Role:      RoleBuyer,
		Signature: firstSig, // Same signature bytes
		PublicKey: []byte("buyer-pubkey"),
		SignedAt:  time.Now(),
		Nonce:     firstNonce,
		PaymentID: paymentID,
	}

	err = em.validateSignatureReplay(idempotentSig, payment)
	if err != nil && !contains(err.Error(), "nonce reuse") {
		// Idempotent submissions should be allowed (returns nil early)
		// But nonce reuse is still caught first
		t.Logf("validateSignatureReplay() idempotent case: %v", err)
	}

	// Test different signature from same signer/role (should fail)
	secondNonce := make([]byte, 32)
	if _, err := rand.Read(secondNonce); err != nil {
		t.Fatalf("Failed to generate nonce: %v", err)
	}

	differentSig := &SignatureData{
		SignerID:  "buyer-1",
		Role:      RoleBuyer,
		Signature: []byte("different-signature"), // Different signature
		PublicKey: []byte("buyer-pubkey"),
		SignedAt:  time.Now(),
		Nonce:     secondNonce, // Different nonce
		PaymentID: paymentID,
	}

	err = em.validateSignatureReplay(differentSig, payment)
	if err == nil {
		t.Error("validateSignatureReplay() expected error for duplicate signer/role with different signature, got nil")
	}
	if err != nil && !contains(err.Error(), "already provided a different signature") {
		t.Errorf("validateSignatureReplay() expected 'already provided a different signature' error, got: %v", err)
	}
}

// TestSignatureReplayProtection_BackwardCompatibility verifies that signatures without nonces work in compatibility mode
func TestSignatureReplayProtection_BackwardCompatibility(t *testing.T) {
	store := NewMemoryStore()
	pw := &Paywall{
		Store:     store,
		HDWallets: make(map[wallet.WalletType]wallet.HDWallet),
		prices:    map[wallet.WalletType]float64{wallet.Bitcoin: 0.001},
	}

	em, err := NewEscrowManager(pw)
	if err != nil {
		t.Fatalf("NewEscrowManager() error = %v", err)
	}

	payment := &Payment{
		ID:              "test-backward-compat",
		MultisigEnabled: true,
		EscrowState:     EscrowFunded,
		Version:         1,
	}

	// Signature without nonce (legacy format)
	legacySig := &SignatureData{
		SignerID:  "buyer-1",
		Role:      RoleBuyer,
		Signature: []byte("legacy-sig"),
		PublicKey: []byte("buyer-pubkey"),
		SignedAt:  time.Now(),
		Nonce:     nil, // No nonce
		PaymentID: "",  // No payment ID binding
	}

	// Should return error but allow in backward compatibility mode
	err = em.validateSignatureReplay(legacySig, payment)
	if err == nil {
		t.Error("validateSignatureReplay() should return error for missing nonce (backward compatibility)")
	}
	if err != nil && !contains(err.Error(), "missing nonce") {
		t.Errorf("validateSignatureReplay() expected missing nonce error, got: %v", err)
	}
}

// TestSignatureReplayProtection_MultipleSigners verifies that different signers can use different nonces
func TestSignatureReplayProtection_MultipleSigners(t *testing.T) {
	store := NewMemoryStore()
	pw := &Paywall{
		Store:     store,
		HDWallets: make(map[wallet.WalletType]wallet.HDWallet),
		prices:    map[wallet.WalletType]float64{wallet.Bitcoin: 0.001},
	}

	em, err := NewEscrowManager(pw)
	if err != nil {
		t.Fatalf("NewEscrowManager() error = %v", err)
	}

	paymentID := "test-multi-signer"

	buyerNonce := make([]byte, 32)
	if _, err := rand.Read(buyerNonce); err != nil {
		t.Fatalf("Failed to generate buyer nonce: %v", err)
	}

	payment := &Payment{
		ID:              paymentID,
		MultisigEnabled: true,
		EscrowState:     EscrowFunded,
		Version:         1,
		Signatures: map[wallet.WalletType][]SignatureData{
			wallet.Bitcoin: {
				{
					SignerID:  "buyer-1",
					Role:      RoleBuyer,
					Signature: []byte("buyer-sig"),
					PublicKey: []byte("buyer-pubkey"),
					SignedAt:  time.Now(),
					Nonce:     buyerNonce,
					PaymentID: paymentID,
				},
			},
		},
	}

	if err := store.CreatePayment(payment); err != nil {
		t.Fatalf("Failed to create payment: %v", err)
	}

	// Seller signature with different nonce (should succeed)
	sellerNonce := make([]byte, 32)
	if _, err := rand.Read(sellerNonce); err != nil {
		t.Fatalf("Failed to generate seller nonce: %v", err)
	}

	sellerSig := &SignatureData{
		SignerID:  "seller-1",
		Role:      RoleSeller,
		Signature: []byte("seller-sig"),
		PublicKey: []byte("seller-pubkey"),
		SignedAt:  time.Now(),
		Nonce:     sellerNonce, // Different nonce than buyer
		PaymentID: paymentID,
	}

	err = em.validateSignatureReplay(sellerSig, payment)
	if err != nil {
		t.Errorf("validateSignatureReplay() unexpected error for different signer: %v", err)
	}
}