Merge pull request #122 from kvnxiao/prevent-stack-overflow

Add max recursion depth check to avoid stack overflow during statement evaluation
PaddiM8 · Jul 19, 2023 · 28b6e42 · 28b6e42
2 parents 9b55f89 + 46d8801
commit 28b6e42
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 1 deletion.
diff --git a/kalk/src/errors.rs b/kalk/src/errors.rs
@@ -19,6 +19,7 @@ pub enum KalkError {
     InvalidNumberLiteral(String),
     InvalidOperator,
     InvalidUnit,
+    StackOverflow,
     TimedOut,
     VariableReferencesItself,
     PiecewiseConditionsAreFalse,
@@ -61,6 +62,7 @@ impl ToString for KalkError {
             KalkError::InvalidNumberLiteral(x) => format!("Invalid number literal: '{}'.", x),
             KalkError::InvalidOperator => String::from("Invalid operator."),
             KalkError::InvalidUnit => String::from("Invalid unit."),
+            KalkError::StackOverflow => String::from("Operation recursed too deeply."),
             KalkError::TimedOut => String::from("Operation took too long."),
             KalkError::VariableReferencesItself => String::from("Variable references itself."),
             KalkError::PiecewiseConditionsAreFalse => String::from("All the conditions in the piecewise are false."),

diff --git a/kalk/src/interpreter.rs b/kalk/src/interpreter.rs
@@ -9,6 +9,8 @@ use crate::symbol_table::SymbolTable;
 use crate::{as_number_or_zero, numerical};
 use crate::{float, prelude};
 
+const DEFAULT_MAX_RECURSION_DEPTH: u32 = 128;
+
 pub struct Context<'a> {
     pub symbol_table: &'a mut SymbolTable,
     angle_unit: String,
@@ -20,6 +22,8 @@ pub struct Context<'a> {
     #[cfg(not(target_arch = "wasm32"))]
     start_time: std::time::SystemTime,
     is_approximation: bool,
+    recursion_depth: u32,
+    max_recursion_depth: u32,
 }
 
 impl<'a> Context<'a> {
@@ -40,6 +44,8 @@ impl<'a> Context<'a> {
             #[cfg(not(target_arch = "wasm32"))]
             start_time: std::time::SystemTime::now(),
             is_approximation: false,
+            recursion_depth: 0,
+            max_recursion_depth: DEFAULT_MAX_RECURSION_DEPTH,
         }
     }
 
@@ -129,7 +135,10 @@ pub(crate) fn eval_expr(
         Expr::Boolean(value) => Ok(KalkValue::Boolean(*value)),
         Expr::Group(expr) => eval_group_expr(context, expr, unit),
         Expr::FnCall(identifier, expressions) => {
-            eval_fn_call_expr(context, identifier, expressions, unit)
+            context.recursion_depth += 1;
+            let res = eval_fn_call_expr(context, identifier, expressions, unit);
+            context.recursion_depth -= 1;
+            res
         }
         Expr::Piecewise(pieces) => eval_piecewise(context, pieces, unit),
         Expr::Vector(values) => eval_vector(context, values),
@@ -333,6 +342,10 @@ pub(crate) fn eval_fn_call_expr(
     expressions: &[Expr],
     unit: Option<&String>,
 ) -> Result<KalkValue, KalkError> {
+    if context.recursion_depth > context.max_recursion_depth {
+        return Err(KalkError::StackOverflow);
+    }
+
     if identifier.prime_count > 0 {
         context.is_approximation = true;
     }