captchaforge 0.2.34

//! Adversarial training corpus — failure-driven feedback loop
//! storage layer.
//!
//! Every solver attempt that returns failure (whether
//! `Outcome::Failure`, `Outcome::Error`, `Outcome::Timeout`, or
//! `success=false-after-oracle-downgrade`) carries diagnostic
//! information that's gold for the next training round of the
//! captchaforge solvers:
//!
//! - The exact challenge screenshot the VLM mis-classified.
//! - The DOM snapshot showing what the page looked like.
//! - The fingerprint surface the page was probing.
//! - The vendor / detected captcha kind.
//! - The outcome the chain produced + the oracle's classification.
//!
//! [`TrainingCorpus`] is the storage substrate: append-only +
//! seekable, sample-able by vendor / outcome / kind, exportable
//! to a labelled dataset format that downstream training pipelines
//! (the Python ML side — see roadmap C2) consume directly.
//!
//! ## Pure-Rust storage
//!
//! Each sample is one JSONL line in a per-vendor file. Append is
//! atomic per write call; concurrent writers are serialised behind
//! the per-corpus mutex. No external database — operators can
//! `tar` a corpus directory and ship it as the training input.
//!
//! ## Privacy
//!
//! Caller is responsible for redaction. The corpus stores what it
//! receives; if a screenshot contains PII, it gets stored. Use
//! [`TrainingSample::redact_screenshot_to_thumbnail`] before
//! `append` to drop the full image.

#![allow(dead_code)] // module is opt-in; trainer-side wiring lands separately.

use anyhow::{Context, Result};
use serde::{Deserialize, Serialize};
use std::path::{Path, PathBuf};
use std::sync::Mutex;

/// One captured solver attempt for the training feedback loop.
///
/// Every field is optional so a failed solve that didn't manage to
/// capture (say) the screenshot still produces a useful sample —
/// the model can train on partial data, just with lower weight.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TrainingSample {
    /// Stable solver identifier.
    pub solver: String,
    /// Vendor canonical name (`"cloudflare-turnstile"`, `"hcaptcha"`,
    /// custom rule names from the TOML pack).
    pub vendor: String,
    /// What detection said the captcha was.
    pub detected_kind: String,
    /// Origin URL the chain was solving on. Useful for
    /// per-domain patterns; redact in privacy-strict deployments.
    pub url: String,
    /// What the solver returned.
    pub outcome: String,
    /// Confidence in `[0, 1]`. None when the solver didn't report.
    pub confidence: Option<f32>,
    /// Wall-clock duration of the attempt, in milliseconds.
    pub time_ms: u64,
    /// Base64-encoded PNG screenshot of the captcha widget at the
    /// time of solve. Optional — capture is expensive.
    pub screenshot_b64: Option<String>,
    /// DOM snapshot — outerHTML of the captcha-bearing frame.
    /// Optional. Trims at 64KB by `redact_dom_to_summary`.
    pub dom_snapshot: Option<String>,
    /// Verified outcome from the chain's oracle (Advanced /
    /// Recycled / HardBlock / Unknown).
    pub verified_outcome: Option<String>,
    /// Wall-clock unix epoch when the sample was captured.
    pub captured_at_unix: i64,
}

impl TrainingSample {
    /// Replace the full screenshot with a 64×64 thumbnail-equivalent
    /// hash so the sample retains visual signal without storing
    /// identifiable pixels. Stub today (drops the screenshot
    /// entirely); a future thumbnailer lands here.
    pub fn redact_screenshot_to_thumbnail(&mut self) {
        self.screenshot_b64 = None;
    }

    /// Trim the DOM snapshot to the first 8 KB to drop bulk page
    /// content while keeping the captcha-shaped tags. Caller
    /// chooses when to apply (training pipelines often keep the
    /// full snapshot; production telemetry typically trims).
    pub fn redact_dom_to_summary(&mut self) {
        if let Some(d) = self.dom_snapshot.as_mut() {
            d.truncate(8 * 1024);
        }
    }
}

/// On-disk append-only training corpus.
///
/// One directory containing one JSONL file per vendor. Append is
/// atomic per `append()` call. Listing + sampling re-reads the
/// underlying files; use [`TrainingCorpus::load_iter`] for
/// streaming reads.
pub struct TrainingCorpus {
    root: PathBuf,
    write_lock: Mutex<()>,
}

impl TrainingCorpus {
    /// Open or create a corpus rooted at `path`.
    ///
    /// Creates the directory if missing. No further setup needed —
    /// per-vendor files are created on first append.
    pub fn open(path: impl AsRef<Path>) -> Result<Self> {
        let root = path.as_ref().to_path_buf();
        std::fs::create_dir_all(&root)
            .with_context(|| format!("creating training corpus dir {}", root.display()))?;
        Ok(Self {
            root,
            write_lock: Mutex::new(()),
        })
    }

    /// Borrow the root directory.
    pub fn root(&self) -> &Path {
        &self.root
    }

    /// Append one sample to the per-vendor JSONL file. Atomic via
    /// `OpenOptions::append`.
    pub fn append(&self, sample: &TrainingSample) -> Result<()> {
        let _guard = self
            .write_lock
            .lock()
            .map_err(|e| anyhow::anyhow!("training corpus write lock poisoned: {e}"))?;
        let path = self.path_for_vendor(&sample.vendor);
        let mut line = serde_json::to_string(sample).context("serialising training sample")?;
        line.push('\n');
        use std::io::Write;
        let mut file = std::fs::OpenOptions::new()
            .create(true)
            .append(true)
            .open(&path)
            .with_context(|| format!("opening corpus file {}", path.display()))?;
        file.write_all(line.as_bytes())
            .with_context(|| format!("writing to corpus file {}", path.display()))?;
        Ok(())
    }

    /// Load every sample for `vendor`, returning a Vec.
    /// Cheap for typical corpus sizes (<100k samples per vendor);
    /// use [`Self::load_iter`] for streaming reads on larger corpora.
    pub fn load_vendor(&self, vendor: &str) -> Result<Vec<TrainingSample>> {
        let path = self.path_for_vendor(vendor);
        if !path.exists() {
            return Ok(Vec::new());
        }
        let raw = std::fs::read_to_string(&path)
            .with_context(|| format!("reading {}", path.display()))?;
        let mut out = Vec::new();
        for (i, line) in raw.lines().enumerate() {
            if line.trim().is_empty() {
                continue;
            }
            let sample: TrainingSample = serde_json::from_str(line).with_context(|| {
                format!("parsing sample on line {} of {}", i + 1, path.display())
            })?;
            out.push(sample);
        }
        Ok(out)
    }

    /// Streaming iterator over every sample for `vendor`. Each
    /// `next()` reads + parses one JSONL line. Errors during parse
    /// are surfaced as `Some(Err(_))` so consumers can decide
    /// whether to skip or stop.
    pub fn load_iter(&self, vendor: &str) -> Result<TrainingSampleIter> {
        let path = self.path_for_vendor(vendor);
        let file = if path.exists() {
            Some(
                std::fs::File::open(&path)
                    .with_context(|| format!("opening {}", path.display()))?,
            )
        } else {
            None
        };
        Ok(TrainingSampleIter {
            reader: file.map(std::io::BufReader::new),
        })
    }

    /// List every vendor with at least one sample on disk.
    pub fn vendors(&self) -> Result<Vec<String>> {
        let mut out = Vec::new();
        for entry in std::fs::read_dir(&self.root)
            .with_context(|| format!("reading {}", self.root.display()))?
        {
            let entry = entry?;
            let name = entry.file_name();
            let s = name.to_string_lossy();
            if let Some(rest) = s.strip_suffix(".jsonl") {
                out.push(rest.to_string());
            }
        }
        out.sort();
        Ok(out)
    }

    /// Aggregate counts per vendor — useful for "is this corpus
    /// big enough to retrain on?" pre-flight.
    pub fn vendor_counts(&self) -> Result<Vec<(String, usize)>> {
        let mut out = Vec::new();
        for vendor in self.vendors()? {
            let count = self.load_iter(&vendor)?.filter_map(|r| r.ok()).count();
            out.push((vendor, count));
        }
        out.sort();
        Ok(out)
    }

    fn path_for_vendor(&self, vendor: &str) -> PathBuf {
        // Sanitise: replace path separators / null bytes / control
        // chars so the filename is filesystem-safe regardless of
        // what the caller passes.
        let safe: String = vendor
            .chars()
            .map(|c| {
                if c.is_alphanumeric() || matches!(c, '-' | '_') {
                    c
                } else {
                    '_'
                }
            })
            .collect();
        self.root.join(format!("{safe}.jsonl"))
    }
}

/// Streaming sample iterator. Parses one JSONL line per `next()`.
pub struct TrainingSampleIter {
    reader: Option<std::io::BufReader<std::fs::File>>,
}

impl Iterator for TrainingSampleIter {
    type Item = Result<TrainingSample>;

    fn next(&mut self) -> Option<Self::Item> {
        use std::io::BufRead;
        let reader = self.reader.as_mut()?;
        let mut buf = String::new();
        loop {
            buf.clear();
            match reader.read_line(&mut buf) {
                Ok(0) => return None,
                Ok(_) => {
                    if buf.trim().is_empty() {
                        continue;
                    }
                    return Some(
                        serde_json::from_str(buf.trim_end_matches('\n'))
                            .context("parsing training sample"),
                    );
                }
                Err(e) => return Some(Err(anyhow::Error::from(e))),
            }
        }
    }
}

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

    fn sample(vendor: &str, outcome: &str) -> TrainingSample {
        TrainingSample {
            solver: "TestSolver".into(),
            vendor: vendor.into(),
            detected_kind: "Turnstile".into(),
            url: "https://example.com".into(),
            outcome: outcome.into(),
            confidence: Some(0.8),
            time_ms: 1234,
            screenshot_b64: Some("AAAA".into()),
            dom_snapshot: Some("<html></html>".into()),
            verified_outcome: Some("Advanced".into()),
            captured_at_unix: 1_700_000_000,
        }
    }

    #[test]
    fn open_creates_root_directory_when_missing() {
        let tmp = tempdir().unwrap();
        let path = tmp.path().join("nested/subdir");
        let corpus = TrainingCorpus::open(&path).expect("must create dir");
        assert!(path.exists());
        assert_eq!(corpus.root(), path);
    }

    #[test]
    fn append_then_load_round_trips_sample() {
        let tmp = tempdir().unwrap();
        let corpus = TrainingCorpus::open(tmp.path()).unwrap();
        let s = sample("cf-turnstile", "failure");
        corpus.append(&s).unwrap();
        let back = corpus.load_vendor("cf-turnstile").unwrap();
        assert_eq!(back.len(), 1);
        assert_eq!(back[0].vendor, s.vendor);
        assert_eq!(back[0].outcome, s.outcome);
        assert_eq!(back[0].screenshot_b64, s.screenshot_b64);
    }

    #[test]
    fn append_appends_not_overwrites() {
        let tmp = tempdir().unwrap();
        let corpus = TrainingCorpus::open(tmp.path()).unwrap();
        for i in 0..5 {
            let mut s = sample("cf-turnstile", "failure");
            s.time_ms = 1000 + i;
            corpus.append(&s).unwrap();
        }
        let back = corpus.load_vendor("cf-turnstile").unwrap();
        assert_eq!(back.len(), 5);
    }

    #[test]
    fn load_vendor_returns_empty_for_unknown_vendor() {
        let tmp = tempdir().unwrap();
        let corpus = TrainingCorpus::open(tmp.path()).unwrap();
        let back = corpus.load_vendor("never-touched").unwrap();
        assert!(back.is_empty());
    }

    #[test]
    fn load_iter_streams_one_sample_at_a_time() {
        let tmp = tempdir().unwrap();
        let corpus = TrainingCorpus::open(tmp.path()).unwrap();
        for _ in 0..3 {
            corpus.append(&sample("cf-turnstile", "failure")).unwrap();
        }
        let mut iter = corpus.load_iter("cf-turnstile").unwrap();
        let s1 = iter.next().expect("first sample present").unwrap();
        let s2 = iter.next().expect("second sample present").unwrap();
        let s3 = iter.next().expect("third sample present").unwrap();
        assert!(iter.next().is_none(), "iter must be exhausted after 3");
        assert_eq!(s1.vendor, "cf-turnstile");
        assert_eq!(s2.vendor, "cf-turnstile");
        assert_eq!(s3.vendor, "cf-turnstile");
    }

    #[test]
    fn vendors_lists_every_vendor_with_at_least_one_sample() {
        let tmp = tempdir().unwrap();
        let corpus = TrainingCorpus::open(tmp.path()).unwrap();
        corpus.append(&sample("cf-turnstile", "failure")).unwrap();
        corpus.append(&sample("hcaptcha", "failure")).unwrap();
        corpus.append(&sample("hcaptcha", "success")).unwrap();
        let vendors = corpus.vendors().unwrap();
        assert_eq!(vendors, vec!["cf-turnstile", "hcaptcha"]);
    }

    #[test]
    fn vendor_counts_aggregates_per_vendor() {
        let tmp = tempdir().unwrap();
        let corpus = TrainingCorpus::open(tmp.path()).unwrap();
        for _ in 0..3 {
            corpus.append(&sample("cf-turnstile", "failure")).unwrap();
        }
        for _ in 0..7 {
            corpus.append(&sample("hcaptcha", "failure")).unwrap();
        }
        let counts = corpus.vendor_counts().unwrap();
        assert!(counts.contains(&("cf-turnstile".into(), 3)));
        assert!(counts.contains(&("hcaptcha".into(), 7)));
    }

    #[test]
    fn redact_screenshot_drops_base64_payload() {
        let mut s = sample("cf-turnstile", "failure");
        assert!(s.screenshot_b64.is_some());
        s.redact_screenshot_to_thumbnail();
        assert!(s.screenshot_b64.is_none());
    }

    #[test]
    fn redact_dom_truncates_at_eight_kb() {
        let mut s = sample("cf-turnstile", "failure");
        s.dom_snapshot = Some("a".repeat(20 * 1024));
        s.redact_dom_to_summary();
        assert!(s.dom_snapshot.as_ref().unwrap().len() <= 8 * 1024);
    }

    #[test]
    fn path_for_vendor_sanitises_unsafe_chars() {
        let tmp = tempdir().unwrap();
        let corpus = TrainingCorpus::open(tmp.path()).unwrap();
        let path = corpus.path_for_vendor("../../etc/passwd");
        // No `/` or `..` should survive.
        let s = path.to_string_lossy();
        let filename = path.file_name().unwrap().to_string_lossy();
        assert!(!filename.contains('/'));
        assert!(!filename.contains(".."));
        assert!(s.starts_with(tmp.path().to_string_lossy().as_ref()));
    }

    #[test]
    fn append_with_unsafe_vendor_name_writes_to_sanitised_path() {
        let tmp = tempdir().unwrap();
        let corpus = TrainingCorpus::open(tmp.path()).unwrap();
        let mut s = sample("foo/bar:baz", "failure");
        s.vendor = "foo/bar:baz".into();
        corpus.append(&s).unwrap();
        // Vendor-name sanitisation: non-alphanumeric/-_ chars
        // become `_`.
        let back = corpus.load_vendor("foo_bar_baz").unwrap();
        assert_eq!(back.len(), 1);
    }
}