rwadurian/backend/mpc-system/tests/unit/pkg/crypto_test.go

package pkg_test

import (
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rand"
	"testing"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"

	"github.com/rwadurian/mpc-system/pkg/crypto"
)

func TestGenerateRandomBytes(t *testing.T) {
	t.Run("should generate random bytes of correct length", func(t *testing.T) {
		lengths := []int{16, 32, 64, 128}

		for _, length := range lengths {
			bytes, err := crypto.GenerateRandomBytes(length)
			require.NoError(t, err)
			assert.Len(t, bytes, length)
		}
	})

	t.Run("should generate different bytes each time", func(t *testing.T) {
		bytes1, _ := crypto.GenerateRandomBytes(32)
		bytes2, _ := crypto.GenerateRandomBytes(32)

		assert.NotEqual(t, bytes1, bytes2)
	})
}

func TestHashMessage(t *testing.T) {
	t.Run("should hash message consistently", func(t *testing.T) {
		message := []byte("test message")

		hash1 := crypto.HashMessage(message)
		hash2 := crypto.HashMessage(message)

		assert.Equal(t, hash1, hash2)
		assert.Len(t, hash1, 32) // SHA-256 produces 32 bytes
	})

	t.Run("should produce different hashes for different messages", func(t *testing.T) {
		hash1 := crypto.HashMessage([]byte("message1"))
		hash2 := crypto.HashMessage([]byte("message2"))

		assert.NotEqual(t, hash1, hash2)
	})
}

func TestEncryptDecrypt(t *testing.T) {
	t.Run("should encrypt and decrypt data successfully", func(t *testing.T) {
		key := make([]byte, 32)
		rand.Read(key)
		plaintext := []byte("secret data to encrypt")

		ciphertext, err := crypto.Encrypt(key, plaintext)
		require.NoError(t, err)
		assert.NotEqual(t, plaintext, ciphertext)

		decrypted, err := crypto.Decrypt(key, ciphertext)
		require.NoError(t, err)
		assert.Equal(t, plaintext, decrypted)
	})

	t.Run("should fail decryption with wrong key", func(t *testing.T) {
		key1 := make([]byte, 32)
		key2 := make([]byte, 32)
		rand.Read(key1)
		rand.Read(key2)

		plaintext := []byte("secret data")
		ciphertext, _ := crypto.Encrypt(key1, plaintext)

		_, err := crypto.Decrypt(key2, ciphertext)
		assert.Error(t, err)
	})

	t.Run("should produce different ciphertext for same plaintext", func(t *testing.T) {
		key := make([]byte, 32)
		rand.Read(key)
		plaintext := []byte("secret data")

		ciphertext1, _ := crypto.Encrypt(key, plaintext)
		ciphertext2, _ := crypto.Encrypt(key, plaintext)

		// Due to random nonce, ciphertexts should be different
		assert.NotEqual(t, ciphertext1, ciphertext2)
	})
}

func TestDeriveKey(t *testing.T) {
	t.Run("should derive key consistently", func(t *testing.T) {
		secret := []byte("master secret")
		salt := []byte("random salt")

		key1, err := crypto.DeriveKey(secret, salt, 32)
		require.NoError(t, err)

		key2, err := crypto.DeriveKey(secret, salt, 32)
		require.NoError(t, err)

		assert.Equal(t, key1, key2)
		assert.Len(t, key1, 32)
	})

	t.Run("should derive different keys with different salts", func(t *testing.T) {
		secret := []byte("master secret")

		key1, _ := crypto.DeriveKey(secret, []byte("salt1"), 32)
		key2, _ := crypto.DeriveKey(secret, []byte("salt2"), 32)

		assert.NotEqual(t, key1, key2)
	})
}

func TestSignAndVerify(t *testing.T) {
	t.Run("should sign and verify successfully", func(t *testing.T) {
		privateKey, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
		require.NoError(t, err)

		message := []byte("message to sign")
		signature, err := crypto.SignMessage(privateKey, message)
		require.NoError(t, err)
		assert.NotEmpty(t, signature)

		// Hash the message for verification (SignMessage internally hashes)
		messageHash := crypto.HashMessage(message)
		valid := crypto.VerifySignature(&privateKey.PublicKey, messageHash, signature)
		assert.True(t, valid)
	})

	t.Run("should fail verification with wrong message", func(t *testing.T) {
		privateKey, _ := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)

		signature, _ := crypto.SignMessage(privateKey, []byte("original message"))

		wrongHash := crypto.HashMessage([]byte("different message"))
		valid := crypto.VerifySignature(&privateKey.PublicKey, wrongHash, signature)
		assert.False(t, valid)
	})

	t.Run("should fail verification with wrong public key", func(t *testing.T) {
		privateKey1, _ := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
		privateKey2, _ := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)

		message := []byte("message")
		signature, _ := crypto.SignMessage(privateKey1, message)

		messageHash := crypto.HashMessage(message)
		valid := crypto.VerifySignature(&privateKey2.PublicKey, messageHash, signature)
		assert.False(t, valid)
	})
}

func TestEncodeDecodeHex(t *testing.T) {
	t.Run("should encode and decode hex successfully", func(t *testing.T) {
		original := []byte("test data")

		encoded := crypto.EncodeToHex(original)
		assert.NotEmpty(t, encoded)

		decoded, err := crypto.DecodeFromHex(encoded)
		require.NoError(t, err)
		assert.Equal(t, original, decoded)
	})

	t.Run("should fail to decode invalid hex", func(t *testing.T) {
		_, err := crypto.DecodeFromHex("invalid-hex-string!")
		assert.Error(t, err)
	})
}

func TestPublicKeyMarshaling(t *testing.T) {
	t.Run("should marshal and unmarshal public key", func(t *testing.T) {
		privateKey, _ := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)

		encoded := crypto.MarshalPublicKey(&privateKey.PublicKey)
		assert.NotEmpty(t, encoded)

		decoded, err := crypto.ParsePublicKey(encoded)
		require.NoError(t, err)

		// Verify keys are equal by comparing coordinates
		assert.Equal(t, privateKey.PublicKey.X.Bytes(), decoded.X.Bytes())
		assert.Equal(t, privateKey.PublicKey.Y.Bytes(), decoded.Y.Bytes())
	})
}

func TestCompareBytes(t *testing.T) {
	t.Run("should return true for equal byte slices", func(t *testing.T) {
		a := []byte("test data")
		b := []byte("test data")

		assert.True(t, crypto.CompareBytes(a, b))
	})

	t.Run("should return false for different byte slices", func(t *testing.T) {
		a := []byte("test data 1")
		b := []byte("test data 2")

		assert.False(t, crypto.CompareBytes(a, b))
	})

	t.Run("should return false for different length byte slices", func(t *testing.T) {
		a := []byte("short")
		b := []byte("longer string")

		assert.False(t, crypto.CompareBytes(a, b))
	})
}