From 3ee339eb6d449630bad251db5e0aa40a0d28c17c Mon Sep 17 00:00:00 2001
From: Euclid Ye <yezhizhenjiakang@gmail.com>
Date: Wed, 18 Jun 2025 22:43:07 +0800
Subject: [PATCH] script::webdriver_handler: Fully implement
 `get_known_element` (#37532)

1. `get_known_element`: Refactor to follow same step of spec (which
reduces unnecessary search) and implement previously missing step 4.2:
If node is stale return error with error code stale element reference.
An element is stale if its node document is not the active document or
if it is not connected.

2. Refactor `find_node_by_unique_id_in_document` to make it not check
error and really return a `Option<DomRoot<Node>>` as the name suggests.
This will greatly reduce duplication when implement [get a known shadow
root](https://w3c.github.io/webdriver/#dfn-get-a-known-shadow-root) soon

Testing: All WebDriver Conformance test after removing two problematic
commits in #37520

Signed-off-by: Euclid Ye <yezhizhenjiakang@gmail.com>
---
 components/script/dom/window.rs         |  9 ++---
 components/script/webdriver_handlers.rs | 47 ++++++++++++++++---------
 2 files changed, 32 insertions(+), 24 deletions(-)

diff --git a/components/script/dom/window.rs b/components/script/dom/window.rs
index ebf0891c0e1..dc3606fa1aa 100644
--- a/components/script/dom/window.rs
+++ b/components/script/dom/window.rs
@@ -1439,14 +1439,11 @@ impl WindowMethods<crate::DomTypeHolder> for Window {
     }
 
     fn WebdriverElement(&self, id: DOMString) -> Option<DomRoot<Element>> {
-        find_node_by_unique_id_in_document(&self.Document(), id.into())
-            .ok()
-            .and_then(Root::downcast)
+        find_node_by_unique_id_in_document(&self.Document(), id.into()).and_then(Root::downcast)
     }
 
     fn WebdriverFrame(&self, id: DOMString) -> Option<DomRoot<Element>> {
         find_node_by_unique_id_in_document(&self.Document(), id.into())
-            .ok()
             .and_then(Root::downcast::<HTMLIFrameElement>)
             .map(Root::upcast::<Element>)
     }
@@ -1457,9 +1454,7 @@ impl WindowMethods<crate::DomTypeHolder> for Window {
     }
 
     fn WebdriverShadowRoot(&self, id: DOMString) -> Option<DomRoot<ShadowRoot>> {
-        find_node_by_unique_id_in_document(&self.Document(), id.into())
-            .ok()
-            .and_then(Root::downcast)
+        find_node_by_unique_id_in_document(&self.Document(), id.into()).and_then(Root::downcast)
     }
 
     // https://drafts.csswg.org/cssom/#dom-window-getcomputedstyle
diff --git a/components/script/webdriver_handlers.rs b/components/script/webdriver_handlers.rs
index e31a79a49e7..2f127d73b79 100644
--- a/components/script/webdriver_handlers.rs
+++ b/components/script/webdriver_handlers.rs
@@ -80,35 +80,48 @@ fn get_known_element(
     let doc = documents
         .find_document(pipeline)
         .expect("webdriver_handlers::Document should exists");
+    // Step 1. If not node reference is known with session, session's current browsing context,
+    // and reference return error with error code no such element.
+    if !ScriptThread::has_node_id(pipeline, &node_id) {
+        // If the node is known, but not found in the document, it is stale.
+        return Err(ErrorStatus::NoSuchElement);
+    }
+    // Step 2.Let node be the result of get a node with session,
+    // session's current browsing context, and reference.
+    let node = find_node_by_unique_id_in_document(&doc, node_id);
+
     // Step 3. If node is not null and node does not implement Element
     // return error with error code no such element.
-    find_node_by_unique_id_in_document(&doc, node_id).and_then(|node| {
-        node.downcast::<Element>()
-            .map(DomRoot::from_ref)
-            .ok_or(ErrorStatus::NoSuchElement)
-    })
+    if let Some(ref node) = node {
+        if !node.is::<Element>() {
+            return Err(ErrorStatus::NoSuchElement);
+        }
+    }
+    // Step 4.1. If node is null return error with error code stale element reference.
+    if node.is_none() {
+        return Err(ErrorStatus::StaleElementReference);
+    }
+    // Step 4.2. If node is stale return error with error code stale element reference.
+    // An element is stale if its node document is not the active document
+    // or if it is not connected.
+    let element = DomRoot::from_ref(node.unwrap().downcast::<Element>().unwrap());
+    if !element.owner_document().is_active() || !element.is_connected() {
+        return Err(ErrorStatus::StaleElementReference);
+    }
+    // Step 5. Return success with data node.
+    Ok(element)
 }
 
 // This is also used by `dom/window.rs`
 pub(crate) fn find_node_by_unique_id_in_document(
     document: &Document,
     node_id: String,
-) -> Result<DomRoot<Node>, ErrorStatus> {
+) -> Option<DomRoot<Node>> {
     let pipeline = document.window().pipeline_id();
-    match document
+    document
         .upcast::<Node>()
         .traverse_preorder(ShadowIncluding::Yes)
         .find(|node| node.unique_id(pipeline) == node_id)
-    {
-        Some(node) => Ok(node),
-        None => {
-            if ScriptThread::has_node_id(pipeline, &node_id) {
-                Err(ErrorStatus::StaleElementReference)
-            } else {
-                Err(ErrorStatus::NoSuchElement)
-            }
-        },
-    }
 }
 
 /// <https://w3c.github.io/webdriver/#dfn-link-text-selector>