mev-beta/pkg/security/security_test.go

package security

import (
	"encoding/json"
	"fmt"
	"math/big"
	"strings"
	"testing"
	"time"

	"github.com/ethereum/go-ethereum/common"
	"github.com/ethereum/go-ethereum/core/types"

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

// newTestLogger creates a simple test logger
func newTestLogger() *logger.Logger {
	return logger.New("info", "text", "")
}

// FuzzRPCResponseParser tests RPC response parsing with malformed inputs
func FuzzRPCResponseParser(f *testing.F) {
	// Add seed corpus with valid RPC responses
	validResponses := []string{
		`{"jsonrpc":"2.0","id":1,"result":"0x1"}`,
		`{"jsonrpc":"2.0","id":2,"result":{"blockNumber":"0x1b4","hash":"0x..."}}`,
		`{"jsonrpc":"2.0","id":3,"error":{"code":-32000,"message":"insufficient funds"}}`,
		`{"jsonrpc":"2.0","id":4,"result":null}`,
		`{"jsonrpc":"2.0","id":5,"result":[]}`,
	}

	for _, response := range validResponses {
		f.Add([]byte(response))
	}

	f.Fuzz(func(t *testing.T, data []byte) {
		// Test that RPC response parsing doesn't panic
		defer func() {
			if r := recover(); r != nil {
				t.Errorf("Panic on RPC input: %v\nInput: %q", r, string(data))
			}
		}()

		// Test JSON parsing
		var result interface{}
		_ = json.Unmarshal(data, &result)

		// Test with InputValidator
		validator := NewInputValidator(42161) // Arbitrum chain ID
		_ = validator.ValidateRPCResponse(data)
	})
}

// FuzzTransactionSigning tests transaction signing with various inputs
func FuzzTransactionSigning(f *testing.F) {
	// Setup key manager for testing
	config := &KeyManagerConfig{
		KeystorePath:    "test_keystore",
		EncryptionKey:   "test_encryption_key_for_fuzzing_32chars",
		SessionTimeout:  time.Hour,
		AuditLogPath:    "",
		MaxSigningRate:  1000,
		KeyRotationDays: 30,
	}

	testLogger := newTestLogger()
	km, err := newKeyManagerForTesting(config, testLogger)
	if err != nil {
		f.Skip("Failed to create key manager for fuzzing")
	}

	// Generate test key
	testKeyAddr, err := km.GenerateKey("test", KeyPermissions{
		CanSign:     true,
		CanTransfer: true,
	})
	if err != nil {
		f.Skip("Failed to generate test key")
	}

	// Seed corpus with valid transaction data
	validTxData := [][]byte{
		{0x02}, // EIP-1559 transaction type
		{0x01}, // EIP-2930 transaction type
		{0x00}, // Legacy transaction type
	}

	for _, data := range validTxData {
		f.Add(data)
	}

	f.Fuzz(func(t *testing.T, data []byte) {
		defer func() {
			if r := recover(); r != nil {
				t.Errorf("Panic in transaction signing: %v\nInput: %x", r, data)
			}
		}()

		// Try to create transaction from fuzzed data
		if len(data) == 0 {
			return
		}

		// Create a basic transaction for signing tests
		tx := types.NewTransaction(
			0,                               // nonce
			common.HexToAddress("0x1234"),   // to
			big.NewInt(1000000000000000000), // value (1 ETH)
			21000,                           // gas limit
			big.NewInt(20000000000),         // gas price (20 gwei)
			data,                            // data
		)

		// Test signing
		request := &SigningRequest{
			Transaction:  tx,
			From:         testKeyAddr,
			Purpose:      "fuzz_test",
			ChainID:      big.NewInt(42161),
			UrgencyLevel: 1,
		}

		_, _ = km.SignTransaction(request)
	})
}

// FuzzKeyValidation tests key validation with various encryption keys
func FuzzKeyValidation(f *testing.F) {
	// Seed with common weak keys
	weakKeys := []string{
		"test123",
		"password",
		"12345678901234567890123456789012",
		"aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa",
		"test_encryption_key_default_config",
	}

	for _, key := range weakKeys {
		f.Add(key)
	}

	f.Fuzz(func(t *testing.T, encryptionKey string) {
		defer func() {
			if r := recover(); r != nil {
				t.Errorf("Panic in key validation: %v\nKey: %q", r, encryptionKey)
			}
		}()

		config := &KeyManagerConfig{
			EncryptionKey: encryptionKey,
			KeystorePath:  "test_keystore",
		}

		// This should not panic, even with invalid keys
		err := validateProductionConfig(config)

		// Check for expected security rejections
		if strings.Contains(strings.ToLower(encryptionKey), "test") ||
			strings.Contains(strings.ToLower(encryptionKey), "default") ||
			len(encryptionKey) < 32 {
			if err == nil {
				t.Errorf("Expected validation error for weak key: %q", encryptionKey)
			}
		}
	})
}

// FuzzInputValidator tests input validation with malicious inputs
func FuzzInputValidator(f *testing.F) {
	validator := NewInputValidator(42161)

	// Seed with various address formats
	addresses := []string{
		"0x1234567890123456789012345678901234567890",
		"0x0000000000000000000000000000000000000000",
		"0xffffffffffffffffffffffffffffffffffffffff",
		"0x",
		"",
		"not_an_address",
	}

	for _, addr := range addresses {
		f.Add(addr)
	}

	f.Fuzz(func(t *testing.T, addressStr string) {
		defer func() {
			if r := recover(); r != nil {
				t.Errorf("Panic in address validation: %v\nAddress: %q", r, addressStr)
			}
		}()

		// Test RPC response validation
		rpcData := []byte(fmt.Sprintf(`{"jsonrpc":"2.0","id":1,"result":"%s"}`, addressStr))
		_ = validator.ValidateRPCResponse(rpcData)

		// Test amount validation if it looks like a number
		if len(addressStr) > 0 && addressStr[0] >= '0' && addressStr[0] <= '9' {
			amount := new(big.Int)
			amount.SetString(addressStr, 10)
			// Test basic amount validation logic
			if amount.Sign() < 0 {
				// Negative amounts should be rejected
			}
		}
	})
}

// TestConcurrentKeyAccess tests concurrent access to key manager
func TestConcurrentKeyAccess(t *testing.T) {
	config := &KeyManagerConfig{
		KeystorePath:    "test_concurrent_keystore",
		EncryptionKey:   "concurrent_test_encryption_key_32c",
		SessionTimeout:  time.Hour,
		MaxSigningRate:  1000,
		KeyRotationDays: 30,
	}

	testLogger := newTestLogger()
	km, err := newKeyManagerForTesting(config, testLogger)
	if err != nil {
		t.Fatalf("Failed to create key manager: %v", err)
	}

	// Generate test key
	testKeyAddr, err := km.GenerateKey("concurrent_test", KeyPermissions{
		CanSign:     true,
		CanTransfer: true,
	})
	if err != nil {
		t.Fatalf("Failed to generate test key: %v", err)
	}

	// Test concurrent signing
	const numGoroutines = 100
	const signingsPerGoroutine = 10

	results := make(chan error, numGoroutines*signingsPerGoroutine)

	for i := 0; i < numGoroutines; i++ {
		go func(workerID int) {
			for j := 0; j < signingsPerGoroutine; j++ {
				tx := types.NewTransaction(
					uint64(workerID*signingsPerGoroutine+j),
					common.HexToAddress("0x1234"),
					big.NewInt(1000000000000000000),
					21000,
					big.NewInt(20000000000),
					[]byte(fmt.Sprintf("worker_%d_tx_%d", workerID, j)),
				)

				request := &SigningRequest{
					Transaction:  tx,
					From:         testKeyAddr,
					Purpose:      fmt.Sprintf("concurrent_test_%d_%d", workerID, j),
					ChainID:      big.NewInt(42161),
					UrgencyLevel: 1,
				}

				_, err := km.SignTransaction(request)
				results <- err
			}
		}(i)
	}

	// Collect results
	for i := 0; i < numGoroutines*signingsPerGoroutine; i++ {
		if err := <-results; err != nil {
			t.Errorf("Concurrent signing failed: %v", err)
		}
	}
}

// TestSecurityMetrics tests security metrics collection
func TestSecurityMetrics(t *testing.T) {
	validator := NewInputValidator(42161)

	// Test metrics for various validation scenarios
	testCases := []struct {
		name        string
		testFunc    func() error
		expectError bool
	}{
		{
			name: "valid_rpc_response",
			testFunc: func() error {
				return validator.ValidateRPCResponse([]byte(`{"jsonrpc":"2.0","id":1,"result":"0x1"}`))
			},
			expectError: false,
		},
		{
			name: "invalid_rpc_response",
			testFunc: func() error {
				return validator.ValidateRPCResponse([]byte(`invalid json`))
			},
			expectError: true,
		},
		{
			name: "empty_rpc_response",
			testFunc: func() error {
				return validator.ValidateRPCResponse([]byte{})
			},
			expectError: true,
		},
		{
			name: "oversized_rpc_response",
			testFunc: func() error {
				largeData := make([]byte, 11*1024*1024) // 11MB
				return validator.ValidateRPCResponse(largeData)
			},
			expectError: true,
		},
	}

	for _, tc := range testCases {
		t.Run(tc.name, func(t *testing.T) {
			err := tc.testFunc()

			if tc.expectError && err == nil {
				t.Errorf("Expected error for %s, but got none", tc.name)
			}

			if !tc.expectError && err != nil {
				t.Errorf("Unexpected error for %s: %v", tc.name, err)
			}
		})
	}
}

// BenchmarkSecurityOperations benchmarks critical security operations
func BenchmarkSecurityOperations(b *testing.B) {
	config := &KeyManagerConfig{
		KeystorePath:    "benchmark_keystore",
		EncryptionKey:   "benchmark_encryption_key_32chars",
		SessionTimeout:  time.Hour,
		MaxSigningRate:  10000,
		KeyRotationDays: 30,
	}

	testLogger := newTestLogger()
	km, err := newKeyManagerForTesting(config, testLogger)
	if err != nil {
		b.Fatalf("Failed to create key manager: %v", err)
	}

	testKeyAddr, err := km.GenerateKey("benchmark_test", KeyPermissions{
		CanSign:     true,
		CanTransfer: true,
	})
	if err != nil {
		b.Fatalf("Failed to generate test key: %v", err)
	}

	tx := types.NewTransaction(
		0,
		common.HexToAddress("0x1234"),
		big.NewInt(1000000000000000000),
		21000,
		big.NewInt(20000000000),
		[]byte("benchmark_data"),
	)

	b.Run("SignTransaction", func(b *testing.B) {
		for i := 0; i < b.N; i++ {
			request := &SigningRequest{
				Transaction:  tx,
				From:         testKeyAddr,
				Purpose:      fmt.Sprintf("benchmark_%d", i),
				ChainID:      big.NewInt(42161),
				UrgencyLevel: 1,
			}

			_, err := km.SignTransaction(request)
			if err != nil {
				b.Fatalf("Signing failed: %v", err)
			}
		}
	})

	validator := NewInputValidator(42161)
	b.Run("ValidateRPCResponse", func(b *testing.B) {
		testData := []byte(`{"jsonrpc":"2.0","id":1,"result":"0x1"}`)
		for i := 0; i < b.N; i++ {
			_ = validator.ValidateRPCResponse(testData)
		}
	})
}

// Additional helper for RPC response validation
func (iv *InputValidator) ValidateRPCResponse(data []byte) error {
	if len(data) == 0 {
		return fmt.Errorf("empty RPC response")
	}

	if len(data) > 10*1024*1024 { // 10MB limit
		return fmt.Errorf("RPC response too large: %d bytes", len(data))
	}

	// Check for valid JSON
	var result interface{}
	if err := json.Unmarshal(data, &result); err != nil {
		return fmt.Errorf("invalid JSON in RPC response: %w", err)
	}

	// Check for common RPC response structure
	if resultMap, ok := result.(map[string]interface{}); ok {
		if jsonrpc, exists := resultMap["jsonrpc"]; exists {
			if jsonrpcStr, ok := jsonrpc.(string); !ok || jsonrpcStr != "2.0" {
				return fmt.Errorf("invalid JSON-RPC version")
			}
		}
	}

	return nil
}