diff --git a/src/js_runtime.rs b/src/js_runtime.rs index 86f0e0f..b3d5f67 100644 --- a/src/js_runtime.rs +++ b/src/js_runtime.rs @@ -1,8 +1,21 @@ -use crate::config::JsRuntimeConfig; +mod script_termination_reason; + +use crate::{ + config::JsRuntimeConfig, js_runtime::script_termination_reason::ScriptTerminationReason, +}; use anyhow::{bail, Context}; use deno_core::{serde_v8, v8, PollEventLoopOptions, RuntimeOptions}; use serde::{Deserialize, Serialize}; -use std::time::{Duration, Instant}; +use std::{ + sync::{ + atomic::{AtomicBool, AtomicUsize, Ordering}, + Arc, + }, + time::{Duration, Instant}, +}; + +/// Defines a maximum interval on which script is checked for timeout. +const SCRIPT_TIMEOUT_CHECK_INTERVAL: Duration = Duration::from_secs(2); /// An abstraction over the V8/Deno runtime that allows any utilities to execute custom user /// JavaScript scripts. @@ -38,13 +51,30 @@ impl JsRuntime { ) -> Result<(R, Duration), anyhow::Error> { let now = Instant::now(); + let termination_reason = + Arc::new(AtomicUsize::new(ScriptTerminationReason::Unknown as usize)); + let timeout_token = Arc::new(AtomicBool::new(false)); let isolate_handle = self.inner_runtime.v8_isolate().thread_safe_handle(); + + // Track memory usage and terminate execution if threshold is exceeded. + let isolate_handle_clone = isolate_handle.clone(); + let termination_reason_clone = termination_reason.clone(); + let timeout_token_clone = timeout_token.clone(); self.inner_runtime .add_near_heap_limit_callback(move |current_value, _| { log::error!( "Approaching the memory limit of ({current_value}), terminating execution." ); - isolate_handle.terminate_execution(); + + // Define termination reason and terminate execution. + isolate_handle_clone.terminate_execution(); + + timeout_token_clone.swap(true, Ordering::Relaxed); + termination_reason_clone.store( + ScriptTerminationReason::MemoryLimit as usize, + Ordering::Relaxed, + ); + // Give the runtime enough heap to terminate without crashing the process. 5 * current_value }); @@ -65,22 +95,77 @@ impl JsRuntime { .set(scope, context_key.into(), context_value); } + // Track the time the script takes to execute, and terminate execution if threshold is exceeded. + let termination_timeout = self.max_user_script_execution_time; + let termination_reason_clone = termination_reason.clone(); + let timeout_token_clone = timeout_token.clone(); + std::thread::spawn(move || { + let now = Instant::now(); + loop { + // If task is cancelled, return immediately. + if timeout_token_clone.load(Ordering::Relaxed) { + return; + } + + // Otherwise, terminate execution if time is out, or sleep for max `SCRIPT_TIMEOUT_CHECK_INTERVAL`. + let Some(time_left) = termination_timeout.checked_sub(now.elapsed()) else { + termination_reason_clone.store( + ScriptTerminationReason::TimeLimit as usize, + Ordering::Relaxed, + ); + isolate_handle.terminate_execution(); + return; + }; + + std::thread::sleep(std::cmp::min(time_left, SCRIPT_TIMEOUT_CHECK_INTERVAL)); + } + }); + // Retrieve the result `Promise`. - let promise = self + let script_result_promise = self .inner_runtime - .execute_script("<anon>", js_code.into().into())?; + .execute_script("<anon>", js_code.into().into()) + .map_err(|err| { + match ScriptTerminationReason::from(termination_reason.load(Ordering::Relaxed)) { + ScriptTerminationReason::MemoryLimit => { + err.context("Script exceeded memory limit.") + } + ScriptTerminationReason::TimeLimit => { + err.context("Script exceeded time limit.") + } + ScriptTerminationReason::Unknown => err, + } + })?; // Wait for the promise to resolve. - let resolve = self.inner_runtime.resolve(promise); - let out = tokio::time::timeout( - self.max_user_script_execution_time, - self.inner_runtime - .with_event_loop_promise(resolve, PollEventLoopOptions::default()), - ) - .await??; + let resolve = self.inner_runtime.resolve(script_result_promise); + let script_result = self + .inner_runtime + .with_event_loop_promise(resolve, PollEventLoopOptions::default()) + .await + .map_err(|err| { + timeout_token.swap(true, Ordering::Relaxed); + match ScriptTerminationReason::from(termination_reason.load(Ordering::Relaxed)) { + ScriptTerminationReason::MemoryLimit => { + err.context("Script exceeded memory limit.") + } + ScriptTerminationReason::TimeLimit => { + err.context("Script exceeded time limit.") + } + ScriptTerminationReason::Unknown => err, + } + })?; + + // If execution was terminated due to timeout, but script managed to complete execution nevertheless, cancel + // termination. + timeout_token.swap(true, Ordering::Relaxed); + let isolate = self.inner_runtime.v8_isolate(); + if isolate.is_execution_terminating() { + isolate.cancel_terminate_execution(); + } let scope = &mut self.inner_runtime.handle_scope(); - let local = v8::Local::new(scope, out); + let local = v8::Local::new(scope, script_result); serde_v8::from_v8(scope, local) .map(|result| (result, now.elapsed())) .with_context(|| "Error deserializing script result") @@ -93,7 +178,7 @@ pub mod tests { use deno_core::error::JsError; use serde::{Deserialize, Serialize}; - #[tokio::test] + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn can_execute_scripts() -> anyhow::Result<()> { let config = JsRuntimeConfig { max_heap_size_bytes: 10 * 1024 * 1024, @@ -147,46 +232,88 @@ pub mod tests { "Uncaught (in promise) Error: Uh oh." ); - // Limit execution time. + Ok(()) + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn can_limit_execution_time() -> anyhow::Result<()> { + let config = JsRuntimeConfig { + max_heap_size_bytes: 10 * 1024 * 1024, + max_user_script_execution_time: std::time::Duration::from_secs(5), + }; + + let mut runtime = JsRuntime::new(&config); + + // Limit execution time (async). let result = runtime .execute_script::<String>( r#" -(async () => {{ - return new Promise((resolve) => { - Deno.core.queueTimer( - Deno.core.getTimerDepth() + 1, - false, - 10 * 1000, - () => resolve("Done") + (async () => {{ + return new Promise((resolve) => { + Deno.core.queueTimer( + Deno.core.getTimerDepth() + 1, + false, + 10 * 1000, + () => resolve("Done") + ); + }); + }})(); + "#, + None::<()>, + ) + .await + .unwrap_err(); + assert_eq!( + format!("{result}"), + "Script exceeded time limit.".to_string() ); - }); -}})(); -"#, + + // Limit execution time (sync). + let result = runtime + .execute_script::<String>( + r#" + (() => {{ + while (true) {} + }})(); + "#, None::<()>, ) .await - .unwrap_err() - .downcast::<tokio::time::error::Elapsed>()?; - assert_eq!(format!("{result}"), "deadline has elapsed".to_string()); + .unwrap_err(); + assert_eq!( + format!("{result}"), + "Script exceeded time limit.".to_string() + ); + + Ok(()) + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn can_limit_execution_memory() -> anyhow::Result<()> { + let config = JsRuntimeConfig { + max_heap_size_bytes: 10 * 1024 * 1024, + max_user_script_execution_time: std::time::Duration::from_secs(5), + }; + + let mut runtime = JsRuntime::new(&config); // Limit memory usage. let result = runtime .execute_script::<String>( r#" -(async () => {{ - let s = ""; - while(true) { s += "Hello"; } - return "Done"; -}})(); -"#, + (async () => {{ + let s = ""; + while(true) { s += "Hello"; } + return "Done"; + }})(); + "#, None::<()>, ) .await - .unwrap_err() - .downcast::<JsError>()?; + .unwrap_err(); assert_eq!( - result.exception_message, - "Uncaught Error: execution terminated" + format!("{result}"), + "Script exceeded memory limit.".to_string() ); Ok(()) diff --git a/src/js_runtime/script_termination_reason.rs b/src/js_runtime/script_termination_reason.rs new file mode 100644 index 0000000..ed970f3 --- /dev/null +++ b/src/js_runtime/script_termination_reason.rs @@ -0,0 +1,45 @@ +/// Defines the reason why a script was terminated. +#[derive(Debug, Clone, Copy, PartialEq)] +pub enum ScriptTerminationReason { + /// The script was terminated because it hit the memory limit. + MemoryLimit = 0, + /// The script was terminated because it hit the time limit. + TimeLimit = 1, + /// The script was terminated for an unknown reason. + Unknown = 2, +} + +impl From<usize> for ScriptTerminationReason { + fn from(value: usize) -> Self { + match value { + 0 => Self::MemoryLimit, + 1 => Self::TimeLimit, + _ => Self::Unknown, + } + } +} + +#[cfg(test)] +mod tests { + use super::ScriptTerminationReason; + + #[test] + fn conversion() { + assert_eq!( + ScriptTerminationReason::MemoryLimit, + ScriptTerminationReason::from(0) + ); + assert_eq!( + ScriptTerminationReason::TimeLimit, + ScriptTerminationReason::from(1) + ); + assert_eq!( + ScriptTerminationReason::Unknown, + ScriptTerminationReason::from(2) + ); + assert_eq!( + ScriptTerminationReason::Unknown, + ScriptTerminationReason::from(100500) + ); + } +}