[jwe] Work with non X25519 ECDH encryption (#1442)

* WIP code

* more WIP

* [jwe] Work with non X25519 ECDH keys

* go fmt

* omit non-pointer test

https://github.com/lestrrat-go/jwx/actions/runs/16959725199/job/48069344266

This test is not _that_ important.

* reduce code duplication

* typo

* Update Changes

Author: lestrrat

Changes

@@ -4,6 +4,10 @@ Changes
 v3 has many incompatibilities with v2. To see the full list of differences between
 v2 and v3, please read the Changes-v3.md file (https://github.com/lestrrat-go/jwx/blob/develop/v3/Changes-v3.md)
 
+v3.0.11 UNRELEASED
+  * [jwe] Previously, ecdh.PrivateKey/ecdh.PublicKey were not properly handled
+    when used for encryption, which has been fixed.
+
 v3.0.10 04 Aug 2025
   * [jws/jwsbb] Add `jwsbb.ErrHeaderNotFound()` to return the same error type as when
     a non-existent header is requested. via `HeaderGetXXX()` functions. Previously, this

internal/keyconv/keyconv.go

@@ -5,8 +5,10 @@ import (
 	"crypto/ecdh"
 	"crypto/ecdsa"
 	"crypto/ed25519"
+	"crypto/elliptic"
 	"crypto/rsa"
 	"fmt"
+	"math/big"
 
 	"github.com/lestrrat-go/blackmagic"
 	"github.com/lestrrat-go/jwx/v3/jwk"
@@ -263,3 +265,90 @@ func ECDHPublicKey(dst, src any) error {
 
 	return blackmagic.AssignIfCompatible(dst, pubECDH)
 }
+
+// ecdhCurveToElliptic maps ECDH curves to elliptic curves
+func ecdhCurveToElliptic(ecdhCurve ecdh.Curve) (elliptic.Curve, error) {
+	switch ecdhCurve {
+	case ecdh.P256():
+		return elliptic.P256(), nil
+	case ecdh.P384():
+		return elliptic.P384(), nil
+	case ecdh.P521():
+		return elliptic.P521(), nil
+	default:
+		return nil, fmt.Errorf(`keyconv: unsupported ECDH curve: %v`, ecdhCurve)
+	}
+}
+
+// ecdhPublicKeyToECDSA converts an ECDH public key to an ECDSA public key
+func ecdhPublicKeyToECDSA(ecdhPubKey *ecdh.PublicKey) (*ecdsa.PublicKey, error) {
+	curve, err := ecdhCurveToElliptic(ecdhPubKey.Curve())
+	if err != nil {
+		return nil, err
+	}
+
+	pubBytes := ecdhPubKey.Bytes()
+
+	// Parse the uncompressed point format (0x04 prefix + X + Y coordinates)
+	if len(pubBytes) == 0 || pubBytes[0] != 0x04 {
+		return nil, fmt.Errorf(`keyconv: invalid ECDH public key format`)
+	}
+
+	keyLen := (len(pubBytes) - 1) / 2
+	if len(pubBytes) != 1+2*keyLen {
+		return nil, fmt.Errorf(`keyconv: invalid ECDH public key length`)
+	}
+
+	x := new(big.Int).SetBytes(pubBytes[1 : 1+keyLen])
+	y := new(big.Int).SetBytes(pubBytes[1+keyLen:])
+
+	return &ecdsa.PublicKey{
+		Curve: curve,
+		X:     x,
+		Y:     y,
+	}, nil
+}
+
+func ECDHToECDSA(dst, src any) error {
+	// convert ecdh.PublicKey to ecdsa.PublicKey, ecdh.PrivateKey to ecdsa.PrivateKey
+
+	// First, handle value types by converting to pointers
+	switch s := src.(type) {
+	case ecdh.PrivateKey:
+		src = &s
+	case ecdh.PublicKey:
+		src = &s
+	}
+
+	var privBytes []byte
+	var pubkey *ecdh.PublicKey
+	// Now handle the actual conversion with pointer types
+	switch src := src.(type) {
+	case *ecdh.PrivateKey:
+		pubkey = src.PublicKey()
+		privBytes = src.Bytes()
+	case *ecdh.PublicKey:
+		pubkey = src
+	default:
+		return fmt.Errorf(`keyconv: expected ecdh.PrivateKey, *ecdh.PrivateKey, ecdh.PublicKey, or *ecdh.PublicKey, got %T`, src)
+	}
+
+	// convert the public key
+	ecdsaPubKey, err := ecdhPublicKeyToECDSA(pubkey)
+	if err != nil {
+		return fmt.Errorf(`keyconv.ECDHToECDSA: failed to convert ECDH public key to ECDSA public key: %w`, err)
+	}
+
+	// return if we were being asked to convert *ecdh.PublicKey
+	if privBytes == nil {
+		return blackmagic.AssignIfCompatible(dst, ecdsaPubKey)
+	}
+
+	// Then create the private key with the public key embedded
+	ecdsaPrivKey := &ecdsa.PrivateKey{
+		D:         new(big.Int).SetBytes(privBytes),
+		PublicKey: *ecdsaPubKey,
+	}
+
+	return blackmagic.AssignIfCompatible(dst, ecdsaPrivKey)
+}

internal/keyconv/keyconv_test.go

@@ -1,7 +1,9 @@
 package keyconv_test
 
 import (
+	"crypto/ecdh"
 	"crypto/ecdsa"
+	"crypto/rand"
 	"crypto/rsa"
 	"testing"
 
@@ -205,3 +207,126 @@ func TestKeyconv(t *testing.T) {
 		})
 	})
 }
+
+func TestECDHToECDSA(t *testing.T) {
+	curves := []struct {
+		name      string
+		ecdhCurve ecdh.Curve
+		jwaAlg    jwa.EllipticCurveAlgorithm
+	}{
+		{"P256", ecdh.P256(), jwa.P256()},
+		{"P384", ecdh.P384(), jwa.P384()},
+		{"P521", ecdh.P521(), jwa.P521()},
+	}
+
+	for _, curve := range curves {
+		t.Run(curve.name, func(t *testing.T) {
+			// Generate an ECDSA key for comparison
+			ecdsaKey, err := jwxtest.GenerateEcdsaKey(curve.jwaAlg)
+			require.NoError(t, err, `ecdsa.GenerateKey should succeed`)
+
+			// Convert ECDSA key to ECDH key
+			ecdhPrivKey, err := ecdsaKey.ECDH()
+			require.NoError(t, err, `ECDSA to ECDH conversion should succeed`)
+
+			ecdhPubKey := ecdhPrivKey.PublicKey()
+
+			t.Run("PrivateKey", func(t *testing.T) {
+				testcases := []struct {
+					name  string
+					src   any
+					error bool
+				}{
+					{"*ecdh.PrivateKey", ecdhPrivKey, false},
+					{"invalid type", "not a key", true},
+				}
+
+				for _, tc := range testcases {
+					t.Run(tc.name, func(t *testing.T) {
+						var dst *ecdsa.PrivateKey
+						err := keyconv.ECDHToECDSA(&dst, tc.src)
+
+						if tc.error {
+							require.Error(t, err, `ECDHToECDSA should fail for invalid input`)
+						} else {
+							require.NoError(t, err, `ECDHToECDSA should succeed`)
+							require.NotNil(t, dst, `destination should not be nil`)
+
+							// Verify the converted key has the same curve
+							require.Equal(t, ecdsaKey.Curve, dst.Curve, `curves should match`)
+
+							// Verify the private key values match
+							require.Equal(t, ecdsaKey.D, dst.D, `private key values should match`)
+
+							// Verify the public key coordinates match
+							require.Equal(t, ecdsaKey.PublicKey.X, dst.PublicKey.X, `X coordinates should match`)
+							require.Equal(t, ecdsaKey.PublicKey.Y, dst.PublicKey.Y, `Y coordinates should match`)
+						}
+					})
+				}
+			})
+
+			t.Run("PublicKey", func(t *testing.T) {
+				testcases := []struct {
+					name  string
+					src   any
+					error bool
+				}{
+					{"*ecdh.PublicKey", ecdhPubKey, false},
+					{"ecdh.PublicKey", *ecdhPubKey, false},
+					{"invalid type", "not a key", true},
+				}
+
+				for _, tc := range testcases {
+					t.Run(tc.name, func(t *testing.T) {
+						var dst *ecdsa.PublicKey
+						err := keyconv.ECDHToECDSA(&dst, tc.src)
+
+						if tc.error {
+							require.Error(t, err, `ECDHToECDSA should fail for invalid input`)
+						} else {
+							require.NoError(t, err, `ECDHToECDSA should succeed`)
+							require.NotNil(t, dst, `destination should not be nil`)
+
+							// Verify the converted key has the same curve
+							require.Equal(t, ecdsaKey.PublicKey.Curve, dst.Curve, `curves should match`)
+
+							// Verify the public key coordinates match
+							require.Equal(t, ecdsaKey.PublicKey.X, dst.X, `X coordinates should match`)
+							require.Equal(t, ecdsaKey.PublicKey.Y, dst.Y, `Y coordinates should match`)
+						}
+					})
+				}
+			})
+
+			t.Run("RoundTrip", func(t *testing.T) {
+				// Test that ECDSA -> ECDH -> ECDSA produces the same key
+				var convertedPrivKey *ecdsa.PrivateKey
+				err := keyconv.ECDHToECDSA(&convertedPrivKey, ecdhPrivKey)
+				require.NoError(t, err, `ECDHToECDSA should succeed`)
+
+				var convertedPubKey *ecdsa.PublicKey
+				err = keyconv.ECDHToECDSA(&convertedPubKey, ecdhPubKey)
+				require.NoError(t, err, `ECDHToECDSA should succeed`)
+
+				// Verify the keys are equivalent
+				require.Equal(t, ecdsaKey.D, convertedPrivKey.D, `private key values should match`)
+				require.Equal(t, ecdsaKey.PublicKey.X, convertedPrivKey.PublicKey.X, `private key X coordinates should match`)
+				require.Equal(t, ecdsaKey.PublicKey.Y, convertedPrivKey.PublicKey.Y, `private key Y coordinates should match`)
+				require.Equal(t, ecdsaKey.PublicKey.X, convertedPubKey.X, `public key X coordinates should match`)
+				require.Equal(t, ecdsaKey.PublicKey.Y, convertedPubKey.Y, `public key Y coordinates should match`)
+			})
+		})
+	}
+
+	t.Run("UnsupportedCurve", func(t *testing.T) {
+		// Create a mock ECDH key with X25519 curve (not supported for ECDSA)
+		x25519Key, err := ecdh.X25519().GenerateKey(rand.Reader)
+		require.NoError(t, err, `X25519 key generation should succeed`)
+
+		var dst *ecdsa.PrivateKey
+		err = keyconv.ECDHToECDSA(&dst, x25519Key)
+		require.Error(t, err, `ECDHToECDSA should fail for unsupported curve`)
+		require.Contains(t, err.Error(), "unsupported ECDH curve", `error should mention unsupported curve`)
+	})
+}

jwe/encrypt.go

@@ -96,23 +96,48 @@ func (e *encrypter) EncryptKey(cek []byte) (keygen.ByteSource, error) {
 			keyToUse = e.pubkey
 		}
 
-		// Handle ecdsa.PublicKey by value - convert to pointer
-		if pk, ok := keyToUse.(ecdsa.PublicKey); ok {
-			keyToUse = &pk
+		switch key := keyToUse.(type) {
+		case *ecdsa.PublicKey:
+			// no op
+		case ecdsa.PublicKey:
+			keyToUse = &key
+		case *ecdsa.PrivateKey:
+			keyToUse = &key.PublicKey
+		case ecdsa.PrivateKey:
+			keyToUse = &key.PublicKey
+		case *ecdh.PublicKey:
+			// no op
+		case ecdh.PublicKey:
+			keyToUse = &key
+		case ecdh.PrivateKey:
+			keyToUse = key.PublicKey()
+		case *ecdh.PrivateKey:
+			keyToUse = key.PublicKey()
 		}
 
 		// Determine key type and call appropriate function
+		switch key := keyToUse.(type) {
+		case *ecdh.PublicKey:
+			if key.Curve() == ecdh.X25519() {
+				if !keywrap {
+					return jwebb.KeyEncryptECDHESX25519(cek, e.keyalg.String(), e.apu, e.apv, key, keysize, e.ctalg.String())
+				}
+				return jwebb.KeyEncryptECDHESKeyWrapX25519(cek, e.keyalg.String(), e.apu, e.apv, key, keysize, e.ctalg.String())
+			}
+
+			var ecdsaKey *ecdsa.PublicKey
+			if err := keyconv.ECDHToECDSA(&ecdsaKey, key); err != nil {
+				return nil, fmt.Errorf(`encrypt: failed to convert ECDH public key to ECDSA: %w`, err)
+			}
+			keyToUse = ecdsaKey
+		}
+
 		switch key := keyToUse.(type) {
 		case *ecdsa.PublicKey:
 			if !keywrap {
 				return jwebb.KeyEncryptECDHESECDSA(cek, e.keyalg.String(), e.apu, e.apv, key, keysize, e.ctalg.String())
 			}
 			return jwebb.KeyEncryptECDHESKeyWrapECDSA(cek, e.keyalg.String(), e.apu, e.apv, key, keysize, e.ctalg.String())
-		case *ecdh.PublicKey:
-			if !keywrap {
-				return jwebb.KeyEncryptECDHESX25519(cek, e.keyalg.String(), e.apu, e.apv, key, keysize, e.ctalg.String())
-			}
-			return jwebb.KeyEncryptECDHESKeyWrapX25519(cek, e.keyalg.String(), e.apu, e.apv, key, keysize, e.ctalg.String())
 		default:
 			return nil, fmt.Errorf(`encrypt: unsupported key type for ECDH-ES: %T`, keyToUse)
 		}

jwe/jwe_test.go

@@ -1008,3 +1008,18 @@ func BenchmarkParseCompat(b *testing.B) {
 		}
 	}
 }
+
+func TestGH1434(t *testing.T) {
+	// Test if we can use ecdh.PrivateKey to encrypt and decrypt
+
+	key, err := ecdh.P256().GenerateKey(rand.Reader)
+	require.NoError(t, err, `ecdh.P256().GenerateKey should succeed`)
+
+	const payload = `hello, world!`
+	encrypted, err := jwe.Encrypt([]byte(payload), jwe.WithKey(jwa.ECDH_ES_A256KW(), key))
+	require.NoError(t, err, `jwe.Encrypt should succeed`)
+
+	decrypted, err := jwe.Decrypt(encrypted, jwe.WithKey(jwa.ECDH_ES_A256KW(), key))
+	require.NoError(t, err, `jwe.Decrypt should succeed`)
+	require.Equal(t, []byte(payload), decrypted, `decrypted payload should match original payload`)
+}