captchaforge 0.2.39

//! Thompson-sampling multi-armed bandit over solver methods.
//!
//! `CaptchaSolverChain` runs strategies in a STATIC priority order
//! (`BehavioralCaptchaSolver` first, then `VlmCaptchaSolver`, then
//! audio, etc.). The `PatternStore` re-orders eligible solvers
//! per-domain using an EMA over historical success. Both work fine
//! at low traffic but converge slowly when the vendor changes
//! behaviour (e.g., reCAPTCHA suddenly throws every visitor into
//! the audio path).
//!
//! `Bandit` solves the explore/exploit tradeoff properly:
//!
//! - One arm per `(captcha_type, solver_method)` pair.
//! - Beta(α, β) posterior per arm; α counts successes, β failures.
//! - On every solve, sample θ_arm ~ Beta(α_arm, β_arm) for each
//!   candidate, pick argmax. Naturally explores low-data arms while
//!   exploiting high-data winners.
//! - Bayesian regret bounds: O(√(T log T)) per arm, which is
//!   provably optimal for the multi-armed bandit problem.
//!
//! No public captcha-solver crate ships a bandit chooser. The
//! closest prior art is `linkedin/photon-ml`'s online-learning
//! library (but that's for ad-ranking, not bot-bypass).
//!
//! Persisted to disk per-process via [`Bandit::save_to`] +
//! [`Bandit::load_from`] so a long-running service doesn't lose
//! its priors on restart.

use rand::Rng;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::path::Path;
use std::sync::RwLock;

use crate::solver::{CaptchaType, SolveMethod};

/// One arm of the bandit: a `(captcha_type, solver_method)` pair
/// with Beta posterior.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Arm {
    pub captcha_type: String,
    pub solver_method: String,
    /// Beta α parameter (successes + 1).
    pub alpha: f64,
    /// Beta β parameter (failures + 1).
    pub beta: f64,
}

impl Arm {
    pub fn new(captcha_type: String, solver_method: String) -> Self {
        Self {
            captcha_type,
            solver_method,
            alpha: 1.0,
            beta: 1.0,
        }
    }

    /// Posterior mean (for diagnostics / rendering).
    pub fn posterior_mean(&self) -> f64 {
        self.alpha / (self.alpha + self.beta)
    }

    /// Posterior variance (for diagnostics).
    pub fn posterior_variance(&self) -> f64 {
        let a = self.alpha;
        let b = self.beta;
        let denom = (a + b).powi(2) * (a + b + 1.0);
        (a * b) / denom
    }

    /// Update on a solve outcome. `succeeded=true` increments α.
    pub fn observe(&mut self, succeeded: bool) {
        if succeeded {
            self.alpha += 1.0;
        } else {
            self.beta += 1.0;
        }
    }
}

/// Bandit chooser. Internally a HashMap of arms keyed by
/// `(captcha_type, solver_method)` strings; thread-safe via RwLock.
pub struct Bandit {
    arms: RwLock<HashMap<(String, String), Arm>>,
    /// Decay factor applied per arm per observation outside its
    /// rolling-window quota, keeps the priors fresh when a vendor
    /// changes behaviour. `1.0` = no decay; sensible default 0.999.
    decay: f64,
}

impl Bandit {
    /// Construct an empty bandit. Arms are created lazily on first
    /// `choose` / `observe` call.
    ///
    /// # Example
    ///
    /// ```rust
    /// use captchaforge::solver::bandit::Bandit;
    /// use captchaforge::solver::{CaptchaType, SolveMethod};
    ///
    /// let bandit = Bandit::new();
    /// assert!(bandit.is_empty());
    /// bandit.observe(
    ///     &CaptchaType::CloudflareTurnstile,
    ///     &SolveMethod::AutoPass,
    ///     true,
    /// );
    /// assert_eq!(bandit.len(), 1);
    /// ```
    pub fn new() -> Self {
        Self {
            arms: RwLock::new(HashMap::new()),
            decay: 1.0,
        }
    }

    /// Configure exponential decay per observation. Lower → more
    /// reactive to recent observations; higher → more stable.
    /// Must be in [0.5, 1.0].
    pub fn with_decay(mut self, decay: f64) -> Self {
        self.decay = decay.clamp(0.5, 1.0);
        self
    }

    /// Sample a θ from each arm's Beta posterior and pick the
    /// arm with highest θ among `candidates`. Returns the chosen
    /// solver_method. If no arms exist yet for the captcha_type,
    /// returns the first candidate.
    pub fn choose(
        &self,
        captcha_type: &CaptchaType,
        candidates: &[SolveMethod],
    ) -> Option<SolveMethod> {
        if candidates.is_empty() {
            return None;
        }
        let arms = self.arms.read().unwrap_or_else(|e| e.into_inner());
        let ct = render_captcha_type(captcha_type);
        let mut rng = rand::thread_rng();
        let mut best: Option<(f64, SolveMethod)> = None;
        for cand in candidates {
            let sm = render_solve_method(cand);
            let arm = arms
                .get(&(ct.clone(), sm.clone()))
                .cloned()
                .unwrap_or_else(|| Arm::new(ct.clone(), sm.clone()));
            let theta = sample_beta(arm.alpha, arm.beta, &mut rng);
            if best.as_ref().map(|(b, _)| theta > *b).unwrap_or(true) {
                best = Some((theta, cand.clone()));
            }
        }
        best.map(|(_, m)| m)
    }

    /// Observe a solve outcome and update the arm's posterior.
    pub fn observe(
        &self,
        captcha_type: &CaptchaType,
        solver_method: &SolveMethod,
        succeeded: bool,
    ) {
        let ct = render_captcha_type(captcha_type);
        let sm = render_solve_method(solver_method);
        let mut arms = self.arms.write().unwrap_or_else(|e| e.into_inner());
        // Decay all other arms slightly so priors stay fresh.
        if self.decay < 1.0 {
            for arm in arms.values_mut() {
                if arm.alpha + arm.beta > 2.0 {
                    arm.alpha = ((arm.alpha - 1.0) * self.decay) + 1.0;
                    arm.beta = ((arm.beta - 1.0) * self.decay) + 1.0;
                }
            }
        }
        let arm = arms
            .entry((ct.clone(), sm.clone()))
            .or_insert_with(|| Arm::new(ct, sm));
        arm.observe(succeeded);
    }

    /// Snapshot every arm's state.
    pub fn snapshot(&self) -> Vec<Arm> {
        let arms = self.arms.read().unwrap_or_else(|e| e.into_inner());
        arms.values().cloned().collect()
    }

    /// Persist the bandit to a file at `path` as JSON.
    pub fn save_to(&self, path: &Path) -> std::io::Result<()> {
        let snap = self.snapshot();
        let s = serde_json::to_string(&snap)
            .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))?;
        std::fs::write(path, s)
    }

    /// Load arms from a JSON file. Adds to existing state (does
    /// not clear).
    pub fn load_from(&self, path: &Path) -> std::io::Result<usize> {
        let body = std::fs::read_to_string(path)?;
        let snap: Vec<Arm> = serde_json::from_str(&body)
            .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))?;
        let count = snap.len();
        let mut arms = self.arms.write().unwrap_or_else(|e| e.into_inner());
        for arm in snap {
            arms.insert((arm.captcha_type.clone(), arm.solver_method.clone()), arm);
        }
        Ok(count)
    }

    /// Number of distinct arms tracked.
    pub fn len(&self) -> usize {
        self.arms.read().unwrap_or_else(|e| e.into_inner()).len()
    }

    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }
}

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

/// Sample x ~ Beta(α, β) via the Gamma-ratio method:
///   X = G1 / (G1 + G2)
/// where G1 ~ Gamma(α, 1), G2 ~ Gamma(β, 1).
///
/// We use the Marsaglia & Tsang squeeze for Gamma sampling, which
/// is the same algorithm Cargo `rand_distr` uses internally.
fn sample_beta(alpha: f64, beta: f64, rng: &mut impl Rng) -> f64 {
    let g1 = sample_gamma(alpha, rng);
    let g2 = sample_gamma(beta, rng);
    if g1 + g2 == 0.0 {
        return 0.5;
    }
    g1 / (g1 + g2)
}

fn sample_gamma(shape: f64, rng: &mut impl Rng) -> f64 {
    if shape < 1.0 {
        // Use the boost trick: Gamma(shape) = Gamma(shape+1) * U^(1/shape).
        let g = sample_gamma(shape + 1.0, rng);
        let u: f64 = rng.gen();
        return g * u.powf(1.0 / shape);
    }
    // Marsaglia & Tsang algorithm.
    let d = shape - 1.0 / 3.0;
    let c = 1.0 / (9.0 * d).sqrt();
    loop {
        let mut x;
        let mut v;
        loop {
            x = sample_standard_normal(rng);
            v = 1.0 + c * x;
            if v > 0.0 {
                break;
            }
        }
        v = v * v * v;
        let u: f64 = rng.gen();
        if u < 1.0 - 0.0331 * (x * x * x * x) {
            return d * v;
        }
        if u.ln() < 0.5 * x * x + d * (1.0 - v + v.ln()) {
            return d * v;
        }
    }
}

fn sample_standard_normal(rng: &mut impl Rng) -> f64 {
    // Box–Muller.
    let u1: f64 = rng.gen::<f64>().max(1e-12);
    let u2: f64 = rng.gen();
    let r = (-2.0 * u1.ln()).sqrt();
    let theta = 2.0 * std::f64::consts::PI * u2;
    r * theta.cos()
}

/// Stable string form for CaptchaType. Used as the bandit map key
/// so persistence is deterministic across versions.
fn render_captcha_type(c: &CaptchaType) -> String {
    match c {
        CaptchaType::Custom(name) => {
            // Hash custom names so the bandit doesn't have unbounded
            // cardinality (operator-supplied labels would blow up
            // arm count).
            use sha2::{Digest, Sha256};
            let mut h = Sha256::new();
            h.update(name.as_bytes());
            let d = h.finalize();
            format!("Custom:{:02x}{:02x}{:02x}{:02x}", d[0], d[1], d[2], d[3])
        }
        other => format!("{other:?}"),
    }
}

fn render_solve_method(m: &SolveMethod) -> String {
    format!("{m:?}")
}

#[cfg(test)]
#[path = "bandit/tests.rs"]
mod tests;