From cc83d11cf01a1bd66ba390498d9d63991b4577e9 Mon Sep 17 00:00:00 2001
From: bobbyphilip <bobbyphilip@gmail.com>
Date: Tue, 26 May 2020 21:08:34 +0200
Subject: [PATCH] Refactor JWTEncryptionImpl to remove redundant code paths
 (#249)

* Refactor JWTEncryptionImpl to remove redundant code paths
Updated the ExceptionMessage to match the one declared in the API

* Updated the encrypt() API singature to throw the right Exception

* Corrected the java doc
---
 .../io/smallrye/jwt/build/JwtEncryption.java  |  6 +--
 .../jwt/build/impl/JwtEncryptionImpl.java     | 46 +++++++------------
 2 files changed, 19 insertions(+), 33 deletions(-)

diff --git a/implementation/src/main/java/io/smallrye/jwt/build/JwtEncryption.java b/implementation/src/main/java/io/smallrye/jwt/build/JwtEncryption.java
index 1f8728d1..e66c0f45 100644
--- a/implementation/src/main/java/io/smallrye/jwt/build/JwtEncryption.java
+++ b/implementation/src/main/java/io/smallrye/jwt/build/JwtEncryption.java
@@ -53,9 +53,9 @@ public interface JwtEncryption {
      * unless different ones have been set with {@code JwtEncryptionBuilder}.
      * A key of size 2048 bits or larger MUST be used with the 'RSA-OAEP' and 'RSA-OAEP-256' algorithms.
      *
-     * @return signed JWT token
-     * @throws JwtSignatureException the exception if the signing operation has failed
+     * @return encrypted JWT token
+     * @throws JwtEncryptionException the exception if the encryption operation has failed
      */
-    String encrypt() throws JwtSignatureException;
+    String encrypt() throws JwtEncryptionException;
 
 }
diff --git a/implementation/src/main/java/io/smallrye/jwt/build/impl/JwtEncryptionImpl.java b/implementation/src/main/java/io/smallrye/jwt/build/impl/JwtEncryptionImpl.java
index c58f3ff5..f5aa3650 100644
--- a/implementation/src/main/java/io/smallrye/jwt/build/impl/JwtEncryptionImpl.java
+++ b/implementation/src/main/java/io/smallrye/jwt/build/impl/JwtEncryptionImpl.java
@@ -64,15 +64,15 @@ public String encrypt(String keyLocation) throws JwtEncryptionException {
         } catch (Exception ex) {
             throw new JwtEncryptionException(ex);
         }
-        return key instanceof PublicKey ? encryptInternal((PublicKey) key) : encryptInternal((SecretKey) key);
+        return encryptInternal(key);
     }
 
     /**
      * {@inheritDoc}
      */
     @Override
-    public String encrypt() throws JwtSignatureException {
-        return encryptInternal(getKeyEncryptionKeyFromConfig((String) headers.get("kid")));
+    public String encrypt() throws JwtEncryptionException {
+        return encrypt(readKeyLocationFromConfig());
     }
 
     /**
@@ -147,42 +147,28 @@ private String getKeyEncryptionAlgorithm(Key keyEncryptionKey) {
         if ("dir".equals(alg)) {
             throw ImplMessages.msg.directContentEncryptionUnsupported();
         }
-
-        if (keyEncryptionKey instanceof RSAPublicKey) {
-            if (alg == null) {
-                return KeyEncryptionAlgorithm.RSA_OAEP_256.getAlgorithm();
-            } else {
-                return alg;
-            }
-        } else if (keyEncryptionKey instanceof SecretKey) {
-            if (alg == null) {
-                return KeyEncryptionAlgorithm.A256KW.getAlgorithm();
-            } else {
-                return alg;
-            }
-        } else if (keyEncryptionKey instanceof ECPublicKey) {
-            if (alg == null) {
-                return KeyEncryptionAlgorithm.ECDH_ES_A256KW.getAlgorithm();
-            } else {
-                return alg;
+        if (alg == null) {
+            if (keyEncryptionKey instanceof RSAPublicKey) {
+                alg = KeyEncryptionAlgorithm.RSA_OAEP_256.getAlgorithm();
+            } else if (keyEncryptionKey instanceof SecretKey) {
+                alg = KeyEncryptionAlgorithm.A256KW.getAlgorithm();
+            } else if (keyEncryptionKey instanceof ECPublicKey) {
+                alg = KeyEncryptionAlgorithm.ECDH_ES_A256KW.getAlgorithm();
             }
         }
-
-        throw ImplMessages.msg.unsupportedKeyEncryptionAlgorithm(keyEncryptionKey.getAlgorithm());
+        if (alg == null) {
+            throw ImplMessages.msg.unsupportedKeyEncryptionAlgorithm(keyEncryptionKey.getAlgorithm());
+        }
+        return alg;
     }
 
     private String getContentEncryptionAlgorithm() {
         return headers.containsKey("enc") ? headers.get("enc").toString() : ContentEncryptionAlgorithm.A256GCM.name();
     }
 
-    static Key getKeyEncryptionKeyFromConfig(String kid) {
+    private static String readKeyLocationFromConfig() {
         try {
-            String keyLocation = ConfigProvider.getConfig().getValue("smallrye.jwt.encrypt.key-location", String.class);
-            try {
-                return KeyUtils.readEncryptionKey(keyLocation, kid);
-            } catch (Exception ex) {
-                throw ImplMessages.msg.encryptionKeyNotFound(keyLocation);
-            }
+            return ConfigProvider.getConfig().getValue("smallrye.jwt.encrypt.key-location", String.class);
         } catch (NoSuchElementException ex) {
             throw ImplMessages.msg.keyLocationPropertyEmpty();
         }