/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.apache.solr.security.jwt;

import java.io.IOException;
import java.lang.invoke.MethodHandles;
import java.security.Key;
import java.util.ArrayList;
import java.util.Collection;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Locale;
import java.util.Map;
import java.util.Set;
import javax.net.ssl.SSLHandshakeException;
import org.apache.solr.common.SolrException;
import org.jose4j.jwk.HttpsJwks;
import org.jose4j.jwk.JsonWebKey;
import org.jose4j.jwk.VerificationJwkSelector;
import org.jose4j.jws.JsonWebSignature;
import org.jose4j.jwt.JwtClaims;
import org.jose4j.jwt.MalformedClaimException;
import org.jose4j.jwt.consumer.InvalidJwtException;
import org.jose4j.jwx.JsonWebStructure;
import org.jose4j.keys.resolvers.VerificationKeyResolver;
import org.jose4j.lang.JoseException;
import org.jose4j.lang.UnresolvableKeyException;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * Resolves jws signature verification keys from a set of {@link JWTIssuerConfig} objects, which may
 * represent any valid configuration in Solr's security.json, i.e. static list of JWKs or keys
 * retrieved from HTTPs JWK endpoints.
 *
 * <p>This implementation maintains a map of issuers, each with its own list of {@link JsonWebKey},
 * and resolves correct key from correct issuer similar to HttpsJwksVerificationKeyResolver. If
 * issuer claim is not required, we will select the first IssuerConfig if there is exactly one such
 * config.
 *
 * <p>If a key is not found, and issuer is backed by HTTPsJWKs, we attempt one cache refresh before
 * failing.
 */
public class JWTVerificationkeyResolver implements VerificationKeyResolver {
  private static final Logger log = LoggerFactory.getLogger(MethodHandles.lookup().lookupClass());

  private VerificationJwkSelector verificationJwkSelector = new VerificationJwkSelector();

  private Map<String, JWTIssuerConfig> issuerConfigs = new HashMap<>();
  private final boolean requireIssuer;

  /**
   * Resolves key from a JWKs from one or more IssuerConfigs
   *
   * @param issuerConfigs Collection of configuration objects for the issuer(s)
   * @param requireIssuer if true, will require 'iss' claim on jws
   */
  public JWTVerificationkeyResolver(
      Collection<JWTIssuerConfig> issuerConfigs, boolean requireIssuer) {
    this.requireIssuer = requireIssuer;
    issuerConfigs.forEach(
        ic -> {
          this.issuerConfigs.put(ic.getIss(), ic);
        });
  }

  @Override
  public Key resolveKey(JsonWebSignature jws, List<JsonWebStructure> nestingContext)
      throws UnresolvableKeyException {
    JsonWebKey theChosenOne;
    List<JsonWebKey> jsonWebKeys = new ArrayList<>();

    String keysSource = "N/A";
    try {
      String tokenIssuer = JwtClaims.parse(jws.getUnverifiedPayload()).getIssuer();
      JWTIssuerConfig issuerConfig;
      if (tokenIssuer == null) {
        if (requireIssuer) {
          throw new UnresolvableKeyException("Token does not contain required issuer claim");
        } else if (issuerConfigs.size() == 1) {
          issuerConfig = issuerConfigs.values().iterator().next();
        } else {
          throw new SolrException(
              SolrException.ErrorCode.SERVER_ERROR,
              "Signature verifiction not supported for multiple issuers without 'iss' claim in token.");
        }
      } else {
        issuerConfig = issuerConfigs.get(tokenIssuer);
        if (issuerConfig == null) {
          if (issuerConfigs.size() > 1) {
            throw new UnresolvableKeyException(
                "No issuers configured for iss='" + tokenIssuer + "', cannot validate signature");
          } else if (issuerConfigs.size() == 1) {
            issuerConfig = issuerConfigs.values().iterator().next();
            log.debug(
                "No issuer matching token's iss claim, but exactly one configured, selecting that one");
          } else {
            throw new SolrException(
                SolrException.ErrorCode.SERVER_ERROR,
                "Signature verifiction failed due to no configured issuer with id " + tokenIssuer);
          }
        }
      }

      // Add all keys into a leader list
      if (issuerConfig.usesHttpsJwk()) {
        keysSource = "[" + String.join(", ", issuerConfig.getJwksUrls()) + "]";
        for (HttpsJwks hjwks : issuerConfig.getHttpsJwks()) {
          try {
            jsonWebKeys.addAll(hjwks.getJsonWebKeys());
          } catch (SSLHandshakeException e) {
            throw new SolrException(
                SolrException.ErrorCode.SERVER_ERROR,
                "Failed to connect with "
                    + hjwks.getLocation()
                    + ", do you have the correct SSL certificate configured?",
                e);
          }
        }
      } else {
        keysSource = "static list of keys in security.json";
        jsonWebKeys.addAll(issuerConfig.getJsonWebKeySet().getJsonWebKeys());
      }

      theChosenOne = verificationJwkSelector.select(jws, jsonWebKeys);
      if (theChosenOne == null && issuerConfig.usesHttpsJwk()) {
        if (log.isDebugEnabled()) {
          log.debug(
              "Refreshing JWKs from all {} locations, as no suitable verification key for JWS w/ header {} was found in {}",
              issuerConfig.getHttpsJwks().size(),
              jws.getHeaders().getFullHeaderAsJsonString(),
              jsonWebKeys);
        }

        jsonWebKeys.clear();
        for (HttpsJwks hjwks : issuerConfig.getHttpsJwks()) {
          hjwks.refresh();
          jsonWebKeys.addAll(hjwks.getJsonWebKeys());
        }
        theChosenOne = verificationJwkSelector.select(jws, jsonWebKeys);
      }
    } catch (JoseException | IOException | InvalidJwtException | MalformedClaimException e) {
      String msg =
          String.format(
              Locale.ROOT,
              "Unable to find a suitable verification key for JWS w/ header %s due to an unexpected exception (%s) "
                  + "while obtaining or using keys from source %s",
              jws.getHeaders().getFullHeaderAsJsonString(),
              e,
              keysSource);
      throw new UnresolvableKeyException(msg, e);
    }

    if (theChosenOne == null) {
      String msg =
          String.format(
              Locale.ROOT,
              "Unable to find a suitable verification key for JWS w/ header %s from %d keys from source %s",
              jws.getHeaders().getFullHeaderAsJsonString(),
              jsonWebKeys.size(),
              keysSource);
      throw new UnresolvableKeyException(msg);
    }

    return theChosenOne.getKey();
  }

  Set<JWTIssuerConfig> getIssuerConfigs() {
    return new HashSet<>(issuerConfigs.values());
  }
}

@@ -4067,24 +4067,25 @@
 ture verific
+a
 tion not sup
@@ -4841,16 +4841,17 @@
  verific
+a
 tion fai

Fix typo in JWTVerificationkeyResolver exception message - verifiction -> verification (#823)