Merge pull request #762 from DataDog/julien/K9VULN-8052

[K9VULN-8052] Add new helper functions

Author: juli1

crates/static-analysis-kernel/src/analysis/ddsa_lib/js/ddsa.js

@@ -2,10 +2,10 @@
 // This product includes software developed at Datadog (https://www.datadoghq.com/).
 // Copyright 2024 Datadog, Inc.
 
-import { _findTaintFlows, transpose, vertexId } from "ext:ddsa_lib/flow/graph";
-import { MethodFlow } from "ext:ddsa_lib/flow/java";
-import { SEALED_EMPTY_ARRAY } from "ext:ddsa_lib/utility";
-import { TreeSitterFieldChildNode } from "ext:ddsa_lib/ts_node";
+import {_findTaintFlows, transpose, vertexId} from "ext:ddsa_lib/flow/graph";
+import {MethodFlow} from "ext:ddsa_lib/flow/java";
+import {SEALED_EMPTY_ARRAY} from "ext:ddsa_lib/utility";
+import {TreeSitterFieldChildNode} from "ext:ddsa_lib/ts_node";
 
 const { op_digraph_adjacency_list_to_dot, op_ts_node_named_children, op_ts_node_parent } = Deno.core.ops;
 
@@ -42,6 +42,50 @@ export class DDSA {
         return children;
     }
 
+    /**
+     * Returns an ancestor node that satisfies a given condition, or undefined if none could be found.
+     * @param {TreeSitterNode | TreeSitterFieldChildNode} node
+     * @param {function(TreeSitterNode | TreeSitterFieldChildNode): boolean} predicate
+     * @returns {TreeSitterNode | TreeSitterFieldChildNode | undefined}
+     */
+    findAncestor(node, predicate) {
+        if (node === undefined) {
+            return undefined;
+        }
+        let n = ddsa.getParent(node);
+        while (n !== undefined) {
+            if (predicate(n)) {
+                return n;
+            }
+            n = ddsa.getParent(n);
+        }
+        return undefined;
+    }
+
+    /**
+     * Returns a descendant node that that satisfies a given condition, or undefined if none could be found.
+     * Descendants are searched depth-first: the first child and all its descendants, then the second child and all its descendants, and so on.
+     * @param {TreeSitterNode | TreeSitterFieldChildNode} node
+     * @param {function(TreeSitterNode | TreeSitterFieldChildNode): boolean} predicate
+     * @returns {TreeSitterNode | TreeSitterFieldChildNode | undefined}
+     */
+    findDescendant(node, predicate) {
+        if (predicate === undefined) {
+            return undefined;
+        }
+
+        for (const c of ddsa.getChildren(node)) {
+            if (predicate(c)) {
+                return c;
+            }
+            const r = ddsa.findDescendant(c, predicate)
+            if (r !== undefined) {
+                return r;
+            }
+        }
+        return undefined;
+    }
+
     /**
      * Fetches and returns the provided node's parent in the tree-sitter tree.
      * If the node is the root node of the tree, `undefined` will be returned.

crates/static-analysis-kernel/src/analysis/ddsa_lib/js/ddsa.rs

@@ -1,6 +1,6 @@
 // Unless explicitly stated otherwise all files in this repository are licensed under the Apache License, Version 2.0.
 // This product includes software developed at Datadog (https://www.datadoghq.com/).
-// Copyright 2024 Datadog, Inc.
+// Copyright 2024-2025 Datadog, Inc.
 
 // (NB: There is no need for any Rust business logic here, as `ddsa.js` is purely the user-facing JavaScript API)
 
@@ -12,14 +12,16 @@ mod tests {
     use crate::analysis::ddsa_lib::test_utils::{
         cfg_test_v8, js_class_eq, js_instance_eq, shorthand_execute_rule,
     };
-    use crate::analysis::ddsa_lib::JsRuntime;
     use crate::analysis::tree_sitter::get_tree_sitter_language;
+    use crate::model::common::Language::Python;
 
     #[test]
     fn js_properties_canary() {
         let expected = &[
             // Methods
             "getChildren",
+            "findAncestor",
+            "findDescendant",
             "getParent",
             "getTaintSinks",
             "getTaintSources",
@@ -35,6 +37,48 @@ mod tests {
         compat_helper_op_ts_node_named_children("ddsa.getChildren(node)")
     }
 
+    #[test]
+    fn test_find_descendant() {
+        let mut rt = cfg_test_v8().new_runtime();
+        let text = "print(foo)";
+        let ts_query = "(module)@module";
+        let rule_code = r#"
+function isIdentifier(n) { return n.cstType === "identifier"; }
+
+function visit(query, filename, code) {{
+  const n = query.captures.module;
+  const c = ddsa.findDescendant(n, isIdentifier);
+  console.log(c.text);
+}}
+"#;
+
+        // Then execute the rule that fetches the children of the node.
+        let res =
+            shorthand_execute_rule(&mut rt, Python, ts_query, &rule_code, text, None).unwrap();
+        assert_eq!(res.console_lines[0], "print");
+    }
+
+    #[test]
+    fn test_find_ancestor() {
+        let mut rt = cfg_test_v8().new_runtime();
+        let text = "print(\"foo\")";
+        let ts_query = "(string_content)@sc";
+        let rule_code = r#"
+function isModule(n) { return n.cstType === "module"; }
+
+function visit(query, filename, code) {
+  const n = query.captures.sc;
+  const c = ddsa.findAncestor(n, isModule);
+  console.log(c.text);
+}
+"#;
+
+        // Then execute the rule that fetches the children of the node.
+        let res =
+            shorthand_execute_rule(&mut rt, Python, ts_query, &rule_code, text, None).unwrap();
+        assert_eq!(res.console_lines[0], "print(\"foo\")");
+    }
+
     /// Stella syntax can get named children
     #[test]
     fn op_ts_node_named_children_stella_compat() {