ecdhkw_test.go

package ecdhkw_test

import (
	"bytes"
	"crypto/ecdh"
	"crypto/rand"
	"encoding/hex"
	"testing"

	"github.com/filecoin-project/go-fee/aeskw"
	"github.com/filecoin-project/go-fee/ecdhkw"
	"github.com/stretchr/testify/require"
)

func newRecipient(t *testing.T) *ecdh.PrivateKey {
	t.Helper()
	priv, err := ecdh.X25519().GenerateKey(rand.Reader)
	require.NoError(t, err, "generate recipient key")
	return priv
}

func mustDecode(t *testing.T, s string) []byte {
	t.Helper()
	b, err := hex.DecodeString(s)
	require.NoErrorf(t, err, "bad hex %q", s)
	return b
}

// AC: "I can wrap a CEK to an X25519 public key and unwrap it with the
// corresponding private key, recovering the original CEK." Checked across the
// valid CEK sizes (16/24/32).
func TestWrapUnwrapRoundTrip(t *testing.T) {
	recipient := newRecipient(t)
	for _, size := range []int{16, 24, 32} {
		cek := make([]byte, size)
		_, err := rand.Read(cek)
		require.NoError(t, err, "rand cek")

		w, err := ecdhkw.Wrap(recipient.PublicKey(), cek)
		require.NoErrorf(t, err, "Wrap(%d-byte CEK)", size)
		require.Equal(t, ecdh.X25519(), w.EphemeralPublicKey.Curve(), "ephemeral key is not X25519")
		require.Lenf(t, w.WrappedCEK, size+8, "wrapped CEK length")

		got, err := ecdhkw.Unwrap(recipient, w)
		require.NoErrorf(t, err, "Unwrap(%d-byte CEK)", size)
		require.Equalf(t, cek, got, "round-trip mismatch for %d-byte CEK", size)
	}
}

// AC: "When I attempt to unwrap with the wrong private key, unwrap returns an
// error." The wrong key derives a different KEK, so AES-KW's integrity check
// fails; the error wraps aeskw.ErrIntegrity.
func TestUnwrapWrongPrivateKey(t *testing.T) {
	recipient := newRecipient(t)
	wrongKey := newRecipient(t)
	cek := bytes.Repeat([]byte{0x5A}, 32)

	w, err := ecdhkw.Wrap(recipient.PublicKey(), cek)
	require.NoError(t, err, "Wrap")

	got, err := ecdhkw.Unwrap(wrongKey, w)
	require.ErrorIs(t, err, aeskw.ErrIntegrity, "Unwrap with wrong key should wrap aeskw.ErrIntegrity")
	require.Nil(t, got)
}

// AC: "When I wrap the same CEK twice to the same public key, the two wrapped
// outputs differ — the ephemeral sender key is fresh each time." Both the
// ephemeral public key and the wrapped bytes must differ, and both copies must
// still unwrap to the original CEK.
func TestWrapFreshEphemeralPerCall(t *testing.T) {
	recipient := newRecipient(t)
	cek := bytes.Repeat([]byte{0xC3}, 32)

	first, err := ecdhkw.Wrap(recipient.PublicKey(), cek)
	require.NoError(t, err, "first Wrap")
	second, err := ecdhkw.Wrap(recipient.PublicKey(), cek)
	require.NoError(t, err, "second Wrap")

	require.NotEqual(t, first.EphemeralPublicKey.Bytes(), second.EphemeralPublicKey.Bytes(),
		"ephemeral public keys are identical across wraps; expected a fresh key each time")
	require.NotEqual(t, first.WrappedCEK, second.WrappedCEK,
		"wrapped CEKs are identical across wraps; expected different output")

	// Both independently recover the same plaintext CEK.
	for i, w := range []*ecdhkw.Wrapped{first, second} {
		got, err := ecdhkw.Unwrap(recipient, w)
		require.NoErrorf(t, err, "Unwrap copy %d", i)
		require.Equalf(t, cek, got, "Unwrap copy %d mismatch", i)
	}
}

// A wrap is bound to its ephemeral key: swapping in a different ephemeral
// public key (even a valid one) breaks the derivation and fails the unwrap.
func TestUnwrapTamperedEphemeral(t *testing.T) {
	recipient := newRecipient(t)
	cek := bytes.Repeat([]byte{0x11}, 32)

	w, err := ecdhkw.Wrap(recipient.PublicKey(), cek)
	require.NoError(t, err, "Wrap")
	other, err := ecdh.X25519().GenerateKey(rand.Reader)
	require.NoError(t, err, "generate other key")
	w.EphemeralPublicKey = other.PublicKey()

	_, err = ecdhkw.Unwrap(recipient, w)
	require.ErrorIs(t, err, aeskw.ErrIntegrity, "Unwrap with swapped ephemeral key")
}

// A low-order ephemeral point would force the ECDH shared secret into a small
// subgroup; crypto/ecdh rejects it, and Unwrap must surface that as an error
// rather than deriving a KEK from a degenerate secret.
func TestUnwrapLowOrderEphemeral(t *testing.T) {
	recipient := newRecipient(t)
	// The all-zero u-coordinate is a classic X25519 low-order point. It is a
	// valid 32-byte public key to construct, but ECDH against it fails.
	lowOrder, err := ecdh.X25519().NewPublicKey(make([]byte, 32))
	require.NoError(t, err, "construct low-order point")
	w := &ecdhkw.Wrapped{EphemeralPublicKey: lowOrder, WrappedCEK: make([]byte, 40)}
	_, err = ecdhkw.Unwrap(recipient, w)
	require.Error(t, err, "Unwrap with low-order ephemeral point should fail")
}

// A wrap is bound to its wrapped bytes: flipping any bit fails the unwrap.
func TestUnwrapTamperedCEK(t *testing.T) {
	recipient := newRecipient(t)
	cek := bytes.Repeat([]byte{0x22}, 32)

	w, err := ecdhkw.Wrap(recipient.PublicKey(), cek)
	require.NoError(t, err, "Wrap")
	for i := range w.WrappedCEK {
		tampered := &ecdhkw.Wrapped{
			EphemeralPublicKey: w.EphemeralPublicKey,
			WrappedCEK:         bytes.Clone(w.WrappedCEK),
		}
		tampered.WrappedCEK[i] ^= 0x01
		_, err = ecdhkw.Unwrap(recipient, tampered)
		require.ErrorIsf(t, err, aeskw.ErrIntegrity, "tampering wrapped byte %d not detected", i)
	}
}

func TestWrapInputValidation(t *testing.T) {
	recipient := newRecipient(t)
	p256, err := ecdh.P256().GenerateKey(rand.Reader)
	require.NoError(t, err, "generate P256 key")

	tests := []struct {
		name string
		pub  *ecdh.PublicKey
		cek  []byte
	}{
		{"nil public key", nil, make([]byte, 32)},
		{"wrong curve", p256.PublicKey(), make([]byte, 32)},
		{"cek too short", recipient.PublicKey(), make([]byte, 8)},
		{"cek not block-aligned", recipient.PublicKey(), make([]byte, 20)},
		{"nil cek", recipient.PublicKey(), nil},
	}
	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			_, err := ecdhkw.Wrap(tc.pub, tc.cek)
			require.Error(t, err)
		})
	}
}

func TestUnwrapInputValidation(t *testing.T) {
	recipient := newRecipient(t)
	p256, err := ecdh.P256().GenerateKey(rand.Reader)
	require.NoError(t, err, "generate P256 key")
	valid, err := ecdhkw.Wrap(recipient.PublicKey(), make([]byte, 32))
	require.NoError(t, err, "Wrap")

	t.Run("nil private key", func(t *testing.T) {
		_, err := ecdhkw.Unwrap(nil, valid)
		require.Error(t, err)
	})
	t.Run("nil wrapped", func(t *testing.T) {
		_, err := ecdhkw.Unwrap(recipient, nil)
		require.Error(t, err)
	})
	t.Run("nil ephemeral key", func(t *testing.T) {
		_, err := ecdhkw.Unwrap(recipient, &ecdhkw.Wrapped{WrappedCEK: valid.WrappedCEK})
		require.Error(t, err)
	})
	t.Run("ephemeral wrong curve", func(t *testing.T) {
		w := &ecdhkw.Wrapped{EphemeralPublicKey: p256.PublicKey(), WrappedCEK: valid.WrappedCEK}
		_, err := ecdhkw.Unwrap(recipient, w)
		require.Error(t, err)
	})
}

// TestKnownAnswerVector pins the ECDH-ES+A256KW construction to a fixed
// decryption vector: a known recipient private key, ephemeral public key, and
// wrapped CEK must Unwrap to a known CEK. This is the gold vector for the
// cross-implementation tests (foc-encryption / FIL-473). Encryption is
// nondeterministic (a fresh ephemeral key per wrap), so the shared anchor is
// the decrypt direction — which any conforming implementation must reproduce
// without an injection seam. It exercises the full path end to end: X25519
// ECDH, the COSE Concat-KDF context (kdf.go), and AES-KW unwrap.
//
// All values use the standard RFC 9053 §5.2 context (AlgorithmID = A256KW,
// keyDataLength = 256, empty PartyU/PartyV/protected) documented in kdf.go.
func TestKnownAnswerVector(t *testing.T) {
	const (
		recipientPrivHex = "0102030405060708090a0b0c0d0e0f101112131415161718191a1b1c1d1e1f20"
		ephemeralPubHex  = "605a725d2a4adfeeb1a29e17edd621c1b7593ee8cdbc44ac6c4ab6e2f805d23c"
		wrappedCEKHex    = "1433cdb050fc4ab1ccb616c395a81908c001fcfa7fb865366a3ef8db0af45c7c4305ae7de7080007"
		cekHex           = "00112233445566778899aabbccddeeff000102030405060708090a0b0c0d0e0f"

		// Derived from recipientPrivHex; recorded so a cross-impl test starts
		// from the same recipient key.
		wantRecipientPubHex = "07a37cbc142093c8b755dc1b10e86cb426374ad16aa853ed0bdfc0b2b86d1c7c"
	)

	priv, err := ecdh.X25519().NewPrivateKey(mustDecode(t, recipientPrivHex))
	require.NoError(t, err, "recipient private key")
	require.Equal(t, wantRecipientPubHex, hex.EncodeToString(priv.PublicKey().Bytes()), "recipient public key")

	ephemeralPub, err := ecdh.X25519().NewPublicKey(mustDecode(t, ephemeralPubHex))
	require.NoError(t, err, "ephemeral public key")

	w := &ecdhkw.Wrapped{EphemeralPublicKey: ephemeralPub, WrappedCEK: mustDecode(t, wrappedCEKHex)}
	got, err := ecdhkw.Unwrap(priv, w)
	require.NoError(t, err, "unwrap")
	require.Equal(t, cekHex, hex.EncodeToString(got), "unwrapped CEK")
}