//! Library path <-> ComfyUI path mapping (#192). //! //! ImageApi sees `/media/nas/Media/Pictures/a.jpg`; ComfyUI (on Windows) //! sees `Z:\Media\Pictures\a.jpg` or `\\nas\Media\Pictures\a.jpg`. Each //! mapping is a prefix pair; the longest matching prefix wins and the //! separators are converted. Used in both directions: library -> ComfyUI on //! submit, and ComfyUI -> library to recognise the path `approve` returns. //! With no mappings configured, paths pass through unchanged. use std::path::{Path, PathBuf}; use anyhow::{Result, bail}; use super::config::PathMapping; #[derive(Debug, Clone)] struct Entry { /// Library side, `/`-separated, always ending in `/`. from: String, /// ComfyUI side without a trailing separator. to: String, /// Separator ComfyUI's side uses. to_sep: char, } #[derive(Debug, Clone, Default)] pub struct PathMap { entries: Vec, } impl PathMap { pub fn new(mappings: &[PathMapping]) -> Result { let mut entries = Vec::new(); for m in mappings { let from = m.from.replace('\\', "/"); let to = m.to.trim(); if from.trim_matches('/').is_empty() || to.is_empty() { bail!("path_map entries need non-empty from and to ({m:?})"); } let to_sep = if to.contains('\\') { '\\' } else { '/' }; entries.push(Entry { from: format!("{}/", from.trim_end_matches('/')), to: to.trim_end_matches(['/', '\\']).to_string(), to_sep, }); } Ok(PathMap { entries }) } pub fn is_identity(&self) -> bool { self.entries.is_empty() } /// The path ComfyUI should open for a library file. `None` when mappings /// exist but none covers the file. pub fn to_comfy(&self, abs: &Path) -> Option { let text = abs.to_str()?.replace('\\', "/"); if self.entries.is_empty() { return Some(abs.to_string_lossy().into_owned()); } let best = self .entries .iter() .filter(|e| text.starts_with(e.from.as_str())) .max_by_key(|e| e.from.len())?; let rest = &text[best.from.len()..]; let sep = best.to_sep.to_string(); Some(format!("{}{}{}", best.to, sep, rest.replace('/', &sep))) } /// The library-side absolute path of a path ComfyUI reported. Matching /// is case-insensitive (Windows). `None` when no mapping covers it or it /// tries to climb out with `..`. pub fn from_comfy(&self, comfy: &str) -> Option { let text = comfy.trim().replace('\\', "/"); if self.entries.is_empty() { return Some(PathBuf::from(text)); } let best = self .entries .iter() .filter_map(|e| { let prefix = format!("{}/", e.to.replace('\\', "/")); let head = text.get(..prefix.len())?; head.eq_ignore_ascii_case(&prefix) .then_some((e, prefix.len())) }) .max_by_key(|(_, len)| *len)?; let rest = &text[best.1..]; if rest.split('/').any(|c| c == "..") { return None; } Some(PathBuf::from(format!("{}{}", best.0.from, rest))) } } #[cfg(test)] mod tests { use super::*; fn map(pairs: &[(&str, &str)]) -> PathMap { let m: Vec = pairs .iter() .map(|(f, t)| PathMapping { from: f.to_string(), to: t.to_string(), }) .collect(); PathMap::new(&m).unwrap() } #[test] fn drive_letter_round_trip() { let m = map(&[("/media/nas/", "Z:\\")]); let c = m .to_comfy(Path::new("/media/nas/Media/Pictures/a b.jpg")) .unwrap(); assert_eq!(c, "Z:\\Media\\Pictures\\a b.jpg"); assert_eq!( m.from_comfy(&c).unwrap(), PathBuf::from("/media/nas/Media/Pictures/a b.jpg") ); // Drive letters are case-insensitive. assert_eq!( m.from_comfy("z:\\x\\y.png").unwrap(), PathBuf::from("/media/nas/x/y.png") ); } #[test] fn unc_target_and_missing_trailing_slash() { let m = map(&[("/media/nas", "\\\\nas\\")]); assert_eq!( m.to_comfy(Path::new("/media/nas/Media/a.jpg")).unwrap(), "\\\\nas\\Media\\a.jpg" ); assert_eq!( m.from_comfy("\\\\nas\\Media\\a.jpg").unwrap(), PathBuf::from("/media/nas/Media/a.jpg") ); } #[test] fn longest_prefix_wins_both_ways() { let m = map(&[("/media/nas/", "Z:\\"), ("/media/nas/Media/", "Y:\\")]); assert_eq!( m.to_comfy(Path::new("/media/nas/Media/a.jpg")).unwrap(), "Y:\\a.jpg" ); assert_eq!( m.to_comfy(Path::new("/media/nas/Other/a.jpg")).unwrap(), "Z:\\Other\\a.jpg" ); assert_eq!( m.from_comfy("Y:\\a.jpg").unwrap(), PathBuf::from("/media/nas/Media/a.jpg") ); } #[test] fn uncovered_and_escaping_paths_are_none() { let m = map(&[("/media/nas/", "Z:\\")]); assert!(m.to_comfy(Path::new("/elsewhere/a.jpg")).is_none()); assert!(m.to_comfy(Path::new("/media/nasty/a.jpg")).is_none()); assert!(m.from_comfy("C:\\a.png").is_none()); assert!(m.from_comfy("Z:\\..\\etc\\passwd").is_none()); } #[test] fn no_mappings_is_identity() { let m = PathMap::default(); assert_eq!(m.to_comfy(Path::new("/a/b.jpg")).unwrap(), "/a/b.jpg"); assert_eq!(m.from_comfy("/a/b.png").unwrap(), PathBuf::from("/a/b.png")); } }