From a06d16ef287ef1241b1a3aa215693547d9d9cbd2 Mon Sep 17 00:00:00 2001 From: Lahfir Date: Fri, 5 Jun 2026 19:39:43 -0700 Subject: [PATCH] fix: harden ref fallback resolution --- crates/core/src/refs_store.rs | 30 ++++++++++++++- crates/core/src/refs_store_tests.rs | 52 ++++++++++++++++++++++++++ crates/macos/src/tree/resolve_roots.rs | 46 +++++++++++++++-------- crates/macos/src/tree/resolve_tests.rs | 30 ++++++++++++++- 4 files changed, 140 insertions(+), 18 deletions(-) diff --git a/crates/core/src/refs_store.rs b/crates/core/src/refs_store.rs index 23ee0bb..2c729d2 100644 --- a/crates/core/src/refs_store.rs +++ b/crates/core/src/refs_store.rs @@ -129,7 +129,7 @@ impl RefStore { snapshot_id: &str, ) -> Result, AppError> { let path = Self::snapshot_path_for_base(base_dir, snapshot_id); - let mut file = match std::fs::File::open(&path) { + let mut file = match open_refmap_file(&path) { Ok(file) => file, Err(err) if err.kind() == ErrorKind::NotFound => return Ok(None), Err(err) => return Err(err.into()), @@ -251,6 +251,13 @@ impl RefStore { if path == self.base_dir { continue; } + let name = entry.file_name(); + let Some(name) = name.to_str() else { + continue; + }; + if validate_session_id(name).is_err() { + continue; + } let Ok(file_type) = entry.file_type() else { continue; }; @@ -298,6 +305,27 @@ impl RefStore { } } +fn open_refmap_file(path: &Path) -> std::io::Result { + #[cfg(unix)] + { + use std::os::unix::fs::OpenOptionsExt; + std::fs::OpenOptions::new() + .read(true) + .custom_flags(libc::O_NOFOLLOW) + .open(path) + } + #[cfg(not(unix))] + { + if std::fs::symlink_metadata(path)?.file_type().is_symlink() { + return Err(std::io::Error::new( + ErrorKind::PermissionDenied, + "refmap path must not be a symlink", + )); + } + std::fs::File::open(path) + } +} + #[cfg(test)] #[path = "refs_store_tests.rs"] mod tests; diff --git a/crates/core/src/refs_store_tests.rs b/crates/core/src/refs_store_tests.rs index 84f0733..fbac20c 100644 --- a/crates/core/src/refs_store_tests.rs +++ b/crates/core/src/refs_store_tests.rs @@ -179,6 +179,58 @@ fn duplicate_explicit_snapshot_id_requires_session() { ); } +#[test] +fn discover_skips_invalid_session_names_when_detecting_collisions() { + let _guard = HomeGuard::new(); + let default_store = RefStore::new().unwrap(); + let session_a = RefStore::for_session(Some("agent-a")).unwrap(); + + session_a + .save_snapshot("sdup2", &map_with("Session A")) + .unwrap(); + let invalid_base = default_store.base_dir.join("sessions").join("bad.session"); + let invalid_path = RefStore::snapshot_path_for_base(&invalid_base, "sdup2"); + std::fs::create_dir_all(invalid_path.parent().unwrap()).unwrap(); + std::fs::write( + invalid_path, + map_with("Invalid").serialize_with_size_check().unwrap(), + ) + .unwrap(); + + assert_eq!( + default_store + .load(Some("sdup2")) + .unwrap() + .get("@e1") + .unwrap() + .name + .as_deref(), + Some("Session A") + ); +} + +#[cfg(unix)] +#[test] +fn read_snapshot_rejects_symlinked_refmap() { + let _guard = HomeGuard::new(); + let store = RefStore::new().unwrap(); + + store.save_snapshot("ssym1", &map_with("Original")).unwrap(); + let path = store.snapshot_path("ssym1"); + let target = store.base_dir.join("symlink-target-refmap.json"); + std::fs::write( + target.as_path(), + map_with("Symlinked").serialize_with_size_check().unwrap(), + ) + .unwrap(); + std::fs::remove_file(&path).unwrap(); + std::os::unix::fs::symlink(&target, &path).unwrap(); + + let err = store.load(Some("ssym1")).unwrap_err(); + + assert_eq!(err.code(), "INTERNAL"); +} + #[test] fn save_existing_snapshot_does_not_promote_latest_pointer() { let _guard = HomeGuard::new(); diff --git a/crates/macos/src/tree/resolve_roots.rs b/crates/macos/src/tree/resolve_roots.rs index 1c361b0..c690741 100644 --- a/crates/macos/src/tree/resolve_roots.rs +++ b/crates/macos/src/tree/resolve_roots.rs @@ -190,13 +190,11 @@ fn window_by_title( deadline: Instant, ) -> Option { let source_window_title = source_window_title?; - windows - .iter() - .find(|win| { - prepare_for_read(win, deadline).is_ok() - && copy_string_attr(win, "AXTitle").as_deref() == Some(source_window_title) - }) - .cloned() + let index = unique_matching_index(windows, |win| { + prepare_for_read(win, deadline).is_ok() + && copy_string_attr(win, "AXTitle").as_deref() == Some(source_window_title) + })?; + windows.get(index).cloned() } #[cfg(target_os = "macos")] @@ -208,19 +206,14 @@ fn fallback_source_window_root( if let Some(window) = window_by_title(windows, entry.source_window_title.as_deref(), deadline) { return Some(window); } - if !single_window_fallback_allowed(entry) { + if !sole_source_window_fallback_allowed(entry) { return None; } - let mut candidates = windows.iter().filter(|win| { + let index = unique_matching_index(windows, |win| { prepare_for_read(win, deadline).is_ok() && copy_string_attr(win, "AXRole").as_deref() == Some("AXWindow") - }); - let first = candidates.next()?; - if candidates.next().is_none() { - Some(first.clone()) - } else { - None - } + })?; + windows.get(index).cloned() } #[cfg(target_os = "macos")] @@ -233,6 +226,11 @@ pub(super) fn single_window_fallback_allowed(entry: &RefEntry) -> bool { source_window_scope_required(entry) && entry.bounds_hash.is_some() } +#[cfg(target_os = "macos")] +pub(super) fn sole_source_window_fallback_allowed(entry: &RefEntry) -> bool { + single_window_fallback_allowed(entry) && entry.source_window_title.is_none() +} + #[cfg(target_os = "macos")] pub(super) fn source_window_number(entry: &RefEntry) -> Option { entry @@ -248,3 +246,19 @@ fn prepare_for_read(element: &AXElement, deadline: Instant) -> Result<(), Adapte set_messaging_timeout(element, remaining_before_deadline(deadline)?); Ok(()) } + +#[cfg(target_os = "macos")] +pub(super) fn unique_matching_index( + items: &[T], + mut matches: impl FnMut(&T) -> bool, +) -> Option { + let mut matches = items + .iter() + .enumerate() + .filter_map(|(index, item)| matches(item).then_some(index)); + let first = matches.next()?; + if matches.next().is_some() { + return None; + } + Some(first) +} diff --git a/crates/macos/src/tree/resolve_tests.rs b/crates/macos/src/tree/resolve_tests.rs index a0162a3..ef1a4b8 100644 --- a/crates/macos/src/tree/resolve_tests.rs +++ b/crates/macos/src/tree/resolve_tests.rs @@ -1,5 +1,8 @@ use super::*; -use crate::tree::resolve_roots::{single_window_fallback_allowed, source_window_number}; +use crate::tree::resolve_roots::{ + single_window_fallback_allowed, sole_source_window_fallback_allowed, source_window_number, + unique_matching_index, +}; use agent_desktop_core::adapter::SnapshotSurface; fn entry( @@ -243,6 +246,31 @@ fn single_window_fallback_requires_bounds_hash_not_title() { assert!(!single_window_fallback_allowed(&menu_entry)); } +#[test] +fn sole_window_fallback_requires_missing_title() { + assert!(sole_source_window_fallback_allowed(&entry( + Some(42), + Some("w-10"), + None, + None + ))); + assert!(!sole_source_window_fallback_allowed(&entry( + Some(42), + Some("w-10"), + Some("Documents"), + None + ))); +} + +#[test] +fn unique_matching_index_fails_closed_on_duplicate_matches() { + let values = [1, 2, 3]; + + assert_eq!(unique_matching_index(&values, |value| *value == 2), Some(1)); + assert_eq!(unique_matching_index(&values, |value| *value > 1), None); + assert_eq!(unique_matching_index(&values, |value| *value == 4), None); +} + #[test] fn bounds_hash_keeps_collecting_to_disambiguate_identity_matches() { assert!(!should_stop_collecting(