feat: add query timeout functionality to the server

DataDog · Jan 9, 2025 · dc7d01f · dc7d01f
1 parent 166e94f
commit dc7d01f
Show file tree

Hide file tree

Showing 9 changed files with 69 additions and 20 deletions.
diff --git a/crates/bins/src/bin/datadog-static-analyzer-git-hook.rs b/crates/bins/src/bin/datadog-static-analyzer-git-hook.rs
@@ -43,7 +43,7 @@ use std::io::Write;
 use std::path::PathBuf;
 use std::process::exit;
 use std::sync::Arc;
-use std::time::Instant;
+use std::time::{Duration, Instant};
 use std::{env, fs, io};
 use terminal_emoji::Emoji;
 
@@ -348,6 +348,7 @@ fn main() -> Result<()> {
         .opt_str("rule-timeout-ms")
         .map(|val| {
             val.parse::<u64>()
+                .map(Duration::from_millis)
                 .context("unable to parse `rule-timeout-ms` flag as integer")
         })
         .transpose()?;

diff --git a/crates/bins/src/bin/datadog-static-analyzer.rs b/crates/bins/src/bin/datadog-static-analyzer.rs
@@ -390,6 +390,7 @@ fn main() -> Result<()> {
         .opt_str("rule-timeout-ms")
         .map(|val| {
             val.parse::<u64>()
+                .map(Duration::from_millis)
                 .context("unable to parse `rule-timeout-ms` flag as integer")
         })
         .transpose()?;

diff --git a/crates/bins/src/bin/datadog_static_analyzer_server/cli.rs b/crates/bins/src/bin/datadog_static_analyzer_server/cli.rs
@@ -60,6 +60,13 @@ fn get_opts() -> Options {
         "/tmp/static-analysis-server",
     );
 
+    opts.optopt(
+        "",
+        "rule-timeout-ms",
+        "how long a rule can run before being killed, in milliseconds",
+        "1000",
+    );
+
     opts
 }
 
@@ -92,6 +99,8 @@ pub enum CliError {
     InvalidAddress(String),
     #[error("Invalid log argument: {0}. It must be 'minutely', 'hourly' or 'daily'.")]
     InvalidLogRolling(String),
+    #[error("Invalid timeout argument {0:?}. It must be a number.")]
+    InvalidTimeout(String),
 }
 
 fn try_to_file_appender(
@@ -173,7 +182,19 @@ pub fn prepare_rocket(tx_keep_alive_error: Sender<i32>) -> Result<RocketPreparat
 
     // server state
     tracing::debug!("Preparing the server state and rocket configuration");
-    let mut server_state = ServerState::new(matches.opt_str("s"), matches.opt_present("e"));
+
+    let timeout = matches
+        .opt_str("rule-timeout-ms")
+        .map(|val| {
+            val.parse::<u64>()
+                .map(Duration::from_millis)
+                .map_err(|_| CliError::InvalidTimeout(val))
+        })
+        .transpose()?;
+
+    let mut server_state =
+        ServerState::new(matches.opt_str("s"), matches.opt_present("e"), timeout);
+
     // Create a cache for rules, if requested
     if matches.opt_present("use-rules-cache") {
         RULE_CACHE.set(RuleCache::new()).expect("should init once");

diff --git a/crates/bins/src/bin/datadog_static_analyzer_server/endpoints.rs b/crates/bins/src/bin/datadog_static_analyzer_server/endpoints.rs
@@ -111,10 +111,16 @@ fn languages(span: TraceSpan) -> Value {
 }
 
 #[rocket::post("/analyze", format = "application/json", data = "<request>")]
-async fn analyze(span: TraceSpan, request: Json<AnalysisRequest<ServerRule>>) -> Value {
+async fn analyze(
+    span: TraceSpan,
+    state: &State<ServerState>,
+    request: Json<AnalysisRequest<ServerRule>>,
+) -> Value {
     let _entered = span.enter();
 
-    rocket::tokio::task::spawn_blocking(|| {
+    let timeout = state.rule_timeout_ms;
+
+    rocket::tokio::task::spawn_blocking(move || {
         let pool = RAYON_POOL.get().expect("pool should have been created");
         pool.scope_fifo(|_| {
             thread_local! {
@@ -128,7 +134,7 @@ async fn analyze(span: TraceSpan, request: Json<AnalysisRequest<ServerRule>>) ->
             });
             let request = request.into_inner();
             let (rule_responses, errors) =
-                match cached_analysis_request(runtime_ref, request, RULE_CACHE.get()) {
+                match cached_analysis_request(runtime_ref, request, timeout, RULE_CACHE.get()) {
                     Ok(resp) => (resp, vec![]),
                     Err(err) => (vec![], vec![err.to_string()]),
                 };

diff --git a/crates/bins/src/bin/datadog_static_analyzer_server/rule_cache.rs b/crates/bins/src/bin/datadog_static_analyzer_server/rule_cache.rs
@@ -5,6 +5,7 @@ use server::model::analysis_request::{AnalysisRequest, ServerRule};
 use server::model::analysis_response::RuleResponse;
 use server::request::process_analysis_request;
 use std::sync::Arc;
+use std::time::Duration;
 
 /// A [`ServerRule`] with its corresponding pre-compiled [`RuleInternal`].
 #[derive(Debug)]
@@ -82,6 +83,7 @@ impl RuleCache {
 pub(crate) fn cached_analysis_request(
     runtime: &mut JsRuntime,
     request: AnalysisRequest<ServerRule>,
+    timeout: Option<Duration>,
     cache: Option<&RuleCache>,
 ) -> Result<Vec<RuleResponse>, &'static str> {
     if let Some(cache) = cache {
@@ -113,7 +115,7 @@ pub(crate) fn cached_analysis_request(
             configuration_base64: request.configuration_base64,
             options: request.options,
         };
-        process_analysis_request(req_with_compiled, runtime)
+        process_analysis_request(req_with_compiled, runtime, timeout)
     } else {
         let mut rule_internals = Vec::with_capacity(request.rules.len());
         for server_rule in request.rules {
@@ -131,7 +133,7 @@ pub(crate) fn cached_analysis_request(
             configuration_base64: request.configuration_base64,
             options: request.options,
         };
-        process_analysis_request(req_with_internal, runtime)
+        process_analysis_request(req_with_internal, runtime, timeout)
     }
 }
 
@@ -233,7 +235,7 @@ function visit(captures) {
         for test_case in [None, Some(&RuleCache::new())] {
             let mut rt = v8.new_runtime();
             let rule_responses =
-                cached_analysis_request(&mut rt, req_v1.clone(), test_case).unwrap();
+                cached_analysis_request(&mut rt, req_v1.clone(), None, test_case).unwrap();
             assert_eq!(rule_responses[0].violations[0].0.message, "rule_v1");
             if let Some(cache) = test_case {
                 assert_eq!(
@@ -243,7 +245,7 @@ function visit(captures) {
             }
 
             let rule_responses =
-                cached_analysis_request(&mut rt, req_v2.clone(), test_case).unwrap();
+                cached_analysis_request(&mut rt, req_v2.clone(), None, test_case).unwrap();
             assert_eq!(rule_responses[0].violations[0].0.message, "rule_v2");
             if let Some(cache) = test_case {
                 assert_eq!(

diff --git a/crates/bins/src/bin/datadog_static_analyzer_server/state.rs b/crates/bins/src/bin/datadog_static_analyzer_server/state.rs
@@ -1,21 +1,30 @@
 use super::utils::get_current_timestamp_ms;
-use std::sync::{Arc, RwLock};
+use std::{
+    sync::{Arc, RwLock},
+    time::Duration,
+};
 
 #[derive(Clone)]
 pub struct ServerState {
     pub last_ping_request_timestamp_ms: Arc<RwLock<u128>>,
     pub static_directory: Option<String>,
     pub is_shutdown_enabled: bool,
     pub is_keepalive_enabled: bool,
+    pub rule_timeout_ms: Option<Duration>,
 }
 
 impl ServerState {
-    pub fn new(static_directory: Option<String>, is_shutdown_enabled: bool) -> Self {
+    pub fn new(
+        static_directory: Option<String>,
+        is_shutdown_enabled: bool,
+        timeout: Option<Duration>,
+    ) -> Self {
         Self {
             last_ping_request_timestamp_ms: Arc::new(RwLock::new(get_current_timestamp_ms())),
             static_directory,
             is_shutdown_enabled,
             is_keepalive_enabled: false,
+            rule_timeout_ms: timeout,
         }
     }
 }
diff --git a/crates/common/src/analysis_options.rs b/crates/common/src/analysis_options.rs
@@ -1,3 +1,5 @@
+use std::time::Duration;
+
 use serde::{Deserialize, Serialize};
 
 // Used internally to pass options to the analysis
@@ -6,7 +8,7 @@ pub struct AnalysisOptions {
     pub log_output: bool,
     pub use_debug: bool,
     pub ignore_generated_files: bool,
-    pub timeout: Option<u64>,
+    pub timeout: Option<Duration>,
 }
 
 impl Default for AnalysisOptions {

diff --git a/crates/static-analysis-kernel/src/analysis/analyze.rs b/crates/static-analysis-kernel/src/analysis/analyze.rs
@@ -165,11 +165,7 @@ 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)
-    };
+    let timeout = analysis_option.timeout.or(Some(RULE_EXECUTION_TIMEOUT));
 
     rules
         .into_iter()

diff --git a/crates/static-analysis-server/src/request.rs b/crates/static-analysis-server/src/request.rs
@@ -14,11 +14,13 @@ use kernel::rule_config::RuleConfigProvider;
 use kernel::utils::decode_base64_string;
 use std::borrow::Borrow;
 use std::sync::Arc;
+use std::time::Duration;
 
 #[tracing::instrument(skip_all)]
 pub fn process_analysis_request<T: Borrow<RuleInternal>>(
     request: AnalysisRequest<T>,
     runtime: &mut JsRuntime,
+    timeout: Option<Duration>,
 ) -> Result<Vec<RuleResponse>, &'static str> {
     tracing::debug!("Processing analysis request");
 
@@ -97,7 +99,7 @@ pub fn process_analysis_request<T: Borrow<RuleInternal>>(
                 .map(|o| o.log_output.unwrap_or(false))
                 .unwrap_or(false),
             ignore_generated_files: false,
-            timeout: None,
+            timeout,
         },
     );
 
@@ -125,6 +127,8 @@ pub fn process_analysis_request<T: Borrow<RuleInternal>>(
 
 #[cfg(test)]
 mod tests {
+    use std::time::Duration;
+
     use super::{AnalysisRequest, RuleResponse};
     use crate::constants::{
         ERROR_CHECKSUM_MISMATCH, ERROR_CODE_LANGUAGE_MISMATCH, ERROR_CODE_NOT_BASE64,
@@ -145,6 +149,13 @@ mod tests {
     /// without requiring a `ServerRule` -> `RuleInternal` conversion or an explicitly-created [`JsRuntime`].
     pub fn shorthand_process_req(
         request: AnalysisRequest<ServerRule>,
+    ) -> Result<Vec<RuleResponse>, &'static str> {
+        shorthand_process_req_with_timeout(request, None)
+    }
+
+    pub fn shorthand_process_req_with_timeout(
+        request: AnalysisRequest<ServerRule>,
+        timeout: Option<Duration>,
     ) -> Result<Vec<RuleResponse>, &'static str> {
         let v8 = ddsa_lib::test_utils::cfg_test_v8();
         let mut runtime = v8.new_runtime();
@@ -163,7 +174,7 @@ mod tests {
             configuration_base64: request.configuration_base64,
             options: request.options,
         };
-        super::process_analysis_request(req_with_internal, &mut runtime)
+        super::process_analysis_request(req_with_internal, &mut runtime, timeout)
     }
 
     /// A sample JavaScript rule and tree-sitter query used to assert violation behavior.
@@ -708,7 +719,7 @@ rulesets:
             filename: "myfile.js".to_string(),
             language: Language::JavaScript,
             file_encoding: "utf-8".to_string(),
-            code_base64: encode_base64_string("function foo() { const baz = 1; }=".repeat(10000)),
+            code_base64: encode_base64_string("function foo() { const baz = 1; }".repeat(10000)),
             configuration_base64: None,
             options: Some(AnalysisRequestOptions {
                 use_tree_sitter: None,