From eb0a32f47f0bbebf5504d88229d5991e3c9fd93c Mon Sep 17 00:00:00 2001 From: mogurt Date: Mon, 28 Sep 2026 14:52:28 +0800 Subject: [PATCH 1/8] fix(security): sanitize filenames, match hosts exactly, redact tokens in logs - Add safety module: sanitize remote filenames to a bare basename (blocks ../ and \ traversal, drive prefixes, illegal chars, reserved names); create_download / check_file_exists reject unsafe names. - Match civitai.com / huggingface.co / hf-mirror.com by exact host or subdomain instead of substring, in both backend and frontend, so tokens are never sent to look-alike domains. - Redact token-like query params from URLs in backend logs and forwarded aria2 output. --- src-tauri/src/aria2/process.rs | 3 + src-tauri/src/commands/download.rs | 10 +- src-tauri/src/lib.rs | 1 + src-tauri/src/safety.rs | 173 ++++++++++++++++++++++++ src-tauri/src/url_parser/civitai.rs | 5 +- src-tauri/src/url_parser/generic.rs | 5 +- src-tauri/src/url_parser/huggingface.rs | 5 +- src-tauri/src/url_parser/mod.rs | 13 +- src/stores/taskStore.ts | 18 ++- 9 files changed, 216 insertions(+), 17 deletions(-) create mode 100644 src-tauri/src/safety.rs diff --git a/src-tauri/src/aria2/process.rs b/src-tauri/src/aria2/process.rs index fc4d92d..7fd05ef 100644 --- a/src-tauri/src/aria2/process.rs +++ b/src-tauri/src/aria2/process.rs @@ -156,6 +156,9 @@ where if line.is_empty() { continue; } + // aria2 echoes request URIs, which may carry API tokens. + let line = crate::safety::redact_text(line); + let line = line.as_str(); let lowered = line.to_ascii_lowercase(); let level = if lowered.contains("error") { diff --git a/src-tauri/src/commands/download.rs b/src-tauri/src/commands/download.rs index f27f776..5736d64 100644 --- a/src-tauri/src/commands/download.rs +++ b/src-tauri/src/commands/download.rs @@ -1,5 +1,6 @@ use crate::aria2::rpc::Aria2Rpc; use crate::model_type::rules::{suggest_model_type, UserRule}; +use crate::safety::{redact_url, validate_filename}; use crate::url_parser; use crate::url_parser::ParseResult; use log::info; @@ -58,6 +59,7 @@ pub async fn suggest_type( #[tauri::command] pub async fn check_file_exists(dir: String, filename: String) -> Result { + validate_filename(&filename)?; let path = Path::new(&dir).join(&filename); Ok(path.exists()) } @@ -70,9 +72,15 @@ pub async fn create_download( filename: String, headers: Option>, ) -> Result { + validate_filename(&filename)?; let rpc = get_rpc(&rpc_state).await?; - info!("Creating download: {} -> {}/{}", url, dir, filename); + info!( + "Creating download: {} -> {}/{}", + redact_url(&url), + dir, + filename + ); let gid = rpc.add_uri(vec![url], &dir, &filename, headers).await?; diff --git a/src-tauri/src/lib.rs b/src-tauri/src/lib.rs index 8994768..8402a4b 100644 --- a/src-tauri/src/lib.rs +++ b/src-tauri/src/lib.rs @@ -2,6 +2,7 @@ mod aria2; mod commands; mod db; mod model_type; +mod safety; mod url_parser; use aria2::process::{find_available_port, resolve_aria2_path, Aria2Process}; diff --git a/src-tauri/src/safety.rs b/src-tauri/src/safety.rs new file mode 100644 index 0000000..cf02a79 --- /dev/null +++ b/src-tauri/src/safety.rs @@ -0,0 +1,173 @@ +//! Helpers for handling untrusted input: remote filenames, hosts, and URLs that +//! may carry credentials. + +use url::Url; + +const WINDOWS_RESERVED_NAMES: [&str; 22] = [ + "CON", "PRN", "AUX", "NUL", "COM1", "COM2", "COM3", "COM4", "COM5", "COM6", "COM7", "COM8", + "COM9", "LPT1", "LPT2", "LPT3", "LPT4", "LPT5", "LPT6", "LPT7", "LPT8", "LPT9", +]; + +const SENSITIVE_QUERY_KEYS: [&str; 5] = ["token", "access_token", "api_key", "apikey", "key"]; + +/// Turns an untrusted (remote-supplied) name into a plain file name that cannot +/// escape the download directory. Returns `None` if nothing usable remains. +pub fn sanitize_filename(raw: &str) -> Option { + let base = raw.rsplit(['/', '\\']).next().unwrap_or(raw); + + let cleaned: String = base + .chars() + .map(|c| match c { + '<' | '>' | ':' | '"' | '|' | '?' | '*' => '_', + c if c.is_control() => '_', + c => c, + }) + .collect(); + + // Windows silently strips trailing dots and spaces. + let cleaned = cleaned.trim().trim_end_matches(['.', ' ']).to_string(); + + if cleaned.is_empty() || cleaned.chars().all(|c| c == '.' || c == '_') { + return None; + } + + let stem = cleaned.split('.').next().unwrap_or("").to_ascii_uppercase(); + if WINDOWS_RESERVED_NAMES.contains(&stem.as_str()) { + return Some(format!("_{}", cleaned)); + } + + Some(cleaned) +} + +/// Validates a file name coming from the frontend. Unlike [`sanitize_filename`] +/// this rejects instead of rewriting, so the stored name always matches the file on disk. +pub fn validate_filename(name: &str) -> Result<(), String> { + match sanitize_filename(name) { + Some(clean) if clean == name => Ok(()), + _ => Err(format!("Invalid file name: {}", name)), + } +} + +/// True if `host` is `domain` itself or one of its subdomains. +pub fn host_matches(host: &str, domain: &str) -> bool { + let host = host.trim_end_matches('.').to_ascii_lowercase(); + host == domain || host.ends_with(&format!(".{}", domain)) +} + +/// Masks credential-looking query parameters so a URL is safe to log. +pub fn redact_url(raw: &str) -> String { + let Ok(mut url) = Url::parse(raw) else { + return raw.to_string(); + }; + if url.query().is_none() { + return raw.to_string(); + } + + let pairs: Vec<(String, String)> = url + .query_pairs() + .map(|(k, v)| { + let value = if SENSITIVE_QUERY_KEYS.contains(&k.to_ascii_lowercase().as_str()) { + "***".to_string() + } else { + v.into_owned() + }; + (k.into_owned(), value) + }) + .collect(); + + url.query_pairs_mut().clear().extend_pairs(pairs); + url.to_string() +} + +/// Redacts any URLs embedded in a free-form log line (e.g. aria2 output). +pub fn redact_text(text: &str) -> String { + text.split(' ') + .map(|part| { + if part.contains("://") && part.contains('?') { + let start = part.find("http").unwrap_or(0); + format!("{}{}", &part[..start], redact_url(&part[start..])) + } else { + part.to_string() + } + }) + .collect::>() + .join(" ") +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn sanitize_strips_traversal_and_separators() { + assert_eq!( + sanitize_filename("../../Startup/a.bat").as_deref(), + Some("a.bat") + ); + assert_eq!( + sanitize_filename("..\\..\\evil.dll").as_deref(), + Some("evil.dll") + ); + assert_eq!( + sanitize_filename("C:\\Windows\\x.dll").as_deref(), + Some("x.dll") + ); + assert_eq!(sanitize_filename("/etc/passwd").as_deref(), Some("passwd")); + } + + #[test] + fn sanitize_rejects_empty_and_dot_names() { + assert_eq!(sanitize_filename(""), None); + assert_eq!(sanitize_filename(".."), None); + assert_eq!(sanitize_filename("foo/.."), None); + assert_eq!(sanitize_filename("dir/"), None); + } + + #[test] + fn sanitize_replaces_illegal_chars_and_reserved_names() { + assert_eq!( + sanitize_filename("a:b?c*.safetensors").as_deref(), + Some("a_b_c_.safetensors") + ); + assert_eq!(sanitize_filename("model. ").as_deref(), Some("model")); + assert_eq!(sanitize_filename("con.txt").as_deref(), Some("_con.txt")); + assert_eq!( + sanitize_filename("my model.safetensors").as_deref(), + Some("my model.safetensors") + ); + } + + #[test] + fn validate_only_accepts_clean_names() { + assert!(validate_filename("flux1-dev.safetensors").is_ok()); + assert!(validate_filename("../x.safetensors").is_err()); + assert!(validate_filename("C:\\x.safetensors").is_err()); + assert!(validate_filename("").is_err()); + } + + #[test] + fn host_matching_is_exact() { + assert!(host_matches("civitai.com", "civitai.com")); + assert!(host_matches("www.civitai.com", "civitai.com")); + assert!(host_matches("CIVITAI.COM.", "civitai.com")); + assert!(!host_matches("civitai.com.evil.net", "civitai.com")); + assert!(!host_matches("evilcivitai.com", "civitai.com")); + } + + #[test] + fn redacts_tokens_in_urls_and_text() { + let redacted = + redact_url("https://civitai.com/api/download/models/1?type=Model&token=abc123"); + assert!(!redacted.contains("abc123")); + assert!(redacted.contains("type=Model")); + + let line = redact_text("Download aborted. URI=https://civitai.com/x?token=secret done"); + assert!(!line.contains("secret")); + assert!(line.starts_with("Download aborted. URI=https://civitai.com/x?token=")); + + assert_eq!( + redact_url("https://example.com/a.bin"), + "https://example.com/a.bin" + ); + } +} diff --git a/src-tauri/src/url_parser/civitai.rs b/src-tauri/src/url_parser/civitai.rs index 39620a3..284e894 100644 --- a/src-tauri/src/url_parser/civitai.rs +++ b/src-tauri/src/url_parser/civitai.rs @@ -9,7 +9,10 @@ pub async fn parse( proxy: Option<&str>, token: Option<&str>, ) -> Result { - info!("Parsing Civitai URL: {}", raw_url); + info!( + "Parsing Civitai URL: {}", + crate::safety::redact_url(raw_url) + ); let mut headers = HeaderMap::new(); if let Some(t) = token { diff --git a/src-tauri/src/url_parser/generic.rs b/src-tauri/src/url_parser/generic.rs index 0b372a8..a46adb0 100644 --- a/src-tauri/src/url_parser/generic.rs +++ b/src-tauri/src/url_parser/generic.rs @@ -3,7 +3,10 @@ use log::info; use url::Url; pub async fn parse(raw_url: &str, proxy: Option<&str>) -> Result { - info!("Parsing generic URL: {}", raw_url); + info!( + "Parsing generic URL: {}", + crate::safety::redact_url(raw_url) + ); let url = Url::parse(raw_url).map_err(|e| format!("Invalid URL: {}", e))?; diff --git a/src-tauri/src/url_parser/huggingface.rs b/src-tauri/src/url_parser/huggingface.rs index a7cae15..711b0f8 100644 --- a/src-tauri/src/url_parser/huggingface.rs +++ b/src-tauri/src/url_parser/huggingface.rs @@ -8,7 +8,10 @@ pub async fn parse( proxy: Option<&str>, token: Option<&str>, ) -> Result { - info!("Parsing HuggingFace URL: {}", raw_url); + info!( + "Parsing HuggingFace URL: {}", + crate::safety::redact_url(raw_url) + ); let url = Url::parse(raw_url).map_err(|e| format!("Invalid URL: {}", e))?; let path = url.path(); diff --git a/src-tauri/src/url_parser/mod.rs b/src-tauri/src/url_parser/mod.rs index ba0461e..ec6bf7c 100644 --- a/src-tauri/src/url_parser/mod.rs +++ b/src-tauri/src/url_parser/mod.rs @@ -2,6 +2,7 @@ pub mod civitai; pub mod generic; pub mod huggingface; +use crate::safety::{host_matches, sanitize_filename}; use serde::{Deserialize, Serialize}; use url::Url; @@ -23,11 +24,17 @@ pub async fn parse_url( let url = Url::parse(raw_url).map_err(|e| format!("Invalid URL: {}", e))?; let host = url.host_str().unwrap_or(""); - if host.contains("civitai.com") { + let mut result = if host_matches(host, "civitai.com") { civitai::parse(raw_url, proxy, civitai_token).await - } else if host.contains("huggingface.co") || host.contains("hf-mirror.com") { + } else if host_matches(host, "huggingface.co") || host_matches(host, "hf-mirror.com") { huggingface::parse(raw_url, proxy, huggingface_token).await } else { generic::parse(raw_url, proxy).await - } + }?; + + // Filenames come from remote URLs / headers / metadata and must never + // be able to point outside the target directory. + result.filename = + sanitize_filename(&result.filename).unwrap_or_else(|| "unknown_model".to_string()); + Ok(result) } diff --git a/src/stores/taskStore.ts b/src/stores/taskStore.ts index 4f25590..9d3926a 100644 --- a/src/stores/taskStore.ts +++ b/src/stores/taskStore.ts @@ -43,23 +43,21 @@ function matchesTaskGid(taskGid: string | null | undefined, reportedGid: string) return normalizedTaskGid === normalizedReportedGid || normalizedTaskGid.startsWith(normalizedReportedGid); } -function isHuggingFaceUrl(url: string): boolean { +function hostMatches(url: string, domains: string[]): boolean { try { - const parsed = new URL(url); - const host = parsed.host.toLowerCase(); - return host.includes("huggingface.co") || host.includes("hf-mirror.com"); + const host = new URL(url).hostname.toLowerCase().replace(/\.$/, ""); + return domains.some((domain) => host === domain || host.endsWith(`.${domain}`)); } catch { return false; } } +function isHuggingFaceUrl(url: string): boolean { + return hostMatches(url, ["huggingface.co", "hf-mirror.com"]); +} + function isCivitaiUrl(url: string): boolean { - try { - const parsed = new URL(url); - return parsed.host.toLowerCase().includes("civitai.com"); - } catch { - return false; - } + return hostMatches(url, ["civitai.com"]); } async function getDb(): Promise { From bb70948222e3948c27038274c13b96514d4cfad1 Mon Sep 17 00:00:00 2001 From: mogurt Date: Mon, 28 Sep 2026 14:52:29 +0800 Subject: [PATCH 2/8] fix(macos): locate bundled aria2 sidecar and stop hardcoding "\\" paths - Look for the aria2c sidecar next to the main executable (where Tauri places externalBin, e.g. Contents/MacOS) before resource_dir, and warn when falling back to PATH. - Add joinPath/getModelBaseDir helpers that follow the base path's separator style; replace hardcoded backslash joins in TaskInput, Search and Settings. - Workflow page now uses the same default `/models` base dir as the other pages instead of reporting every model missing. --- src-tauri/src/aria2/process.rs | 24 +++++++++++++++++++----- src/components/TaskInput.tsx | 11 +++++------ src/lib/utils.ts | 13 +++++++++++++ src/pages/Search.tsx | 21 ++++++++------------- src/pages/Settings.tsx | 12 +++++------- src/pages/Workflow.tsx | 3 ++- 6 files changed, 52 insertions(+), 32 deletions(-) diff --git a/src-tauri/src/aria2/process.rs b/src-tauri/src/aria2/process.rs index 7fd05ef..c194aa2 100644 --- a/src-tauri/src/aria2/process.rs +++ b/src-tauri/src/aria2/process.rs @@ -210,11 +210,21 @@ pub fn resolve_aria2_path(app_handle: &tauri::AppHandle) -> Result Result { if (!baseDir) { @@ -72,15 +71,15 @@ export default function TaskInput() { setSelectedSubSubdir(""); return; } - const fullPath = `${baseDir}\\${selectedSubdir}`; + const fullPath = joinPath(baseDir, selectedSubdir); api.listSubdirs(fullPath).then(setSubSubdirs).catch(() => setSubSubdirs([])); }, [baseDir, selectedSubdir]); const targetDir = baseDir ? selectedSubSubdir - ? `${baseDir}\\${selectedSubdir}\\${selectedSubSubdir}` + ? joinPath(baseDir, selectedSubdir, selectedSubSubdir) : selectedSubdir && selectedSubdir !== "__manual__" - ? `${baseDir}\\${selectedSubdir}` + ? joinPath(baseDir, selectedSubdir) : "" : ""; diff --git a/src/lib/utils.ts b/src/lib/utils.ts index bd0c391..edfa6ae 100644 --- a/src/lib/utils.ts +++ b/src/lib/utils.ts @@ -4,3 +4,16 @@ import { twMerge } from "tailwind-merge" export function cn(...inputs: ClassValue[]) { return twMerge(clsx(inputs)) } + +/** Joins path segments using the separator style already present in the base path. */ +export function joinPath(base: string, ...parts: string[]): string { + const sep = base.includes("\\") && !base.includes("/") ? "\\" : "/"; + const segments = parts.filter(Boolean).map((p) => p.replace(/^[\/]+|[\/]+$/g, "")); + return [base.replace(/[\/]+$/, ""), ...segments].join(sep); +} + +/** The configured model base dir, or `/models` when unset. */ +export function getModelBaseDir(settings: { model_base_dir: string; comfyui_root: string }): string { + if (settings.model_base_dir) return settings.model_base_dir; + return settings.comfyui_root ? joinPath(settings.comfyui_root, "models") : ""; +} diff --git a/src/pages/Search.tsx b/src/pages/Search.tsx index 991d058..84e42a5 100644 --- a/src/pages/Search.tsx +++ b/src/pages/Search.tsx @@ -38,6 +38,7 @@ import { useSearchStore, type SortOption, type SourceFilter } from "@/stores/sea import { useSettingsStore } from "@/stores/settingsStore"; import { useTaskStore } from "@/stores/taskStore"; import * as api from "@/lib/api"; +import { getModelBaseDir, joinPath } from "@/lib/utils"; import { extractPath } from "@/lib/api"; import { open } from "@tauri-apps/plugin-dialog"; import { open as shellOpen } from "@tauri-apps/plugin-shell"; @@ -148,9 +149,7 @@ export default function Search({ const debounceRef = useRef>(undefined); const scrollRef = useRef(null); - const baseDir = - settings.model_base_dir || - (settings.comfyui_root ? `${settings.comfyui_root}\\models` : ""); + const baseDir = getModelBaseDir(settings); const handleInputChange = (value: string) => { setQuery(value); @@ -735,9 +734,7 @@ function FileRow({ const rules = useSettingsStore((s) => s.rules); const addLog = useTaskStore((s) => s.addLog); - const baseDir = - settings.model_base_dir || - (settings.comfyui_root ? `${settings.comfyui_root}\\models` : ""); + const baseDir = getModelBaseDir(settings); const handleClick = useCallback(async () => { setPreparing(true); @@ -839,9 +836,7 @@ function DownloadConfirmDialog({ const startDownload = useTaskStore((s) => s.startDownload); const addLog = useTaskStore((s) => s.addLog); - const baseDir = - settings.model_base_dir || - (settings.comfyui_root ? `${settings.comfyui_root}\\models` : ""); + const baseDir = getModelBaseDir(settings); useEffect(() => { if (!pending) return; @@ -852,7 +847,7 @@ function DownloadConfirmDialog({ setSubSubdirs([]); if (initial && baseDir) { - api.listSubdirs(`${baseDir}\\${initial}`) + api.listSubdirs(joinPath(baseDir, initial)) .then(setSubSubdirs) .catch(() => setSubSubdirs([])); } @@ -864,7 +859,7 @@ function DownloadConfirmDialog({ setSelectedSubSubdir(""); return; } - api.listSubdirs(`${baseDir}\\${selectedSubdir}`) + api.listSubdirs(joinPath(baseDir, selectedSubdir)) .then(setSubSubdirs) .catch(() => setSubSubdirs([])); }, [selectedSubdir, baseDir]); @@ -880,9 +875,9 @@ function DownloadConfirmDialog({ const resolvedDir = isManual ? manualDir : selectedSubSubdir - ? `${baseDir}\\${selectedSubdir}\\${selectedSubSubdir}` + ? joinPath(baseDir, selectedSubdir, selectedSubSubdir) : selectedSubdir - ? `${baseDir}\\${selectedSubdir}` + ? joinPath(baseDir, selectedSubdir) : ""; const handlePickDir = async () => { diff --git a/src/pages/Settings.tsx b/src/pages/Settings.tsx index b4cb5b7..5e443e4 100644 --- a/src/pages/Settings.tsx +++ b/src/pages/Settings.tsx @@ -21,7 +21,7 @@ import * as api from "@/lib/api"; import { extractPath } from "@/lib/api"; import type { ModelType } from "@/lib/types"; import type { ThemeMode } from "@/lib/theme"; -import { cn } from "@/lib/utils"; +import { cn, getModelBaseDir, joinPath } from "@/lib/utils"; import { useI18n } from "@/lib/i18n"; const MODEL_TYPES: ModelType[] = [ @@ -121,9 +121,7 @@ export default function Settings() { setNewKeyword(""); }; - const effectiveBaseDir = - settings.model_base_dir || - (settings.comfyui_root ? `${settings.comfyui_root}\\models` : ""); + const effectiveBaseDir = getModelBaseDir(settings); useEffect(() => { if (effectiveBaseDir) { @@ -135,7 +133,7 @@ export default function Settings() { const handleApplyDerivedBaseDir = async () => { if (settings.comfyui_root) { - await updateSetting("model_base_dir", `${settings.comfyui_root}\\models`); + await updateSetting("model_base_dir", joinPath(settings.comfyui_root, "models")); } }; @@ -371,7 +369,7 @@ export default function Settings() { className="w-96" placeholder={ settings.comfyui_root - ? `Auto: ${settings.comfyui_root}\\models` + ? `Auto: ${joinPath(settings.comfyui_root, "models")}` : t("settings.placeholder.modelBase") } /> @@ -383,7 +381,7 @@ export default function Settings() { {!settings.model_base_dir && settings.comfyui_root && (
- {t("settings.autoDetectedBaseDir", { path: `${settings.comfyui_root}\\models` })} + {t("settings.autoDetectedBaseDir", { path: joinPath(settings.comfyui_root, "models") })}