AI Flow: Reverse Engineering

//! `serve`: the flow window. A small local HTTP server (127.0.0.1 only, std library, one thread per
//! request) that serves the flow page, `/state` (the steps, their states and the recorded views) and
//! the files under reverse-flow/ (views, photos, evidence). The page polls /state and shows the newest
//! view as it is recorded, so the window follows the run; its history list replays the process.
use crate::flow::{self, Flow};
use serde_json::{json, Value};
use std::io::{BufRead, BufReader, Write};
use std::net::{TcpListener, TcpStream};
use std::path::{Component, Path, PathBuf};

const PAGE: &str = include_str!("../assets/flow.html");
/// AI Flow's timeline page, copied unchanged from adom/adom-aiflow tools/timeline (MIT, Adom Industries).
const TIMELINE: &str = include_str!("../assets/timeline.html");
const TIMELINE_ICON: &str = include_str!("../assets/timeline-favicon.svg");

pub fn state(flow_dir: &Path, fl: &Flow) -> Value {
    let cur = flow::current(fl, flow_dir).map(|s| s.name.clone());
    let steps: Vec<Value> = fl.steps.iter().map(|s| {
        let done = s.done.iter().all(|d| flow::exists(flow_dir, d));
        let stale: Vec<String> = s.done.iter().flat_map(|d| flow::stale(flow_dir, d)).collect();
        let st = if done { "done" } else if Some(&s.name) == cur.as_ref() { "now" } else if !stale.is_empty() { "stale" } else { "todo" };
        json!({"name": s.name, "state": st, "stale": stale, "do": s.todo})
    }).collect();
    let st = crate::items::read_json(&flow_dir.join("state.json")).unwrap_or(json!({}));
    let views = crate::items::read_json(&flow_dir.join("views.json")).unwrap_or(json!({"entries": []}));
    let evidence: Vec<String> = std::fs::read_dir(flow_dir.join("evidence")).map(|rd| {
        let mut v: Vec<String> = rd.flatten().map(|e| format!("evidence/{}", e.file_name().to_string_lossy())).filter(|n| is_image(n)).collect();
        v.sort(); v
    }).unwrap_or_default();
    let findings = crate::findings::all(flow_dir);
    json!({"flow": fl.name, "project": flow_dir.parent().map(|p| p.display().to_string()), "current": cur, "steps": steps,
           "titles": fl.steps.iter().map(|s| (s.name.clone(), json!(crate::timeline::title(&s.name)))).collect::<serde_json::Map<_, _>>(),
           "photos": st["photos"], "kicad": st["kicad"], "views": views["entries"], "evidence": evidence, "findings": findings})
}

fn is_image(n: &str) -> bool { let l = n.to_lowercase(); [".jpg", ".jpeg", ".png", ".webp", ".svg"].iter().any(|x| l.ends_with(x)) }

fn ctype(p: &Path) -> &'static str {
    match p.extension().and_then(|x| x.to_str()).map(|x| x.to_lowercase()).as_deref() {
        Some("html") => "text/html; charset=utf-8", Some("json") | Some("jsonl") => "application/json", Some("svg") => "image/svg+xml",
        Some("png") => "image/png", Some("jpg") | Some("jpeg") => "image/jpeg", Some("webp") => "image/webp", Some("md") => "text/plain; charset=utf-8",
        _ => "application/octet-stream",
    }
}

fn send(s: &mut TcpStream, code: &str, ct: &str, body: &[u8]) {
    let _ = write!(s, "HTTP/1.1 {code}\r\nContent-Type: {ct}\r\nContent-Length: {}\r\nCache-Control: no-store\r\nConnection: close\r\n\r\n", body.len());
    let _ = s.write_all(body);
}

fn decode(p: &str) -> String {
    let b = p.as_bytes();
    let mut out = vec![];
    let mut i = 0;
    while i < b.len() {
        if b[i] == b'%' && i + 2 < b.len() { if let Ok(v) = u8::from_str_radix(&p[i + 1..i + 3], 16) { out.push(v); i += 3; continue; } }
        out.push(b[i]); i += 1;
    }
    String::from_utf8_lossy(&out).to_string()
}

/// A request path under reverse-flow/, refusing anything that climbs out of it.
pub fn safe(flow_dir: &Path, rel: &str) -> Option<PathBuf> {
    let r = Path::new(rel);
    if r.components().any(|c| !matches!(c, Component::Normal(_))) { return None; }
    let p = flow_dir.join(r);
    p.is_file().then_some(p)
}

fn handle(mut s: TcpStream, flow_dir: &Path, fl: &Flow) {
    let mut line = String::new();
    if BufReader::new(&s).read_line(&mut line).is_err() { return; }
    let path = line.split_whitespace().nth(1).unwrap_or("/").split('?').next().unwrap_or("/").to_string();
    match path.as_str() {
        "/" | "/index.html" => send(&mut s, "200 OK", "text/html; charset=utf-8", PAGE.as_bytes()),
        "/state" => send(&mut s, "200 OK", "application/json", state(flow_dir, fl).to_string().as_bytes()),
        "/timeline/" | "/timeline/index.html" => send(&mut s, "200 OK", "text/html; charset=utf-8", TIMELINE.as_bytes()),
        "/timeline/favicon.svg" => send(&mut s, "200 OK", "image/svg+xml", TIMELINE_ICON.as_bytes()),
        "/timeline/timeline.json" => { let f = crate::findings::all(flow_dir); send(&mut s, "200 OK", "application/json", crate::timeline::timeline(flow_dir, fl, &f).to_string().as_bytes()) }
        "/timeline/presentation.json" => { let f = crate::findings::all(flow_dir); send(&mut s, "200 OK", "application/json", crate::timeline::presentation(flow_dir, fl, &f).to_string().as_bytes()) }
        "/health" => send(&mut s, "200 OK", "application/json", br#"{"ok":true}"#),
        p if p.starts_with("/files/") || p.starts_with("/timeline/files/") => match safe(flow_dir, &decode(p.split_once("/files/").map(|x| x.1).unwrap_or(""))) {
            Some(f) => match std::fs::read(&f) { Ok(b) => send(&mut s, "200 OK", ctype(&f), &b), Err(_) => send(&mut s, "500 Error", "text/plain", b"read failed") },
            None => send(&mut s, "404 Not Found", "text/plain", b"not found"),
        },
        _ => send(&mut s, "404 Not Found", "text/plain", b"not found"),
    }
}

pub fn serve(flow_dir: PathBuf, fl: Flow, port: u16) -> std::io::Result<()> {
    let l = TcpListener::bind(("127.0.0.1", port))?;
    let fl = std::sync::Arc::new(fl);
    let dir = std::sync::Arc::new(flow_dir);
    for c in l.incoming().flatten() {
        let (fl, dir) = (fl.clone(), dir.clone());
        std::thread::spawn(move || handle(c, &dir, &fl));
    }
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    #[test]
    fn paths_stay_inside() {
        let d = std::env::temp_dir().join(format!("srv-test-{}", std::process::id()));
        std::fs::create_dir_all(d.join("views")).unwrap();
        std::fs::write(d.join("views/a.html"), "x").unwrap();
        assert!(safe(&d, "views/a.html").is_some());
        assert!(safe(&d, "../etc/passwd").is_none());
        assert!(safe(&d, "/etc/passwd").is_none());
        assert_eq!(decode("views/a%20b.html"), "views/a b.html");
        std::fs::remove_dir_all(&d).ok();
    }
}