Merge pull request #555 from DataDog/amaan.qureshi/STAL-2846

[STAL-2846] feat: add a timeout flag, and use it in the rule execution

Cargo.toml

@@ -39,4 +39,4 @@ sha2 = "0.10.7"
 num_cpus = "1.15.0"
 tracing = "0.1.40"
 uuid = { version = "1.6.1", features = ["v4"] }
-tree-sitter = "0.22.6"
+tree-sitter = "0.24.4"

crates/bins/src/bin/datadog-static-analyzer-git-hook.rs

@@ -159,6 +159,12 @@ fn main() -> Result<()> {
     opts.optflag("t", "include-testing-rules", "include testing rules");
     opts.optflag("", "secrets", "enable secrets detection (BETA)");
     opts.optflag("", "static-analysis", "enable static-analysis");
+    opts.optopt(
+        "",
+        "rule-timeout-ms",
+        "how long a rule can run before being killed, in milliseconds",
+        "1000",
+    );
 
     let matches = match opts.parse(&args[1..]) {
         Ok(m) => m,
@@ -329,10 +335,19 @@ fn main() -> Result<()> {
         print_configuration(&configuration);
     }
 
+    let timeout = matches
+        .opt_str("rule-timeout-ms")
+        .map(|val| {
+            val.parse::<u64>()
+                .context("unable to parse `rule-timeout-ms` flag as integer")
+        })
+        .transpose()?;
+
     let analysis_options = AnalysisOptions {
         log_output: true,
         use_debug,
         ignore_generated_files,
+        timeout,
     };
 
     if should_verify_checksum {

crates/bins/src/bin/datadog-static-analyzer-test-ruleset.rs

@@ -25,6 +25,7 @@ fn test_rule(rule: &Rule, test: &RuleTest) -> Result<String> {
         log_output: true,
         use_debug: true,
         ignore_generated_files: false,
+        timeout: None,
     };
     let rules = vec![rule_internal];
     let analyze_result = analyze(

crates/bins/src/bin/datadog-static-analyzer.rs

@@ -131,6 +131,12 @@ fn main() -> Result<()> {
     );
     // TODO (JF): Remove this when releasing 0.3.8
     opts.optflag("", "ddsa-runtime", "(deprecated)");
+    opts.optopt(
+        "",
+        "rule-timeout-ms",
+        "how long a rule can run before being killed, in milliseconds",
+        "1000",
+    );
 
     let matches = match opts.parse(&args[1..]) {
         Ok(m) => m,
@@ -368,10 +374,19 @@ fn main() -> Result<()> {
     if matches.opt_present("ddsa-runtime") {
         println!("[WARNING] the --ddsa-runtime flag is deprecated and will be removed in the next version");
     }
+    let timeout = matches
+        .opt_str("rule-timeout-ms")
+        .map(|val| {
+            val.parse::<u64>()
+                .context("unable to parse `rule-timeout-ms` flag as integer")
+        })
+        .transpose()?;
+
     let analysis_options = AnalysisOptions {
         log_output: true,
         use_debug,
         ignore_generated_files,
+        timeout,
     };
 
     if should_verify_checksum {

crates/common/src/analysis_options.rs

@@ -6,6 +6,7 @@ pub struct AnalysisOptions {
     pub log_output: bool,
     pub use_debug: bool,
     pub ignore_generated_files: bool,
+    pub timeout: Option<u64>,
 }
 
 impl Default for AnalysisOptions {
@@ -14,6 +15,7 @@ impl Default for AnalysisOptions {
             log_output: false,
             use_debug: false,
             ignore_generated_files: true,
+            timeout: None,
         }
     }
 }

crates/static-analysis-kernel/Cargo.toml

@@ -22,6 +22,7 @@ globset = "0.4.14"
 graphviz-rust = "0.9.0"
 sequence_trie = "0.3.6"
 serde_yaml = "0.9.21"
+streaming-iterator = "0.1.9"
 thiserror = "1.0.59"
 
 [build-dependencies]

crates/static-analysis-kernel/src/analysis/analyze.rs

@@ -16,8 +16,10 @@ use std::collections::HashMap;
 use std::sync::Arc;
 use std::time::{Duration, Instant};
 
-/// The duration an individual execution of `v8` may run before it will be forcefully halted.
-const JAVASCRIPT_EXECUTION_TIMEOUT: Duration = Duration::from_millis(5000);
+/// The duration an individual execution of a rule may run before it will be forcefully halted.
+/// This includes the time it takes for the tree-sitter query to collect its matches, as well as
+/// the time it takes for the JavaScript rule to execute.
+const RULE_EXECUTION_TIMEOUT: Duration = Duration::from_millis(2000);
 
 thread_local! {
     /// A thread-local `JsRuntime`
@@ -200,6 +202,12 @@ where
     let tree = Arc::new(tree);
     let cst_parsing_time = now.elapsed();
 
+    let timeout = if let Some(timeout) = analysis_option.timeout {
+        Some(Duration::from_millis(timeout))
+    } else {
+        Some(RULE_EXECUTION_TIMEOUT)
+    };
+
     rules
         .into_iter()
         .filter(|rule| rule_config.rule_is_enabled(&rule.borrow().name))
@@ -215,7 +223,7 @@ where
                 filename,
                 rule,
                 &rule_config.get_arguments(&rule.name),
-                Some(JAVASCRIPT_EXECUTION_TIMEOUT),
+                timeout,
             );
 
             // NOTE: This is a translation layer to map Result<T, E> to a `RuleResult` struct.
@@ -258,7 +266,8 @@ where
                 Err(err) => {
                     let r_f = format!("{}:{}", rule.name, filename);
                     let (err_kind, execution_error) = match err {
-                        DDSAJsRuntimeError::JavaScriptTimeout { timeout } => {
+                        DDSAJsRuntimeError::JavaScriptTimeout { timeout }
+                        | DDSAJsRuntimeError::TreeSitterTimeout { timeout } => {
                             if analysis_option.use_debug {
                                 eprintln!(
                                     "rule:file {} TIMED OUT ({} ms)",
@@ -1216,6 +1225,7 @@ function visit(node, filename, code) {
             log_output: true,
             use_debug: false,
             ignore_generated_files: false,
+            timeout: None,
         };
         let rule_config_provider = RuleConfigProvider::from_config(
             &parse_config_file(

crates/static-analysis-kernel/src/analysis/ddsa_lib/bridge/query_match.rs

@@ -126,7 +126,7 @@ const ghi = 'hello' + ' world';
         let query = TSQuery::try_new(&tree.language(), query).unwrap();
         let matches = query
             .cursor()
-            .matches(tree.root_node(), text)
+            .matches(tree.root_node(), text, None)
             .collect::<Vec<_>>();
         assert!(query_match_bridge.is_empty());
         assert!(ts_node_bridge.is_empty());
@@ -152,7 +152,7 @@ const alpha = 'bravo';
         let tree = get_tree(text, &Language::JavaScript).unwrap();
         let matches = query
             .cursor()
-            .matches(tree.root_node(), text)
+            .matches(tree.root_node(), text, None)
             .collect::<Vec<_>>();
         query_match_bridge.set_data(scope, matches, &mut ts_node_bridge);
         assert_eq!(get_node_id_at_idx(&query_match_bridge, 0), 3);

crates/static-analysis-kernel/src/analysis/ddsa_lib/common.rs

@@ -15,6 +15,8 @@ pub type NodeId = u32;
 pub enum DDSAJsRuntimeError {
     #[error("{error}")]
     Execution { error: JsError },
+    #[error("Tree-sitter query execution timeout")]
+    TreeSitterTimeout { timeout: Duration },
     #[error("JavaScript execution timeout")]
     JavaScriptTimeout { timeout: Duration },
     #[error("expected `{name}` to exist within the v8 context")]

crates/static-analysis-kernel/src/analysis/ddsa_lib/context/file_go.rs

@@ -8,6 +8,7 @@ use crate::analysis::tree_sitter::get_tree_sitter_language;
 use crate::model::common::Language;
 use deno_core::v8;
 use deno_core::v8::HandleScope;
+use streaming_iterator::StreamingIterator;
 
 /// Structure for the file context that is specific to Go.
 #[derive(Debug)]
@@ -44,8 +45,9 @@ impl FileContextGo {
         // the second capture is the name of the package.
 
         let mut query_cursor = tree_sitter::QueryCursor::new();
-        let query_result = query_cursor.matches(&self.ts_query, tree.root_node(), code.as_bytes());
-        for query_match in query_result {
+        let mut query_result =
+            query_cursor.matches(&self.ts_query, tree.root_node(), code.as_bytes());
+        while let Some(query_match) = query_result.next() {
             let mut package_name: Option<&str> = None;
             let mut package_alias: Option<&str> = None;
 

crates/static-analysis-kernel/src/analysis/ddsa_lib/context/file_js.rs

@@ -4,6 +4,7 @@
 
 use deno_core::v8;
 use deno_core::v8::HandleScope;
+use streaming_iterator::StreamingIterator;
 
 use crate::analysis::ddsa_lib::common::{Class, DDSAJsRuntimeError};
 use crate::analysis::ddsa_lib::js::JSPackageImport;
@@ -165,8 +166,7 @@ impl FileContextJavaScript {
 
         let query_result = query_cursor.matches(&self.query, tree.root_node(), code.as_bytes());
         let imports = query_result
-            .into_iter()
-            .filter_map(|query_match| {
+            .filter_map_deref(|query_match| {
                 let mut name = None;
                 let mut imported_from = None;
                 let mut has_default = false;

crates/static-analysis-kernel/src/analysis/ddsa_lib/context/file_tf.rs

@@ -3,6 +3,7 @@
 // Copyright 2024 Datadog, Inc.
 
 use deno_core::v8::{self, HandleScope};
+use streaming_iterator::StreamingIterator;
 
 use crate::analysis::ddsa_lib::common::{Class, DDSAJsRuntimeError};
 use crate::analysis::ddsa_lib::js;
@@ -55,8 +56,7 @@ impl FileContextTerraform {
         let query_result = query_cursor.matches(&self.query, tree.root_node(), code.as_bytes());
 
         let resources = query_result
-            .into_iter()
-            .map(|query_match| {
+            .map_deref(|query_match| {
                 let mut resource_type = None;
                 let mut resource_name = None;
 

crates/static-analysis-kernel/src/analysis/ddsa_lib/context/root.rs

@@ -118,7 +118,10 @@ impl RootContext {
         /////
 
         let mut depth = 0;
-        while let Some(child) = current_node.child_containing_descendant(node) {
+        while let Some(child) = current_node.child_with_descendant(node) {
+            if child == node {
+                break;
+            }
             // Cache the relationship discovered as a side effect of traversing down from the root.
             // Note that because the `child_containing_descendant` API filter its output, we
             // only end up caching the path directly from the root to this node.

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

@@ -106,6 +106,7 @@ mod tests {
     use crate::model::common::Language;
     use deno_core::v8;
     use std::marker::PhantomData;
+    use streaming_iterator::StreamingIterator;
 
     #[test]
     fn js_properties_canary() {