lazyBoy/crates/api/src/attachments.rs

//! Chat attachments: bytes go to the model once, then we drop them.
//!
//! The chat row only stores name/mime/size. If the bot's computer needs a copy
//! (open a PDF, upload a spreadsheet), we write it under `inbox/` on the bind-
//! mounted home and delete anything older than [`INBOX_TTL`].

use std::path::{Path, PathBuf};
use std::time::Duration;

use base64::Engine;
use lazyboy_contracts::SessionAttachment;
use lazyboy_control::resolve_bot_workspace_path;
use rig_core::completion::message::{ImageDetail, ImageMediaType, UserContent};
use serde::Serialize;
use serde_json::{Value, json};

use crate::computer;
use crate::db::{Actor, parse_mode};
use crate::state::AppState;

pub const INBOX_TTL: Duration = Duration::from_secs(2 * 60 * 60);
pub const MAX_COUNT: usize = 4;
pub const MAX_BYTES: usize = 10 * 1024 * 1024;
pub const MAX_TOTAL_BYTES: usize = 20 * 1024 * 1024;
const TEXT_INLINE_CHARS: usize = 80_000;
const NAME_MAX: usize = 80;

pub type IncomingAttachment = SessionAttachment;

#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct StoredAttachment {
    pub kind: &'static str,
    pub name: String,
    pub mime_type: String,
    pub size: usize,
    pub path: String,
}

#[derive(Debug, Clone)]
pub struct DecodedAttachment {
    pub name: String,
    pub stored_name: String,
    pub mime_type: String,
    pub bytes: Vec<u8>,
}

pub fn decode_incoming(items: &[IncomingAttachment]) -> Result<Vec<DecodedAttachment>, String> {
    if items.len() > MAX_COUNT {
        return Err(format!("最多 {MAX_COUNT} 個附件"));
    }
    let mut out = Vec::with_capacity(items.len());
    let mut total = 0usize;
    for item in items {
        let name = safe_name(&item.name)?;
        let mime = normalize_mime(&item.mime_type, &name);
        if !allowed_mime(&mime) {
            return Err(format!("不支援的檔案類型：{name}"));
        }
        let bytes = decode_base64(&item.content)?;
        if bytes.is_empty() {
            return Err(format!("空檔案：{name}"));
        }
        if bytes.len() > MAX_BYTES {
            return Err(format!("{name} 超過 10 MB"));
        }
        total = total.saturating_add(bytes.len());
        if total > MAX_TOTAL_BYTES {
            return Err("附件合計超過 20 MB".into());
        }
        let (stem, ext) = split_ext(&name);
        let stored_name = format!("{stem}-{}{ext}", uuid::Uuid::new_v4());
        out.push(DecodedAttachment {
            name,
            stored_name,
            mime_type: mime,
            bytes,
        });
    }
    Ok(out)
}

pub fn stored_blocks(files: &[DecodedAttachment]) -> Vec<Value> {
    files
        .iter()
        .map(|file| {
            json!({
                "kind": if is_image(&file.mime_type) { "image" } else { "file" },
                "name": file.name,
                "mimeType": file.mime_type,
                "size": file.bytes.len(),
                "path": format!("inbox/{}", file.stored_name),
            })
        })
        .collect()
}

pub fn caption_for_title(text: &str, files: &[DecodedAttachment]) -> String {
    let trimmed = text.trim();
    if !trimmed.is_empty() {
        return trimmed.to_string();
    }
    match files.first() {
        Some(file) => format!("附件 {}", file.name),
        None => String::new(),
    }
}

/// Write copies into each bot's `inbox/` so the desktop can open them.
pub async fn stage_for_bots(
    state: &AppState,
    actor: &Actor,
    bot_ids: &[String],
    files: &[DecodedAttachment],
) -> Result<(), String> {
    if files.is_empty() {
        return Ok(());
    }
    for bot_id in bot_ids {
        let Some(dir) = inbox_dir_for_bot(state, actor, bot_id).await? else {
            continue;
        };
        let root = PathBuf::from(&state.data_dir).join("homes");
        let relative = dir
            .strip_prefix(&root)
            .map_err(|_| "invalid inbox root")?
            .to_path_buf();
        let files = files.to_vec();
        tokio::task::spawn_blocking(move || -> Result<(), String> {
            use std::io::Write;
            let root = cap_std::fs::Dir::open_ambient_dir(root, cap_std::ambient_authority())
                .map_err(|e| e.to_string())?;
            root.create_dir_all(&relative).map_err(|e| e.to_string())?;
            let inbox = root.open_dir(&relative).map_err(|e| e.to_string())?;
            for file in files {
                let mut opts = cap_std::fs::OpenOptions::new();
                opts.write(true).create_new(true);
                let mut output = inbox
                    .open_with(&file.stored_name, &opts)
                    .map_err(|e| e.to_string())?;
                output.write_all(&file.bytes).map_err(|e| e.to_string())?;
            }
            Ok(())
        })
        .await
        .map_err(|e| e.to_string())??;
    }
    Ok(())
}

pub async fn llm_parts(
    state: &AppState,
    actor: &Actor,
    bot_id: &str,
    blocks: &[Value],
    vision: bool,
) -> Vec<UserContent> {
    let files = load_from_blocks(state, actor, bot_id, blocks).await;
    if files.is_empty() && blocks.is_empty() {
        return Vec::new();
    }
    let mut parts = Vec::new();
    let mut notes = Vec::new();
    if !blocks.is_empty() {
        notes.push(format!(
            "User attached {} file(s). Chat history does not keep the bytes. Copies live in inbox/ on this computer and expire after 2 hours.",
            blocks.len()
        ));
    }
    for (stored, bytes) in files {
        notes.push(format!(
            "- {} ({}, {} bytes) at {}",
            stored.name, stored.mime_type, stored.size, stored.path
        ));
        if is_image(&stored.mime_type) {
            if vision {
                if let Some(bytes) = bytes {
                    parts.push(UserContent::image_base64(
                        base64::engine::general_purpose::STANDARD.encode(bytes),
                        Some(image_media(&stored.mime_type)),
                        Some(ImageDetail::High),
                    ));
                } else {
                    notes.push("  (image missing from inbox; it may have expired)".into());
                }
            }
        } else if is_text_mime(&stored.mime_type) {
            if let Some(bytes) = bytes
                && let Some(text) = utf8_preview(&bytes)
            {
                parts.push(UserContent::text(format!(
                    "Contents of {}:\n```\n{text}\n```",
                    stored.name
                )));
            }
        } else {
            notes.push(
                "  Open this with open_path or the desktop if you need the original file.".into(),
            );
        }
    }
    if !notes.is_empty() {
        parts.insert(0, UserContent::text(notes.join("\n")));
    }
    parts
}

pub async fn sweep_all_inboxes(data_dir: &str) {
    let homes = PathBuf::from(data_dir).join("homes");
    let _ = tokio::task::spawn_blocking(move || {
        let Ok(root) = cap_std::fs::Dir::open_ambient_dir(homes, cap_std::ambient_authority())
        else {
            return;
        };
        let Ok(entries) = root.entries() else {
            return;
        };
        for entry in entries.flatten() {
            let Ok(home) = root.open_dir(entry.file_name()) else {
                continue;
            };
            sweep_cap_inbox(&home, "inbox");
            if let Ok(bots) = home.open_dir("bots")
                && let Ok(entries) = bots.entries()
            {
                for bot in entries.flatten() {
                    if let Ok(dir) = bots.open_dir(bot.file_name()) {
                        sweep_cap_inbox(&dir, "inbox");
                    }
                }
            }
        }
    })
    .await;
}

fn sweep_cap_inbox(root: &cap_std::fs::Dir, path: &str) {
    let Ok(dir) = root.open_dir(path) else {
        return;
    };
    let Ok(entries) = dir.entries() else {
        return;
    };
    for entry in entries.flatten() {
        let Ok(meta) = entry.metadata() else {
            continue;
        };
        if meta.is_file()
            && meta
                .modified()
                .ok()
                .and_then(|t| t.into_std().elapsed().ok())
                .is_some_and(|age| age > INBOX_TTL)
        {
            let _ = dir.remove_file(entry.file_name());
        }
    }
}

async fn inbox_dir_for_bot(
    state: &AppState,
    actor: &Actor,
    bot_id: &str,
) -> Result<Option<PathBuf>, String> {
    let Some(bot) = state
        .db
        .get_bot(actor, bot_id)
        .await
        .map_err(|error| error.to_string())?
    else {
        return Ok(None);
    };
    let Some(computer_id) = bot.computer_id.as_deref() else {
        return Ok(None);
    };
    let Some(computer) = state
        .db
        .get_computer(computer_id)
        .await
        .map_err(|error| error.to_string())?
    else {
        return Ok(None);
    };
    let home = computer::home_path(&state.data_dir, &computer.home_key);
    let relative = resolve_bot_workspace_path(parse_mode(&computer.scope), bot_id, "inbox")
        .unwrap_or_else(|_| "inbox".into());
    Ok(Some(home.join(relative)))
}

async fn load_from_blocks(
    state: &AppState,
    actor: &Actor,
    bot_id: &str,
    blocks: &[Value],
) -> Vec<(StoredAttachment, Option<Vec<u8>>)> {
    let Some(dir) = inbox_dir_for_bot(state, actor, bot_id).await.ok().flatten() else {
        return blocks
            .iter()
            .filter_map(parse_stored)
            .map(|file| (file, None))
            .collect();
    };
    let mut out = Vec::new();
    for block in blocks {
        let Some(file) = parse_stored(block) else {
            continue;
        };
        let root = PathBuf::from(&state.data_dir).join("homes");
        let stored = file.path.strip_prefix("inbox/").unwrap_or(&file.name);
        let relative = dir.strip_prefix(&root).ok().map(|p| p.join(stored));
        let bytes = if let Some(relative) = relative {
            tokio::task::spawn_blocking(move || {
                use std::io::Read;
                let root =
                    cap_std::fs::Dir::open_ambient_dir(root, cap_std::ambient_authority()).ok()?;
                let input = root.open(relative).ok()?;
                if !input.metadata().ok()?.is_file() {
                    return None;
                }
                let mut bytes = Vec::new();
                input
                    .take(MAX_BYTES as u64 + 1)
                    .read_to_end(&mut bytes)
                    .ok()?;
                (bytes.len() <= MAX_BYTES).then_some(bytes)
            })
            .await
            .ok()
            .flatten()
        } else {
            None
        };
        out.push((file, bytes));
    }
    out
}

fn parse_stored(value: &Value) -> Option<StoredAttachment> {
    let kind = value.get("kind").and_then(Value::as_str)?;
    if kind != "file" && kind != "image" {
        return None;
    }
    let name = value.get("name").and_then(Value::as_str)?;
    if name.is_empty()
        || name == "."
        || name == ".."
        || name.contains(['/', '\\'])
        || name.chars().any(char::is_control)
    {
        return None;
    }
    let path = value.get("path").and_then(Value::as_str).unwrap_or("");
    if !path.is_empty() {
        let stored = path.strip_prefix("inbox/")?;
        if stored.is_empty() || stored == "." || stored == ".." || stored.contains(['/', '\\']) {
            return None;
        }
    }
    Some(StoredAttachment {
        kind: if kind == "image" { "image" } else { "file" },
        name: name.to_string(),
        mime_type: value
            .get("mimeType")
            .and_then(Value::as_str)
            .unwrap_or("application/octet-stream")
            .to_string(),
        size: value.get("size").and_then(Value::as_u64).unwrap_or(0) as usize,
        path: value
            .get("path")
            .and_then(Value::as_str)
            .unwrap_or("")
            .to_string(),
    })
}

fn split_ext(name: &str) -> (String, String) {
    match name.rfind('.') {
        Some(index) if index > 0 => (name[..index].to_string(), name[index..].to_string()),
        _ => (name.to_string(), String::new()),
    }
}

pub fn safe_name(name: &str) -> Result<String, String> {
    let file = Path::new(name)
        .file_name()
        .and_then(|value| value.to_str())
        .unwrap_or("")
        .trim();
    let mut out = String::new();
    for ch in file.chars() {
        if ch.is_control() || "/\\:".contains(ch) {
            continue;
        }
        out.push(ch);
        if out.chars().count() >= NAME_MAX {
            break;
        }
    }
    let out = out.trim_matches('.').trim().to_string();
    if out.is_empty() || out == "." || out == ".." {
        return Err("檔名無效".into());
    }
    Ok(out)
}

fn decode_base64(value: &str) -> Result<Vec<u8>, String> {
    let trimmed = value.trim();
    if trimmed.is_empty() {
        return Err("附件內容是空的".into());
    }
    base64::engine::general_purpose::STANDARD
        .decode(trimmed)
        .or_else(|_| base64::engine::general_purpose::STANDARD.decode(trimmed.replace('\n', "")))
        .map_err(|_| "附件不是有效的 base64".into())
}

fn normalize_mime(mime: &str, name: &str) -> String {
    let mime = mime.trim().to_ascii_lowercase();
    if !mime.is_empty() && mime != "application/octet-stream" {
        return mime.split(';').next().unwrap_or(&mime).trim().to_string();
    }
    match Path::new(name)
        .extension()
        .and_then(|ext| ext.to_str())
        .unwrap_or("")
        .to_ascii_lowercase()
        .as_str()
    {
        "png" => "image/png",
        "jpg" | "jpeg" => "image/jpeg",
        "gif" => "image/gif",
        "webp" => "image/webp",
        "pdf" => "application/pdf",
        "txt" | "log" => "text/plain",
        "md" => "text/markdown",
        "csv" => "text/csv",
        "json" => "application/json",
        "html" | "htm" => "text/html",
        "xml" => "application/xml",
        "docx" => "application/vnd.openxmlformats-officedocument.wordprocessingml.document",
        "xlsx" => "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
        "pptx" => "application/vnd.openxmlformats-officedocument.presentationml.presentation",
        _ => "application/octet-stream",
    }
    .into()
}

fn allowed_mime(mime: &str) -> bool {
    is_image(mime)
        || is_text_mime(mime)
        || matches!(
            mime,
            "application/pdf"
                | "application/xml"
                | "application/vnd.openxmlformats-officedocument.wordprocessingml.document"
                | "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"
                | "application/vnd.openxmlformats-officedocument.presentationml.presentation"
        )
}

fn is_image(mime: &str) -> bool {
    matches!(
        mime,
        "image/jpeg" | "image/png" | "image/gif" | "image/webp"
    )
}

fn is_text_mime(mime: &str) -> bool {
    mime.starts_with("text/") || matches!(mime, "application/json" | "application/xml")
}

fn image_media(mime: &str) -> ImageMediaType {
    if mime == "image/jpeg" {
        ImageMediaType::JPEG
    } else {
        ImageMediaType::PNG
    }
}

fn utf8_preview(bytes: &[u8]) -> Option<String> {
    if bytes.contains(&0) {
        return None;
    }
    let text = std::str::from_utf8(bytes).ok()?;
    let preview: String = text.chars().take(TEXT_INLINE_CHARS).collect();
    if text.chars().count() > TEXT_INLINE_CHARS {
        Some(format!("{preview}\n…(truncated)"))
    } else {
        Some(preview)
    }
}

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

    #[test]
    fn rejects_path_in_name() {
        assert_eq!(safe_name("../../etc/passwd").unwrap(), "passwd");
        assert!(safe_name("...").is_err());
    }

    #[test]
    fn guesses_png_mime() {
        assert_eq!(normalize_mime("", "shot.PNG"), "image/png");
        assert!(allowed_mime("image/png"));
        assert!(!allowed_mime("application/x-msdownload"));
    }

    #[test]
    fn stored_blocks_drop_bytes() {
        let files = [DecodedAttachment {
            name: "a.png".into(),
            stored_name: "a.png".into(),
            mime_type: "image/png".into(),
            bytes: vec![1, 2, 3],
        }];
        let blocks = stored_blocks(&files);
        assert_eq!(blocks[0]["kind"], "image");
        assert_eq!(blocks[0]["path"], "inbox/a.png");
        assert!(blocks[0].get("content").is_none());
    }

    #[test]
    fn decode_limits_count() {
        let item = IncomingAttachment {
            name: "a.txt".into(),
            mime_type: "text/plain".into(),
            content: base64::engine::general_purpose::STANDARD.encode("hi"),
        };
        let many = vec![item; 5];
        assert!(decode_incoming(&many).is_err());
    }
}

#[cfg(test)]
mod boundary_tests {
    use super::*;
    #[test]
    fn stored_metadata_cannot_escape_inbox() {
        for name in ["../secret", "/etc/passwd", "..", "x\\secret"] {
            assert!(parse_stored(&json!({"kind":"file","name":name})).is_none());
        }
    }
    #[cfg(unix)]
    #[test]
    fn symlink_from_computer_home_cannot_read_host_secret() {
        let base = std::env::temp_dir().join(format!("lazyboy-boundary-{}", uuid::Uuid::new_v4()));
        std::fs::create_dir_all(base.join("homes")).unwrap();
        std::fs::write(base.join("secret"), b"private").unwrap();
        std::os::unix::fs::symlink(base.join("secret"), base.join("homes/link")).unwrap();
        let root =
            cap_std::fs::Dir::open_ambient_dir(base.join("homes"), cap_std::ambient_authority())
                .unwrap();
        assert!(root.read("link").is_err());
        assert!(root.write("link", b"overwritten").is_err());
        assert_eq!(std::fs::read(base.join("secret")).unwrap(), b"private");
        drop(root);
        std::fs::remove_dir_all(base).unwrap();
    }
}