//! Verifying atproto service auth (javascript/src/serviceauth): an `Authorization: Bearer <JWT>` a
//! caller's PDS minted with com.atproto.server.getServiceAuth, signed by the caller's #atproto key.

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

use base64::Engine;
use base64::engine::general_purpose::URL_SAFE_NO_PAD;
use serde_json::Value;

use crate::BoxFut;
use crate::crypto::unbase58;
use crate::wasm::sha256_hex;

/// Longest a token may still be valid for, and the clock skew allowed on its iat.
pub const MAX_TTL_S: u64 = 3600;
pub const SKEW_S: u64 = 60;

/// Who may call: the JWT's aud must be one of `audiences` (the service's DIDs); `issuers` empty lets
/// any account in, else only those DIDs (an entry ending in `*` is a prefix: `did:plc:*`, one DID method).
pub struct Auth {
    pub audiences: Vec<String>,
    pub issuers: Vec<String>,
}

/// A DID -> its DID document.
pub type Resolve = Arc<dyn Fn(String) -> BoxFut<'static, Result<Value, String>> + Send + Sync>;

/// Tokens a Replay remembers at once, and how long and how many DID documents `cached` keeps.
pub const REPLAY_MAX: usize = 10_000;
pub const DOC_TTL: Duration = Duration::from_secs(300);
pub const DOC_MAX: usize = 1000;

/// The tokens `verify` accepted, each until its exp: a token is good for one request (a replay of
/// it, by whoever saw it on the way, is refused). Keyed by iss + jti, or the token's sha256 when it
/// has no jti.
pub struct Replay {
    seen: Mutex<HashMap<String, u64>>,
    max: usize,
}

impl Default for Replay {
    fn default() -> Self {
        Self::new(REPLAY_MAX)
    }
}

impl Replay {
    pub fn new(max: usize) -> Self {
        Self { seen: Mutex::default(), max: max.max(1) }
    }

    fn admit(&self, key: String, exp: u64, now_s: u64) -> Result<(), String> {
        let mut seen = self.seen.lock().unwrap();
        if seen.get(&key).is_some_and(|e| *e > now_s) {
            return Err("the token was already used".into());
        }
        if seen.len() >= self.max {
            seen.retain(|_, e| *e > now_s);
        }
        // Still full: `max` live tokens at once. Evicting one would let it be replayed, so the new
        // token is refused until some expire.
        // ponytail: one global cap, so one issuer can fill it; per-iss caps if that happens.
        if seen.len() >= self.max && !seen.contains_key(&key) {
            return Err("busy: too many tokens in use; try again later".into());
        }
        seen.insert(key, exp);
        Ok(())
    }
}

/// `resolve` with its documents kept for `ttl`, at most `max` of them (expired ones go first, then
/// the oldest). Failures are not kept.
/// ponytail: a key rotated within ttl fails the issuer's new tokens until the entry expires; refetch
/// on a bad signature if that matters.
pub fn cached(resolve: Resolve, ttl: Duration, max: usize) -> Resolve {
    let docs: Arc<Mutex<HashMap<String, (Instant, Value)>>> = Arc::default();
    Arc::new(move |did: String| {
        let (docs, resolve) = (docs.clone(), resolve.clone());
        Box::pin(async move {
            if let Some((at, doc)) = docs.lock().unwrap().get(&did)
                && at.elapsed() < ttl
            {
                return Ok(doc.clone());
            }
            let doc = resolve(did.clone()).await?;
            let mut docs = docs.lock().unwrap();
            if docs.len() >= max && !docs.contains_key(&did) {
                docs.retain(|_, (at, _)| at.elapsed() < ttl);
                if docs.len() >= max
                    && let Some(oldest) = docs.iter().min_by_key(|(_, (at, _))| *at).map(|(k, _)| k.clone())
                {
                    docs.remove(&oldest);
                }
            }
            docs.insert(did, (Instant::now(), doc.clone()));
            Ok(doc)
        })
    })
}

// The #atproto key of a DID document: (multicodec prefix, compressed key).
fn atproto_key(doc: &Value) -> Result<([u8; 2], Vec<u8>), String> {
    let multibase = doc["verificationMethod"]
        .as_array()
        .into_iter()
        .flatten()
        .find(|m| m["id"].as_str().is_some_and(|id| id.ends_with("#atproto")))
        .and_then(|m| m["publicKeyMultibase"].as_str())
        .ok_or("the issuer's DID document has no #atproto key")?;
    let bytes = unbase58(
        multibase
            .strip_prefix('z')
            .ok_or("publicKeyMultibase is not base58btc")?,
    )?;
    if bytes.len() < 3 {
        return Err("publicKeyMultibase is too short".into());
    }
    Ok(([bytes[0], bytes[1]], bytes[2..].to_vec()))
}

fn check_signature(
    alg: &str,
    prefix: [u8; 2],
    key: &[u8],
    message: &[u8],
    signature: &[u8],
) -> Result<(), String> {
    use k256::ecdsa::signature::Verifier;
    let bad = |e: &dyn std::fmt::Display| format!("bad signature: {e}");
    match (alg, prefix) {
        ("ES256K", [0xe7, 0x01]) => {
            let sig = k256::ecdsa::Signature::from_slice(signature).map_err(|e| bad(&e))?;
            if sig.normalize_s().is_some() {
                return Err("bad signature: high S".into());
            }
            let vk = k256::ecdsa::VerifyingKey::from_sec1_bytes(key).map_err(|e| bad(&e))?;
            vk.verify(message, &sig).map_err(|e| bad(&e))
        }
        ("ES256", [0x80, 0x24]) => {
            let sig = p256::ecdsa::Signature::from_slice(signature).map_err(|e| bad(&e))?;
            if sig.normalize_s().is_some() {
                return Err("bad signature: high S".into());
            }
            let vk = p256::ecdsa::VerifyingKey::from_sec1_bytes(key).map_err(|e| bad(&e))?;
            vk.verify(message, &sig).map_err(|e| bad(&e))
        }
        _ => Err(format!("alg {alg} does not match the issuer's key type")),
    }
}

/// The issuer DID of a valid `Authorization: Bearer <service-auth JWT>` for method `lxm`: aud one of
/// `auth.audiences` (or that DID with a #fragment), iat no later than now + SKEW_S, exp in the future
/// but at most MAX_TTL_S away, iss allowed, signed by the issuer's #atproto key, and not seen by
/// `replay` before (a caller mints one token per request).
pub async fn verify(
    auth: &Auth,
    authorization: Option<&str>,
    lxm: &str,
    now_s: u64,
    resolve: &Resolve,
    replay: &Replay,
) -> Result<String, String> {
    let token = authorization
        .and_then(|h| h.strip_prefix("Bearer "))
        .ok_or("Bearer service-auth token required")?
        .trim();
    let [h, p, s] = token.split('.').collect::<Vec<_>>()[..] else {
        return Err("the token is not a JWT".into());
    };
    let part = |x: &str| -> Result<Value, String> {
        serde_json::from_slice(
            &URL_SAFE_NO_PAD
                .decode(x)
                .map_err(|_| "the token is not base64url")?,
        )
        .map_err(|_| "the token is not JSON".into())
    };
    let (header, claims) = (part(h)?, part(p)?);
    let alg = header["alg"].as_str().unwrap_or_default();
    let claim = |k: &str| claims[k].as_str().unwrap_or_default();
    let iss = claim("iss").split('#').next().unwrap_or_default().to_string();
    if claim("lxm") != lxm {
        return Err(format!("the token is for {:?}, not {lxm}", claim("lxm")));
    }
    let aud = claim("aud");
    if !auth
        .audiences
        .iter()
        .any(|a| a == aud || a == aud.split('#').next().unwrap_or_default())
    {
        return Err(format!("the token's aud {aud:?} is not this service"));
    }
    let exp = claims["exp"].as_u64().unwrap_or(0);
    if exp <= now_s {
        return Err("the token has expired".into());
    }
    if exp > now_s + MAX_TTL_S {
        return Err(format!("the token's exp is more than {MAX_TTL_S} s away"));
    }
    if claims["iat"].as_u64().is_none_or(|iat| iat > now_s + SKEW_S) {
        return Err("the token's iat is missing or in the future".into());
    }
    if !auth.issuers.is_empty() && !auth.issuers.iter().any(|i| i == &iss || i.strip_suffix('*').is_some_and(|p| iss.starts_with(p))) {
        return Err(format!("issuer {iss} is not allowed"));
    }
    let (prefix, key) = atproto_key(&resolve(iss.clone()).await?)?;
    let signature = URL_SAFE_NO_PAD
        .decode(s)
        .map_err(|_| "the signature is not base64url")?;
    check_signature(alg, prefix, &key, format!("{h}.{p}").as_bytes(), &signature)?;
    let key = match claim("jti") {
        "" => format!("sha256 {}", sha256_hex(token.as_bytes())),
        jti => format!("{iss} {jti}"),
    };
    replay.admit(key, exp, now_s)?;
    Ok(iss)
}

/// did:plc via plc.directory, did:web via /.well-known/did.json, each document `cached` (DOC_TTL,
/// DOC_MAX).
/// ponytail: unguarded fetch of the issuer's document; hardening is the open SSRF item (one hardened DID resolver).
pub fn resolver(client: reqwest::Client) -> Resolve {
    cached(fetching(client), DOC_TTL, DOC_MAX)
}

fn fetching(client: reqwest::Client) -> Resolve {
    Arc::new(move |did: String| {
        let client = client.clone();
        Box::pin(async move {
            let url = match did.strip_prefix("did:web:") {
                Some(host) => format!("https://{host}/.well-known/did.json"),
                None if did.starts_with("did:plc:") => format!("https://plc.directory/{did}"),
                None => return Err(format!("cannot resolve {did}")),
            };
            let res = client
                .get(&url)
                .send()
                .await
                .map_err(|e| format!("resolving {did}: {e}"))?;
            if !res.status().is_success() {
                return Err(format!("resolving {did}: {}", res.status()));
            }
            res.json().await.map_err(|e| format!("resolving {did}: {e}"))
        })
    })
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::crypto::{DidKey, base58};
    use k256::ecdsa::signature::Signer;
    use serde_json::json;

    const ISS: &str = "did:plc:caller0000000000000000000";
    const AUD: &str = "did:key:zDnaeService";
    const LXM: &str = "com.example.launch";
    const NOW: u64 = 1_800_000_000;

    fn jwt(alg: &str, claims: Value, sign: impl Fn(&[u8]) -> Vec<u8>) -> String {
        let enc = |v: Value| URL_SAFE_NO_PAD.encode(v.to_string());
        let signing = format!("{}.{}", enc(json!({ "typ": "JWT", "alg": alg })), enc(claims));
        format!("{signing}.{}", URL_SAFE_NO_PAD.encode(sign(signing.as_bytes())))
    }

    fn k256_key() -> k256::ecdsa::SigningKey {
        k256::ecdsa::SigningKey::from_slice(&[7u8; 32]).unwrap()
    }

    fn doc_for(prefix: [u8; 2], compressed: &[u8]) -> Value {
        json!({ "verificationMethod": [{ "id": "#atproto", "publicKeyMultibase": format!("z{}", base58(&[&prefix[..], compressed].concat())) }] })
    }

    fn resolve_to(doc: Value) -> Resolve {
        Arc::new(move |_did| {
            let doc = doc.clone();
            Box::pin(async move { Ok(doc) })
        })
    }

    fn auth() -> Auth {
        Auth {
            audiences: vec![AUD.into(), "did:web:svc.example".into()],
            issuers: vec![ISS.into()],
        }
    }

    fn claims(lxm: &str) -> Value {
        json!({ "iss": ISS, "aud": AUD, "lxm": lxm, "iat": NOW, "exp": NOW + 60, "jti": "x" })
    }

    fn es256k(claims: Value) -> String {
        let key = k256_key();
        jwt("ES256K", claims, |m| {
            let sig: k256::ecdsa::Signature = key.sign(m);
            sig.to_bytes().to_vec()
        })
    }

    fn k256_doc() -> Value {
        doc_for(
            [0xe7, 0x01],
            k256_key().verifying_key().to_encoded_point(true).as_bytes(),
        )
    }

    #[test]
    fn base58_round_trip() {
        for bytes in [vec![], vec![0, 0, 1], vec![0xe7, 0x01, 2, 255, 0, 7], (0..40).collect::<Vec<u8>>()] {
            assert_eq!(unbase58(&base58(&bytes)).unwrap(), bytes);
        }
    }

    #[tokio::test]
    async fn es256k_token_from_an_allowed_issuer_passes() {
        let header = format!("Bearer {}", es256k(claims(LXM)));
        assert_eq!(
            verify(&auth(), Some(&header), LXM, NOW, &resolve_to(k256_doc()), &Replay::default())
                .await
                .unwrap(),
            ISS
        );
        let mut c = claims(LXM);
        c["aud"] = json!("did:web:svc.example#svc");
        let header = format!("Bearer {}", es256k(c));
        assert!(
            verify(&auth(), Some(&header), LXM, NOW, &resolve_to(k256_doc()), &Replay::default())
                .await
                .is_ok()
        );
    }

    #[tokio::test]
    async fn es256_token_passes() {
        let key = DidKey::new(None).unwrap();
        let token = key.service_auth(AUD, LXM, None);
        let mb = key.did.strip_prefix("did:key:").unwrap();
        let doc = json!({ "verificationMethod": [{ "id": "did:key:x#atproto", "publicKeyMultibase": mb }] });
        let auth = Auth {
            audiences: vec![AUD.into()],
            issuers: vec![key.did.clone()],
        };
        let now = crate::now_ms() / 1000;
        let header = format!("Bearer {token}");
        assert_eq!(
            verify(&auth, Some(&header), LXM, now, &resolve_to(doc), &Replay::default())
                .await
                .unwrap(),
            key.did
        );
    }

    #[tokio::test]
    async fn bad_tokens_are_refused() {
        let r = resolve_to(k256_doc());
        let check = |header: Option<String>, lxm: &'static str| {
            let r = r.clone();
            async move {
                verify(&auth(), header.as_deref(), lxm, NOW, &r, &Replay::default())
                    .await
                    .unwrap_err()
            }
        };
        let bearer = |t: String| Some(format!("Bearer {t}"));
        let with = |key: &str, value: Value| {
            let mut c = claims(LXM);
            c[key] = value;
            bearer(es256k(c))
        };
        assert!(check(None, LXM).await.contains("Bearer"));
        assert!(check(Some("Basic x".into()), LXM).await.contains("Bearer"));
        assert!(check(bearer("a.b".into()), LXM).await.contains("not a JWT"));
        assert!(
            check(bearer(es256k(claims("other.method"))), LXM)
                .await
                .contains("not com.example.launch")
        );
        assert!(
            check(with("aud", json!("did:web:someone.else")), LXM)
                .await
                .contains("is not this service")
        );
        assert!(check(with("exp", json!(NOW)), LXM).await.contains("expired"));
        // exp far ahead (a long-lived token) and iat missing or ahead of the clock
        assert!(
            check(with("exp", json!(NOW + MAX_TTL_S + 1)), LXM)
                .await
                .contains("more than 3600 s away")
        );
        assert!(check(with("iat", Value::Null), LXM).await.contains("iat"));
        assert!(
            check(with("iat", json!(NOW + SKEW_S + 1)), LXM)
                .await
                .contains("iat")
        );
        let skewed = with("iat", json!(NOW + SKEW_S));
        assert!(verify(&auth(), skewed.as_deref(), LXM, NOW, &r, &Replay::default()).await.is_ok());
        assert!(
            check(with("iss", json!("did:plc:stranger")), LXM)
                .await
                .contains("not allowed")
        );
        let other = k256::ecdsa::SigningKey::from_slice(&[9u8; 32]).unwrap();
        let forged = jwt("ES256K", claims(LXM), |m| {
            let sig: k256::ecdsa::Signature = other.sign(m);
            sig.to_bytes().to_vec()
        });
        assert!(check(bearer(forged), LXM).await.contains("bad signature"));
        // a valid signature with S flipped high (atproto requires low S)
        let key = k256_key();
        let high = jwt("ES256K", claims(LXM), |m| {
            let sig: k256::ecdsa::Signature = key.sign(m);
            let (r, s) = sig.split_scalars();
            k256::ecdsa::Signature::from_scalars(r.to_bytes(), (-*s).to_bytes())
                .unwrap()
                .to_bytes()
                .to_vec()
        });
        assert!(check(bearer(high), LXM).await.contains("high S"));
        let header = |alg: &str| URL_SAFE_NO_PAD.encode(json!({ "typ": "JWT", "alg": alg }).to_string());
        let wrong = es256k(claims(LXM)).replacen(&header("ES256K"), &header("ES256"), 1);
        assert!(
            check(bearer(wrong), LXM)
                .await
                .contains("does not match the issuer")
        );
    }

    #[tokio::test]
    async fn a_token_is_good_once_until_it_expires() {
        let (r, replay) = (resolve_to(k256_doc()), Replay::default());
        let once = |c: Value| format!("Bearer {}", es256k(c));
        let header = once(claims(LXM));
        assert!(verify(&auth(), Some(&header), LXM, NOW, &r, &replay).await.is_ok());
        let again = verify(&auth(), Some(&header), LXM, NOW + 1, &r, &replay).await.unwrap_err();
        assert_eq!(again, "the token was already used");
        // the same jti from the issuer is the same token, whatever else changed
        let mut c = claims(LXM);
        c["exp"] = json!(NOW + 30);
        assert!(verify(&auth(), Some(&once(c)), LXM, NOW, &r, &replay).await.is_err());
        // another jti passes; a token without one is keyed by its hash
        let mut c = claims(LXM);
        c["jti"] = json!("y");
        assert!(verify(&auth(), Some(&once(c)), LXM, NOW, &r, &replay).await.is_ok());
        let mut c = claims(LXM);
        c.as_object_mut().unwrap().remove("jti");
        let bare = once(c);
        assert!(verify(&auth(), Some(&bare), LXM, NOW, &r, &replay).await.is_ok());
        assert!(verify(&auth(), Some(&bare), LXM, NOW, &r, &replay).await.is_err());
        assert_eq!(replay.seen.lock().unwrap().len(), 3, "x, y and the bare token");
    }

    #[test]
    fn replay_is_bounded_and_never_evicts_a_live_token() {
        let replay = Replay::new(2);
        // a legitimate token, then an attacker's two long-lived ones fill the cache
        replay.admit("a".into(), 60, 0).unwrap();
        replay.admit("x".into(), 3600, 0).unwrap();
        assert_eq!(replay.admit("y".into(), 3600, 0).unwrap_err(), "busy: too many tokens in use; try again later");
        assert_eq!(replay.admit("a".into(), 60, 1).unwrap_err(), "the token was already used", "a is still remembered");
        // once a expires it is pruned and room is made
        replay.admit("b".into(), 120, 61).unwrap();
        let seen = replay.seen.lock().unwrap();
        assert!(seen.contains_key("x") && seen.contains_key("b") && !seen.contains_key("a"), "{seen:?}");
        drop(seen);
        // past its exp a key may come back (its token is refused as expired anyway), even when full
        replay.admit("b".into(), 200, 121).unwrap();
    }

    #[tokio::test]
    async fn did_documents_are_cached_for_a_while_and_bounded() {
        let fetched = Arc::new(Mutex::new(Vec::<String>::new()));
        let log = fetched.clone();
        let inner: Resolve = Arc::new(move |did: String| {
            log.lock().unwrap().push(did.clone());
            Box::pin(async move { if did == "did:plc:bad" { Err("nope".to_string()) } else { Ok(json!({ "id": did })) } })
        });
        let resolve = cached(inner, Duration::from_millis(80), 2);
        for did in ["did:plc:a", "did:plc:a", "did:plc:b", "did:plc:a"] {
            assert_eq!(resolve(did.into()).await.unwrap()["id"], did);
        }
        assert_eq!(*fetched.lock().unwrap(), ["did:plc:a", "did:plc:b"]);
        // a failure is not kept; a third document pushes out the oldest
        assert!(resolve("did:plc:bad".into()).await.is_err());
        assert!(resolve("did:plc:bad".into()).await.is_err());
        resolve("did:plc:c".into()).await.unwrap();
        resolve("did:plc:a".into()).await.unwrap();
        assert_eq!(fetched.lock().unwrap()[2..], ["did:plc:bad", "did:plc:bad", "did:plc:c", "did:plc:a"]);
        // past the ttl it is fetched again
        tokio::time::sleep(Duration::from_millis(100)).await;
        resolve("did:plc:c".into()).await.unwrap();
        assert_eq!(fetched.lock().unwrap().last().unwrap(), "did:plc:c");
        assert_eq!(fetched.lock().unwrap().len(), 7);
    }
}