Improve keyprovider reliability

lib/src/main/java/com/auth0/jwt/JWTCreator.java

@@ -4,12 +4,14 @@
 import com.auth0.jwt.exceptions.JWTCreationException;
 import com.auth0.jwt.exceptions.SignatureGenerationException;
 import com.auth0.jwt.impl.*;
+import com.auth0.jwt.interfaces.PrivateKeyDetail;
 import com.fasterxml.jackson.core.JsonProcessingException;
 import com.fasterxml.jackson.databind.MapperFeature;
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.fasterxml.jackson.databind.module.SimpleModule;
 
 import java.nio.charset.StandardCharsets;
+import java.security.PrivateKey;
 import java.time.Instant;
 import java.util.*;
 import java.util.Map.Entry;
@@ -547,11 +549,13 @@ public String sign(Algorithm algorithm) throws IllegalArgumentException, JWTCrea
             if (!headerClaims.containsKey(PublicClaims.TYPE)) {
                 headerClaims.put(PublicClaims.TYPE, "JWT");
             }
-            String signingKeyId = algorithm.getSigningKeyId();
-            if (signingKeyId != null) {
-                withKeyId(signingKeyId);
+            PrivateKeyDetail<?> privateKeyDetail = algorithm.getPrivateKeyDetails();
+            String keyId = privateKeyDetail == null ? null : privateKeyDetail.getPrivateKeyId();
+            PrivateKey key = privateKeyDetail == null ? null : privateKeyDetail.getPrivateKey();
+            if (keyId != null) {
+                withKeyId(privateKeyDetail.getPrivateKeyId());
             }
-            return new JWTCreator(algorithm, headerClaims, payloadClaims).sign();
+            return new JWTCreator(algorithm, headerClaims, payloadClaims).sign(key);
         }
 
         private void assertNonNull(String name) {
@@ -565,14 +569,14 @@ private void addClaim(String name, Object value) {
         }
     }
 
-    private String sign() throws SignatureGenerationException {
+    private String sign(PrivateKey privateKey) throws SignatureGenerationException {
         String header = Base64.getUrlEncoder().withoutPadding()
                 .encodeToString(headerJson.getBytes(StandardCharsets.UTF_8));
         String payload = Base64.getUrlEncoder().withoutPadding()
                 .encodeToString(payloadJson.getBytes(StandardCharsets.UTF_8));
 
         byte[] signatureBytes = algorithm.sign(header.getBytes(StandardCharsets.UTF_8),
-                payload.getBytes(StandardCharsets.UTF_8));
+                payload.getBytes(StandardCharsets.UTF_8), privateKey);
         String signature = Base64.getUrlEncoder().withoutPadding().encodeToString((signatureBytes));
 
         return String.format("%s.%s.%s", header, payload, signature);

lib/src/main/java/com/auth0/jwt/algorithms/Algorithm.java

@@ -4,8 +4,10 @@
 import com.auth0.jwt.exceptions.SignatureVerificationException;
 import com.auth0.jwt.interfaces.DecodedJWT;
 import com.auth0.jwt.interfaces.ECDSAKeyProvider;
+import com.auth0.jwt.interfaces.PrivateKeyDetail;
 import com.auth0.jwt.interfaces.RSAKeyProvider;
 
+import java.security.PrivateKey;
 import java.security.interfaces.*;
 
 /**
@@ -318,7 +320,7 @@ protected Algorithm(String name, String description) {
      *
      * @return the Key Id that identifies the Signing Key or null if it's not specified.
      */
-    public String getSigningKeyId() {
+    public PrivateKeyDetail<?> getPrivateKeyDetails() {
         return null;
     }
 
@@ -367,14 +369,32 @@ public String toString() {
      * @throws SignatureGenerationException if the Key is invalid.
      */
     public byte[] sign(byte[] headerBytes, byte[] payloadBytes) throws SignatureGenerationException {
+        PrivateKey privateKey = getPrivateKeyDetails() == null ? null : getPrivateKeyDetails().getPrivateKey();
+        return sign(headerBytes, payloadBytes, privateKey);
+    }
+
+    /**
+     * Sign the given content using this Algorithm instance.
+     *
+     * @param headerBytes  an array of bytes representing the base64 encoded header content
+     *                     to be verified against the signature.
+     * @param payloadBytes an array of bytes representing the base64 encoded payload content
+     *                     to be verified against the signature.
+     * @param privateKey   A private key with which the content has to be signed
+     *
+     * @return the signature in a base64 encoded array of bytes
+     * @throws SignatureGenerationException if the Key is invalid.
+     */
+    public byte[] sign(byte[] headerBytes, byte[] payloadBytes, PrivateKey privateKey)
+            throws SignatureGenerationException {
         // default implementation; keep around until sign(byte[]) method is removed
         byte[] contentBytes = new byte[headerBytes.length + 1 + payloadBytes.length];
 
         System.arraycopy(headerBytes, 0, contentBytes, 0, headerBytes.length);
         contentBytes[headerBytes.length] = (byte) '.';
         System.arraycopy(payloadBytes, 0, contentBytes, headerBytes.length + 1, payloadBytes.length);
 
-        return sign(contentBytes);
+        return sign(contentBytes, privateKey);
     }
 
     /**
@@ -386,7 +406,22 @@ public byte[] sign(byte[] headerBytes, byte[] payloadBytes) throws SignatureGene
      * @return the signature in a base64 encoded array of bytes
      * @throws SignatureGenerationException if the Key is invalid.
      */
+    public byte[] sign(byte[] contentBytes) throws SignatureGenerationException {
+        PrivateKey privateKey = getPrivateKeyDetails() == null ? null : getPrivateKeyDetails().getPrivateKey();
+        return sign(contentBytes, privateKey);
+    }
 
-    public abstract byte[] sign(byte[] contentBytes) throws SignatureGenerationException;
+    /**
+     * Sign the given content using this Algorithm instance.
+     * To get the correct JWT Signature, ensure the content is in the format {HEADER}.{PAYLOAD}
+     *
+     * @param contentBytes an array of bytes representing the base64 encoded content
+     *                     to be verified against the signature.
+     * @param privateKey   A private key with which the content has to be signed
+     *
+     * @return the signature in a base64 encoded array of bytes
+     * @throws SignatureGenerationException if the Key is invalid.
+     */
+    public abstract byte[] sign(byte[] contentBytes, PrivateKey privateKey) throws SignatureGenerationException;
 
 }

lib/src/main/java/com/auth0/jwt/algorithms/CryptoHelper.java

@@ -112,33 +112,6 @@ boolean verifySignatureFor(
         return s.verify(signatureBytes);
     }
 
-    /**
-     * Create signature for JWT header and payload using a private key.
-     *
-     * @param algorithm    algorithm name.
-     * @param privateKey   the private key to use for signing.
-     * @param headerBytes  JWT header.
-     * @param payloadBytes JWT payload.
-     * @return the signature bytes.
-     * @throws NoSuchAlgorithmException if the algorithm is not supported.
-     * @throws InvalidKeyException      if the given key is inappropriate for initializing the specified algorithm.
-     * @throws SignatureException       if this signature object is not initialized properly
-     *                                  or if this signature algorithm is unable to process the input data provided.
-     */
-    byte[] createSignatureFor(
-            String algorithm,
-            PrivateKey privateKey,
-            byte[] headerBytes,
-            byte[] payloadBytes
-    ) throws NoSuchAlgorithmException, InvalidKeyException, SignatureException {
-        final Signature s = Signature.getInstance(algorithm);
-        s.initSign(privateKey);
-        s.update(headerBytes);
-        s.update(JWT_PART_SEPARATOR);
-        s.update(payloadBytes);
-        return s.sign();
-    }
-
     /**
      * Create signature for JWT header and payload.
      *

lib/src/main/java/com/auth0/jwt/algorithms/ECDSAAlgorithm.java

@@ -4,9 +4,11 @@
 import com.auth0.jwt.exceptions.SignatureVerificationException;
 import com.auth0.jwt.interfaces.DecodedJWT;
 import com.auth0.jwt.interfaces.ECDSAKeyProvider;
+import com.auth0.jwt.interfaces.PrivateKeyDetail;
 
 import java.security.InvalidKeyException;
 import java.security.NoSuchAlgorithmException;
+import java.security.PrivateKey;
 import java.security.SignatureException;
 import java.security.interfaces.ECPrivateKey;
 import java.security.interfaces.ECPublicKey;
@@ -61,23 +63,8 @@ public void verify(DecodedJWT jwt) throws SignatureVerificationException {
     }
 
     @Override
-    public byte[] sign(byte[] headerBytes, byte[] payloadBytes) throws SignatureGenerationException {
+    public byte[] sign(byte[] contentBytes, PrivateKey privateKey) throws SignatureGenerationException {
         try {
-            ECPrivateKey privateKey = keyProvider.getPrivateKey();
-            if (privateKey == null) {
-                throw new IllegalStateException("The given Private Key is null.");
-            }
-            byte[] signature = crypto.createSignatureFor(getDescription(), privateKey, headerBytes, payloadBytes);
-            return DERToJOSE(signature);
-        } catch (NoSuchAlgorithmException | SignatureException | InvalidKeyException | IllegalStateException e) {
-            throw new SignatureGenerationException(this, e);
-        }
-    }
-
-    @Override
-    public byte[] sign(byte[] contentBytes) throws SignatureGenerationException {
-        try {
-            ECPrivateKey privateKey = keyProvider.getPrivateKey();
             if (privateKey == null) {
                 throw new IllegalStateException("The given Private Key is null.");
             }
@@ -89,8 +76,8 @@ public byte[] sign(byte[] contentBytes) throws SignatureGenerationException {
     }
 
     @Override
-    public String getSigningKeyId() {
-        return keyProvider.getPrivateKeyId();
+    public PrivateKeyDetail<ECPrivateKey> getPrivateKeyDetails() {
+        return keyProvider.getPrivateKeyDetails();
     }
 
     //Visible for testing

lib/src/main/java/com/auth0/jwt/algorithms/HMACAlgorithm.java

@@ -7,6 +7,7 @@
 import java.nio.charset.StandardCharsets;
 import java.security.InvalidKeyException;
 import java.security.NoSuchAlgorithmException;
+import java.security.PrivateKey;
 import java.util.Arrays;
 import java.util.Base64;
 
@@ -62,16 +63,7 @@ public void verify(DecodedJWT jwt) throws SignatureVerificationException {
     }
 
     @Override
-    public byte[] sign(byte[] headerBytes, byte[] payloadBytes) throws SignatureGenerationException {
-        try {
-            return crypto.createSignatureFor(getDescription(), secret, headerBytes, payloadBytes);
-        } catch (NoSuchAlgorithmException | InvalidKeyException e) {
-            throw new SignatureGenerationException(this, e);
-        }
-    }
-
-    @Override
-    public byte[] sign(byte[] contentBytes) throws SignatureGenerationException {
+    public byte[] sign(byte[] contentBytes, PrivateKey privateKey) throws SignatureGenerationException {
         try {
             return crypto.createSignatureFor(getDescription(), secret, contentBytes);
         } catch (NoSuchAlgorithmException | InvalidKeyException e) {

lib/src/main/java/com/auth0/jwt/algorithms/NoneAlgorithm.java

@@ -3,6 +3,8 @@
 import com.auth0.jwt.exceptions.SignatureGenerationException;
 import com.auth0.jwt.exceptions.SignatureVerificationException;
 import com.auth0.jwt.interfaces.DecodedJWT;
+
+import java.security.PrivateKey;
 import java.util.Base64;
 
 class NoneAlgorithm extends Algorithm {
@@ -33,4 +35,9 @@ public byte[] sign(byte[] headerBytes, byte[] payloadBytes) throws SignatureGene
     public byte[] sign(byte[] contentBytes) throws SignatureGenerationException {
         return new byte[0];
     }
+
+    @Override
+    public byte[] sign(byte[] contentBytes, PrivateKey privateKey) throws SignatureGenerationException {
+        return new byte[0];
+    }
 }

lib/src/main/java/com/auth0/jwt/algorithms/RSAAlgorithm.java

@@ -3,11 +3,13 @@
 import com.auth0.jwt.exceptions.SignatureGenerationException;
 import com.auth0.jwt.exceptions.SignatureVerificationException;
 import com.auth0.jwt.interfaces.DecodedJWT;
+import com.auth0.jwt.interfaces.PrivateKeyDetail;
 import com.auth0.jwt.interfaces.RSAKeyProvider;
 
 import java.nio.charset.StandardCharsets;
 import java.security.InvalidKeyException;
 import java.security.NoSuchAlgorithmException;
+import java.security.PrivateKey;
 import java.security.SignatureException;
 import java.security.interfaces.RSAPrivateKey;
 import java.security.interfaces.RSAPublicKey;
@@ -58,22 +60,8 @@ public void verify(DecodedJWT jwt) throws SignatureVerificationException {
     }
 
     @Override
-    public byte[] sign(byte[] headerBytes, byte[] payloadBytes) throws SignatureGenerationException {
+    public byte[] sign(byte[] contentBytes, PrivateKey privateKey) throws SignatureGenerationException {
         try {
-            RSAPrivateKey privateKey = keyProvider.getPrivateKey();
-            if (privateKey == null) {
-                throw new IllegalStateException("The given Private Key is null.");
-            }
-            return crypto.createSignatureFor(getDescription(), privateKey, headerBytes, payloadBytes);
-        } catch (NoSuchAlgorithmException | SignatureException | InvalidKeyException | IllegalStateException e) {
-            throw new SignatureGenerationException(this, e);
-        }
-    }
-
-    @Override
-    public byte[] sign(byte[] contentBytes) throws SignatureGenerationException {
-        try {
-            RSAPrivateKey privateKey = keyProvider.getPrivateKey();
             if (privateKey == null) {
                 throw new IllegalStateException("The given Private Key is null.");
             }
@@ -84,8 +72,8 @@ public byte[] sign(byte[] contentBytes) throws SignatureGenerationException {
     }
 
     @Override
-    public String getSigningKeyId() {
-        return keyProvider.getPrivateKeyId();
+    public PrivateKeyDetail<RSAPrivateKey> getPrivateKeyDetails() {
+        return keyProvider.getPrivateKeyDetails();
     }
 
     //Visible for testing

lib/src/main/java/com/auth0/jwt/interfaces/KeyProvider.java

@@ -24,6 +24,7 @@ interface KeyProvider<U extends PublicKey, R extends PrivateKey> {
      * Getter for the Private Key instance. Used to sign the content on the JWT signing stage.
      *
      * @return the Private Key instance
+     * @deprecated Use {@link KeyProvider#getPrivateKeyDetails()} instead
      */
     R getPrivateKey();
 
@@ -32,6 +33,27 @@ interface KeyProvider<U extends PublicKey, R extends PrivateKey> {
      * This represents the `kid` claim and will be placed in the Header.
      *
      * @return the Key Id that identifies the Private Key or null if it's not specified.
+     * @deprecated Use {@link KeyProvider#getPrivateKeyDetails()} instead
      */
     String getPrivateKeyId();
+
+    /**
+     * Getter for the Private Key instance along with its Id.
+     * Used to sign the content on the JWT signing stage.
+     *
+     * @return the Private Key Details instance
+     */
+    default PrivateKeyDetail<R> getPrivateKeyDetails() {
+        return new PrivateKeyDetail<R>() {
+            @Override
+            public R getPrivateKey() {
+                return KeyProvider.this.getPrivateKey();
+            }
+
+            @Override
+            public String getPrivateKeyId() {
+                return KeyProvider.this.getPrivateKeyId();
+            }
+        };
+    }
 }

lib/src/main/java/com/auth0/jwt/interfaces/PrivateKeyDetail.java

@@ -0,0 +1,25 @@
+package com.auth0.jwt.interfaces;
+
+import java.security.PrivateKey;
+
+/**
+ * Generic representation of Private Key with its Key ID.
+ *
+ * @param <R> the class that represents the Private Key
+ */
+public interface PrivateKeyDetail<R extends PrivateKey> {
+    /**
+     * Getter for the Private Key instance. Used to sign the content on the JWT signing stage.
+     *
+     * @return the Private Key instance
+     */
+    R getPrivateKey();
+
+    /**
+     * Getter for the Id of the Private Key used to sign the tokens.
+     * This represents the `kid` claim and will be placed in the Header.
+     *
+     * @return the Key Id that identifies the Private Key or null if it's not specified.
+     */
+    String getPrivateKeyId();
+}