searchi/tools/desktop-capture/backend/src/main.rs

//! Independent screenshot test backend, not a production HTTP service.
use anyhow::{bail, ensure, Context, Result};
use clap::Parser;
use searchi_core::{SearchIndex, SearchOptions};
use searchi_service::{FfmpegTools, HybridSearch, MediaLimits};
use serde::Deserialize;
use serde_json::{json, Value};
use sha2::{Digest, Sha256};
use std::{
    collections::BTreeMap,
    fs,
    io::Read,
    path::{Component, Path, PathBuf},
};
use tiny_http::{Header, Method, Request, Response, Server, StatusCode};

#[derive(Parser)]
struct Args {
    /// Dedicated empty/generated corpus directory. Never a user document root.
    #[arg(long)]
    root: PathBuf,
    #[arg(long)]
    models: PathBuf,
    #[arg(long)]
    worker: PathBuf,
    #[arg(long)]
    db: PathBuf,
    #[arg(long)]
    dist: PathBuf,
    #[arg(long)]
    ffmpeg: PathBuf,
    #[arg(long)]
    ffprobe: PathBuf,
}
#[derive(Default, Deserialize)]
#[serde(deny_unknown_fields)]
struct Filters {
    code_only: Option<bool>,
    folder: Option<PathBuf>,
    extension: Option<String>,
    min_size: Option<u64>,
    max_size: Option<u64>,
    min_duration_ms: Option<u32>,
    max_duration_ms: Option<u32>,
    modified_after_unix_ms: Option<i64>,
    modified_before_unix_ms: Option<i64>,
}
impl Filters {
    fn options(self, root: &Path) -> Result<SearchOptions> {
        if let Some(folder) = &self.folder {
            ensure!(
                folder.canonicalize()?.starts_with(root),
                "Capture folder filter must stay within the fixture root"
            );
        }
        Ok(SearchOptions {
            code_only: self.code_only.unwrap_or(false),
            folder: self.folder,
            extension: self.extension,
            min_size: self.min_size,
            max_size: self.max_size,
            min_duration_ms: self.min_duration_ms,
            max_duration_ms: self.max_duration_ms,
            modified_after_unix_ms: self.modified_after_unix_ms,
            modified_before_unix_ms: self.modified_before_unix_ms,
        })
    }
}
struct App {
    root: PathBuf,
    db: PathBuf,
    dist: PathBuf,
    models: PathBuf,
    token: String,
    lexical: SearchIndex,
    hybrid: HybridSearch,
    media: BTreeMap<String, PathBuf>,
}

fn main() -> Result<()> {
    let args = Args::parse();
    std::env::set_var("SEARCHI_FORCE_CPU", "1");
    let root = prepare_fixture(&args.root)?;
    if let Some(parent) = args.db.parent() {
        fs::create_dir_all(parent)?;
    }
    let mut lexical = SearchIndex::open(&args.db)?;
    validate_database(&args.db, &root)?;
    lexical.index_folder(&root)?;
    lexical.prewarm()?;
    let tools = FfmpegTools::new(args.ffmpeg, args.ffprobe)?;
    let mut hybrid = HybridSearch::open(&args.db, &args.models, &args.worker, 1)?;
    let status = HybridSearch::embedding_status(&args.db)?;
    if status.stale || status.generation.is_none() {
        let report = hybrid.index_embeddings_with_media(&tools, &MediaLimits::default())?;
        ensure!(
            report.media_failures == 0,
            "Capture media indexing failed: {:?}",
            report.failures
        );
        println!(
            "Indexed real fixture: {} passages, {} vectors, {} images",
            report.passages, report.chunks, report.media_files
        );
    }
    hybrid.warmup()?;
    let dist = args.dist.canonicalize()?;
    ensure!(
        dist.join("index.html").is_file(),
        "Build the frontend before capture"
    );
    let media = BTreeMap::from([
        ("lake".into(), root.join("Mountain Lake.jpg")),
        ("pier".into(), root.join("Night Pier.jpg")),
    ]);
    let app = App {
        root,
        db: args.db,
        dist,
        models: args.models,
        token: uuid::Uuid::new_v4().to_string(),
        lexical,
        hybrid,
        media,
    };
    let server = Server::http(("127.0.0.1", 1431)).map_err(|e| anyhow::anyhow!(e.to_string()))?;
    println!("READY http://127.0.0.1:1431/?capture=1 (real CPU engine, local fixture only)");
    serve(server, app)
}

fn valid_host(host: &str) -> bool {
    matches!(host, "127.0.0.1:1431" | "localhost:1431")
}
fn permitted(headers: &[(String, String)], post: bool, token: &str) -> bool {
    let get = |key: &str| {
        headers
            .iter()
            .find(|(name, _)| name.eq_ignore_ascii_case(key))
            .map(|(_, value)| value.as_str())
    };
    let Some(host) = get("host") else {
        return false;
    };
    if !valid_host(host) {
        return false;
    }
    if let Some(origin) = get("origin") {
        if origin != format!("http://{host}") {
            return false;
        }
    }
    !post
        || (get("origin") == Some(format!("http://{host}").as_str())
            && get("x-searchi-capture-token") == Some(token))
}
fn serve(server: Server, mut app: App) -> Result<()> {
    for mut request in server.incoming_requests() {
        let headers = request
            .headers()
            .iter()
            .map(|h| (h.field.as_str().to_string(), h.value.as_str().to_owned()))
            .collect::<Vec<_>>();
        let post = request.method() == &Method::Post;
        if !permitted(&headers, post, &app.token) {
            reply(
                request,
                403,
                "application/json",
                br#"{"message":"Capture request denied"}"#.to_vec(),
            )?;
            continue;
        }
        let path = request.url().split('?').next().unwrap_or("").to_string();
        let outcome: Result<(u16, String, Vec<u8>)> = (|| {
            if request.method() == &Method::Get && path == "/capture-info" {
                return Ok((
                    200,
                    "application/json".into(),
                    serde_json::to_vec(
                        &json!({"mode":"readme-capture","token":app.token,"ready":true,"worker_info":app.hybrid.worker_info(),"rust_build_profile":if cfg!(debug_assertions) {"debug"} else {"release"}}),
                    )?,
                ));
            }
            if post {
                let command = path
                    .strip_prefix("/capture-api/")
                    .context("Unknown capture route")?;
                ensure!(
                    request.body_length().is_none_or(|n| n <= 65536),
                    "Capture request too large"
                );
                let mut body = Vec::new();
                request.as_reader().take(65537).read_to_end(&mut body)?;
                ensure!(body.len() <= 65536, "Capture request too large");
                let args: Value = serde_json::from_slice(&body)?;
                return Ok((
                    200,
                    "application/json".into(),
                    serde_json::to_vec(&app.command(command, args)?)?,
                ));
            }
            ensure!(request.method() == &Method::Get, "Unsupported method");
            if let Some(id) = path.strip_prefix("/capture-media/") {
                let file = app.media.get(id).context("Unknown fixture media ID")?;
                ensure!(
                    app.lexical.contains_path(&file.to_string_lossy())?,
                    "Media no longer indexed"
                );
                return Ok((200, "image/jpeg".into(), bounded_read(file)?));
            }
            let file = static_path(&app.dist, &path)?;
            let mime = mime_guess::from_path(&file)
                .first_or_octet_stream()
                .to_string();
            Ok((200, mime, bounded_read(&file)?))
        })();
        match outcome {
            Ok((status, mime, body)) => reply(request, status, &mime, body)?,
            Err(error) => reply(
                request,
                400,
                "application/json",
                serde_json::to_vec(&json!({"message":format!("{error:#}")}))?,
            )?,
        }
    }
    Ok(())
}
fn reply(request: Request, status: u16, mime: &str, body: Vec<u8>) -> Result<()> {
    let response = Response::from_data(body)
        .with_status_code(StatusCode(status))
        .with_header(
            Header::from_bytes("Content-Type", mime)
                .map_err(|_| anyhow::anyhow!("Invalid content type"))?,
        )
        .with_header(Header::from_bytes("Cache-Control", "no-store").unwrap())
        .with_header(Header::from_bytes("X-Content-Type-Options", "nosniff").unwrap());
    request.respond(response)?;
    Ok(())
}
fn bounded_read(path: &Path) -> Result<Vec<u8>> {
    ensure!(path.is_file(), "Capture resource must be a regular file");
    let mut bytes = Vec::new();
    fs::File::open(path)?
        .take(20 * 1024 * 1024 + 1)
        .read_to_end(&mut bytes)?;
    ensure!(
        bytes.len() <= 20 * 1024 * 1024,
        "Capture resource exceeds20MiB"
    );
    Ok(bytes)
}
fn static_path(dist: &Path, url: &str) -> Result<PathBuf> {
    ensure!(
        !url.contains('%') && !url.contains('\\'),
        "Encoded or Windows paths are not capture routes"
    );
    let relative = Path::new(url.trim_start_matches('/'));
    ensure!(
        relative
            .components()
            .all(|c| matches!(c, Component::Normal(_))),
        "Unsafe static resource path"
    );
    let target = if relative.as_os_str().is_empty() {
        dist.join("index.html")
    } else {
        dist.join(relative)
    };
    let target = target.canonicalize()?;
    ensure!(
        target.starts_with(dist),
        "Static resource escapes build directory"
    );
    Ok(target)
}
impl App {
    fn command(&mut self, name: &str, args: Value) -> Result<Value> {
        match name {
            "desktop_status" => Ok(
                json!({"settings":{"roots":[self.root],"cpu_threads":1},"indexing":false,"shortcut_status":"Ctrl K · reproducible CPU capture"}),
            ),
            "stats" => Ok(serde_json::to_value(self.lexical.stats()?)?),
            "semantic_status" => {
                let status = HybridSearch::embedding_status(&self.db)?;
                Ok(
                    json!({"model":"Google EmbeddingGemma 2","ready":true,"model_directory":self.models,"worker_ready":true,"downloading":false,"embedding":false,"device_policy":"cpu","error":null,"downloaded_bytes":864676480u64,"total_bytes":864676480u64,"generation":status.generation,"chunks":status.chunks,"stale":status.stale,"pending":status.pending,"warming":false,"worker_info":self.hybrid.worker_info(),"media_tools_ready":true,"media_setup_error":null,"media_failures":status.failures}),
                )
            }
            "search_measured" | "hybrid_search_measured" => {
                let query = args["query"].as_str().context("Query required")?;
                ensure!(query.len() <= 4096, "Query too large");
                let filters: Filters = if args["filters"].is_null() {
                    Filters::default()
                } else {
                    serde_json::from_value(args["filters"].clone())?
                };
                let options = filters.options(&self.root)?;
                let debug = args["debugPerformance"].as_bool().unwrap_or(false);
                if name == "search_measured" {
                    Ok(serde_json::to_value(
                        self.lexical
                            .search_with_options_measured(query, 50, &options, debug)?,
                    )?)
                } else {
                    Ok(serde_json::to_value(
                        self.hybrid.search_measured(query, 50, &options, debug)?,
                    )?)
                }
            }
            "prepare_media_preview" => {
                let path = args["path"].as_str().context("Media path required")?;
                let (id, _) = self
                    .media
                    .iter()
                    .find(|(_, file)| file.to_string_lossy() == path)
                    .context("Only known indexed fixture photos may be previewed")?;
                ensure!(self.lexical.contains_path(path)?, "Photo no longer indexed");
                Ok(
                    json!({"path":format!("/capture-media/{id}"),"kind":"image","mime":"image/jpeg"}),
                )
            }
            _ => bail!("This capture harness supports read-only fixture commands"),
        }
    }
}

fn validate_database(db: &Path, root: &Path) -> Result<()> {
    let connection = rusqlite::Connection::open(db)?;
    let mut query = connection.prepare("SELECT path FROM files")?;
    for path in query.query_map([], |row| row.get::<_, String>(0))? {
        let path = PathBuf::from(path?);
        ensure!(
            path.parent() == Some(root),
            "Capture database contains files outside dedicated fixture root"
        );
        ensure!(
            fixture_names().contains(&path.file_name().unwrap().to_string_lossy().as_ref()),
            "Capture database contains unrelated files"
        );
    }
    Ok(())
}
fn fixture_names() -> Vec<&'static str> {
    vec![
        "Supplier Agreement.pdf",
        "Vendor Onboarding.md",
        "Q3 Procurement.md",
        "Travel Notes.md",
        "Garden Notes.md",
        "Mountain Lake.jpg",
        "Night Pier.jpg",
    ]
}
fn prepare_fixture(root: &Path) -> Result<PathBuf> {
    fs::create_dir_all(root)?;
    let root = root.canonicalize()?;
    for entry in fs::read_dir(&root)? {
        let entry = entry?;
        ensure!(
            entry.file_type()?.is_file()
                && fixture_names().contains(&entry.file_name().to_string_lossy().as_ref()),
            "Use an empty dedicated capture root, not user files"
        );
    }
    let texts=[("Vendor Onboarding.md","# Vendor onboarding\nCollect supplier tax details, banking information and insurance documents. The purchasing team completes compliance review before signing a new partnership."),("Q3 Procurement.md","# Q3 procurement review\nWe will review supplier delivery times, packaging costs and the annual purchasing budget. Present the comparison of existing vendors at Thursday's planning meeting."),("Travel Notes.md","# Travel plans\nTake the train to Berlin for the museum weekend. Reserve a room near the station and pack a raincoat."),("Garden Notes.md","# Garden notes\nTomatoes need sunlight and regular watering. Plant herbs in pots on the balcony and add compost in spring.")];
    for (name, text) in texts {
        write_known(&root.join(name), text.as_bytes())?;
    }
    supplier_pdf(&root.join("Supplier Agreement.pdf"))?;
    let repo = Path::new(env!("CARGO_MANIFEST_DIR"))
        .join("../../..")
        .canonicalize()?;
    let manifest: Value =
        serde_json::from_slice(&fs::read(repo.join("fixtures/media/fixture.json"))?)?;
    for (id, name) in [
        ("image03", "Mountain Lake.jpg"),
        ("image05", "Night Pier.jpg"),
    ] {
        let asset = manifest["assets"]
            .as_array()
            .context("Missing licensed assets")?
            .iter()
            .find(|a| a["id"] == id)
            .context("Missing fixed licensed photo")?;
        let source = repo
            .join("fixtures/media")
            .join(asset["path"].as_str().unwrap());
        let bytes = bounded_read(&source)?;
        ensure!(
            asset["sha256"].as_str() == Some(format!("{:x}", Sha256::digest(&bytes)).as_str()),
            "Licensed photo checksum changed"
        );
        write_known(&root.join(name), &bytes)?;
    }
    Ok(root)
}
fn write_known(path: &Path, bytes: &[u8]) -> Result<()> {
    if path.exists() {
        ensure!(
            fs::read(path)? == bytes,
            "Capture fixture file differs; use a new dedicated root"
        );
    } else {
        fs::write(path, bytes)?;
    }
    Ok(())
}
fn supplier_pdf(path: &Path) -> Result<()> {
    use lopdf::{
        content::{Content, Operation},
        dictionary, Document, Object, Stream,
    };
    let mut document = Document::with_version("1.5");
    let pages_id = document.new_object_id();
    let font=document.add_object(dictionary!{"Type"=>"Font","Subtype"=>"Type1","BaseFont"=>"Helvetica","Encoding"=>"WinAnsiEncoding"});
    let resources = document.add_object(dictionary! {"Font"=>dictionary!{"F1"=>font}});
    let mut pages = Vec::new();
    for lines in [
        vec![
            "Supplier Agreement",
            "The purchasing team and supplier enter a partnership.",
            "Orders and delivery schedules are confirmed in writing.",
        ],
        vec![
            "Termination and cancellation",
            "The supplier partnership may be ended with sixty days",
            "written notice. Send the cancellation letter to the vendor.",
            "Return unused materials and settle outstanding invoices.",
        ],
    ]
    .iter()
    {
        let mut operations = vec![
            Operation::new("BT", vec![]),
            Operation::new("Tf", vec![Object::Name(b"F1".to_vec()), 12.into()]),
            Operation::new("Td", vec![50.into(), 740.into()]),
        ];
        for line in lines {
            operations.push(Operation::new(
                "Tj",
                vec![Object::string_literal(format!("{line} "))],
            ));
            operations.push(Operation::new("Td", vec![0.into(), (-20).into()]));
        }
        operations.push(Operation::new("ET", vec![]));
        let stream = document.add_object(Stream::new(
            dictionary! {},
            Content { operations }.encode()?,
        ));
        let page =
            document.add_object(dictionary! {"Type"=>"Page","Parent"=>pages_id,"Contents"=>stream});
        pages.push(Object::Reference(page));
    }
    document.objects.insert(pages_id,Object::Dictionary(dictionary!{"Type"=>"Pages","Kids"=>pages,"Count"=>2,"Resources"=>resources,"MediaBox"=>vec![0.into(),0.into(),595.into(),842.into()]}));
    let catalog = document.add_object(dictionary! {"Type"=>"Catalog","Pages"=>pages_id});
    document.trailer.set("Root", catalog);
    document.compress();
    let mut bytes = Vec::new();
    document.save_to(&mut bytes)?;
    write_known(path, &bytes)?;
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    #[test]
    fn fixture_reuses_unchanged_files_and_real_pdf_page_two() -> Result<()> {
        let temporary = tempfile::tempdir()?;
        let root = prepare_fixture(&temporary.path().join("corpus"))?;
        let db = temporary.path().join("index.db");
        let mut index = SearchIndex::open(&db)?;
        validate_database(&db, &root)?;
        let first = index.index_folder(&root)?;
        assert_eq!(first.discovered, 7);
        let hits = index.search("sixty days", 10)?;
        assert_eq!(hits.len(), 1);
        assert_eq!(hits[0].filename, "Supplier Agreement.pdf");
        assert_eq!(hits[0].source.as_ref().unwrap().number, 2);
        prepare_fixture(&root)?;
        let reused = index.index_folder(&root)?;
        assert_eq!(reused.unchanged, 7);
        assert_eq!(reused.indexed, 0);
        let foreign = temporary.path().join("foreign");
        fs::create_dir(&foreign)?;
        fs::write(foreign.join("secret.txt"), "private")?;
        index.index_folder(&foreign)?;
        assert!(validate_database(&db, &root).is_err());
        Ok(())
    }
    #[test]
    fn origin_host_and_token_all_required_for_post() {
        let headers = vec![
            ("Host".into(), "127.0.0.1:1431".into()),
            ("Origin".into(), "http://127.0.0.1:1431".into()),
            ("X-Searchi-Capture-Token".into(), "secret".into()),
        ];
        assert!(permitted(&headers, true, "secret"));
        assert!(!permitted(&headers, true, "wrong"));
        assert!(!permitted(&headers[..1], true, "secret"));
        let mut attacker = headers.clone();
        attacker[1].1 = "https://attacker.example".into();
        assert!(!permitted(&attacker, true, "secret"));
        attacker[0].1 = "attacker.example:1431".into();
        assert!(!permitted(&attacker, false, "secret"));
    }
    #[test]
    fn static_routes_reject_traversal_and_unrelated_root_files() -> Result<()> {
        let root = tempfile::tempdir()?;
        fs::write(root.path().join("index.html"), "real frontend")?;
        let canonical = root.path().canonicalize()?;
        assert!(static_path(&canonical, "/").is_ok());
        for path in ["/../secret", "/%2e%2e/secret", "/C:\\secret"] {
            assert!(static_path(&canonical, path).is_err());
        }
        fs::write(root.path().join("user-secret.txt"), "private")?;
        assert!(prepare_fixture(root.path()).is_err());
        Ok(())
    }
}