mev-beta/pkg/security/keymanager_test.go

package security

import (
	"crypto/ecdsa"
	"math/big"
	"testing"
	"time"

	"github.com/ethereum/go-ethereum/common"
	"github.com/ethereum/go-ethereum/core/types"
	"github.com/ethereum/go-ethereum/crypto"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"

	"github.com/fraktal/mev-beta/internal/logger"
)

// TestNewKeyManager tests the creation of a new KeyManager
func TestNewKeyManager(t *testing.T) {
	// Test with valid configuration
	config := &KeyManagerConfig{
		KeystorePath:  "/tmp/test_keystore",
		EncryptionKey: "test_encryption_key_very_long_and_secure_for_testing",
	}

	log := logger.New("info", "text", "")
	km, err := newKeyManagerForTesting(config, log)

	require.NoError(t, err)
	assert.NotNil(t, km)
	assert.NotNil(t, km.keystore)
	assert.NotNil(t, km.keys)
	assert.NotNil(t, km.encryptionKey)
	assert.Equal(t, config, km.config)

	// Test with nil configuration (should use defaults with test encryption key)
	defaultConfig := &KeyManagerConfig{
		KeystorePath:  "/tmp/test_default_keystore",
		EncryptionKey: "default_test_encryption_key_very_long_and_secure_32chars",
	}
	km2, err := newKeyManagerForTesting(defaultConfig, log)
	require.NoError(t, err)
	assert.NotNil(t, km2)
	assert.NotNil(t, km2.config)
	assert.NotEmpty(t, km2.config.KeystorePath)
}

// TestNewKeyManagerInvalidConfig tests error cases for KeyManager creation
func TestNewKeyManagerInvalidConfig(t *testing.T) {
	log := logger.New("info", "text", "")

	// Test with empty encryption key
	config := &KeyManagerConfig{
		KeystorePath:  "/tmp/test_keystore",
		EncryptionKey: "",
	}

	km, err := newKeyManagerForTesting(config, log)
	assert.Error(t, err)
	assert.Nil(t, km)
	assert.Contains(t, err.Error(), "encryption key cannot be empty")

	// Test with short encryption key
	config = &KeyManagerConfig{
		KeystorePath:  "/tmp/test_keystore",
		EncryptionKey: "short",
	}

	km, err = newKeyManagerForTesting(config, log)
	assert.Error(t, err)
	assert.Nil(t, km)
	assert.Contains(t, err.Error(), "encryption key must be at least 32 characters")

	// Test with empty keystore path
	config = &KeyManagerConfig{
		KeystorePath:  "",
		EncryptionKey: "test_encryption_key_very_long_and_secure_for_testing",
	}

	km, err = newKeyManagerForTesting(config, log)
	assert.Error(t, err)
	assert.Nil(t, km)
	assert.Contains(t, err.Error(), "keystore path cannot be empty")
}

// TestGenerateKey tests key generation functionality
func TestGenerateKey(t *testing.T) {
	config := &KeyManagerConfig{
		KeystorePath:  "/tmp/test_keystore_generate",
		EncryptionKey: "test_encryption_key_very_long_and_secure_for_testing",
	}

	log := logger.New("info", "text", "")
	km, err := newKeyManagerForTesting(config, log)
	require.NoError(t, err)

	// Test generating a trading key
	permissions := KeyPermissions{
		CanSign:        true,
		CanTransfer:    true,
		MaxTransferWei: big.NewInt(1000000000000000000), // 1 ETH
	}

	address, err := km.GenerateKey("trading", permissions)
	require.NoError(t, err)
	assert.NotEqual(t, common.Address{}, address)

	// Verify the key exists
	keyInfo, err := km.GetKeyInfo(address)
	require.NoError(t, err)
	assert.Equal(t, address, keyInfo.Address)
	assert.Equal(t, "trading", keyInfo.KeyType)
	assert.Equal(t, permissions, keyInfo.Permissions)
	assert.WithinDuration(t, time.Now(), keyInfo.CreatedAt, time.Second)
	assert.WithinDuration(t, time.Now(), keyInfo.GetLastUsed(), time.Second)
	assert.Equal(t, int64(0), keyInfo.GetUsageCount())

	// Test generating an emergency key (should have expiration)
	emergencyAddress, err := km.GenerateKey("emergency", permissions)
	require.NoError(t, err)
	assert.NotEqual(t, common.Address{}, emergencyAddress)

	emergencyKeyInfo, err := km.GetKeyInfo(emergencyAddress)
	require.NoError(t, err)
	assert.NotNil(t, emergencyKeyInfo.ExpiresAt)
	assert.True(t, emergencyKeyInfo.ExpiresAt.After(time.Now()))
}

// TestImportKey tests key import functionality
func TestImportKey(t *testing.T) {
	config := &KeyManagerConfig{
		KeystorePath:  "/tmp/test_keystore_import",
		EncryptionKey: "test_encryption_key_very_long_and_secure_for_testing",
	}

	log := logger.New("info", "text", "")
	km, err := newKeyManagerForTesting(config, log)
	require.NoError(t, err)

	// Generate a test private key
	privateKey, err := crypto.GenerateKey()
	require.NoError(t, err)
	privateKeyHex := common.Bytes2Hex(crypto.FromECDSA(privateKey))

	// Import the key
	permissions := KeyPermissions{
		CanSign:        true,
		CanTransfer:    false,
		MaxTransferWei: nil,
	}

	address, err := km.ImportKey(privateKeyHex, "test", permissions)
	require.NoError(t, err)
	assert.NotEqual(t, common.Address{}, address)

	// Verify the imported key
	keyInfo, err := km.GetKeyInfo(address)
	require.NoError(t, err)
	assert.Equal(t, address, keyInfo.Address)
	assert.Equal(t, "test", keyInfo.KeyType)
	assert.Equal(t, permissions, keyInfo.Permissions)

	// Test importing invalid key
	_, err = km.ImportKey("invalid_private_key", "test", permissions)
	assert.Error(t, err)
	assert.Contains(t, err.Error(), "invalid private key")

	// Test importing duplicate key
	_, err = km.ImportKey(privateKeyHex, "duplicate", permissions)
	assert.Error(t, err)
	assert.Contains(t, err.Error(), "key already exists")
}

// TestListKeys tests key listing functionality
func TestListKeys(t *testing.T) {
	config := &KeyManagerConfig{
		KeystorePath:  "/tmp/test_keystore_list",
		EncryptionKey: "test_encryption_key_very_long_and_secure_for_testing",
	}

	log := logger.New("info", "text", "")
	km, err := newKeyManagerForTesting(config, log)
	require.NoError(t, err)

	// Initially should be empty
	keys := km.ListKeys()
	assert.Empty(t, keys)

	// Generate a few keys
	permissions := KeyPermissions{CanSign: true}
	addr1, err := km.GenerateKey("test1", permissions)
	require.NoError(t, err)

	addr2, err := km.GenerateKey("test2", permissions)
	require.NoError(t, err)

	// Check that both keys are listed
	keys = km.ListKeys()
	assert.Len(t, keys, 2)
	assert.Contains(t, keys, addr1)
	assert.Contains(t, keys, addr2)
}

// TestGetKeyInfo tests key information retrieval
func TestGetKeyInfo(t *testing.T) {
	config := &KeyManagerConfig{
		KeystorePath:  "/tmp/test_keystore_info",
		EncryptionKey: "test_encryption_key_very_long_and_secure_for_testing",
	}

	log := logger.New("info", "text", "")
	km, err := newKeyManagerForTesting(config, log)
	require.NoError(t, err)

	// Generate a key
	permissions := KeyPermissions{CanSign: true, CanTransfer: true}
	address, err := km.GenerateKey("test", permissions)
	require.NoError(t, err)

	// Get key info
	keyInfo, err := km.GetKeyInfo(address)
	require.NoError(t, err)
	assert.Equal(t, address, keyInfo.Address)
	assert.Equal(t, "test", keyInfo.KeyType)
	assert.Equal(t, permissions, keyInfo.Permissions)
	// EncryptedKey should be nil in the returned info for security
	assert.Nil(t, keyInfo.EncryptedKey)

	// Test getting non-existent key
	nonExistentAddr := common.HexToAddress("0x1234567890123456789012345678901234567890")
	_, err = km.GetKeyInfo(nonExistentAddr)
	assert.Error(t, err)
	assert.Contains(t, err.Error(), "key not found")
}

// TestEncryptDecryptPrivateKey tests the encryption/decryption functionality
func TestEncryptDecryptPrivateKey(t *testing.T) {
	config := &KeyManagerConfig{
		KeystorePath:  "/tmp/test_keystore_encrypt",
		EncryptionKey: "test_encryption_key_very_long_and_secure_for_testing",
	}

	log := logger.New("info", "text", "")
	km, err := newKeyManagerForTesting(config, log)
	require.NoError(t, err)

	// Generate a test private key
	privateKey, err := crypto.GenerateKey()
	require.NoError(t, err)

	// Test encryption
	encryptedKey, err := km.encryptPrivateKey(privateKey)
	require.NoError(t, err)
	assert.NotNil(t, encryptedKey)
	assert.NotEmpty(t, encryptedKey)

	// Test decryption
	decryptedKey, err := km.decryptPrivateKey(encryptedKey)
	require.NoError(t, err)
	assert.NotNil(t, decryptedKey)

	// Verify the keys are the same
	assert.Equal(t, crypto.PubkeyToAddress(privateKey.PublicKey), crypto.PubkeyToAddress(decryptedKey.PublicKey))
	assert.Equal(t, crypto.FromECDSA(privateKey), crypto.FromECDSA(decryptedKey))

	// Test decryption with invalid data
	_, err = km.decryptPrivateKey([]byte("x")) // Very short data to trigger "encrypted key too short"
	assert.Error(t, err)
	assert.Contains(t, err.Error(), "encrypted key too short")
}

// TestRotateKey tests key rotation functionality
func TestRotateKey(t *testing.T) {
	config := &KeyManagerConfig{
		KeystorePath:  "/tmp/test_keystore_rotate",
		EncryptionKey: "test_encryption_key_very_long_and_secure_for_testing",
	}

	log := logger.New("info", "text", "")
	km, err := newKeyManagerForTesting(config, log)
	require.NoError(t, err)

	// Generate an original key
	permissions := KeyPermissions{CanSign: true, CanTransfer: true}
	originalAddr, err := km.GenerateKey("test", permissions)
	require.NoError(t, err)

	// Rotate the key
	newAddr, err := km.RotateKey(originalAddr)
	require.NoError(t, err)
	assert.NotEqual(t, originalAddr, newAddr)

	// Check that the original key still exists but has restricted permissions
	originalInfo, err := km.GetKeyInfo(originalAddr)
	require.NoError(t, err)
	assert.False(t, originalInfo.Permissions.CanSign)
	assert.False(t, originalInfo.Permissions.CanTransfer)

	// Check that the new key has the same permissions
	newInfo, err := km.GetKeyInfo(newAddr)
	require.NoError(t, err)
	assert.Equal(t, permissions, newInfo.Permissions)
	assert.True(t, newInfo.Permissions.CanSign)
	assert.True(t, newInfo.Permissions.CanTransfer)

	// Test rotating non-existent key
	nonExistentAddr := common.HexToAddress("0x1234567890123456789012345678901234567890")
	_, err = km.RotateKey(nonExistentAddr)
	assert.Error(t, err)
	assert.Contains(t, err.Error(), "key not found")
}

// TestSignTransaction tests transaction signing with various scenarios
func TestSignTransaction(t *testing.T) {
	config := &KeyManagerConfig{
		KeystorePath:  "/tmp/test_keystore_sign",
		EncryptionKey: "test_encryption_key_very_long_and_secure_for_testing",
	}

	log := logger.New("info", "text", "")
	km, err := newKeyManagerForTesting(config, log)
	require.NoError(t, err)

	// Generate a key with signing permissions
	permissions := KeyPermissions{
		CanSign:        true,
		CanTransfer:    true,
		MaxTransferWei: big.NewInt(1000000000000000000), // 1 ETH (safe int64 value)
	}
	signerAddr, err := km.GenerateKey("signer", permissions)
	require.NoError(t, err)

	// Create a test transaction using Arbitrum chain ID (EIP-155 transaction)
	chainID := big.NewInt(42161) // Arbitrum One

	// Create transaction data for EIP-155 transaction
	toAddr := common.HexToAddress("0x1234567890123456789012345678901234567890")
	value := big.NewInt(1000000000000000000) // 1 ETH
	gasLimit := uint64(21000)
	gasPrice := big.NewInt(20000000000) // 20 Gwei
	nonce := uint64(0)

	// Create DynamicFeeTx (EIP-1559) which properly handles chain ID
	tx := types.NewTx(&types.DynamicFeeTx{
		ChainID:   chainID,
		Nonce:     nonce,
		To:        &toAddr,
		Value:     value,
		Gas:       gasLimit,
		GasFeeCap: gasPrice,
		GasTipCap: big.NewInt(1000000000), // 1 Gwei tip
		Data:      nil,
	})

	// Create signing request
	request := &SigningRequest{
		Transaction:  tx,
		ChainID:      chainID,
		From:         signerAddr,
		Purpose:      "Test transaction",
		UrgencyLevel: 3,
	}

	// Sign the transaction
	result, err := km.SignTransaction(request)
	require.NoError(t, err)
	assert.NotNil(t, result)
	assert.NotNil(t, result.SignedTx)
	assert.NotNil(t, result.Signature)
	assert.NotEmpty(t, result.AuditID)
	assert.WithinDuration(t, time.Now(), result.SignedAt, time.Second)
	assert.Equal(t, signerAddr, result.KeyUsed)

	// Verify the signature is valid
	signedTx := result.SignedTx
	// Use appropriate signer based on transaction type
	var signer types.Signer
	switch signedTx.Type() {
	case types.LegacyTxType:
		signer = types.NewEIP155Signer(chainID)
	case types.DynamicFeeTxType:
		signer = types.NewLondonSigner(chainID)
	default:
		t.Fatalf("Unsupported transaction type: %d", signedTx.Type())
	}
	from, err := types.Sender(signer, signedTx)
	require.NoError(t, err)
	assert.Equal(t, signerAddr, from)

	// Test signing with non-existent key
	nonExistentAddr := common.HexToAddress("0x1234567890123456789012345678901234567890")
	request.From = nonExistentAddr
	_, err = km.SignTransaction(request)
	assert.Error(t, err)
	assert.Contains(t, err.Error(), "key not found")

	// Test signing with key that can't sign
	km2, err := newKeyManagerForTesting(config, log)
	require.NoError(t, err)

	noSignPermissions := KeyPermissions{
		CanSign:        false,
		CanTransfer:    true,
		MaxTransferWei: big.NewInt(1000000000000000000), // 1 ETH (safe int64 value)
	}
	noSignAddr, err := km2.GenerateKey("no_sign", noSignPermissions)
	require.NoError(t, err)

	request.From = noSignAddr
	_, err = km2.SignTransaction(request)
	assert.Error(t, err)
	assert.Contains(t, err.Error(), "not permitted to sign")
}

// TestSignTransactionTransferLimits tests transfer limits during signing
func TestSignTransactionTransferLimits(t *testing.T) {
	config := &KeyManagerConfig{
		KeystorePath:  "/tmp/test_keystore_limits",
		EncryptionKey: "test_encryption_key_very_long_and_secure_for_testing",
	}

	log := logger.New("info", "text", "")
	km, err := newKeyManagerForTesting(config, log)
	require.NoError(t, err)

	// Generate a key with limited transfer permissions
	maxTransfer := big.NewInt(1000000000000000000) // 1 ETH
	permissions := KeyPermissions{
		CanSign:        true,
		CanTransfer:    true,
		MaxTransferWei: maxTransfer,
	}
	signerAddr, err := km.GenerateKey("limited_signer", permissions)
	require.NoError(t, err)

	// Create a transaction that exceeds the limit
	chainID := big.NewInt(1)
	excessiveTx := types.NewTransaction(0, common.Address{}, big.NewInt(2000000000000000000), 21000, big.NewInt(20000000000), nil) // 2 ETH

	request := &SigningRequest{
		Transaction:  excessiveTx,
		ChainID:      chainID,
		From:         signerAddr,
		Purpose:      "Excessive transfer",
		UrgencyLevel: 3,
	}

	_, err = km.SignTransaction(request)
	assert.Error(t, err)
	assert.Contains(t, err.Error(), "transfer amount exceeds limit")
}

// TestDeriveEncryptionKey tests the key derivation function
func TestDeriveEncryptionKey(t *testing.T) {
	// Test with valid master key
	masterKey := "test_encryption_key_very_long_and_secure_for_testing"
	key, err := deriveEncryptionKey(masterKey)
	require.NoError(t, err)
	assert.NotNil(t, key)
	assert.Len(t, key, 32) // Should be 32 bytes for AES-256

	// Test with different master key (should produce different result)
	differentKey := "different_test_encryption_key_very_long_and_secure_for_testing"
	key2, err := deriveEncryptionKey(differentKey)
	require.NoError(t, err)
	assert.NotEqual(t, key, key2)

	// Test with empty master key
	_, err = deriveEncryptionKey("")
	assert.Error(t, err)
}

// TestValidateConfig tests the configuration validation function
func TestValidateConfig(t *testing.T) {
	// Test with valid config
	validConfig := &KeyManagerConfig{
		KeystorePath:  "/tmp/test",
		EncryptionKey: "test_encryption_key_very_long_and_secure_for_testing",
	}
	err := validateConfig(validConfig)
	assert.NoError(t, err)

	// Test with empty encryption key
	emptyKeyConfig := &KeyManagerConfig{
		KeystorePath:  "/tmp/test",
		EncryptionKey: "",
	}
	err = validateConfig(emptyKeyConfig)
	assert.Error(t, err)
	assert.Contains(t, err.Error(), "encryption key cannot be empty")

	// Test with short encryption key
	shortKeyConfig := &KeyManagerConfig{
		KeystorePath:  "/tmp/test",
		EncryptionKey: "short",
	}
	err = validateConfig(shortKeyConfig)
	assert.Error(t, err)
	assert.Contains(t, err.Error(), "encryption key must be at least 32 characters")

	// Test with empty keystore path
	emptyPathConfig := &KeyManagerConfig{
		KeystorePath:  "",
		EncryptionKey: "test_encryption_key_very_long_and_secure_for_testing",
	}
	err = validateConfig(emptyPathConfig)
	assert.Error(t, err)
	assert.Contains(t, err.Error(), "keystore path cannot be empty")
}

// TestClearPrivateKey tests the private key clearing function
func TestClearPrivateKey(t *testing.T) {
	// Generate a test private key
	privateKey, err := crypto.GenerateKey()
	require.NoError(t, err)

	// Store original D value for comparison
	originalD := new(big.Int).Set(privateKey.D)

	// Clear the private key
	clearPrivateKey(privateKey)

	// Verify the D value has been cleared (nil or zero)
	if privateKey.D != nil {
		assert.Zero(t, privateKey.D.Sign())
	} else {
		assert.Nil(t, privateKey.D)
	}
	assert.NotEqual(t, originalD, privateKey.D)

	// Test with nil private key (should not panic)
	clearPrivateKey(nil)
}

// TestGenerateAuditID tests the audit ID generation function
func TestGenerateAuditID(t *testing.T) {
	id1 := generateAuditID()
	id2 := generateAuditID()

	// Both should be non-empty and different
	assert.NotEmpty(t, id1)
	assert.NotEmpty(t, id2)
	assert.NotEqual(t, id1, id2)

	// Should be a valid hex string
	hash1 := common.HexToHash(id1)
	assert.NotEqual(t, hash1, common.Hash{})

	hash2 := common.HexToHash(id2)
	assert.NotEqual(t, hash2, common.Hash{})
}

// TestCalculateRiskScore tests the risk score calculation function
func TestCalculateRiskScore(t *testing.T) {
	// Test failed operations (high risk)
	score := calculateRiskScore("TRANSACTION_SIGNED", false)
	assert.Equal(t, 8, score)

	// Test successful transaction signing (low-medium risk)
	score = calculateRiskScore("TRANSACTION_SIGNED", true)
	assert.Equal(t, 3, score)

	// Test key generation (medium risk)
	score = calculateRiskScore("KEY_GENERATED", true)
	assert.Equal(t, 5, score)

	// Test key import (medium risk)
	score = calculateRiskScore("KEY_IMPORTED", true)
	assert.Equal(t, 5, score)

	// Test key rotation (medium risk)
	score = calculateRiskScore("KEY_ROTATED", true)
	assert.Equal(t, 4, score)

	// Test default (low risk)
	score = calculateRiskScore("UNKNOWN_OPERATION", true)
	assert.Equal(t, 2, score)
}

// TestKeyPermissions tests the KeyPermissions struct
func TestKeyPermissions(t *testing.T) {
	// Test creating permissions with max transfer limit
	maxTransfer := big.NewInt(1000000000000000000) // 1 ETH
	permissions := KeyPermissions{
		CanSign:        true,
		CanTransfer:    true,
		MaxTransferWei: maxTransfer,
		AllowedContracts: []string{
			"0x1234567890123456789012345678901234567890",
			"0x0987654321098765432109876543210987654321",
		},
		RequireConfirm: true,
	}

	assert.True(t, permissions.CanSign)
	assert.True(t, permissions.CanTransfer)
	assert.Equal(t, maxTransfer, permissions.MaxTransferWei)
	assert.Len(t, permissions.AllowedContracts, 2)
	assert.True(t, permissions.RequireConfirm)
}

// BenchmarkKeyGeneration benchmarks key generation performance
func BenchmarkKeyGeneration(b *testing.B) {
	config := &KeyManagerConfig{
		KeystorePath:  "/tmp/benchmark_keystore",
		EncryptionKey: "test_encryption_key_very_long_and_secure_for_testing",
	}

	log := logger.New("info", "text", "")
	km, err := newKeyManagerForTesting(config, log)
	require.NoError(b, err)

	permissions := KeyPermissions{CanSign: true}

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		_, err := km.GenerateKey("benchmark", permissions)
		if err != nil {
			b.Fatal(err)
		}
	}
}

// BenchmarkTransactionSigning benchmarks transaction signing performance
func BenchmarkTransactionSigning(b *testing.B) {
	config := &KeyManagerConfig{
		KeystorePath:  "/tmp/benchmark_signing",
		EncryptionKey: "test_encryption_key_very_long_and_secure_for_testing",
	}

	log := logger.New("info", "text", "")
	km, err := newKeyManagerForTesting(config, log)
	require.NoError(b, err)

	permissions := KeyPermissions{CanSign: true, CanTransfer: true}
	signerAddr, err := km.GenerateKey("benchmark_signer", permissions)
	require.NoError(b, err)

	chainID := big.NewInt(1)
	tx := types.NewTransaction(0, common.Address{}, big.NewInt(1000000000000000000), 21000, big.NewInt(20000000000), nil)

	request := &SigningRequest{
		Transaction: tx,
		ChainID:     chainID,
		From:        signerAddr,
		Purpose:     "Benchmark transaction",
	}

	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		_, err := km.SignTransaction(request)
		if err != nil {
			b.Fatal(err)
		}
	}
}

// ENHANCED: Unit tests for memory clearing verification
func TestMemoryClearing(t *testing.T) {
	t.Run("TestSecureClearBigInt", func(t *testing.T) {
		// Create a big.Int with sensitive data
		sensitiveValue := big.NewInt(0)
		sensitiveValue.SetString("123456789012345678901234567890123456789012345678901234567890", 10)

		// Capture the original bits for verification
		originalBits := make([]big.Word, len(sensitiveValue.Bits()))
		copy(originalBits, sensitiveValue.Bits())

		// Ensure we have actual data to clear
		require.True(t, len(originalBits) > 0, "Test requires non-zero big.Int")

		// Clear the sensitive value
		secureClearBigInt(sensitiveValue)

		// Verify all bits are zeroed
		clearedBits := sensitiveValue.Bits()
		for i, bit := range clearedBits {
			assert.Equal(t, big.Word(0), bit, "Bit %d should be zero after clearing", i)
		}

		// Verify the value is actually zero
		assert.True(t, sensitiveValue.Cmp(big.NewInt(0)) == 0, "BigInt should be zero after clearing")
	})

	t.Run("TestSecureClearBytes", func(t *testing.T) {
		// Create sensitive byte data
		sensitiveData := []byte("This is very sensitive private key data that should be cleared")
		originalData := make([]byte, len(sensitiveData))
		copy(originalData, sensitiveData)

		// Verify we have data to clear
		require.True(t, len(sensitiveData) > 0, "Test requires non-empty byte slice")

		// Clear the sensitive data
		secureClearBytes(sensitiveData)

		// Verify all bytes are zeroed
		for i, b := range sensitiveData {
			assert.Equal(t, byte(0), b, "Byte %d should be zero after clearing", i)
		}

		// Verify the data was actually changed
		assert.NotEqual(t, originalData, sensitiveData, "Data should be different after clearing")
	})

	t.Run("TestClearPrivateKey", func(t *testing.T) {
		// Generate a test private key
		privateKey, err := crypto.GenerateKey()
		require.NoError(t, err)

		// Store original values for verification
		originalD := new(big.Int).Set(privateKey.D)
		originalX := new(big.Int).Set(privateKey.PublicKey.X)
		originalY := new(big.Int).Set(privateKey.PublicKey.Y)

		// Verify we have actual key material
		require.True(t, originalD.Cmp(big.NewInt(0)) != 0, "Private key D should not be zero")
		require.True(t, originalX.Cmp(big.NewInt(0)) != 0, "Public key X should not be zero")
		require.True(t, originalY.Cmp(big.NewInt(0)) != 0, "Public key Y should not be zero")

		// Clear the private key
		clearPrivateKey(privateKey)

		// Verify all components are nil or zero
		assert.Nil(t, privateKey.D, "Private key D should be nil after clearing")
		assert.Nil(t, privateKey.PublicKey.X, "Public key X should be nil after clearing")
		assert.Nil(t, privateKey.PublicKey.Y, "Public key Y should be nil after clearing")
		assert.Nil(t, privateKey.PublicKey.Curve, "Curve should be nil after clearing")
	})
}

// ENHANCED: Test memory usage monitoring
func TestKeyMemoryMetrics(t *testing.T) {
	config := &KeyManagerConfig{
		KeystorePath:      "/tmp/test_keystore_metrics",
		EncryptionKey:     "test_encryption_key_very_long_and_secure_for_testing",
		BackupEnabled:     false,
		MaxFailedAttempts: 3,
		LockoutDuration:   5 * time.Minute,
	}

	log := logger.New("info", "text", "")
	km, err := newKeyManagerForTesting(config, log)
	require.NoError(t, err)

	// Get initial metrics
	initialMetrics := km.GetMemoryMetrics()
	assert.NotNil(t, initialMetrics)
	assert.Equal(t, 0, initialMetrics.ActiveKeys)
	assert.Greater(t, initialMetrics.MemoryUsageBytes, int64(0))

	// Generate some keys
	permissions := KeyPermissions{
		CanSign:        true,
		CanTransfer:    true,
		MaxTransferWei: big.NewInt(1000000000000000000),
	}

	addr1, err := km.GenerateKey("test", permissions)
	require.NoError(t, err)

	// Check metrics after adding a key
	metricsAfterKey := km.GetMemoryMetrics()
	assert.Equal(t, 1, metricsAfterKey.ActiveKeys)

	// Test memory protection wrapper
	err = withMemoryProtection(func() error {
		_, err := km.GenerateKey("test2", permissions)
		return err
	})
	require.NoError(t, err)

	// Check final metrics
	finalMetrics := km.GetMemoryMetrics()
	assert.Equal(t, 2, finalMetrics.ActiveKeys)

	// Note: No cleanup method available, keys remain for test duration
	_ = addr1 // Silence unused variable warning
}

// ENHANCED: Benchmark memory clearing performance
func BenchmarkMemoryClearing(b *testing.B) {
	b.Run("BenchmarkSecureClearBigInt", func(b *testing.B) {
		// Create test big.Int values
		values := make([]*big.Int, b.N)
		for i := 0; i < b.N; i++ {
			values[i] = big.NewInt(0)
			values[i].SetString("123456789012345678901234567890123456789012345678901234567890", 10)
		}

		b.ResetTimer()
		for i := 0; i < b.N; i++ {
			secureClearBigInt(values[i])
		}
	})

	b.Run("BenchmarkSecureClearBytes", func(b *testing.B) {
		// Create test byte slices
		testData := make([][]byte, b.N)
		for i := 0; i < b.N; i++ {
			testData[i] = make([]byte, 64) // 64 bytes like a private key
			for j := range testData[i] {
				testData[i][j] = byte(j % 256)
			}
		}

		b.ResetTimer()
		for i := 0; i < b.N; i++ {
			secureClearBytes(testData[i])
		}
	})

	b.Run("BenchmarkClearPrivateKey", func(b *testing.B) {
		// Generate test private keys
		keys := make([]*ecdsa.PrivateKey, b.N)
		for i := 0; i < b.N; i++ {
			key, err := crypto.GenerateKey()
			if err != nil {
				b.Fatal(err)
			}
			keys[i] = key
		}

		b.ResetTimer()
		for i := 0; i < b.N; i++ {
			clearPrivateKey(keys[i])
		}
	})
}