Support AWS ALB token verification (#707)

* Support AWS ALB token verification

* Cache AWS ALB keys

* Add properties to control the key cache

doc/modules/ROOT/pages/configuration.adoc

@@ -42,9 +42,12 @@ SmallRye JWT supports many properties which can be used to customize the token p
 |smallrye.jwt.verify.key.location|NONE|Location of the verification key which can point to both public and secret keys. Secret keys can only be in the JWK format. Note that 'mp.jwt.verify.publickey.location' will be ignored if this property is set.
 |smallrye.jwt.verify.algorithm|`RS256`|Signature algorithm. Set it to `ES256` to support the Elliptic Curve signature algorithm. This property is deprecated, use `mp.jwt.verify.publickey.algorithm`.
 |smallrye.jwt.verify.key-format|`ANY`|Set this property to a specific key format such as `PEM_KEY`, `PEM_CERTIFICATE`, `JWK` or `JWK_BASE64URL` to optimize the way the verification key is loaded.
+|smallrye.jwt.verify.key-provider|`DEFAULT`|By default, PEM, JWK or JWK key sets can be read from the local file system or fetched from URIs as required by MicroProfile JWT specification. Set this property to `AWS_ALB` to support an AWS Application Load Balancer verification key resolution.
 |smallrye.jwt.verify.relax-key-validation|false|Relax the validation of the verification keys, setting this property to `true` will allow public RSA keys with the length less than 2048 bit.
 |smallrye.jwt.verify.certificate-thumbprint|false|If this property is enabled then a signed token must contain either 'x5t' or 'x5t#S256' X509Certificate thumbprint headers. Verification keys can only be in JWK or PEM Certificate key formats in this case. JWK keys must have a 'x5c' (Base64-encoded X509Certificate) property set.
 |smallrye.jwt.token.header|`Authorization`|Set this property if another header such as `Cookie` is used to pass the token. This property is deprecated, use `mp.jwt.token.header`.
+|smallrye.jwt.key-cache-size|`100`|Key cache size. Use this property, as well as `smallrye.jwt.key-cache-time-to-live`, to control the key cache when a key provider such as `AWS_ALB` is configured with `smallrye.jwt.verify.key-provider=AWS_ALB` for resolving the keys dynamically.
+|smallrye.jwt.key-cache-time-to-live|`10`|Key cache entry time-to-live in minutes. Use this property, as well as `smallrye.jwt.key-cache-size`,  to control the key cache when a key provider such as `AWS_ALB` is configured with `smallrye.jwt.verify.key-provider=AWS_ALB` for resolving the keys dynamically.
 |smallrye.jwt.token.cookie|none|Name of the cookie containing a token. This property will be effective only if  `smallrye.jwt.token.header` is set to `Cookie`. This property is deprecated, use `mp.jwt.token.cookie`.
 |smallrye.jwt.always-check-authorization|false|Set this property to `true` for `Authorization` header be checked even if the `smallrye.jwt.token.header` is set to `Cookie` but no cookie with a `smallrye.jwt.token.cookie` name exists.
 |smallrye.jwt.token.schemes|`Bearer`|Comma-separated list containing an alternative single or multiple schemes, for example, `DPoP`.

implementation/jwt-auth/src/main/java/io/smallrye/jwt/KeyProvider.java

@@ -0,0 +1,20 @@
+package io.smallrye.jwt;
+
+/**
+ * Well-known key providers.
+ */
+public enum KeyProvider {
+    /**
+     * AWS Application Load Balancer.
+     *
+     * Verification key in PEM format is fetched from the URI which is created by
+     * adding the current token `kid` (key identifier) header value to the AWS ALB URI.
+     */
+    AWS_ALB,
+
+    /**
+     * Verification key is resolved as required by the MP JWT specification:
+     * PEM or JWK key or JWK key set can be read from the local file system or fetched from URIs.
+     */
+    DEFAULT
+}

implementation/jwt-auth/src/main/java/io/smallrye/jwt/auth/principal/AbstractKeyLocationResolver.java

@@ -314,7 +314,7 @@ protected Key getSecretKeyFromJwk(JsonWebKey jwk) {
         return null;
     }
 
-    protected X509Certificate loadPEMCertificate(String content) {
+    protected static X509Certificate loadPEMCertificate(String content) {
         PrincipalLogging.log.checkKeyContentIsBase64EncodedPEMCertificate();
         X509Certificate cert = null;
         try {

implementation/jwt-auth/src/main/java/io/smallrye/jwt/auth/principal/AwsAlbKeyResolver.java

@@ -0,0 +1,171 @@
+package io.smallrye.jwt.auth.principal;
+
+import java.io.IOException;
+import java.security.Key;
+import java.time.Duration;
+import java.util.HashMap;
+import java.util.Iterator;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.atomic.AtomicInteger;
+
+import org.jose4j.http.Get;
+import org.jose4j.http.SimpleGet;
+import org.jose4j.http.SimpleResponse;
+import org.jose4j.jwk.JsonWebKey;
+import org.jose4j.jws.JsonWebSignature;
+import org.jose4j.jwx.JsonWebStructure;
+import org.jose4j.keys.resolvers.VerificationKeyResolver;
+import org.jose4j.lang.UnresolvableKeyException;
+
+import io.smallrye.jwt.auth.principal.AbstractKeyLocationResolver.TrustAllHostnameVerifier;
+import io.smallrye.jwt.auth.principal.AbstractKeyLocationResolver.TrustedHostsHostnameVerifier;
+import io.smallrye.jwt.util.KeyUtils;
+import io.smallrye.jwt.util.ResourceUtils;
+
+public class AwsAlbKeyResolver implements VerificationKeyResolver {
+    private JWTAuthContextInfo authContextInfo;
+    private long cacheTimeToLive;
+    private Map<String, CacheEntry> keys = new HashMap<>();
+    private AtomicInteger size = new AtomicInteger();
+
+    public AwsAlbKeyResolver(JWTAuthContextInfo authContextInfo) throws UnresolvableKeyException {
+        if (authContextInfo.getPublicKeyLocation() == null) {
+            throw PrincipalMessages.msg.nullKeyLocation();
+        }
+        this.authContextInfo = authContextInfo;
+        this.cacheTimeToLive = Duration.ofMinutes(authContextInfo.getKeyCacheTimeToLive()).toMillis();
+    }
+
+    @Override
+    public Key resolveKey(JsonWebSignature jws, List<JsonWebStructure> nestingContext) throws UnresolvableKeyException {
+        String kid = jws.getHeaders().getStringHeaderValue(JsonWebKey.KEY_ID_PARAMETER);
+        verifyKid(kid);
+
+        CacheEntry entry = findValidCacheEntry(kid);
+        if (entry != null) {
+            return entry.key;
+        } else if (prepareSpaceForNewCacheEntry()) {
+            entry = new CacheEntry(retrieveKey(kid));
+            keys.put(kid, entry);
+            return entry.key;
+        } else {
+            return retrieveKey(kid);
+        }
+
+    }
+
+    protected Key retrieveKey(String kid) throws UnresolvableKeyException {
+        String keyLocation = authContextInfo.getPublicKeyLocation() + "/" + kid;
+        SimpleResponse simpleResponse = null;
+        try {
+            simpleResponse = getHttpGet().get(keyLocation);
+        } catch (IOException ex) {
+            AbstractKeyLocationResolver.reportLoadKeyException(null, keyLocation, ex);
+        }
+        String keyContent = simpleResponse.getBody();
+        try {
+            return KeyUtils.decodePublicKey(keyContent, authContextInfo.getSignatureAlgorithm());
+        } catch (Exception e) {
+            AbstractKeyLocationResolver.reportUnresolvableKeyException(keyContent, keyLocation);
+        }
+        return null;
+    }
+
+    protected SimpleGet getHttpGet() throws UnresolvableKeyException {
+        Get httpGet = new Get();
+        if (authContextInfo.isTlsTrustAll()) {
+            httpGet.setHostnameVerifier(new TrustAllHostnameVerifier());
+        } else if (authContextInfo.getTlsTrustedHosts() != null) {
+            httpGet.setHostnameVerifier(new TrustedHostsHostnameVerifier(authContextInfo.getTlsTrustedHosts()));
+        }
+        if (authContextInfo.getTlsCertificate() != null) {
+            httpGet.setTrustedCertificates(
+                    AbstractKeyLocationResolver.loadPEMCertificate(authContextInfo.getTlsCertificate()));
+        } else if (authContextInfo.getTlsCertificatePath() != null) {
+            httpGet.setTrustedCertificates(AbstractKeyLocationResolver.loadPEMCertificate(
+                    readKeyContent(authContextInfo.getTlsCertificatePath())));
+        }
+        return httpGet;
+    }
+
+    protected String readKeyContent(String keyLocation) throws UnresolvableKeyException {
+        try {
+            String content = ResourceUtils.readResource(keyLocation);
+            if (content == null) {
+                throw PrincipalMessages.msg.resourceNotFound(keyLocation);
+            }
+            return content;
+        } catch (IOException ex) {
+            AbstractKeyLocationResolver.reportLoadKeyException(null, keyLocation, ex);
+            return null;
+        }
+    }
+
+    private void verifyKid(String kid) throws UnresolvableKeyException {
+        if (kid == null) {
+            throw PrincipalMessages.msg.nullKeyIdentifier();
+        }
+        String expectedKid = authContextInfo.getTokenKeyId();
+        if (expectedKid != null && !kid.equals(expectedKid)) {
+            PrincipalLogging.log.invalidTokenKidHeader(kid, expectedKid);
+            throw PrincipalMessages.msg.invalidTokenKid();
+        }
+    }
+
+    private void removeInvalidEntries() {
+        long now = now();
+        for (Iterator<Map.Entry<String, CacheEntry>> it = keys.entrySet().iterator(); it.hasNext();) {
+            Map.Entry<String, CacheEntry> next = it.next();
+            if (isEntryExpired(next.getValue(), now)) {
+                it.remove();
+                size.decrementAndGet();
+            }
+        }
+    }
+
+    private boolean prepareSpaceForNewCacheEntry() {
+        int currentSize;
+        do {
+            currentSize = size.get();
+            if (currentSize == authContextInfo.getKeyCacheSize()) {
+                removeInvalidEntries();
+                if (currentSize == authContextInfo.getKeyCacheSize()) {
+                    return false;
+                }
+            }
+        } while (!size.compareAndSet(currentSize, currentSize + 1));
+        return true;
+    }
+
+    private CacheEntry findValidCacheEntry(String kid) {
+        CacheEntry entry = keys.get(kid);
+        if (entry != null) {
+            long now = now();
+            if (isEntryExpired(entry, now)) {
+                // Entry has expired, remote introspection will be required
+                entry = null;
+                keys.remove(kid);
+                size.decrementAndGet();
+            }
+        }
+        return entry;
+    }
+
+    private boolean isEntryExpired(CacheEntry entry, long now) {
+        return entry.createdTime + cacheTimeToLive < now;
+    }
+
+    private static long now() {
+        return System.currentTimeMillis();
+    }
+
+    private static class CacheEntry {
+        volatile Key key;
+        long createdTime = System.currentTimeMillis();
+
+        public CacheEntry(Key key) {
+            this.key = key;
+        }
+    }
+}

implementation/jwt-auth/src/main/java/io/smallrye/jwt/auth/principal/DefaultJWTTokenParser.java

@@ -36,6 +36,7 @@
 import org.jose4j.lang.JoseException;
 import org.jose4j.lang.UnresolvableKeyException;
 
+import io.smallrye.jwt.KeyProvider;
 import io.smallrye.jwt.algorithm.KeyEncryptionAlgorithm;
 
 /**
@@ -257,9 +258,13 @@ protected VerificationKeyResolver getVerificationKeyResolver(JWTAuthContextInfo
         if (keyResolver == null) {
             synchronized (this) {
                 if (keyResolver == null)
-                    keyResolver = authContextInfo.isVerifyCertificateThumbprint()
-                            ? new X509KeyLocationResolver(authContextInfo)
-                            : new KeyLocationResolver(authContextInfo);
+                    if (KeyProvider.AWS_ALB == authContextInfo.getKeyProvider()) {
+                        keyResolver = new AwsAlbKeyResolver(authContextInfo);
+                    } else {
+                        keyResolver = authContextInfo.isVerifyCertificateThumbprint()
+                                ? new X509KeyLocationResolver(authContextInfo)
+                                : new KeyLocationResolver(authContextInfo);
+                    }
             }
         }
         return keyResolver;

implementation/jwt-auth/src/main/java/io/smallrye/jwt/auth/principal/JWTAuthContextInfo.java

@@ -27,6 +27,7 @@
 import javax.crypto.SecretKey;
 
 import io.smallrye.jwt.KeyFormat;
+import io.smallrye.jwt.KeyProvider;
 import io.smallrye.jwt.algorithm.KeyEncryptionAlgorithm;
 import io.smallrye.jwt.algorithm.SignatureAlgorithm;
 
@@ -47,7 +48,7 @@ public class JWTAuthContextInfo {
     private String publicKeyContent;
     private String decryptionKeyLocation;
     private String decryptionKeyContent;
-    private Integer jwksRefreshInterval;
+    private Integer jwksRefreshInterval = 60;
     private int forcedJwksRefreshInterval = 30;
     private String tokenHeader = "Authorization";
     private String tokenCookie;
@@ -64,6 +65,7 @@ public class JWTAuthContextInfo {
     private Set<KeyEncryptionAlgorithm> keyEncryptionAlgorithm = new HashSet<>(Arrays.asList(KeyEncryptionAlgorithm.RSA_OAEP,
             KeyEncryptionAlgorithm.RSA_OAEP_256));
     private KeyFormat keyFormat = KeyFormat.ANY;
+    private KeyProvider keyProvider = KeyProvider.DEFAULT;
     private Set<String> expectedAudience;
     private String groupsSeparator = " ";
     private Set<String> requiredClaims;
@@ -75,6 +77,8 @@ public class JWTAuthContextInfo {
     private boolean tlsTrustAll;
     private String httpProxyHost;
     private int httpProxyPort;
+    private int keyCacheSize = 100;
+    private int keyCacheTimeToLive = 10;
 
     public JWTAuthContextInfo() {
     }
@@ -129,6 +133,9 @@ public JWTAuthContextInfo(JWTAuthContextInfo orig) {
         this.signatureAlgorithm = orig.getSignatureAlgorithm();
         this.keyEncryptionAlgorithm = orig.getKeyEncryptionAlgorithm();
         this.keyFormat = orig.getKeyFormat();
+        this.keyProvider = orig.getKeyProvider();
+        this.keyCacheSize = orig.getKeyCacheSize();
+        this.keyCacheTimeToLive = orig.getKeyCacheTimeToLive();
         this.expectedAudience = orig.getExpectedAudience();
         this.groupsSeparator = orig.getGroupsSeparator();
         this.requiredClaims = orig.getRequiredClaims();
@@ -378,6 +385,14 @@ public void setKeyFormat(KeyFormat keyFormat) {
         this.keyFormat = keyFormat;
     }
 
+    public KeyProvider getKeyProvider() {
+        return keyProvider;
+    }
+
+    public void setKeyProvider(KeyProvider keyProvider) {
+        this.keyProvider = keyProvider;
+    }
+
     public boolean isAlwaysCheckAuthorization() {
         return alwaysCheckAuthorization;
     }
@@ -424,6 +439,9 @@ public String toString() {
                 ", signatureAlgorithm=" + signatureAlgorithm +
                 ", keyEncryptionAlgorithm=" + keyEncryptionAlgorithm +
                 ", keyFormat=" + keyFormat +
+                ", keyProvider=" + keyProvider +
+                ", keyCacheSize=" + keyCacheSize +
+                ", keyCacheTimeToLive=" + keyCacheTimeToLive +
                 ", expectedAudience=" + expectedAudience +
                 ", groupsSeparator='" + groupsSeparator + '\'' +
                 ", relaxVerificationKeyValidation=" + relaxVerificationKeyValidation +
@@ -515,4 +533,20 @@ public int getClockSkew() {
     public void setClockSkew(int clockSkew) {
         this.clockSkew = clockSkew;
     }
+
+    public int getKeyCacheTimeToLive() {
+        return keyCacheTimeToLive;
+    }
+
+    public void setKeyCacheTimeToLive(int keyCacheTimeToLive) {
+        this.keyCacheTimeToLive = keyCacheTimeToLive;
+    }
+
+    public int getKeyCacheSize() {
+        return keyCacheSize;
+    }
+
+    public void setKeyCacheSize(int keyCacheSize) {
+        this.keyCacheSize = keyCacheSize;
+    }
 }

implementation/jwt-auth/src/main/java/io/smallrye/jwt/auth/principal/PrincipalMessages.java

@@ -74,4 +74,10 @@ interface PrincipalMessages {
 
     @Message(id = 7018, value = "The token age has exceeded %d seconds")
     ParseException tokenAgeExceeded(long tokenAge);
+
+    @Message(id = 7019, value = "Required key location is null")
+    UnresolvableKeyException nullKeyLocation();
+
+    @Message(id = 7020, value = "Required key identifier is null")
+    UnresolvableKeyException nullKeyIdentifier();
 }