/*
 * 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.kafka.common.security.oauthbearer.internals.secured;

import org.apache.kafka.common.security.oauthbearer.BrokerJwtValidator;
import org.apache.kafka.common.utils.Time;

import org.jose4j.jwk.HttpsJwks;
import org.jose4j.jwk.JsonWebKey;
import org.jose4j.lang.JoseException;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.util.Collections;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.locks.ReadWriteLock;
import java.util.concurrent.locks.ReentrantReadWriteLock;

/**
 * Implementation of {@link HttpsJwks} that will periodically refresh the JWKS cache to reduce or
 * even prevent HTTP/HTTPS traffic in the hot path of validation. It is assumed that it's
 * possible to receive a JWT that contains a <code>kid</code> that points to yet-unknown JWK,
 * thus requiring a connection to the OAuth/OIDC provider to be made. Hopefully, in practice,
 * keys are made available for some amount of time before they're used within JWTs.
 * <p>
 * This instance is created and provided to the
 * {@link org.jose4j.keys.resolvers.HttpsJwksVerificationKeyResolver} that is used when using
 * an HTTP-/HTTPS-based {@link org.jose4j.keys.resolvers.VerificationKeyResolver}, which is then
 * provided to the {@link BrokerJwtValidator} to use in validating the signature of
 * a JWT.
 *
 * @see org.jose4j.keys.resolvers.HttpsJwksVerificationKeyResolver
 * @see org.jose4j.keys.resolvers.VerificationKeyResolver
 * @see BrokerJwtValidator
 */
public final class RefreshingHttpsJwks implements OAuthBearerConfigurable {

    private static final Logger log = LoggerFactory.getLogger(RefreshingHttpsJwks.class);

    private static final int MISSING_KEY_ID_CACHE_MAX_ENTRIES = 16;

    static final long MISSING_KEY_ID_CACHE_IN_FLIGHT_MS = 60000;

    static final int MISSING_KEY_ID_MAX_KEY_LENGTH = 1000;

    private static final int SHUTDOWN_TIMEOUT = 10;

    private static final TimeUnit SHUTDOWN_TIME_UNIT = TimeUnit.SECONDS;

    /**
     * {@link HttpsJwks} does the actual work of contacting the OAuth/OIDC endpoint to get the
     * JWKS. In some cases, the call to {@link HttpsJwks#getJsonWebKeys()} will trigger a call
     * to {@link HttpsJwks#refresh()} which will block the current thread in network I/O. We cache
     * the JWKS ourselves (see {@link #jsonWebKeys}) to avoid the network I/O.
     * <p>
     * We want to be very careful where we use the {@link HttpsJwks} instance so that we don't
     * perform any operation (directly or indirectly) that could cause blocking. This is because
     * the JWKS logic is part of the larger authentication logic which operates on Kafka's network
     * thread. It's OK to execute {@link HttpsJwks#getJsonWebKeys()} (which calls
     * {@link HttpsJwks#refresh()}) from within {@link #init()} as that method is called only at
     * startup, and we can afford the blocking hit there.
     */

    private final HttpsJwks httpsJwks;

    private final ScheduledExecutorService executorService;

    private final Time time;

    private final long refreshMs;

    private final long refreshRetryBackoffMs;

    private final long refreshRetryBackoffMaxMs;

    /**
     * Protects {@link #missingKeyIds} and {@link #jsonWebKeys}.
     */

    private final ReadWriteLock refreshLock = new ReentrantReadWriteLock();

    private final Map<String, Long> missingKeyIds;

    /**
     * Flag to prevent concurrent refresh invocations.
     */

    private final AtomicBoolean refreshInProgressFlag = new AtomicBoolean(false);

    /**
     * As mentioned in the comments for {@link #httpsJwks}, we cache the JWKS ourselves so that
     * we can return the list immediately without any network I/O. They are only cached within
     * calls to {@link #refresh()}.
     */

    private List<JsonWebKey> jsonWebKeys;

    private boolean isInitialized;

    /**
     * Creates a <code>RefreshingHttpsJwks</code>. It should only be used for testing to pass in a mock executor
     * service. Otherwise the constructor below should be used.
     */

    // VisibleForTesting
    RefreshingHttpsJwks(Time time,
                        HttpsJwks httpsJwks,
                        long refreshMs,
                        long refreshRetryBackoffMs,
                        long refreshRetryBackoffMaxMs,
                        ScheduledExecutorService executorService) {
        if (refreshMs <= 0)
            throw new IllegalArgumentException("JWKS validation key refresh configuration value retryWaitMs value must be positive");

        this.httpsJwks = httpsJwks;
        this.time = time;
        this.refreshMs = refreshMs;
        this.refreshRetryBackoffMs = refreshRetryBackoffMs;
        this.refreshRetryBackoffMaxMs = refreshRetryBackoffMaxMs;
        this.executorService = executorService;
        this.missingKeyIds = new LinkedHashMap<>(MISSING_KEY_ID_CACHE_MAX_ENTRIES, .75f, true) {
            @Override
            protected boolean removeEldestEntry(Map.Entry<String, Long> eldest) {
                return this.size() > MISSING_KEY_ID_CACHE_MAX_ENTRIES;
            }
        };
    }

    /**
     * Creates a <code>RefreshingHttpsJwks</code> that will be used by the
     * {@link RefreshingHttpsJwksVerificationKeyResolver} to resolve new key IDs in JWTs.
     *
     * @param time                     {@link Time} instance
     * @param httpsJwks                {@link HttpsJwks} instance from which to retrieve the JWKS
     *                                 based on the OAuth/OIDC standard
     * @param refreshMs                The number of milliseconds between refresh passes to connect
     *                                 to the OAuth/OIDC JWKS endpoint to retrieve the latest set
     * @param refreshRetryBackoffMs    Time for delay after initial failed attempt to retrieve JWKS
     * @param refreshRetryBackoffMaxMs Maximum time to retrieve JWKS
     */

    public RefreshingHttpsJwks(Time time,
                               HttpsJwks httpsJwks,
                               long refreshMs,
                               long refreshRetryBackoffMs,
                               long refreshRetryBackoffMaxMs) {
        this(time, httpsJwks, refreshMs, refreshRetryBackoffMs, refreshRetryBackoffMaxMs, Executors.newSingleThreadScheduledExecutor());
    }

    public void init() throws IOException {
        try {
            log.debug("init started");

            List<JsonWebKey> localJWKs;

            try {
                localJWKs = httpsJwks.getJsonWebKeys();
            } catch (JoseException e) {
                throw new IOException("Could not refresh JWKS", e);
            }

            try {
                refreshLock.writeLock().lock();
                jsonWebKeys = Collections.unmodifiableList(localJWKs);
            } finally {
                refreshLock.writeLock().unlock();
            }

            // Since we just grabbed the keys (which will have invoked a HttpsJwks.refresh()
            // internally), we can delay our first invocation by refreshMs.
            //
            // Note: we refer to this as a _scheduled_ refresh.
            executorService.scheduleAtFixedRate(this::refresh,
                    refreshMs,
                    refreshMs,
                    TimeUnit.MILLISECONDS);

            log.info("JWKS validation key refresh thread started with a refresh interval of {} ms", refreshMs);
        } finally {
            isInitialized = true;

            log.debug("init completed");
        }
    }

    @Override
    public void close() {
        try {
            log.debug("close started");

            try {
                log.debug("JWKS validation key refresh thread shutting down");
                executorService.shutdown();

                if (!executorService.awaitTermination(SHUTDOWN_TIMEOUT, SHUTDOWN_TIME_UNIT)) {
                    log.warn("JWKS validation key refresh thread termination did not end after {} {}",
                            SHUTDOWN_TIMEOUT, SHUTDOWN_TIME_UNIT);
                }
            } catch (InterruptedException e) {
                log.warn("JWKS validation key refresh thread error during close", e);
            }
        } finally {
            log.debug("close completed");
        }
    }

    /**
     * Our implementation avoids the blocking call within {@link HttpsJwks#refresh()} that is
     * sometimes called internal to {@link HttpsJwks#getJsonWebKeys()}. We want to avoid any
     * blocking I/O as this code is running in the authentication path on the Kafka network thread.
     * <p>
     * The list may be stale up to {@link #refreshMs}.
     *
     * @return {@link List} of {@link JsonWebKey} instances
     * @throws JoseException Thrown if a problem is encountered parsing the JSON content into JWKs
     * @throws IOException   Thrown f a problem is encountered making the HTTP request
     */

    public List<JsonWebKey> getJsonWebKeys() throws JoseException, IOException {
        if (!isInitialized)
            throw new IllegalStateException("Please call init() first");

        try {
            refreshLock.readLock().lock();
            return jsonWebKeys;
        } finally {
            refreshLock.readLock().unlock();
        }
    }

    public String getLocation() {
        return httpsJwks.getLocation();
    }

    /**
     * <p>
     * <code>refresh</code> is an internal method that will refresh the JWKS cache and is
     * invoked in one of two ways:
     *
     * <ol>
     *     <li>Scheduled</li>
     *     <li>Expedited</li>
     * </ol>
     * </p>
     *
     * <p>
     * The <i>scheduled</i> refresh is scheduled in {@link #init()} and runs every
     * {@link #refreshMs} milliseconds. An <i>expedited</i> refresh is performed when an
     * incoming JWT refers to a key ID that isn't in our JWKS cache ({@link #jsonWebKeys})
     * and we try to perform a refresh sooner than the next scheduled refresh.
     * </p>
     */

    private void refresh() {
        if (!refreshInProgressFlag.compareAndSet(false, true)) {
            log.debug("OAuth JWKS refresh is already in progress; ignoring concurrent refresh");
            return;
        }

        try {
            log.info("OAuth JWKS refresh of {} starting", httpsJwks.getLocation());
            Retry<List<JsonWebKey>> retry = new Retry<>(refreshRetryBackoffMs, refreshRetryBackoffMaxMs);
            List<JsonWebKey> localJWKs = retry.execute(() -> {
                try {
                    log.debug("JWKS validation key calling refresh of {} starting", httpsJwks.getLocation());
                    // Call the *actual* refresh implementation that will more than likely issue
                    // HTTP(S) calls over the network.
                    httpsJwks.refresh();
                    List<JsonWebKey> jwks = httpsJwks.getJsonWebKeys();
                    log.debug("JWKS validation key refresh of {} complete", httpsJwks.getLocation());
                    return jwks;
                } catch (Exception e) {
                    throw new ExecutionException(e);
                }
            });

            try {
                refreshLock.writeLock().lock();

                for (JsonWebKey jwk : localJWKs)
                    missingKeyIds.remove(jwk.getKeyId());

                jsonWebKeys = Collections.unmodifiableList(localJWKs);
            } finally {
                refreshLock.writeLock().unlock();
            }

            log.info("OAuth JWKS refresh of {} complete", httpsJwks.getLocation());
        } catch (ExecutionException e) {
            log.warn("OAuth JWKS refresh of {} encountered an error; not updating local JWKS cache", httpsJwks.getLocation(), e);
        } finally {
            refreshInProgressFlag.set(false);
        }
    }

    /**
     * <p>
     * <code>maybeExpediteRefresh</code> is a public method that will trigger a refresh of
     * the JWKS cache if all of the following conditions are met:
     *
     * <ul>
     *     <li>The given <code>keyId</code> parameter is &lte; the
     *     {@link #MISSING_KEY_ID_MAX_KEY_LENGTH}</li>
     *     <li>The key isn't in the process of being expedited already</li>
     * </ul>
     *
     * <p>
     * This <i>expedited</i> refresh is scheduled immediately.
     * </p>
     *
     * @param keyId JWT key ID
     * @return <code>true</code> if an expedited refresh was scheduled, <code>false</code> otherwise
     */

    public boolean maybeExpediteRefresh(String keyId) {
        if (keyId.length() > MISSING_KEY_ID_MAX_KEY_LENGTH) {
            // Although there's no limit on the length of the key ID, they're generally
            // "reasonably" short. If we have a very long key ID length, we're going to assume
            // the JWT is malformed, and we will not actually try to resolve the key.
            //
            // In this case, let's prevent blowing out our memory in two ways:
            //
            //     1. Don't try to resolve the key as the large ID will sit in our cache
            //     2. Report the issue in the logs but include only the first N characters
            int actualLength = keyId.length();
            String trimmedKeyId = keyId.substring(0, MISSING_KEY_ID_MAX_KEY_LENGTH);
            String snippet = String.format("%s (trimmed to first %d characters out of %d total)", trimmedKeyId, MISSING_KEY_ID_MAX_KEY_LENGTH, actualLength);
            log.warn("Key ID {} was too long to cache", snippet);
            return false;
        } else {
            try {
                refreshLock.writeLock().lock();

                Long nextCheckTime = missingKeyIds.get(keyId);
                long currTime = time.milliseconds();
                log.debug("For key ID {}, nextCheckTime: {}, currTime: {}", keyId, nextCheckTime, currTime);

                if (nextCheckTime == null || nextCheckTime <= currTime) {
                    // If there's no entry in the missing key ID cache for the incoming key ID,
                    // or it has expired, schedule a refresh ASAP.
                    nextCheckTime = currTime + MISSING_KEY_ID_CACHE_IN_FLIGHT_MS;
                    missingKeyIds.put(keyId, nextCheckTime);
                    executorService.schedule(this::refresh, 0, TimeUnit.MILLISECONDS);
                    return true;
                } else {
                    return false;
                }
            } finally {
                refreshLock.writeLock().unlock();
            }
        }
    }
}
