techne-sdk 0.3.2

//! Service-auth for a service of its own: verify the atproto service-auth JWTs addressed to one
//! service DID (`net`).
//!
//! Ported from regenOS `regenos_appview::auth::service_auth` and the in-process half of
//! `regenos_appview::auth::jti_cache` (P20.7 G4); the Postgres `jti` cache stays in the platform (a
//! caller that needs one implements [`JtiCache`]).
//!
//! `iss` = the caller's DID, `aud` = this service's DID, `lxm` = the method called, a short `exp`, a
//! `jti`. The pipeline: the I/O-free [`precheck_claims`] (so junk tokens cost no DID lookup), resolve
//! `iss` to its `#atproto` key, [`verify_service_jwt`], then the `jti` replay cache.
//!
//! Resolved keys are cached for [`KEY_CACHE_TTL`], so a caller polling does not hit the PLC directory
//! on every call, and a token that fails verification against a cached key is checked once more
//! against a freshly resolved one (a rotated key works at once), without evicting the cached key
//! first. A directory that cannot answer a first lookup is [`AuthError::Unavailable`] (HTTP 503, worth
//! retrying), not a refusal.
//!
//! ## The replay cache fails open, on purpose
//! Replay protection is defense-in-depth; the primary controls are `exp` + `aud` + `lxm`. A
//! [`JtiCache`] that cannot reach its store answers "fresh" (the trait has no error channel), and the
//! short `exp` bounds the exposure.

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

use async_trait::async_trait;
use k256::ecdsa::VerifyingKey;

use crate::did::{DidKeyResolver, DidResolveError};
use crate::service_jwt::{precheck_claims, verify_service_jwt};

/// Seconds since the Unix epoch (0 if the clock is before it).
pub fn now_unix() -> u64 {
    SystemTime::now().duration_since(UNIX_EPOCH).map(|d| d.as_secs()).unwrap_or(0)
}

/// Record-and-check a `jti` for replay. The contract is "first use wins":
///   * returns `true`  ⇒ this `jti` was NOT seen before (fresh) — proceed.
///   * returns `false` ⇒ this `jti` was already recorded within its `exp` window — REPLAY, reject.
///
/// Per the §3.4 fail-OPEN rule, an implementation that cannot reach its backing store returns `true`
/// (treat as fresh) rather than erroring — the trait has no error channel by design.
#[async_trait]
pub trait JtiCache: Send + Sync {
    /// Record `jti` (expiring at `exp_unix`) and report whether it was fresh. Idempotent per `jti`.
    async fn check_and_record(&self, jti: &str, exp_unix: u64) -> bool;
}

/// The default in-process cache — a `Mutex<HashMap>` with lazy eviction at `exp` (sufficient for the
/// single-replica service). A caller that runs several replicas implements [`JtiCache`] over a shared
/// store (the platform keeps a Postgres one).
#[derive(Default)]
pub struct InProcessJtiCache {
    /// `jti` → `exp_unix`. An entry past its `exp` is treated as absent (lazily evicted on the next
    /// touch of that key, and proactively when we observe a now-past entry during a check).
    seen: Mutex<HashMap<String, u64>>,
}

impl InProcessJtiCache {
    /// A fresh empty cache.
    pub fn new() -> Self {
        Self::default()
    }
}

#[async_trait]
impl JtiCache for InProcessJtiCache {
    async fn check_and_record(&self, jti: &str, exp_unix: u64) -> bool {
        let now = now_unix();
        let mut seen = self.seen.lock().unwrap();
        // Lazy eviction: drop everything already past its exp so the map can't grow unbounded.
        seen.retain(|_, &mut exp| exp > now);
        match seen.get(jti) {
            // A live (not-yet-expired) prior record ⇒ replay.
            Some(&exp) if exp > now => false,
            // Either unseen, or a stale record we just cleared — record and report fresh.
            _ => {
                seen.insert(jti.to_string(), exp_unix);
                true
            }
        }
    }
}

/// How long a resolved `#atproto` key is reused before the issuer's DID is resolved again.
pub const KEY_CACHE_TTL: Duration = Duration::from_secs(300);
/// How long a DID that resolved to no usable key is refused without resolving it again, so junk
/// issuers cannot make every request a DID lookup.
pub const NEGATIVE_CACHE_TTL: Duration = Duration::from_secs(60);

/// Why a request was not authenticated.
#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)]
pub enum AuthError {
    /// No token (HTTP 401).
    #[error("authentication required")]
    Missing,
    /// A token that fails verification (HTTP 401).
    #[error("invalid service-auth token")]
    Invalid,
    /// The issuer's DID could not be resolved right now (HTTP 503): the caller should retry.
    #[error("the caller's DID could not be resolved right now")]
    Unavailable,
}

/// Verifies service-auth JWTs addressed to one service DID.
pub struct ServiceAuth {
    service_did: String,
    resolver: Arc<dyn DidKeyResolver>,
    replay: Arc<dyn JtiCache>,
    keys: Mutex<HashMap<String, (VerifyingKey, Instant)>>,
    key_ttl: Duration,
    /// DIDs whose key did not resolve (not "the directory did not answer"), and when.
    unresolvable: Mutex<HashMap<String, Instant>>,
}

impl ServiceAuth {
    /// `service_did` is the `aud` every token must carry.
    pub fn new(service_did: impl Into<String>, resolver: Arc<dyn DidKeyResolver>, replay: Arc<dyn JtiCache>) -> Self {
        Self {
            service_did: service_did.into(),
            resolver,
            replay,
            keys: Mutex::default(),
            key_ttl: KEY_CACHE_TTL,
            unresolvable: Mutex::default(),
        }
    }

    /// Reuse resolved keys for `ttl` instead of [`KEY_CACHE_TTL`] (zero: resolve on every request).
    pub fn with_key_ttl(mut self, ttl: Duration) -> Self {
        self.key_ttl = ttl;
        self
    }

    pub fn service_did(&self) -> &str {
        &self.service_did
    }

    /// The caller's DID, from an `Authorization: Bearer <jwt>` value minted for method `lxm`.
    pub async fn caller(&self, authorization: Option<&str>, lxm: &str) -> Result<String, AuthError> {
        let jwt = authorization
            .and_then(|h| h.strip_prefix("Bearer ").or_else(|| h.strip_prefix("bearer ")))
            .map(str::trim)
            .filter(|t| !t.is_empty())
            .ok_or(AuthError::Missing)?;
        let now = now_unix();
        // Each refusal is logged with its cause before it collapses into one opaque 401.
        let Some(iss) = precheck_claims(jwt, &self.service_did, lxm, now) else {
            tracing::warn!(%lxm, "service-auth refused: malformed, expired, or not minted for this aud and lxm");
            return Err(AuthError::Invalid);
        };
        let (key, cached) = match self.cached_key(&iss) {
            Some(key) => (key, true),
            None => (self.resolve_key(&iss, lxm).await?, false),
        };
        let mut verified = verify_service_jwt(jwt, &key, &self.service_did, lxm, now);
        if verified.is_err() && cached {
            // The cached key may be one the issuer has since rotated away from: resolve once more.
            // The cached key stays until a different one replaces it, so a bad token cannot evict a
            // good reader's key; and since this token already failed against a key the issuer did
            // publish, an unanswering directory or an unchanged key is a refusal, not a 503.
            match self.resolver.atproto_key(&iss).await {
                Ok(fresh) if fresh != key => {
                    self.remember(&iss, fresh);
                    verified = verify_service_jwt(jwt, &fresh, &self.service_did, lxm, now);
                }
                Ok(_) => {}
                Err(error) => tracing::warn!(%iss, %lxm, ?error, "service-auth: re-resolving a cached key failed"),
            }
        }
        let claims = verified.map_err(|error| {
            tracing::warn!(%iss, %lxm, ?error, "service-auth refused: verification failed");
            AuthError::Invalid
        })?;
        if !self.replay.check_and_record(&claims.jti, claims.exp).await {
            tracing::warn!(%iss, %lxm, "service-auth refused: replayed jti");
            return Err(AuthError::Invalid);
        }
        Ok(claims.iss)
    }

    fn cached_key(&self, did: &str) -> Option<VerifyingKey> {
        let keys = self.keys.lock().unwrap_or_else(|e| e.into_inner());
        keys.get(did).filter(|(_, at)| at.elapsed() < self.key_ttl).map(|(key, _)| *key)
    }

    fn remember(&self, did: &str, key: VerifyingKey) {
        let mut keys = self.keys.lock().unwrap_or_else(|e| e.into_inner());
        // Expired entries go when the map is next written, so it stays as small as the set of issuers
        // seen within one TTL.
        keys.retain(|_, (_, at)| at.elapsed() < self.key_ttl);
        keys.insert(did.to_owned(), (key, Instant::now()));
    }

    /// Resolve `did`'s `#atproto` key and cache it. A DID that resolved to no usable key is refused
    /// without a lookup for [`NEGATIVE_CACHE_TTL`]; a directory that did not answer is not cached.
    async fn resolve_key(&self, did: &str, lxm: &str) -> Result<VerifyingKey, AuthError> {
        {
            let mut unresolvable = self.unresolvable.lock().unwrap_or_else(|e| e.into_inner());
            unresolvable.retain(|_, at| at.elapsed() < NEGATIVE_CACHE_TTL);
            if unresolvable.contains_key(did) {
                tracing::warn!(iss = %did, %lxm, "service-auth refused: the issuer's signing key did not resolve (cached)");
                return Err(AuthError::Invalid);
            }
        }
        match self.resolver.atproto_key(did).await {
            Ok(key) => {
                self.remember(did, key);
                Ok(key)
            }
            Err(DidResolveError::Unavailable) => {
                tracing::warn!(iss = %did, %lxm, "service-auth deferred: the issuer's DID directory did not answer");
                Err(AuthError::Unavailable)
            }
            Err(error) => {
                tracing::warn!(iss = %did, %lxm, ?error, "service-auth refused: the issuer's signing key did not resolve");
                self.unresolvable.lock().unwrap_or_else(|e| e.into_inner()).insert(did.to_owned(), Instant::now());
                Err(AuthError::Invalid)
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::did::InMemoryDidKeyResolver;
    use crate::labels::multibase;
    use k256::ecdsa::SigningKey;

    #[tokio::test]
    async fn first_use_is_fresh_replay_is_rejected() {
        let cache = InProcessJtiCache::new();
        let exp = now_unix() + 60;
        assert!(cache.check_and_record("jti-1", exp).await, "first use must be fresh");
        assert!(!cache.check_and_record("jti-1", exp).await, "replay must be rejected");
        // A distinct jti is independent.
        assert!(cache.check_and_record("jti-2", exp).await, "a different jti is fresh");
    }

    #[tokio::test]
    async fn an_expired_record_is_evicted_and_freshly_acceptable() {
        let cache = InProcessJtiCache::new();
        // Record a jti that is ALREADY expired (exp in the past).
        let past = now_unix().saturating_sub(10);
        assert!(cache.check_and_record("old", past).await, "first insert reports fresh");
        // It was already past exp, so a subsequent check evicts it and treats the jti as fresh again.
        assert!(cache.check_and_record("old", now_unix() + 60).await, "expired record must not block");
    }

    const SERVICE: &str = "did:web:rules.example";
    const CALLER: &str = "did:plc:caller000000000000000000";
    const LXM: &str = "pub.techne.draft.commons.runChain";

    fn hex_key(b: u8) -> String {
        hex::encode([b; 32])
    }

    fn token(key: u8, aud: &str, lxm: &str) -> String {
        crate::service_jwt::mint_service_jwt(&hex_key(key), CALLER, aud, lxm, 60, now_unix()).unwrap()
    }

    fn auth_with(resolver: Arc<dyn DidKeyResolver>) -> ServiceAuth {
        ServiceAuth::new(SERVICE, resolver, Arc::new(InProcessJtiCache::new()))
    }

    fn seeded(key: u8) -> Arc<InMemoryDidKeyResolver> {
        let r = InMemoryDidKeyResolver::new();
        r.insert_key(CALLER, multibase(SigningKey::from_slice(&[key; 32]).unwrap().verifying_key()));
        Arc::new(r)
    }

    #[tokio::test]
    async fn a_token_for_this_service_names_its_caller_once() {
        let auth = auth_with(seeded(7));
        assert_eq!(auth.service_did(), SERVICE);
        let jwt = token(7, SERVICE, LXM);
        let bearer = format!("Bearer {jwt}");
        assert_eq!(auth.caller(Some(&bearer), LXM).await, Ok(CALLER.to_owned()));
        assert_eq!(auth.caller(Some(&bearer), LXM).await, Err(AuthError::Invalid), "a replayed jti");
        assert_eq!(auth.caller(Some(&format!("bearer {}", token(7, SERVICE, LXM))), LXM).await, Ok(CALLER.to_owned()));
    }

    #[tokio::test]
    async fn missing_and_wrong_tokens_are_refused() {
        let auth = auth_with(seeded(7));
        assert_eq!(auth.caller(None, LXM).await, Err(AuthError::Missing));
        assert_eq!(auth.caller(Some("Bearer  "), LXM).await, Err(AuthError::Missing));
        assert_eq!(auth.caller(Some("Basic abc"), LXM).await, Err(AuthError::Missing));
        assert_eq!(auth.caller(Some(&format!("Bearer {}", token(7, "did:web:other.example", LXM))), LXM).await, Err(AuthError::Invalid));
        assert_eq!(auth.caller(Some(&format!("Bearer {}", token(7, SERVICE, "x.y.other"))), LXM).await, Err(AuthError::Invalid));
        assert_eq!(auth.caller(Some(&format!("Bearer {}", token(8, SERVICE, LXM))), LXM).await, Err(AuthError::Invalid), "another key");
        assert_eq!(auth.caller(Some("Bearer junk"), LXM).await, Err(AuthError::Invalid));
    }

    /// Answers `Unavailable` until told otherwise, counting lookups.
    struct Flaky {
        inner: Arc<InMemoryDidKeyResolver>,
        down: std::sync::atomic::AtomicBool,
        lookups: std::sync::atomic::AtomicUsize,
    }

    #[async_trait]
    impl DidKeyResolver for Flaky {
        async fn atproto_key(&self, did: &str) -> Result<VerifyingKey, DidResolveError> {
            self.lookups.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
            if self.down.load(std::sync::atomic::Ordering::SeqCst) {
                return Err(DidResolveError::Unavailable);
            }
            self.inner.atproto_key(did).await
        }
        async fn pds_endpoint(&self, _did: &str) -> Result<Option<String>, DidResolveError> {
            Ok(None)
        }
        async fn handle(&self, _did: &str) -> Result<Option<String>, DidResolveError> {
            Ok(None)
        }
    }

    #[tokio::test]
    async fn a_directory_that_does_not_answer_is_unavailable_not_a_refusal() {
        let flaky = Arc::new(Flaky { inner: seeded(7), down: true.into(), lookups: 0.into() });
        let auth = auth_with(flaky.clone());
        assert_eq!(auth.caller(Some(&format!("Bearer {}", token(7, SERVICE, LXM))), LXM).await, Err(AuthError::Unavailable));
        flaky.down.store(false, std::sync::atomic::Ordering::SeqCst);
        assert_eq!(auth.caller(Some(&format!("Bearer {}", token(7, SERVICE, LXM))), LXM).await, Ok(CALLER.to_owned()), "not cached");
        // The key is cached now: a later call does not resolve again.
        let before = flaky.lookups.load(std::sync::atomic::Ordering::SeqCst);
        assert_eq!(auth.caller(Some(&format!("Bearer {}", token(7, SERVICE, LXM))), LXM).await, Ok(CALLER.to_owned()));
        assert_eq!(flaky.lookups.load(std::sync::atomic::Ordering::SeqCst), before);
    }

    #[tokio::test]
    async fn an_unresolvable_issuer_is_refused_without_another_lookup() {
        let flaky = Arc::new(Flaky { inner: Arc::new(InMemoryDidKeyResolver::new()), down: false.into(), lookups: 0.into() });
        let auth = auth_with(flaky.clone());
        for _ in 0..3 {
            assert_eq!(auth.caller(Some(&format!("Bearer {}", token(7, SERVICE, LXM))), LXM).await, Err(AuthError::Invalid));
        }
        assert_eq!(flaky.lookups.load(std::sync::atomic::Ordering::SeqCst), 1, "negative-cached");
    }

    #[tokio::test]
    async fn a_rotated_key_works_at_once() {
        let resolver = seeded(7);
        let auth = auth_with(resolver.clone());
        assert_eq!(auth.caller(Some(&format!("Bearer {}", token(7, SERVICE, LXM))), LXM).await, Ok(CALLER.to_owned()));
        resolver.insert_key(CALLER, multibase(SigningKey::from_slice(&[9; 32]).unwrap().verifying_key()));
        assert_eq!(auth.caller(Some(&format!("Bearer {}", token(9, SERVICE, LXM))), LXM).await, Ok(CALLER.to_owned()));
    }
}