From eccdf88b74e8bb776f43eff7ebc26a2247ef6d04 Mon Sep 17 00:00:00 2001
From: PaddiM8 <bakk@tuta.io>
Date: Thu, 20 Jul 2023 21:06:55 +0200
Subject: [PATCH] Configurable max recursion depth

---
 cli/src/main.rs         |  7 +++++++
 kalk/src/analysis.rs    |  2 +-
 kalk/src/interpreter.rs | 12 +++++++++---
 kalk/src/parser.rs      | 16 +++++++++++++++-
 4 files changed, 32 insertions(+), 5 deletions(-)

diff --git a/cli/src/main.rs b/cli/src/main.rs
index 5028b2b..fe052d5 100644
--- a/cli/src/main.rs
+++ b/cli/src/main.rs
@@ -28,6 +28,10 @@ fn main() {
             Flag::new("angle-unit", FlagType::String)
                 .description("Unit used for angles, either rad or deg. This can also be specified using an environment variable with the name 'ANGLE_UNIT'.")
                 .alias("a"),
+        )
+        .flag(
+            Flag::new("max-recursion-depth", FlagType::Int)
+                .description("The maximum allowed recursion depth. This is used to avoid crashes."),
         );
 
     app.run(args);
@@ -54,6 +58,9 @@ fn default_action(context: &Context) {
     let precision = context
         .int_flag("precision")
         .unwrap_or(output::DEFAULT_PRECISION as isize) as u32;
+    if let Ok(max_recursion_depth) = context.int_flag("max-recursion-depth") {
+        parser_context = parser_context.set_max_recursion_depth(max_recursion_depth as u32);
+    }
 
     if let Some(input_file_path) = get_input_file_by_name("default") {
         load_input_file(&input_file_path, precision, &mut parser_context);
diff --git a/kalk/src/analysis.rs b/kalk/src/analysis.rs
index c23db08..d4f8b43 100644
--- a/kalk/src/analysis.rs
+++ b/kalk/src/analysis.rs
@@ -318,7 +318,7 @@ fn analyse_binary(
             (Expr::Var(identifier), right) => {
                 let right = analyse_expr(context, right)?;
                 analyse_var(context, identifier, None, Some(right))
-            },
+            }
             (left, right) => Ok(Expr::Binary(
                 Box::new(analyse_expr(context, left)?),
                 TokenKind::Power,
diff --git a/kalk/src/interpreter.rs b/kalk/src/interpreter.rs
index 520f062..1d1ebc9 100644
--- a/kalk/src/interpreter.rs
+++ b/kalk/src/interpreter.rs
@@ -9,7 +9,7 @@ use crate::symbol_table::SymbolTable;
 use crate::{as_number_or_zero, numerical};
 use crate::{float, prelude};
 
-const DEFAULT_MAX_RECURSION_DEPTH: u32 = 128;
+pub const DEFAULT_MAX_RECURSION_DEPTH: u32 = 256;
 
 pub struct Context<'a> {
     pub symbol_table: &'a mut SymbolTable,
@@ -49,6 +49,12 @@ impl<'a> Context<'a> {
         }
     }
 
+    pub fn set_max_recursion_depth(mut self, depth: u32) -> Self {
+        self.max_recursion_depth = depth;
+
+        self
+    }
+
     pub fn interpret(
         &mut self,
         statements: Vec<Stmt>,
@@ -791,12 +797,12 @@ fn eval_comprehension(
         }
 
         let condition = eval_expr(context, condition, None);
-        match  condition {
+        match condition {
             Ok(KalkValue::Boolean(boolean)) => {
                 if boolean && vars.len() == 1 {
                     values.push(eval_expr(context, left, None));
                 }
-            },
+            }
             Err(err) => values.push(Err(err)),
             _ => (),
         }
diff --git a/kalk/src/parser.rs b/kalk/src/parser.rs
index 1f8e917..a8eb7dd 100644
--- a/kalk/src/parser.rs
+++ b/kalk/src/parser.rs
@@ -32,6 +32,7 @@ pub struct Context {
     unit_decl_base_unit: Option<String>,
     other_radix: Option<u8>,
     current_stmt_start_pos: usize,
+    max_recursion_depth: Option<u32>,
 }
 
 #[wasm_bindgen]
@@ -48,6 +49,7 @@ impl Context {
             unit_decl_base_unit: None,
             other_radix: None,
             current_stmt_start_pos: 0,
+            max_recursion_depth: None,
         };
 
         parse(&mut context, crate::prelude::INIT).unwrap();
@@ -70,6 +72,12 @@ impl Context {
         self
     }
 
+    pub fn set_max_recursion_depth(mut self, depth: u32) -> Self {
+        self.max_recursion_depth = Some(depth);
+
+        self
+    }
+
     #[wasm_bindgen(js_name = evaluate)]
     #[cfg(not(feature = "rug"))]
     pub fn js_eval(&mut self, input: &str) -> Result<Option<CalculationResult>, JsValue> {
@@ -107,6 +115,11 @@ pub fn eval(
         precision,
         context.timeout.map(|timeout| timeout as u128),
     );
+
+    if let Some(max_recursion_depth) = context.max_recursion_depth {
+        interpreter = interpreter.set_max_recursion_depth(max_recursion_depth);
+    }
+
     let result = interpreter.interpret(statements);
     if let Ok(Some(mut num)) = result {
         num.set_radix(context.other_radix.unwrap_or(10));
@@ -644,7 +657,8 @@ fn parse_identifier(context: &mut Context) -> Result<Expr, KalkError> {
         // vector/group, otherwise it's an expression like sqrt4,
         // which should be parsed as a factor, to allow eg. sqrt2x.
         let mut arguments = if match_token(context, TokenKind::OpenBrace)
-            || match_token(context, TokenKind::OpenParenthesis) {
+            || match_token(context, TokenKind::OpenParenthesis)
+        {
             match parse_primary(context)? {
                 Expr::Vector(arguments) => arguments,
                 Expr::Group(argument) => vec![*argument],