From 98d6bfdb63791f6751d3a60a14252e956d77d19f Mon Sep 17 00:00:00 2001
From: Bojidar Marinov <bojidar.marinov.bg@gmail.com>
Date: Mon, 13 Nov 2023 18:57:27 +0200
Subject: [PATCH] jwe: Support overriding the algorithm when supplying a JWK

Signed-off-by: Bojidar Marinov <bojidar.marinov.bg@gmail.com>
---
 keywrap/jwe/keywrapper_jwe.go      |  6 ++++-
 keywrap/jwe/keywrapper_jwe_test.go | 36 ++++++++++++++++++++++++++++++
 utils/testing.go                   | 13 +++++++++--
 3 files changed, 52 insertions(+), 3 deletions(-)

diff --git a/keywrap/jwe/keywrapper_jwe.go b/keywrap/jwe/keywrapper_jwe.go
index cd2241c..62e03b5 100644
--- a/keywrap/jwe/keywrapper_jwe.go
+++ b/keywrap/jwe/keywrapper_jwe.go
@@ -123,9 +123,13 @@ func addPubKeys(joseRecipients *[]jose.Recipient, pubKeys [][]byte) error {
 		}
 
 		alg := jose.RSA_OAEP
-		switch key.(type) {
+		switch key := key.(type) {
 		case *ecdsa.PublicKey:
 			alg = jose.ECDH_ES_A256KW
+		case *jose.JSONWebKey:
+			if key.Algorithm != "" {
+				alg = jose.KeyAlgorithm(key.Algorithm)
+			}
 		}
 
 		*joseRecipients = append(*joseRecipients, jose.Recipient{
diff --git a/keywrap/jwe/keywrapper_jwe_test.go b/keywrap/jwe/keywrapper_jwe_test.go
index 3beea39..3d5e053 100644
--- a/keywrap/jwe/keywrapper_jwe_test.go
+++ b/keywrap/jwe/keywrapper_jwe_test.go
@@ -70,6 +70,21 @@ func createValidJweCcs() ([]*config.CryptoConfig, error) {
 		return nil, err
 	}
 
+	ecKey, err := utils.CreateECDSAKey(elliptic.P521())
+	if err != nil {
+		return nil, err
+	}
+
+	jweEcPrivKeyJwk, err := jose.JSONWebKey{Key: ecKey, Algorithm: string(jose.ECDH_ES_A256KW)}.MarshalJSON()
+	if err != nil {
+		return nil, err
+	}
+
+	jweEcPubKeyJwk, err := jose.JSONWebKey{Key: &ecKey.PublicKey, Algorithm: string(jose.ECDH_ES_A256KW)}.MarshalJSON()
+	if err != nil {
+		return nil, err
+	}
+
 	validJweCcs := []*config.CryptoConfig{
 		// Key 1
 		{
@@ -226,6 +241,27 @@ func createValidJweCcs() ([]*config.CryptoConfig, error) {
 				},
 			},
 		},
+		// EC Key (JWK format)
+		{
+			EncryptConfig: &config.EncryptConfig{
+				Parameters: map[string][][]byte{
+					"pubkeys": {jweEcPubKeyJwk},
+				},
+				DecryptConfig: config.DecryptConfig{
+					Parameters: map[string][][]byte{
+						"privkeys":           {jweEcPrivKeyJwk},
+						"privkeys-passwords": {oneEmpty},
+					},
+				},
+			},
+
+			DecryptConfig: &config.DecryptConfig{
+				Parameters: map[string][][]byte{
+					"privkeys":           {jweEcPrivKeyJwk},
+					"privkeys-passwords": {oneEmpty},
+				},
+			},
+		},
 	}
 	return validJweCcs, nil
 }
diff --git a/utils/testing.go b/utils/testing.go
index 69bb9d1..050aa88 100644
--- a/utils/testing.go
+++ b/utils/testing.go
@@ -38,6 +38,15 @@ func CreateRSAKey(bits int) (*rsa.PrivateKey, error) {
 	return key, nil
 }
 
+// CreateECDSAKey creates an elliptic curve key for the given curve
+func CreateECDSAKey(curve elliptic.Curve) (*ecdsa.PrivateKey, error) {
+	key, err := ecdsa.GenerateKey(curve, rand.Reader)
+	if err != nil {
+		return nil, fmt.Errorf("ecdsa.GenerateKey failed: %w", err)
+	}
+	return key, nil
+}
+
 // CreateRSATestKey creates an RSA key of the given size and returns
 // the public and private key in PEM or DER format
 func CreateRSATestKey(bits int, password []byte, pemencode bool) ([]byte, []byte, error) {
@@ -85,9 +94,9 @@ func CreateRSATestKey(bits int, password []byte, pemencode bool) ([]byte, []byte
 // CreateECDSATestKey creates and elliptic curve key for the given curve and returns
 // the public and private key in DER format
 func CreateECDSATestKey(curve elliptic.Curve) ([]byte, []byte, error) {
-	key, err := ecdsa.GenerateKey(curve, rand.Reader)
+	key, err := CreateECDSAKey(curve)
 	if err != nil {
-		return nil, nil, fmt.Errorf("ecdsa.GenerateKey failed: %w", err)
+		return nil, nil, err
 	}
 
 	pubData, err := x509.MarshalPKIXPublicKey(&key.PublicKey)