Refactor JWT realm to facilitate new feature changes (#91001)

This is a refactoring PR that restructure how code for the JWT realm is
structured. The goal is to make it ready for incoming new feature related
changes by improve code reuse, modularity, readability and testability. 
This PR should have no functional changes.

x-pack/plugin/security/src/main/java/org/elasticsearch/xpack/security/authc/jwt/JwkSetLoader.java

@@ -0,0 +1,200 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+package org.elasticsearch.xpack.security.authc.jwt;
+
+import com.nimbusds.jose.jwk.JWK;
+
+import org.apache.http.impl.nio.client.CloseableHttpAsyncClient;
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
+import org.elasticsearch.action.ActionListener;
+import org.elasticsearch.action.support.PlainActionFuture;
+import org.elasticsearch.common.Strings;
+import org.elasticsearch.common.hash.MessageDigests;
+import org.elasticsearch.common.util.concurrent.ListenableFuture;
+import org.elasticsearch.core.Nullable;
+import org.elasticsearch.core.Releasable;
+import org.elasticsearch.xpack.core.security.authc.RealmConfig;
+import org.elasticsearch.xpack.core.security.authc.RealmSettings;
+import org.elasticsearch.xpack.core.security.authc.jwt.JwtRealmSettings;
+import org.elasticsearch.xpack.core.ssl.SSLService;
+
+import java.io.IOException;
+import java.net.URI;
+import java.nio.charset.StandardCharsets;
+import java.util.Arrays;
+import java.util.List;
+import java.util.Objects;
+import java.util.concurrent.atomic.AtomicReference;
+
+/**
+ * This class is responsible for loading the JWK set for PKC signature from either a file or URL.
+ * The JWK set is loaded once when the class is instantiated. Subsequent reloading is triggered
+ * by invoking the {@link #reload(ActionListener)} method. The updated JWK set can be retrieved with
+ * the {@link #getContentAndJwksAlgs()} method once loading or reloading is completed.
+ */
+public class JwkSetLoader implements Releasable {
+
+    private static final Logger logger = LogManager.getLogger(JwkSetLoader.class);
+
+    private final AtomicReference<ListenableFuture<Boolean>> reloadFutureRef = new AtomicReference<>();
+    private final RealmConfig realmConfig;
+    private final List<String> allowedJwksAlgsPkc;
+    private final String jwkSetPath;
+    @Nullable
+    private final URI jwkSetPathUri;
+    @Nullable
+    private final CloseableHttpAsyncClient httpClient;
+    private volatile ContentAndJwksAlgs contentAndJwksAlgs = null;
+
+    public JwkSetLoader(final RealmConfig realmConfig, List<String> allowedJwksAlgsPkc, final SSLService sslService) {
+        this.realmConfig = realmConfig;
+        this.allowedJwksAlgsPkc = allowedJwksAlgsPkc;
+        // PKC JWKSet can be URL, file, or not set; only initialize HTTP client if PKC JWKSet is a URL.
+        this.jwkSetPath = realmConfig.getSetting(JwtRealmSettings.PKC_JWKSET_PATH);
+        assert Strings.hasText(this.jwkSetPath);
+        this.jwkSetPathUri = JwtUtil.parseHttpsUri(jwkSetPath);
+        if (this.jwkSetPathUri == null) {
+            this.httpClient = null;
+        } else {
+            this.httpClient = JwtUtil.createHttpClient(realmConfig, sslService);
+        }
+
+        // Any exception during loading requires closing JwkSetLoader's HTTP client to avoid a thread pool leak
+        try {
+            final PlainActionFuture<Boolean> future = new PlainActionFuture<>();
+            reload(future);
+            // ASSUME: Blocking read operations are OK during startup
+            final Boolean isUpdated = future.actionGet();
+            assert isUpdated : "initial reload should have updated the JWK set";
+        } catch (Throwable t) {
+            close();
+            throw t;
+        }
+    }
+
+    /**
+     * Reload the JWK sets, compare to existing JWK sets and update it to the reloaded value if
+     * they are different. The listener is called with false if the reloaded content is the same
+     * as the existing one or true if they are different.
+     */
+    void reload(final ActionListener<Boolean> listener) {
+        final ListenableFuture<Boolean> future = this.getFuture();
+        future.addListener(listener);
+    }
+
+    ContentAndJwksAlgs getContentAndJwksAlgs() {
+        return contentAndJwksAlgs;
+    }
+
+    private ListenableFuture<Boolean> getFuture() {
+        for (;;) {
+            final ListenableFuture<Boolean> existingFuture = this.reloadFutureRef.get();
+            if (existingFuture != null) {
+                return existingFuture;
+            }
+
+            final ListenableFuture<Boolean> newFuture = new ListenableFuture<>();
+            if (this.reloadFutureRef.compareAndSet(null, newFuture)) {
+                loadInternal(ActionListener.runAfter(newFuture, () -> {
+                    final ListenableFuture<Boolean> oldValue = this.reloadFutureRef.getAndSet(null);
+                    assert oldValue == newFuture : "future reference changed unexpectedly";
+                }));
+                return newFuture;
+            }
+            // else, Another thread set the future-ref before us, just try it all again
+        }
+    }
+
+    private void loadInternal(final ActionListener<Boolean> listener) {
+        // PKC JWKSet get contents from local file or remote HTTPS URL
+        if (httpClient == null) {
+            logger.trace("Loading PKC JWKs from path [{}]", jwkSetPath);
+            final byte[] reloadedBytes = JwtUtil.readFileContents(
+                RealmSettings.getFullSettingKey(realmConfig, JwtRealmSettings.PKC_JWKSET_PATH),
+                jwkSetPath,
+                realmConfig.env()
+            );
+            listener.onResponse(handleReloadedContentAndJwksAlgs(reloadedBytes));
+        } else {
+            logger.trace("Loading PKC JWKs from https URI [{}]", jwkSetPathUri);
+            JwtUtil.readUriContents(
+                RealmSettings.getFullSettingKey(realmConfig, JwtRealmSettings.PKC_JWKSET_PATH),
+                jwkSetPathUri,
+                httpClient,
+                listener.map(reloadedBytes -> {
+                    logger.trace("Loaded bytes [{}] from [{}]", reloadedBytes.length, jwkSetPathUri);
+                    return handleReloadedContentAndJwksAlgs(reloadedBytes);
+                })
+            );
+        }
+    }
+
+    private Boolean handleReloadedContentAndJwksAlgs(byte[] bytes) {
+        final ContentAndJwksAlgs newContentAndJwksAlgs = parseContent(bytes);
+        if (contentAndJwksAlgs != null && Arrays.equals(contentAndJwksAlgs.sha256, newContentAndJwksAlgs.sha256)) {
+            logger.debug("No change in reloaded JWK set");
+            return false;
+        } else {
+            logger.debug("Reloaded JWK set is different from the existing set");
+            contentAndJwksAlgs = newContentAndJwksAlgs;
+            return true;
+        }
+    }
+
+    private ContentAndJwksAlgs parseContent(final byte[] jwkSetContentBytesPkc) {
+        final String jwkSetContentsPkc = new String(jwkSetContentBytesPkc, StandardCharsets.UTF_8);
+        final byte[] jwkSetContentsPkcSha256 = JwtUtil.sha256(jwkSetContentsPkc);
+
+        // PKC JWKSet parse contents
+        final List<JWK> jwksPkc = JwkValidateUtil.loadJwksFromJwkSetString(
+            RealmSettings.getFullSettingKey(realmConfig, JwtRealmSettings.PKC_JWKSET_PATH),
+            jwkSetContentsPkc
+        );
+        // Filter JWK(s) vs signature algorithms. Only keep JWKs with a matching alg. Only keep algs with a matching JWK.
+        final JwksAlgs jwksAlgsPkc = JwkValidateUtil.filterJwksAndAlgorithms(jwksPkc, allowedJwksAlgsPkc);
+        logger.info(
+            "Usable PKC: JWKs=[{}] algorithms=[{}] sha256=[{}]",
+            jwksAlgsPkc.jwks().size(),
+            String.join(",", jwksAlgsPkc.algs()),
+            MessageDigests.toHexString(jwkSetContentsPkcSha256)
+        );
+        return new ContentAndJwksAlgs(jwkSetContentsPkcSha256, jwksAlgsPkc);
+    }
+
+    @Override
+    public void close() {
+        if (httpClient != null) {
+            try {
+                httpClient.close();
+            } catch (IOException e) {
+                logger.warn(() -> "Exception closing HTTPS client for realm [" + realmConfig.name() + "]", e);
+            }
+        }
+    }
+
+    // Filtered JWKs and Algs
+    record JwksAlgs(List<JWK> jwks, List<String> algs) {
+        JwksAlgs {
+            Objects.requireNonNull(jwks, "JWKs must not be null");
+            Objects.requireNonNull(algs, "Algs must not be null");
+        }
+
+        boolean isEmpty() {
+            return this.jwks.isEmpty() && this.algs.isEmpty();
+        }
+    }
+
+    // Original PKC JWKSet(for comparison during refresh), and filtered JWKs and Algs
+    record ContentAndJwksAlgs(byte[] sha256, JwksAlgs jwksAlgs) {
+        ContentAndJwksAlgs {
+            Objects.requireNonNull(jwksAlgs, "Filters JWKs and Algs must not be null");
+        }
+    }
+
+}

x-pack/plugin/security/src/main/java/org/elasticsearch/xpack/security/authc/jwt/JwkValidateUtil.java

@@ -40,7 +40,7 @@ public class JwkValidateUtil {
     private static final Logger LOGGER = LogManager.getLogger(JwkValidateUtil.class);
 
     // Static method for unit testing. No need to construct a complete RealmConfig with all settings.
-    static JwtRealm.JwksAlgs filterJwksAndAlgorithms(final List<JWK> jwks, final List<String> algs) throws SettingsException {
+    static JwkSetLoader.JwksAlgs filterJwksAndAlgorithms(final List<JWK> jwks, final List<String> algs) throws SettingsException {
         LOGGER.trace("JWKs [" + jwks.size() + "] and Algorithms [" + String.join(",", algs) + "] before filters.");
 
         final Predicate<JWK> keyUsePredicate = j -> ((j.getKeyUse() == null) || (KeyUse.SIGNATURE.equals(j.getKeyUse())));
@@ -57,7 +57,7 @@ static JwtRealm.JwksAlgs filterJwksAndAlgorithms(final List<JWK> jwks, final Lis
         final List<String> algsFiltered = algs.stream().filter(a -> (jwksFiltered.stream().anyMatch(j -> isMatch(j, a)))).toList();
         LOGGER.trace("Algorithms [" + String.join(",", algsFiltered) + "] after remaining JWKs [" + jwksFiltered.size() + "] filter.");
 
-        return new JwtRealm.JwksAlgs(jwksFiltered, algsFiltered);
+        return new JwkSetLoader.JwksAlgs(jwksFiltered, algsFiltered);
     }
 
     /**

x-pack/plugin/security/src/main/java/org/elasticsearch/xpack/security/authc/jwt/JwtAlgorithmValidator.java

@@ -0,0 +1,43 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+package org.elasticsearch.xpack.security.authc.jwt;
+
+import com.nimbusds.jose.JWSAlgorithm;
+import com.nimbusds.jose.JWSHeader;
+import com.nimbusds.jwt.JWTClaimsSet;
+
+import org.elasticsearch.ElasticsearchSecurityException;
+import org.elasticsearch.common.Strings;
+import org.elasticsearch.rest.RestStatus;
+
+import java.util.List;
+
+public class JwtAlgorithmValidator implements JwtFieldValidator {
+
+    private final List<String> allowedAlgorithms;
+
+    public JwtAlgorithmValidator(List<String> allowedAlgorithms) {
+        this.allowedAlgorithms = allowedAlgorithms;
+    }
+
+    public void validate(JWSHeader jwsHeader, JWTClaimsSet jwtClaimsSet) {
+        final JWSAlgorithm algorithm = jwsHeader.getAlgorithm();
+        if (algorithm == null) {
+            throw new ElasticsearchSecurityException("missing JWT algorithm header", RestStatus.BAD_REQUEST);
+        }
+
+        if (false == allowedAlgorithms.contains(algorithm.getName())) {
+            throw new ElasticsearchSecurityException(
+                "invalid JWT algorithm [{}], allowed algorithms are [{}]",
+                RestStatus.BAD_REQUEST,
+                algorithm,
+                Strings.collectionToCommaDelimitedString(allowedAlgorithms)
+            );
+        }
+    }
+}

x-pack/plugin/security/src/main/java/org/elasticsearch/xpack/security/authc/jwt/JwtAuthenticator.java

@@ -0,0 +1,122 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+package org.elasticsearch.xpack.security.authc.jwt;
+
+import com.nimbusds.jose.JWSHeader;
+import com.nimbusds.jwt.JWTClaimsSet;
+import com.nimbusds.jwt.SignedJWT;
+
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
+import org.elasticsearch.action.ActionListener;
+import org.elasticsearch.common.settings.SecureString;
+import org.elasticsearch.core.Releasable;
+import org.elasticsearch.core.TimeValue;
+import org.elasticsearch.xpack.core.security.authc.RealmConfig;
+import org.elasticsearch.xpack.core.security.authc.jwt.JwtRealmSettings;
+import org.elasticsearch.xpack.core.ssl.SSLService;
+
+import java.text.ParseException;
+import java.time.Clock;
+import java.util.List;
+
+/**
+ * This class performs validations of header, claims and signatures against the incoming {@link JwtAuthenticationToken}.
+ * It returns the {@link JWTClaimsSet} associated to the token if validation is successful.
+ * Note this class does not care about users nor its caching behaviour.
+ */
+public class JwtAuthenticator implements Releasable {
+
+    private static final Logger logger = LogManager.getLogger(JwtAuthenticator.class);
+    private final RealmConfig realmConfig;
+    private final List<JwtFieldValidator> jwtFieldValidators;
+    private final JwtSignatureValidator jwtSignatureValidator;
+
+    public JwtAuthenticator(
+        final RealmConfig realmConfig,
+        final SSLService sslService,
+        final JwtSignatureValidator.PkcJwkSetReloadNotifier reloadNotifier
+    ) {
+        this.realmConfig = realmConfig;
+        final TimeValue allowedClockSkew = realmConfig.getSetting(JwtRealmSettings.ALLOWED_CLOCK_SKEW);
+        final Clock clock = Clock.systemUTC();
+        this.jwtFieldValidators = List.of(
+            JwtTypeValidator.INSTANCE,
+            new JwtStringClaimValidator("iss", List.of(realmConfig.getSetting(JwtRealmSettings.ALLOWED_ISSUER)), true),
+            new JwtStringClaimValidator("aud", realmConfig.getSetting(JwtRealmSettings.ALLOWED_AUDIENCES), false),
+            new JwtAlgorithmValidator(realmConfig.getSetting(JwtRealmSettings.ALLOWED_SIGNATURE_ALGORITHMS)),
+            new JwtDateClaimValidator(clock, "iat", allowedClockSkew, JwtDateClaimValidator.Relationship.BEFORE_NOW, false),
+            new JwtDateClaimValidator(clock, "exp", allowedClockSkew, JwtDateClaimValidator.Relationship.AFTER_NOW, false),
+            new JwtDateClaimValidator(clock, "nbf", allowedClockSkew, JwtDateClaimValidator.Relationship.BEFORE_NOW, true),
+            new JwtDateClaimValidator(clock, "auth_time", allowedClockSkew, JwtDateClaimValidator.Relationship.BEFORE_NOW, true)
+        );
+
+        this.jwtSignatureValidator = new JwtSignatureValidator.DelegatingJwtSignatureValidator(realmConfig, sslService, reloadNotifier);
+    }
+
+    public void authenticate(JwtAuthenticationToken jwtAuthenticationToken, ActionListener<JWTClaimsSet> listener) {
+        final String tokenPrincipal = jwtAuthenticationToken.principal();
+
+        // JWT cache
+        final SecureString serializedJwt = jwtAuthenticationToken.getEndUserSignedJwt();
+        final SignedJWT signedJWT;
+        try {
+            signedJWT = SignedJWT.parse(serializedJwt.toString());
+        } catch (ParseException e) {
+            // TODO: No point to continue to another realm since parsing failed
+            listener.onFailure(e);
+            return;
+        }
+
+        final JWTClaimsSet jwtClaimsSet;
+        try {
+            jwtClaimsSet = signedJWT.getJWTClaimsSet();
+        } catch (ParseException e) {
+            // TODO: No point to continue to another realm since get claimset failed
+            listener.onFailure(e);
+            return;
+        }
+
+        final JWSHeader jwsHeader = signedJWT.getHeader();
+        if (logger.isDebugEnabled()) {
+            logger.debug(
+                "Realm [{}] successfully parsed JWT token [{}] with header [{}] and claimSet [{}]",
+                realmConfig.name(),
+                tokenPrincipal,
+                jwsHeader,
+                jwtClaimsSet
+            );
+        }
+
+        for (JwtFieldValidator jwtFieldValidator : jwtFieldValidators) {
+            try {
+                jwtFieldValidator.validate(jwsHeader, jwtClaimsSet);
+            } catch (Exception e) {
+                listener.onFailure(e);
+                return;
+            }
+        }
+
+        try {
+            jwtSignatureValidator.validate(tokenPrincipal, signedJWT, listener.map(ignored -> jwtClaimsSet));
+        } catch (Exception e) {
+            listener.onFailure(e);
+        }
+    }
+
+    @Override
+    public void close() {
+        jwtSignatureValidator.close();
+    }
+
+    // Package private for testing
+    JwtSignatureValidator.DelegatingJwtSignatureValidator getJwtSignatureValidator() {
+        assert jwtSignatureValidator instanceof JwtSignatureValidator.DelegatingJwtSignatureValidator;
+        return (JwtSignatureValidator.DelegatingJwtSignatureValidator) jwtSignatureValidator;
+    }
+}