diff --git a/crates/obscura-cdp/tests/treewalker_document_order.rs b/crates/obscura-cdp/tests/treewalker_document_order.rs new file mode 100644 index 0000000..2808601 --- /dev/null +++ b/crates/obscura-cdp/tests/treewalker_document_order.rs @@ -0,0 +1,152 @@ +use obscura_cdp::dispatch::{dispatch, CdpContext}; +use obscura_cdp::types::CdpRequest; +use serde_json::{json, Value}; +use tokio::io::{AsyncReadExt, AsyncWriteExt}; +use tokio::net::TcpListener; + +async fn serve_once() -> String { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + tokio::spawn(async move { + let (mut socket, _) = listener.accept().await.unwrap(); + let mut buf = [0u8; 2048]; + let _ = socket.read(&mut buf).await.unwrap(); + let body = ""; + let response = format!( + "HTTP/1.1 200 OK\r\nContent-Type: text/html\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}", + body.len() + ); + let _ = socket.write_all(response.as_bytes()).await; + }); + format!("http://{addr}/") +} + +async fn cdp( + ctx: &mut CdpContext, + id: u64, + method: &str, + params: Value, + session_id: &str, +) -> Value { + let response = dispatch( + &CdpRequest { + id, + method: method.to_string(), + params, + session_id: Some(session_id.to_string()), + }, + ctx, + ) + .await; + assert!( + response.error.is_none(), + "CDP {method} failed: {:?}", + response.error + ); + response.result.unwrap_or_else(|| json!({})) +} + +async fn setup() -> (CdpContext, String) { + std::env::set_var("OBSCURA_ALLOW_PRIVATE_NETWORK", "1"); + let mut ctx = CdpContext::new(); + let page_id = ctx.create_page(); + let session_id = "treewalker-session"; + ctx.sessions.insert(session_id.to_string(), page_id); + cdp( + &mut ctx, + 1, + "Page.navigate", + json!({"url": serve_once().await, "waitUntil": "load"}), + session_id, + ) + .await; + (ctx, session_id.to_string()) +} + +async fn eval(ctx: &mut CdpContext, session_id: &str, expression: &str) -> Value { + cdp( + ctx, + 2, + "Runtime.evaluate", + json!({"expression": expression, "returnByValue": true}), + session_id, + ) + .await +} + +#[tokio::test(flavor = "current_thread")] +async fn next_node_walks_the_whole_subtree_in_document_order() { + let (mut ctx, session_id) = setup().await; + let value = eval( + &mut ctx, + &session_id, + r#"(() => { + document.body.innerHTML = + '

hi

'; + const walker = document.createTreeWalker( + document.getElementById('r'), NodeFilter.SHOW_ELEMENT); + const seen = []; + let node; + while ((node = walker.nextNode())) seen.push(node.tagName); + return JSON.stringify(seen); + })()"#, + ) + .await; + let seen: Vec = + serde_json::from_str(value["result"]["value"].as_str().unwrap()).unwrap(); + assert_eq!(seen, ["SPAN", "P", "SECTION", "A", "B"]); +} + +#[tokio::test(flavor = "current_thread")] +async fn next_node_keeps_searching_after_filtered_leaf_nodes() { + let (mut ctx, session_id) = setup().await; + let value = eval( + &mut ctx, + &session_id, + r#"(() => { + document.body.innerHTML = + '

one

two

three

'; + const walker = document.createTreeWalker( + document.getElementById('r'), NodeFilter.SHOW_TEXT, { + acceptNode(node) { + return node.data === 'two' + ? NodeFilter.FILTER_REJECT + : NodeFilter.FILTER_ACCEPT; + } + }); + const seen = []; + let node; + while ((node = walker.nextNode())) seen.push(node.data); + return JSON.stringify(seen); + })()"#, + ) + .await; + let seen: Vec = + serde_json::from_str(value["result"]["value"].as_str().unwrap()).unwrap(); + assert_eq!(seen, ["one", "three"]); +} + +#[tokio::test(flavor = "current_thread")] +async fn next_node_handles_a_deep_accepted_child_fast_path() { + let (mut ctx, session_id) = setup().await; + let value = eval( + &mut ctx, + &session_id, + r#"(() => { + const root = document.createElement('div'); + let parent = root; + for (let i = 0; i < 5000; i++) { + const child = document.createElement('span'); + parent.appendChild(child); + parent = child; + } + document.body.appendChild(root); + const walker = document.createTreeWalker(root, NodeFilter.SHOW_ELEMENT); + let count = 0; + while (walker.nextNode()) count++; + return count; + })()"#, + ) + .await; + assert_eq!(value["result"]["value"].as_f64(), Some(5000.0)); +} diff --git a/crates/obscura-dom/src/tree.rs b/crates/obscura-dom/src/tree.rs index 4b7369d..725cca7 100644 --- a/crates/obscura-dom/src/tree.rs +++ b/crates/obscura-dom/src/tree.rs @@ -553,6 +553,35 @@ impl DomTree { result } + /// Returns the node after `current` in document order, without leaving the + /// subtree rooted at `root`. + /// + /// Keeping the ancestor climb inside the DOM avoids one JS/native crossing + /// per ancestor when a TreeWalker reaches a deep leaf. + pub fn next_in_subtree(&self, root: NodeId, current: NodeId) -> Option { + let inner = self.inner.borrow(); + let current_node = inner.nodes.get(current.index())?.as_ref()?; + if let Some(child) = current_node.first_child { + return Some(child); + } + + let mut node_id = current; + for _ in 0..=inner.nodes.len() { + if node_id == root { + return None; + } + let node = inner.nodes.get(node_id.index())?.as_ref()?; + if let Some(sibling) = node.next_sibling { + return Some(sibling); + } + node_id = node.parent?; + } + + // Parent cycles are prevented by the mutation APIs. Keep a hard bound + // here as defense in depth for a malformed tree. + None + } + pub fn ancestors(&self, node_id: NodeId) -> Vec { let inner = self.inner.borrow(); let mut result = Vec::new(); @@ -980,6 +1009,24 @@ mod tests { assert_eq!(tree.len(), 1); } + #[test] + fn test_next_in_subtree_follows_document_order_and_stays_within_root() { + let tree = DomTree::new(); + let root = tree.new_node(NodeData::Text { contents: "root".into() }); + let first = tree.new_node(NodeData::Text { contents: "first".into() }); + let nested = tree.new_node(NodeData::Text { contents: "nested".into() }); + let second = tree.new_node(NodeData::Text { contents: "second".into() }); + tree.append_child(tree.document(), root); + tree.append_child(root, first); + tree.append_child(first, nested); + tree.append_child(root, second); + + assert_eq!(tree.next_in_subtree(root, root), Some(first)); + assert_eq!(tree.next_in_subtree(root, first), Some(nested)); + assert_eq!(tree.next_in_subtree(root, nested), Some(second)); + assert_eq!(tree.next_in_subtree(root, second), None); + } + // Builds a chain of `depth` nested
elements under the document and // returns (outermost_div, innermost_div). Depth this large overflows a // recursive tree walk and aborts the process, so it guards the iterative diff --git a/crates/obscura-js/js/bootstrap.js b/crates/obscura-js/js/bootstrap.js index 2fac347..db39028 100644 --- a/crates/obscura-js/js/bootstrap.js +++ b/crates/obscura-js/js/bootstrap.js @@ -2611,18 +2611,10 @@ class Document extends Node { return true; }, nextNode() { - let node = this.currentNode; - let child = node.firstChild; - while (child) { - if (this._accept(child)) { this.currentNode = child; return child; } - if (child.firstChild) { child = child.firstChild; continue; } - if (child.nextSibling) { child = child.nextSibling; continue; } - let parent = child.parentNode; - while (parent && parent !== this.root) { - if (parent.nextSibling) { child = parent.nextSibling; break; } - parent = parent.parentNode; - } - if (!parent || parent === this.root) return null; + let node = _wrap(+_dom("next_in_subtree", this.root._nid, this.currentNode._nid)); + while (node) { + if (this._accept(node)) { this.currentNode = node; return node; } + node = _wrap(+_dom("next_in_subtree", this.root._nid, node._nid)); } return null; }, diff --git a/crates/obscura-js/src/ops.rs b/crates/obscura-js/src/ops.rs index 4815d18..c578ac8 100644 --- a/crates/obscura-js/src/ops.rs +++ b/crates/obscura-js/src/ops.rs @@ -268,6 +268,13 @@ fn op_dom_inner(state: &OpState, cmd: String, arg1: String, arg2: String) -> Str "prev_sibling" => n.prev_sibling, _ => None, }).flatten().map(|id| id.index().to_string()).unwrap_or("-1".into()) } + "next_in_subtree" => { + let root = NodeId::new(arg1.parse::().unwrap_or(0)); + let current = NodeId::new(arg2.parse::().unwrap_or(0)); + dom.next_in_subtree(root, current) + .map(|id| id.index().to_string()) + .unwrap_or("-1".into()) + } "child_nodes" => { let nid = arg1.parse::().unwrap_or(0); let ids: Vec = dom.children(NodeId::new(nid)).iter().map(|id| id.index() as i32).collect();