diff --git a/CHANGELOG.md b/CHANGELOG.md index 6df0d94..94f109f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Fixed +- Bounded AT-SPI traversals now use concurrent indexed child reads, count both + queued references and failed child lookups against their work budgets, and + enforce a snapshot deadline so wide, malformed, or stalled accessibility + trees cannot allocate or retain unbounded work. + ## [0.4.8] - 2026-08-12 ### Fixed diff --git a/src/atspi_tree.rs b/src/atspi_tree.rs index 81c7ec1..c835cdd 100644 --- a/src/atspi_tree.rs +++ b/src/atspi_tree.rs @@ -10,9 +10,11 @@ use atspi::{ // Direct dependency (p2p feature off) — see Cargo.toml for why we bypass // atspi's "connection" re-export. use atspi_connection::AccessibilityConnection; +use futures_util::{stream, StreamExt}; use schemars::JsonSchema; use serde::Serialize; -use std::collections::VecDeque; +use std::{collections::VecDeque, future::Future, time::Duration}; +use tokio::time::timeout; use zbus::{ fdo::DBusProxy, names::{BusName, UniqueName}, @@ -106,6 +108,15 @@ const DEFAULT_SNAPSHOT_MAX_NODES: usize = 1_000; const HARD_SNAPSHOT_MAX_NODES: usize = 2_000; const DEFAULT_SNAPSHOT_MAX_DEPTH: u32 = 32; const HARD_SNAPSHOT_MAX_DEPTH: u32 = 64; +const CHILD_READ_CONCURRENCY: usize = 16; +const SNAPSHOT_TIMEOUT: Duration = Duration::from_secs(10); +const MAX_DISCOVERY_ROOTS: usize = 256; +const ROOT_MATCH_CHILD_LIMIT: usize = 8; +const MAX_DISCOVERY_CHILD_READS: usize = MAX_DISCOVERY_ROOTS * ROOT_MATCH_CHILD_LIMIT; + +fn snapshot_child_read_budgets(max_nodes: usize) -> (usize, usize, usize) { + (MAX_DISCOVERY_ROOTS, MAX_DISCOVERY_CHILD_READS, max_nodes) +} pub(crate) fn snapshot_limits( requested_max_nodes: Option, @@ -121,9 +132,108 @@ pub(crate) fn snapshot_limits( ) } +struct BoundedTraversal { + queue: VecDeque, + attempted: usize, + max_items: usize, +} + +impl BoundedTraversal { + fn new(max_items: usize) -> Self { + Self { + queue: VecDeque::new(), + attempted: 0, + max_items, + } + } + + fn enqueue(&mut self, items: impl IntoIterator) { + self.queue + .extend(items.into_iter().take(self.remaining_capacity())); + } + + fn pop(&mut self) -> Option { + if self.attempted >= self.max_items { + return None; + } + let item = self.queue.pop_front()?; + self.attempted += 1; + Some(item) + } + + fn remaining_capacity(&self) -> usize { + self.max_items + .saturating_sub(self.attempted.saturating_add(self.queue.len())) + } +} + +fn bounded_child_count(reported: i32, limit: usize) -> usize { + usize::try_from(reported).unwrap_or_default().min(limit) +} + +struct IndexedReadBatch { + items: Vec, + attempted: usize, +} + +impl IndexedReadBatch { + fn all_failed(&self) -> bool { + self.attempted > 0 && self.items.is_empty() + } +} + +async fn fetch_indexed_up_to( + reported: i32, + limit: usize, + remaining_attempts: &mut usize, + fetch: F, +) -> IndexedReadBatch +where + F: Fn(i32) -> Fut, + Fut: Future>, +{ + let attempt_count = bounded_child_count(reported, limit).min(*remaining_attempts); + *remaining_attempts = (*remaining_attempts).saturating_sub(attempt_count); + let end_index = i32::try_from(attempt_count).unwrap_or(i32::MAX); + + let items = stream::iter(0..end_index) + .map(fetch) + .buffered(CHILD_READ_CONCURRENCY) + .filter_map(|result| async move { result.ok() }) + .collect() + .await; + + IndexedReadBatch { + items, + attempted: attempt_count, + } +} + +async fn children_up_to( + proxy: &AccessibleProxy<'_>, + limit: usize, + remaining_attempts: &mut usize, +) -> zbus::Result> { + if limit == 0 || *remaining_attempts == 0 { + return Ok(IndexedReadBatch { + items: Vec::new(), + attempted: 0, + }); + } + + let child_count = proxy.child_count().await?; + Ok( + fetch_indexed_up_to(child_count, limit, remaining_attempts, |index| { + proxy.get_child_at_index(index) + }) + .await, + ) +} + pub async fn list_accessible_apps(limit: usize) -> Result> { let conn = connect().await?; - let roots = registry_children(&conn).await?; + let mut remaining_child_reads = limit; + let roots = registry_children(&conn, limit, &mut remaining_child_reads).await?; let dbus = DBusProxy::new(conn.connection()).await.ok(); let mut apps = Vec::new(); @@ -141,38 +251,73 @@ pub async fn snapshot_tree( target_pid: Option, max_nodes: usize, max_depth: u32, +) -> Result> { + let (max_nodes, max_depth) = snapshot_limits(Some(max_nodes), Some(max_depth)); + timeout( + SNAPSHOT_TIMEOUT, + snapshot_tree_inner( + app_name_or_bundle_identifier, + target_pid, + max_nodes, + max_depth, + ), + ) + .await + .context("AT-SPI snapshot exceeded its 10-second deadline")? +} + +async fn snapshot_tree_inner( + app_name_or_bundle_identifier: Option<&str>, + target_pid: Option, + max_nodes: usize, + max_depth: u32, ) -> Result> { let conn = connect().await?; - let roots = registry_children(&conn).await?; - let selected_roots = - select_roots(&conn, roots, app_name_or_bundle_identifier, target_pid).await; + // App discovery is bounded independently so a tiny requested tree still + // finds a target registered after the first accessibility root. + let (mut remaining_registry_reads, mut remaining_filter_reads, mut remaining_traversal_reads) = + snapshot_child_read_budgets(max_nodes); + let roots = + registry_children(&conn, MAX_DISCOVERY_ROOTS, &mut remaining_registry_reads).await?; + let selected_roots = select_roots( + &conn, + roots, + app_name_or_bundle_identifier, + target_pid, + &mut remaining_filter_reads, + ) + .await; let mut nodes = Vec::new(); - let mut queue = VecDeque::new(); - - for object_ref in selected_roots { - queue.push_back((object_ref, 0_u32, None)); - } + let mut traversal = BoundedTraversal::new(max_nodes); - while let Some((object_ref, depth, parent_index)) = queue.pop_front() { - if nodes.len() >= max_nodes { - break; - } + traversal.enqueue( + selected_roots + .into_iter() + .map(|object_ref| (object_ref, 0_u32, None)), + ); + while let Some((object_ref, depth, parent_index)) = traversal.pop() { let Ok(proxy) = open_accessible(&conn, &object_ref).await else { continue; }; let index = nodes.len() as u32; - let child_refs = if depth < max_depth { - proxy.get_children().await.unwrap_or_default() + let remaining = traversal.remaining_capacity(); + let child_refs = if depth < max_depth && remaining > 0 { + children_up_to(&proxy, remaining, &mut remaining_traversal_reads) + .await + .map(|batch| batch.items) + .unwrap_or_default() } else { Vec::new() }; nodes.push(read_node(&proxy, &object_ref, index, parent_index, depth).await); - for child in child_refs { - queue.push_back((child, depth + 1, Some(index))); - } + traversal.enqueue( + child_refs + .into_iter() + .map(|child| (child, depth + 1, Some(index))), + ); } Ok(nodes) @@ -199,21 +344,22 @@ pub async fn focused_element_summary( target_pid: Option, ) -> Result> { let conn = connect().await?; - let roots = registry_children(&conn).await?; - let selected_roots = select_roots(&conn, roots, None, target_pid).await; - let mut visited = 0_usize; - let mut queue = VecDeque::new(); - - for object_ref in selected_roots { - queue.push_back((object_ref, 0_u32)); - } + let mut remaining_registry_reads = MAX_DISCOVERY_ROOTS; + let roots = + registry_children(&conn, MAX_DISCOVERY_ROOTS, &mut remaining_registry_reads).await?; + let mut remaining_filter_reads = MAX_DISCOVERY_CHILD_READS; + let selected_roots = + select_roots(&conn, roots, None, target_pid, &mut remaining_filter_reads).await; + let mut traversal = BoundedTraversal::new(FOCUS_PROBE_MAX_NODES); + let mut remaining_traversal_reads = FOCUS_PROBE_MAX_NODES; - while let Some((object_ref, depth)) = queue.pop_front() { - if visited >= FOCUS_PROBE_MAX_NODES { - break; - } - visited += 1; + traversal.enqueue( + selected_roots + .into_iter() + .map(|object_ref| (object_ref, 0_u32)), + ); + while let Some((object_ref, depth)) = traversal.pop() { let Ok(proxy) = open_accessible(&conn, &object_ref).await else { continue; }; @@ -230,9 +376,12 @@ pub async fn focused_element_summary( })); } if depth < FOCUS_PROBE_MAX_DEPTH { - for child in proxy.get_children().await.unwrap_or_default() { - queue.push_back((child, depth + 1)); - } + let remaining = traversal.remaining_capacity(); + let children = children_up_to(&proxy, remaining, &mut remaining_traversal_reads) + .await + .map(|batch| batch.items) + .unwrap_or_default(); + traversal.enqueue(children.into_iter().map(|child| (child, depth + 1))); } } @@ -341,14 +490,24 @@ async fn open_accessible<'r>( object_ref.as_accessible_proxy(conn.connection()).await } -async fn registry_children(conn: &AccessibilityConnection) -> Result> { +async fn registry_children( + conn: &AccessibilityConnection, + limit: usize, + remaining_child_reads: &mut usize, +) -> Result> { let root = conn .root_accessible_on_registry() .await .context("failed to open AT-SPI registry root")?; - root.get_children() + let batch = children_up_to(&root, limit, remaining_child_reads) .await - .context("failed to read AT-SPI registry children") + .context("failed to read AT-SPI registry children")?; + if batch.all_failed() { + return Err(anyhow!( + "AT-SPI registry reported children, but every indexed child read failed" + )); + } + Ok(batch.items) } async fn select_roots( @@ -356,6 +515,7 @@ async fn select_roots( roots: Vec, app_name_or_bundle_identifier: Option<&str>, target_pid: Option, + remaining_child_reads: &mut usize, ) -> Vec { let needle = app_name_or_bundle_identifier .map(str::trim) @@ -372,7 +532,7 @@ async fn select_roots( for object_ref in remaining { if object_ref_pid(dbus.as_ref(), &object_ref).await == Some(target_pid) { if let Some(needle) = needle.as_deref() { - if root_matches(conn, &object_ref, needle).await { + if root_matches(conn, &object_ref, needle, remaining_child_reads).await { pid_and_filter_matches.push(object_ref); } else { pid_matches.push(object_ref); @@ -401,7 +561,7 @@ async fn select_roots( let mut selected = Vec::new(); for object_ref in remaining { - if root_matches(conn, &object_ref, needle).await { + if root_matches(conn, &object_ref, needle, remaining_child_reads).await { selected.push(object_ref); } } @@ -413,6 +573,7 @@ async fn root_matches( conn: &AccessibilityConnection, object_ref: &ObjectRefOwned, needle: &str, + remaining_child_reads: &mut usize, ) -> bool { let Ok(proxy) = open_accessible(conn, object_ref).await else { return object_ref_id(object_ref) @@ -424,8 +585,11 @@ async fn root_matches( return true; } - let children = proxy.get_children().await.unwrap_or_default(); - for child_ref in children.into_iter().take(8) { + for child_ref in children_up_to(&proxy, ROOT_MATCH_CHILD_LIMIT, remaining_child_reads) + .await + .map(|batch| batch.items) + .unwrap_or_default() + { let Ok(child_proxy) = open_accessible(conn, &child_ref).await else { continue; }; @@ -779,4 +943,74 @@ mod tests { assert_eq!(snapshot_limits(Some(0), Some(0)), (1, 0)); assert_eq!(snapshot_limits(Some(10_000), Some(128)), (2_000, 64)); } + + #[test] + fn app_discovery_budget_is_independent_of_requested_tree_size() { + assert_eq!(snapshot_child_read_budgets(1), (256, 2_048, 1)); + } + + #[test] + fn traversal_attempts_and_queue_share_one_work_budget() { + let mut traversal = BoundedTraversal::new(4); + traversal.enqueue([1]); + + assert_eq!(traversal.pop(), Some(1)); + traversal.enqueue(2..=10_000); + assert_eq!(traversal.queue, VecDeque::from([2, 3, 4])); + + assert_eq!(traversal.pop(), Some(2)); + traversal.enqueue(5..=10_000); + assert_eq!(traversal.queue, VecDeque::from([3, 4])); + assert_eq!(traversal.pop(), Some(3)); + assert_eq!(traversal.pop(), Some(4)); + assert_eq!(traversal.pop(), None); + assert_eq!(traversal.attempted, 4); + } + + #[test] + fn child_count_is_clamped_before_indexed_reads() { + assert_eq!(bounded_child_count(-1, 4), 0); + assert_eq!(bounded_child_count(3, 4), 3); + assert_eq!(bounded_child_count(i32::MAX, 4), 4); + } + + #[tokio::test] + async fn indexed_child_reads_consume_attempts_even_when_one_fails() { + let calls = std::sync::Arc::new(std::sync::Mutex::new(Vec::new())); + let mut remaining_attempts = 3; + let batch = fetch_indexed_up_to(100, 10, &mut remaining_attempts, { + let calls = calls.clone(); + move |index| { + let calls = calls.clone(); + async move { + calls.lock().unwrap().push(index); + if index == 1 { + Err(()) + } else { + Ok(index) + } + } + } + }) + .await; + + assert_eq!(batch.items, vec![0, 2]); + assert_eq!(batch.attempted, 3); + assert!(!batch.all_failed()); + assert_eq!(*calls.lock().unwrap(), vec![0, 1, 2]); + assert_eq!(remaining_attempts, 0); + + let no_children = fetch_indexed_up_to(100, 10, &mut remaining_attempts, |_| async { + Ok::<_, ()>(99) + }) + .await; + assert!(no_children.items.is_empty()); + assert_eq!(no_children.attempted, 0); + assert!(!no_children.all_failed()); + + let mut failed_attempts = 2; + let all_failed = + fetch_indexed_up_to(2, 2, &mut failed_attempts, |_| async { Err::(()) }).await; + assert!(all_failed.all_failed()); + } }