mev-beta/pkg/security/keymanager_private_key_test.go

package security

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

	"github.com/ethereum/go-ethereum/crypto"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func TestEnhancedClearPrivateKey(t *testing.T) {
	// Generate test key
	key, err := ecdsa.GenerateKey(crypto.S256(), rand.Reader)
	require.NoError(t, err)
	require.NotNil(t, key)
	require.NotNil(t, key.D)

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

	// Verify key has valid data before clearing
	assert.True(t, key.D.Sign() != 0)
	assert.True(t, key.PublicKey.X.Sign() != 0)
	assert.True(t, key.PublicKey.Y.Sign() != 0)

	// Clear the key
	clearPrivateKey(key)

	// Verify that the key data is effectively cleared
	assert.Nil(t, key.D, "D should be nil after clearing")
	assert.Nil(t, key.PublicKey.X, "X should be nil after clearing")
	assert.Nil(t, key.PublicKey.Y, "Y should be nil after clearing")
	assert.Nil(t, key.PublicKey.Curve, "Curve should be nil after clearing")

	// Verify original values were actually non-zero
	assert.True(t, originalD.Sign() != 0, "Original D should have been non-zero")
	assert.True(t, originalX.Sign() != 0, "Original X should have been non-zero")
	assert.True(t, originalY.Sign() != 0, "Original Y should have been non-zero")
}

func TestClearPrivateKeyNil(t *testing.T) {
	// Test that clearing a nil key doesn't panic
	clearPrivateKey(nil)
	// Should complete without error
}

func TestClearPrivateKeyPartiallyNil(t *testing.T) {
	// Test key with some nil components
	key := &ecdsa.PrivateKey{}

	// Should not panic with nil components
	clearPrivateKey(key)

	// Test with only D set
	key.D = big.NewInt(12345)
	clearPrivateKey(key)
	assert.Nil(t, key.D)
}

func TestSecureClearBigInt(t *testing.T) {
	tests := []struct {
		name  string
		value *big.Int
	}{
		{
			name:  "small positive value",
			value: big.NewInt(12345),
		},
		{
			name:  "large positive value",
			value: new(big.Int).SetBytes([]byte{0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF}),
		},
		{
			name:  "negative value",
			value: big.NewInt(-9876543210),
		},
		{
			name:  "zero value",
			value: big.NewInt(0),
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			// Create a copy to verify original was non-zero if applicable
			original := new(big.Int).Set(tt.value)

			// Clear the value
			secureClearBigInt(tt.value)

			// Verify it's cleared to zero
			assert.True(t, tt.value.Sign() == 0, "big.Int should be zero after clearing")
			assert.Equal(t, 0, tt.value.Cmp(big.NewInt(0)), "big.Int should equal zero")

			// Verify original wasn't zero (except for zero test case)
			if tt.name != "zero value" {
				assert.True(t, original.Sign() != 0, "Original value should have been non-zero")
			}
		})
	}
}

func TestSecureClearBigIntNil(t *testing.T) {
	// Test that clearing nil doesn't panic
	secureClearBigInt(nil)
	// Should complete without error
}

func TestSecureClearBytes(t *testing.T) {
	tests := []struct {
		name string
		data []byte
	}{
		{
			name: "small byte slice",
			data: []byte{0x01, 0x02, 0x03, 0x04},
		},
		{
			name: "large byte slice",
			data: make([]byte, 1024),
		},
		{
			name: "empty byte slice",
			data: []byte{},
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			// Fill with non-zero data for large slice
			if len(tt.data) > 4 {
				for i := range tt.data {
					tt.data[i] = byte(i % 256)
				}
			}

			// Store original to verify it had data
			original := make([]byte, len(tt.data))
			copy(original, tt.data)

			// Clear the data
			secureClearBytes(tt.data)

			// Verify all bytes are zero
			for i, b := range tt.data {
				assert.Equal(t, byte(0), b, "Byte at index %d should be zero", i)
			}

			// For non-empty slices, verify original had some non-zero data
			if len(original) > 0 && len(original) <= 4 {
				hasNonZero := false
				for _, b := range original {
					if b != 0 {
						hasNonZero = true
						break
					}
				}
				if len(original) > 0 && tt.name != "empty byte slice" {
					assert.True(t, hasNonZero, "Original data should have had non-zero bytes")
				}
			}
		})
	}
}

func TestMemorySecurityIntegration(t *testing.T) {
	// Test the complete workflow of key generation, usage, and clearing

	// Generate multiple keys
	keys := make([]*ecdsa.PrivateKey, 10)
	for i := range keys {
		key, err := ecdsa.GenerateKey(crypto.S256(), rand.Reader)
		require.NoError(t, err)
		keys[i] = key
	}

	// Verify all keys are valid
	for i, key := range keys {
		assert.NotNil(t, key.D, "Key %d D should not be nil", i)
		assert.True(t, key.D.Sign() != 0, "Key %d D should not be zero", i)
	}

	// Clear all keys
	for i, key := range keys {
		clearPrivateKey(key)

		// Verify clearing worked
		assert.Nil(t, key.D, "Key %d D should be nil after clearing", i)
		assert.Nil(t, key.PublicKey.X, "Key %d X should be nil after clearing", i)
		assert.Nil(t, key.PublicKey.Y, "Key %d Y should be nil after clearing", i)
	}
}

func TestConcurrentKeyClearingOperation(t *testing.T) {
	// Test concurrent clearing operations
	const numKeys = 50
	const numWorkers = 10

	keys := make([]*ecdsa.PrivateKey, numKeys)

	// Generate keys
	for i := range keys {
		key, err := ecdsa.GenerateKey(crypto.S256(), rand.Reader)
		require.NoError(t, err)
		keys[i] = key
	}

	// Channel to coordinate workers
	keysChan := make(chan *ecdsa.PrivateKey, numKeys)

	// Send keys to channel
	for _, key := range keys {
		keysChan <- key
	}
	close(keysChan)

	// Start workers to clear keys concurrently
	done := make(chan bool, numWorkers)
	for i := 0; i < numWorkers; i++ {
		go func() {
			defer func() { done <- true }()
			for key := range keysChan {
				clearPrivateKey(key)
			}
		}()
	}

	// Wait for all workers to complete
	for i := 0; i < numWorkers; i++ {
		<-done
	}

	// Verify all keys are cleared
	for i, key := range keys {
		assert.Nil(t, key.D, "Key %d D should be nil after concurrent clearing", i)
		assert.Nil(t, key.PublicKey.X, "Key %d X should be nil after concurrent clearing", i)
		assert.Nil(t, key.PublicKey.Y, "Key %d Y should be nil after concurrent clearing", i)
	}
}