engine_test.go

package metaldca

import (
	"math/bits"
	"testing"
)

// mockTargetKey is the secret key our mock device uses.
const mockTargetKey = 0x42

// mockTraceGenerator simulates a target device executing an algorithm.
// It leaks the Hamming weight of the XOR between the first byte of the payload
// and the secret target key, with some simulated noise.
func mockTraceGenerator(payload []byte) []byte {
	trace := make([]byte, 1)

	// Simulated leakage: S-box lookup or simple XOR
	val := payload[0] ^ mockTargetKey

	// Compute Hamming weight (number of set bits)
	hw := bits.OnesCount8(val)

	// Leak the Hamming weight at trace point 0
	trace[0] = byte(hw)
	
	// Add some artificial trace points to simulate a larger trace window
	for i := 1; i < 10; i++ {
		trace = append(trace, byte(payload[0]^byte(i)))
	}

	return trace
}

func TestEngineCorrelation(t *testing.T) {
	// Initialize the GPU Engine
	// In a real environment, the metallibPath would point to the compiled shader.
	// For testing, we assume the Makefile has built dca_gpu.metallib in the current dir.
	engine, err := NewEngine("./dca_gpu.metallib")
	if err != nil {
		t.Skipf("Skipping GPU test: failed to initialize Metal engine (likely running in headless environment without GPU): %v", err)
	}

	nSamples := 1000 // Number of random traces to generate
	payloadSize := 1 // We only care about recovering 1 byte for this simple test

	// Run the attack
	recoveredKey, confidentBytes, correlations := engine.Correlate(nSamples, payloadSize, mockTraceGenerator)

	if confidentBytes < 1 {
		t.Fatalf("Engine failed to confidently recover any bytes")
	}

	if recoveredKey[0] != mockTargetKey {
		t.Fatalf("Engine recovered incorrect key: got 0x%02x, expected 0x%02x. Correlation: %f", recoveredKey[0], mockTargetKey, correlations[0])
	}

	t.Logf("Successfully recovered target key 0x%02x with correlation %f", recoveredKey[0], correlations[0])
}