captchaforge 0.2.26

//! Per-domain solved-token cache with TTL.
//!
//! Many captcha vendors issue tokens valid for 60–180 seconds.
//! Re-solving the same captcha within that window wastes solver
//! credits, model time, and adds avoidable latency to repeated visits
//! to the same site. The [`TokenCache`] short-circuits the chain when
//! a fresh token already exists for `(domain, captcha_type)`.
//!
//! Default TTL is 60 seconds — conservative enough to keep tokens
//! fresh against most upstream validators. Configure via
//! [`TokenCache::with_ttl`] for sites with longer-lived tokens.
//!
//! Stale-token recovery (a solver tells us its cached token was
//! rejected by the upstream) is intentionally NOT bundled here; it
//! belongs in a separate "solve-result feedback" loop that calls
//! [`TokenCache::invalidate`] on rejection.
use std::collections::HashMap;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Mutex;
use std::time::{Duration, Instant};

use crate::cookies::CapturedCookie;
use crate::solver::CaptchaType;

/// Snapshot of cache effectiveness over the lifetime of a [`TokenCache`].
/// Atomic counters are reset to zero on construction. Read via
/// [`TokenCache::stats`]; production deployments scrape this each
/// minute to plot cache hit rate and detect when TTL needs tuning.
#[derive(Debug, Clone, Copy)]
pub struct CacheStats {
    /// Total `get`/`get_token` calls that returned a fresh entry.
    pub hits: u64,
    /// Total calls that returned `None` because the key was absent.
    pub misses: u64,
    /// Calls that returned `None` because the entry existed but had
    /// expired. A high ratio of `expired_misses : hits` means TTL is
    /// shorter than the typical re-visit window — bump it.
    pub expired_misses: u64,
    /// Total `put`/`put_with_ttl`/`put_full` calls.
    pub puts: u64,
    /// Total `invalidate` calls.
    pub invalidations: u64,
}

impl CacheStats {
    /// Cache hit rate as a fraction in `[0.0, 1.0]`. Returns `None`
    /// when no lookups have occurred yet (avoid 0/0 in dashboards).
    pub fn hit_rate(&self) -> Option<f64> {
        let lookups = self.hits + self.misses + self.expired_misses;
        if lookups == 0 {
            None
        } else {
            Some(self.hits as f64 / lookups as f64)
        }
    }
}

/// One cached token entry. Returned by [`TokenCache::get`].
///
/// Held by the cache; consumers borrow via the accessor methods. The
/// type is `Clone` so callers can take ownership without holding the
/// cache lock.
#[derive(Debug, Clone)]
pub struct CachedToken {
    token: String,
    /// Solver name that produced this token, for telemetry pass-through.
    method_name: &'static str,
    expires_at: Instant,
    /// Browser cookies captured at the moment of the original solve.
    /// Replayed onto the page on cache hit so the WAF/vendor's
    /// trusted session rides along with the captcha token — without
    /// these, a hit returns the token but the next request lands on
    /// a fresh session that immediately re-triggers the captcha.
    cookies: Vec<CapturedCookie>,
}

/// Concurrent cache mapping `(domain, CaptchaType)` to a solved token
/// that hasn't yet expired.
///
/// Backed by a `Mutex<HashMap>` rather than DashMap so the crate keeps
/// its lean dependency surface; reads are still O(1) and the lock is
/// only held during the lookup itself.
pub struct TokenCache {
    inner: Mutex<HashMap<(String, CaptchaType), CachedToken>>,
    default_ttl: Duration,
    hits: AtomicU64,
    misses: AtomicU64,
    expired_misses: AtomicU64,
    puts: AtomicU64,
    invalidations: AtomicU64,
}

impl TokenCache {
    /// New cache with a 60-second default TTL.
    pub fn new() -> Self {
        Self::with_ttl(Duration::from_secs(60))
    }

    /// New cache with the given default TTL applied to every `put`
    /// that doesn't override it.
    pub fn with_ttl(default_ttl: Duration) -> Self {
        Self {
            inner: Mutex::new(HashMap::new()),
            default_ttl,
            hits: AtomicU64::new(0),
            misses: AtomicU64::new(0),
            expired_misses: AtomicU64::new(0),
            puts: AtomicU64::new(0),
            invalidations: AtomicU64::new(0),
        }
    }

    /// Look up a fresh token. Returns `None` for misses AND for
    /// entries past their expiry. Expired entries are not auto-purged
    /// here — that happens lazily on the next `put` for the same key.
    pub fn get(&self, domain: &str, captcha_type: &CaptchaType) -> Option<CachedToken> {
        let map = self.inner.lock().unwrap();
        let key = (domain.to_owned(), captcha_type.clone());
        match map.get(&key) {
            None => {
                self.misses.fetch_add(1, Ordering::Relaxed);
                None
            }
            Some(entry) if entry.expires_at > Instant::now() => {
                self.hits.fetch_add(1, Ordering::Relaxed);
                Some(entry.clone())
            }
            Some(_) => {
                self.expired_misses.fetch_add(1, Ordering::Relaxed);
                None
            }
        }
    }

    /// Convenience: just the token string when present and fresh.
    pub fn get_token(&self, domain: &str, captcha_type: &CaptchaType) -> Option<String> {
        self.get(domain, captcha_type).map(|e| e.token)
    }

    /// Insert or replace a token entry with the default TTL and no
    /// cookies. Kept for back-compat; new code should prefer
    /// [`Self::put_full`] so cache hits replay the captured session.
    pub fn put(
        &self,
        domain: &str,
        captcha_type: &CaptchaType,
        token: String,
        method_name: &'static str,
    ) {
        self.put_with_ttl(domain, captcha_type, token, method_name, self.default_ttl);
    }

    /// Insert or replace with an explicit TTL but no cookies. See
    /// [`Self::put_full`] for the full-fidelity variant.
    pub fn put_with_ttl(
        &self,
        domain: &str,
        captcha_type: &CaptchaType,
        token: String,
        method_name: &'static str,
        ttl: Duration,
    ) {
        self.put_full(domain, captcha_type, token, method_name, ttl, Vec::new());
    }

    /// Insert or replace with explicit TTL + the cookies captured at
    /// the moment of the original solve. Cache hits replay these
    /// cookies via [`crate::cookies::apply_to_page`] so the
    /// WAF/vendor's trusted session is preserved across the cache
    /// boundary — without them a hit returns the token alone and the
    /// next request can immediately re-trigger the captcha.
    pub fn put_full(
        &self,
        domain: &str,
        captcha_type: &CaptchaType,
        token: String,
        method_name: &'static str,
        ttl: Duration,
        cookies: Vec<CapturedCookie>,
    ) {
        let mut map = self.inner.lock().unwrap();
        map.insert(
            (domain.to_owned(), captcha_type.clone()),
            CachedToken {
                token,
                method_name,
                expires_at: Instant::now() + ttl,
                cookies,
            },
        );
        self.puts.fetch_add(1, Ordering::Relaxed);
    }

    /// Drop a single cached token. Used when a downstream consumer
    /// reports the cached token was rejected by the upstream
    /// validator (rotation, blocklist, replay protection).
    pub fn invalidate(&self, domain: &str, captcha_type: &CaptchaType) {
        let mut map = self.inner.lock().unwrap();
        if map
            .remove(&(domain.to_owned(), captcha_type.clone()))
            .is_some()
        {
            self.invalidations.fetch_add(1, Ordering::Relaxed);
        }
    }

    /// Empty the cache. Used in tests and on long-running daemons
    /// that want to reset state between batches.
    pub fn clear(&self) {
        self.inner.lock().unwrap().clear();
    }

    /// Number of entries currently held (including expired). Mostly
    /// useful for diagnostics + tests.
    pub fn len(&self) -> usize {
        self.inner.lock().unwrap().len()
    }

    /// True iff the cache holds zero entries.
    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }

    /// The default TTL applied to `put` calls without an explicit TTL.
    pub fn ttl(&self) -> Duration {
        self.default_ttl
    }

    /// Snapshot of the lifetime hit/miss/expiry counters. Cheap (atomic
    /// loads, no lock acquisition). Production deployments call this
    /// every minute to plot cache effectiveness — a hit rate stuck at
    /// zero means TTL is shorter than the typical re-visit window;
    /// a high `expired_misses : hits` ratio is the same signal.
    pub fn stats(&self) -> CacheStats {
        CacheStats {
            hits: self.hits.load(Ordering::Relaxed),
            misses: self.misses.load(Ordering::Relaxed),
            expired_misses: self.expired_misses.load(Ordering::Relaxed),
            puts: self.puts.load(Ordering::Relaxed),
            invalidations: self.invalidations.load(Ordering::Relaxed),
        }
    }

    /// Reset the lifetime counters to zero. Useful for batched
    /// per-window measurements when a long-running daemon scrapes
    /// stats and wants to start a fresh interval.
    pub fn reset_stats(&self) {
        self.hits.store(0, Ordering::Relaxed);
        self.misses.store(0, Ordering::Relaxed);
        self.expired_misses.store(0, Ordering::Relaxed);
        self.puts.store(0, Ordering::Relaxed);
        self.invalidations.store(0, Ordering::Relaxed);
    }
}

impl CachedToken {
    /// Borrow the cached token string.
    pub fn token(&self) -> &str {
        &self.token
    }

    /// Solver name that produced this entry.
    pub fn method_name(&self) -> &'static str {
        self.method_name
    }

    /// Borrow the cookies captured at the original solve. Empty when
    /// the entry was inserted via [`TokenCache::put`] /
    /// [`TokenCache::put_with_ttl`] (no cookies path).
    pub fn cookies(&self) -> &[CapturedCookie] {
        &self.cookies
    }
}

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

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn miss_returns_none() {
        let c = TokenCache::new();
        assert!(c
            .get("example.com", &CaptchaType::CloudflareTurnstile)
            .is_none());
    }

    #[test]
    fn hit_returns_token_within_ttl() {
        let c = TokenCache::new();
        c.put(
            "example.com",
            &CaptchaType::CloudflareTurnstile,
            "tok123".into(),
            "BehavioralCaptchaSolver",
        );
        let got = c.get_token("example.com", &CaptchaType::CloudflareTurnstile);
        assert_eq!(got.as_deref(), Some("tok123"));
    }

    #[test]
    fn expired_entry_returns_none() {
        let c = TokenCache::with_ttl(Duration::from_millis(1));
        c.put(
            "example.com",
            &CaptchaType::HCaptcha,
            "tok".into(),
            "VlmCaptchaSolver",
        );
        // Sleep just past TTL.
        std::thread::sleep(Duration::from_millis(5));
        assert!(c.get("example.com", &CaptchaType::HCaptcha).is_none());
    }

    #[test]
    fn put_replaces_existing_entry() {
        let c = TokenCache::new();
        c.put(
            "x.test",
            &CaptchaType::HCaptcha,
            "old".into(),
            "VlmCaptchaSolver",
        );
        c.put(
            "x.test",
            &CaptchaType::HCaptcha,
            "new".into(),
            "VlmCaptchaSolver",
        );
        assert_eq!(
            c.get_token("x.test", &CaptchaType::HCaptcha).as_deref(),
            Some("new")
        );
        assert_eq!(c.len(), 1, "replace must not grow the map");
    }

    #[test]
    fn invalidate_drops_only_the_targeted_entry() {
        let c = TokenCache::new();
        c.put("a.test", &CaptchaType::HCaptcha, "ta".into(), "Vlm");
        c.put("b.test", &CaptchaType::HCaptcha, "tb".into(), "Vlm");
        c.invalidate("a.test", &CaptchaType::HCaptcha);
        assert!(c.get("a.test", &CaptchaType::HCaptcha).is_none());
        assert!(c.get("b.test", &CaptchaType::HCaptcha).is_some());
    }

    #[test]
    fn different_captcha_types_on_same_domain_are_isolated() {
        let c = TokenCache::new();
        c.put("d.test", &CaptchaType::HCaptcha, "h".into(), "Vlm");
        c.put(
            "d.test",
            &CaptchaType::CloudflareTurnstile,
            "t".into(),
            "Beh",
        );
        assert_eq!(
            c.get_token("d.test", &CaptchaType::HCaptcha).as_deref(),
            Some("h")
        );
        assert_eq!(
            c.get_token("d.test", &CaptchaType::CloudflareTurnstile)
                .as_deref(),
            Some("t")
        );
    }

    #[test]
    fn put_with_ttl_overrides_default() {
        let c = TokenCache::with_ttl(Duration::from_secs(10));
        c.put_with_ttl(
            "x.test",
            &CaptchaType::HCaptcha,
            "tok".into(),
            "Vlm",
            Duration::from_millis(1),
        );
        std::thread::sleep(Duration::from_millis(5));
        assert!(
            c.get("x.test", &CaptchaType::HCaptcha).is_none(),
            "explicit short TTL wins"
        );
    }

    #[test]
    fn put_with_cookies_round_trips_on_hit() {
        let c = TokenCache::new();
        let cookie = CapturedCookie {
            name: "datadome".into(),
            value: "abc".into(),
            domain: "x.test".into(),
            path: "/".into(),
            expires: None,
            secure: true,
            http_only: true,
            same_site: None,
        };
        c.put_full(
            "x.test",
            &CaptchaType::CloudflareTurnstile,
            "tok".into(),
            "Behavioral",
            Duration::from_secs(60),
            vec![cookie.clone()],
        );
        let entry = c.get("x.test", &CaptchaType::CloudflareTurnstile).unwrap();
        assert_eq!(entry.cookies().len(), 1);
        assert_eq!(entry.cookies()[0].name, "datadome");
    }

    #[test]
    fn put_without_cookies_returns_empty_cookies_on_hit() {
        let c = TokenCache::new();
        c.put("x.test", &CaptchaType::HCaptcha, "t".into(), "Vlm");
        let entry = c.get("x.test", &CaptchaType::HCaptcha).unwrap();
        assert!(entry.cookies().is_empty());
    }

    #[test]
    fn stats_initialise_to_zero() {
        let c = TokenCache::new();
        let s = c.stats();
        assert_eq!(s.hits, 0);
        assert_eq!(s.misses, 0);
        assert_eq!(s.expired_misses, 0);
        assert_eq!(s.puts, 0);
        assert_eq!(s.invalidations, 0);
        assert!(s.hit_rate().is_none(), "no lookups → no rate");
    }

    #[test]
    fn stats_count_each_event_class_separately() {
        let c = TokenCache::new();
        // 1 put, 0 lookups so far.
        c.put("a.test", &CaptchaType::HCaptcha, "t".into(), "Vlm");
        assert_eq!(c.stats().puts, 1);

        // 2 hits.
        let _ = c.get("a.test", &CaptchaType::HCaptcha);
        let _ = c.get("a.test", &CaptchaType::HCaptcha);
        assert_eq!(c.stats().hits, 2);

        // 1 miss.
        let _ = c.get("never.test", &CaptchaType::HCaptcha);
        assert_eq!(c.stats().misses, 1);

        // 1 invalidation.
        c.invalidate("a.test", &CaptchaType::HCaptcha);
        assert_eq!(c.stats().invalidations, 1);

        // hit_rate = 2 / (2+1+0) = 0.667
        let s = c.stats();
        let r = s.hit_rate().unwrap();
        assert!(
            (r - 2.0 / 3.0).abs() < 0.001,
            "expected 0.667 hit rate; got {r}"
        );
    }

    #[test]
    fn stats_track_expired_misses_distinct_from_misses() {
        let c = TokenCache::with_ttl(Duration::from_millis(1));
        c.put("x.test", &CaptchaType::HCaptcha, "t".into(), "Vlm");
        std::thread::sleep(Duration::from_millis(5));
        let _ = c.get("x.test", &CaptchaType::HCaptcha);
        let s = c.stats();
        assert_eq!(s.misses, 0, "key existed → not a plain miss");
        assert_eq!(s.expired_misses, 1, "key existed but expired");
    }

    #[test]
    fn invalidate_on_missing_key_does_not_count() {
        let c = TokenCache::new();
        c.invalidate("missing.test", &CaptchaType::HCaptcha);
        assert_eq!(c.stats().invalidations, 0);
    }

    #[test]
    fn reset_stats_zeroes_counters_but_keeps_entries() {
        let c = TokenCache::new();
        c.put("a.test", &CaptchaType::HCaptcha, "t".into(), "Vlm");
        let _ = c.get("a.test", &CaptchaType::HCaptcha);
        assert_eq!(c.stats().hits, 1);
        c.reset_stats();
        let s = c.stats();
        assert_eq!(s.hits, 0);
        assert_eq!(s.puts, 0);
        // Entry still in cache.
        assert!(c.get("a.test", &CaptchaType::HCaptcha).is_some());
    }

    #[test]
    fn clear_removes_all_entries() {
        let c = TokenCache::new();
        c.put("a.test", &CaptchaType::HCaptcha, "ta".into(), "Vlm");
        c.put("b.test", &CaptchaType::HCaptcha, "tb".into(), "Vlm");
        assert_eq!(c.len(), 2);
        c.clear();
        assert!(c.is_empty());
    }
}