//! Cloudflare Access's signed assertion, `Cf-Access-Jwt-Assertion`. Without it, the proxy
//! sign-in trusts a plain header from any address in `trusted_proxies`, and on Tower that
//! address is the Docker gateway: any container there could send the header and be anyone.
//! Access signs the same identity with keys only Cloudflare holds, so checking that signature
//! takes the network out of the question.

use std::collections::HashMap;
use std::future::Future;
use std::sync::Mutex;
use std::time::{Duration, Instant};

use anyhow::{Context, Result};
use jsonwebtoken::jwk::JwkSet;
use jsonwebtoken::{Algorithm, DecodingKey, Validation, decode, decode_header};

/// Cloudflare rotates its keys every six weeks or so, publishing the new one before using it.
/// A token naming a key not seen yet refetches, but no more often than this, so a stream of
/// made-up key ids cannot turn every request into a request to Cloudflare.
const REFETCH_EVERY: Duration = Duration::from_secs(60);

#[derive(Default)]
pub struct Keys {
    cache: Mutex<Cache>,
}

#[derive(Default)]
struct Cache {
    keys: HashMap<String, DecodingKey>,
    fetched: Option<Instant>,
}

#[derive(serde::Deserialize)]
struct Claims {
    email: Option<String>,
}

impl Keys {
    /// The name the token vouches for, or None: a bad signature, the wrong audience or issuer, an
    /// expired token, a key that cannot be had, or no email in it (a service token has none).
    pub async fn verify(&self, client: &reqwest::Client, team: &str, aud: &str, token: &str) -> Option<String> {
        self.verify_with(team, aud, token, || fetch(client, team)).await
    }

    /// Fill the cache before the first request needs it. A failure is only logged: the next
    /// request tries again, and until one succeeds the proxy sign-in refuses everyone.
    pub async fn prefetch(&self, client: &reqwest::Client, team: &str) {
        match fetch(client, team).await {
            Ok(set) => self.store(set),
            Err(e) => tracing::warn!(error = %format!("{e:#}"), "could not fetch Cloudflare Access's signing keys"),
        }
    }

    async fn verify_with<F, Fut>(&self, team: &str, aud: &str, token: &str, fetch: F) -> Option<String>
    where
        F: FnOnce() -> Fut,
        Fut: Future<Output = Result<JwkSet>>,
    {
        let kid = decode_header(token).ok()?.kid?;
        let key = match self.key(&kid) {
            Some(k) => k,
            None => {
                if !self.may_refetch() {
                    return None;
                }
                match fetch().await {
                    Ok(set) => self.store(set),
                    Err(e) => {
                        tracing::warn!(error = %format!("{e:#}"), "could not fetch Cloudflare Access's signing keys");
                        return None;
                    }
                }
                self.key(&kid)?
            }
        };
        // RS256 only: a token that names HS256 or none is refused here, before its signature
        // is looked at, rather than checked with the public key as if it were a secret.
        let mut v = Validation::new(Algorithm::RS256);
        v.set_audience(&[aud]);
        v.set_issuer(&[format!("https://{team}")]);
        v.validate_nbf = true;
        match decode::<Claims>(token, &key, &v) {
            Ok(data) => crate::auth::name_from_header(&data.claims.email?),
            Err(e) => {
                tracing::warn!(error = %e, "refused a Cloudflare Access token");
                None
            }
        }
    }

    fn key(&self, kid: &str) -> Option<DecodingKey> {
        self.cache.lock().unwrap().keys.get(kid).cloned()
    }

    /// Takes the slot as it answers, so two requests at once do not both fetch.
    fn may_refetch(&self) -> bool {
        let mut c = self.cache.lock().unwrap();
        if c.fetched.is_some_and(|t| t.elapsed() < REFETCH_EVERY) {
            return false;
        }
        c.fetched = Some(Instant::now());
        true
    }

    /// Replaces the whole set, so a key Cloudflare has retired stops being accepted.
    fn store(&self, set: JwkSet) {
        let keys = set
            .keys
            .iter()
            .filter_map(|k| Some((k.common.key_id.clone()?, DecodingKey::from_jwk(k).ok()?)))
            .collect();
        let mut c = self.cache.lock().unwrap();
        c.keys = keys;
        c.fetched = Some(Instant::now());
    }
}

async fn fetch(client: &reqwest::Client, team: &str) -> Result<JwkSet> {
    let url = format!("https://{team}/cdn-cgi/access/certs");
    client
        .get(&url)
        .timeout(Duration::from_secs(10))
        .send()
        .await
        .with_context(|| format!("fetching {url}"))?
        .error_for_status()?
        .json()
        .await
        .context("reading the signing keys")
}

#[cfg(test)]
mod tests {
    use super::*;
    use jsonwebtoken::{EncodingKey, Header, encode, get_current_timestamp};

    const TEAM: &str = "team.cloudflareaccess.com";
    const AUD: &str = "aud-tag";

    fn jwks() -> JwkSet {
        serde_json::from_str(include_str!("../tests/data/access-test.jwks.json")).unwrap()
    }

    fn token(key: &[u8], alg: Algorithm, claims: serde_json::Value) -> String {
        let mut h = Header::new(alg);
        h.kid = Some("k1".into());
        let k = if alg == Algorithm::RS256 { EncodingKey::from_rsa_der(key) } else { EncodingKey::from_secret(key) };
        encode(&h, &claims, &k).unwrap()
    }

    fn claims(aud: &str, exp_in: i64) -> serde_json::Value {
        let now = get_current_timestamp() as i64;
        serde_json::json!({
            "aud": [aud], "iss": format!("https://{TEAM}"), "email": "Rays@SDF1.net",
            "iat": now, "nbf": now, "exp": now + exp_in, "type": "app",
        })
    }

    const SIGNER: &[u8] = include_bytes!("../tests/data/access-test.der");
    const FORGER: &[u8] = include_bytes!("../tests/data/access-forger.der");

    async fn check(keys: &Keys, t: &str) -> Option<String> {
        keys.verify_with(TEAM, AUD, t, || async { Ok(jwks()) }).await
    }

    #[tokio::test]
    async fn only_a_token_cloudflare_signed_for_this_app_signs_anyone_in() {
        let keys = Keys::default();
        keys.store(jwks());
        let ok = token(SIGNER, Algorithm::RS256, claims(AUD, 300));
        assert_eq!(check(&keys, &ok).await.as_deref(), Some("rays@sdf1.net"), "lower-cased like the header");

        let other_app = token(SIGNER, Algorithm::RS256, claims("another-app", 300));
        assert_eq!(check(&keys, &other_app).await, None, "an Access token for another application");

        let expired = token(SIGNER, Algorithm::RS256, claims(AUD, -3600));
        assert_eq!(check(&keys, &expired).await, None, "expired");

        let forged = token(FORGER, Algorithm::RS256, claims(AUD, 300));
        assert_eq!(check(&keys, &forged).await, None, "signed by a key that is not Cloudflare's");

        // HMAC and none, the classic ways to get a token past a verifier that trusts its header.
        let hs = token(b"any secret at all", Algorithm::HS256, claims(AUD, 300));
        assert_eq!(check(&keys, &hs).await, None, "HS256");
        let none = format!("{}.{}.", "eyJhbGciOiJub25lIiwia2lkIjoiazEifQ",
            ok.split('.').nth(1).unwrap());
        assert_eq!(check(&keys, &none).await, None, "alg none");
    }

    #[tokio::test]
    async fn an_unknown_key_refetches_once_a_minute_at_most() {
        let keys = Keys::default();
        let ok = token(SIGNER, Algorithm::RS256, claims(AUD, 300));
        let fetches = std::sync::atomic::AtomicUsize::new(0);
        let count = || { fetches.fetch_add(1, std::sync::atomic::Ordering::SeqCst); async { Ok(jwks()) } };
        assert_eq!(keys.verify_with(TEAM, AUD, &ok, count).await.as_deref(), Some("rays@sdf1.net"),
            "a key not cached yet is fetched");
        assert_eq!(fetches.load(std::sync::atomic::Ordering::SeqCst), 1);

        // A made-up key id straight after: not fetched again.
        let mut h = Header::new(Algorithm::RS256);
        h.kid = Some("nobody".into());
        let stray = encode(&h, &claims(AUD, 300), &EncodingKey::from_rsa_der(SIGNER)).unwrap();
        let count = || { fetches.fetch_add(1, std::sync::atomic::Ordering::SeqCst); async { Ok(jwks()) } };
        assert_eq!(keys.verify_with(TEAM, AUD, &stray, count).await, None);
        assert_eq!(fetches.load(std::sync::atomic::Ordering::SeqCst), 1, "rate-limited");
    }

    #[tokio::test]
    async fn keys_that_cannot_be_fetched_refuse_rather_than_wave_through() {
        let keys = Keys::default();
        let ok = token(SIGNER, Algorithm::RS256, claims(AUD, 300));
        let got = keys.verify_with(TEAM, AUD, &ok, || async { Err(anyhow::anyhow!("offline")) }).await;
        assert_eq!(got, None);
    }
}