feat: add recursion limit to parser

Signed-off-by: Anand Krishnamoorthi <anakrish@microsoft.com>

src/parser.rs

@@ -22,6 +22,11 @@ pub struct Parser<'source> {
     future_keywords: BTreeMap<String, Option<Span>>,
     rego_v1: bool,
 
+    // Tracks current expression/comprehension/query nesting to enforce a recursion limit.
+    expr_depth: usize,
+    max_expr_depth: usize,
+    expr_depth_overflow: bool,
+
     // The index of the last expression that was parsed.
     eidx: u32,
     // The index of the last statement that was parsed.
@@ -31,6 +36,7 @@ pub struct Parser<'source> {
 }
 
 const FUTURE_KEYWORDS: [&str; 4] = ["contains", "every", "if", "in"];
+const DEFAULT_MAX_EXPR_DEPTH: usize = 32;
 
 impl<'source> Parser<'source> {
     pub fn new(source: &'source Source) -> Result<Self> {
@@ -44,6 +50,9 @@ impl<'source> Parser<'source> {
             end: 0,
             future_keywords: BTreeMap::new(),
             rego_v1: false,
+            expr_depth: 0,
+            max_expr_depth: DEFAULT_MAX_EXPR_DEPTH,
+            expr_depth_overflow: false,
             eidx: 0,
             sidx: 0,
             qidx: 0,
@@ -94,6 +103,32 @@ impl<'source> Parser<'source> {
         Ok(())
     }
 
+    fn with_expr_depth<T>(&mut self, f: impl FnOnce(&mut Self) -> Result<T>) -> Result<T> {
+        // Increment expression depth.
+        self.expr_depth = self.expr_depth.saturating_add(1);
+        let current_depth = self.expr_depth;
+
+        // Enforce recursion limit.
+        if self.expr_depth > self.max_expr_depth {
+            self.expr_depth = current_depth.saturating_sub(1);
+            self.expr_depth_overflow = true;
+            bail!(self.tok.1.error(&format!(
+                "expression nesting too deep (>{})",
+                self.max_expr_depth
+            )));
+        }
+
+        let res = f(self);
+
+        // Upon return, ensure that expression depth is still current_depth.
+        if self.expr_depth != current_depth {
+            bail!("internal error: expression depth imbalance");
+        }
+
+        self.expr_depth = current_depth.saturating_sub(1);
+        res
+    }
+
     fn expect(&mut self, text: &str, context: &str) -> Result<()> {
         if self.token_text() == text {
             self.next_token()
@@ -392,7 +427,12 @@ impl<'source> Parser<'source> {
                 span.end = self.end;
                 Ok((term, query))
             }
-            Err(_) if self.end == pos => {
+            Err(err) if self.end == pos => {
+                // Propagate depth overflow error if any.
+                if self.expr_depth_overflow {
+                    return Err(err);
+                }
+
                 // No progress was made in parsing the query.
                 // Restore state and try parsing as set, array or object.
                 *self = state;
@@ -418,7 +458,11 @@ impl<'source> Parser<'source> {
                     eidx: self.next_eidx(),
                 })
             }
-            Err(_) if self.end == pos => {
+            Err(err) if self.end == pos => {
+                // Propagate depth overflow error if any.
+                if self.expr_depth_overflow {
+                    return Err(err);
+                }
                 // No progress was made in parsing comprehension.
                 // Parse as array.
                 let mut items = vec![];
@@ -461,11 +505,20 @@ impl<'source> Parser<'source> {
                 });
             }
             Err(err) if self.end != pos => {
+                // Propagate depth overflow error if any.
+                if self.expr_depth_overflow {
+                    return Err(err);
+                }
                 // Some progress was made parsing the set comprehension.
                 // Report errors.
                 return Err(err);
             }
-            _ => (),
+            Err(err) => {
+                // Propagate depth overflow error if any.
+                if self.expr_depth_overflow {
+                    return Err(err);
+                }
+            }
         }
 
         // It could be a set, object or object comprehension.
@@ -519,11 +572,20 @@ impl<'source> Parser<'source> {
                 });
             }
             Err(err) if self.end != pos => {
+                // Propagate depth overflow error if any.
+                if self.expr_depth_overflow {
+                    return Err(err);
+                }
                 // Some progress was made parsing the object comprehension.
                 // Report errors.
                 return Err(err);
             }
-            _ => (),
+            Err(err) => {
+                // Propagate depth overflow error if any.
+                if self.expr_depth_overflow {
+                    return Err(err);
+                }
+            }
         }
 
         // Parse object
@@ -898,11 +960,17 @@ impl<'source> Parser<'source> {
     }
 
     pub fn parse_expr(&mut self) -> Result<Expr> {
-        #[cfg(feature = "rego-extensions")]
-        return self.parse_or_expr();
+        self.with_expr_depth(|this| {
+            #[cfg(feature = "rego-extensions")]
+            {
+                this.parse_or_expr()
+            }
 
-        #[cfg(not(feature = "rego-extensions"))]
-        return self.parse_membership_expr();
+            #[cfg(not(feature = "rego-extensions"))]
+            {
+                this.parse_membership_expr()
+            }
+        })
     }
 
     #[cfg(feature = "rego-extensions")]
@@ -1925,7 +1993,6 @@ impl<'source> Parser<'source> {
     }
 
     fn parse_target_rule(&mut self) -> Result<Option<String>> {
-        // Check if the current token starts a target rule: __target__
         if self.tok.0 == TokenKind::Ident && self.token_text() == "__target__" {
             // Parse __target__
             self.next_token()?;

tests/parser/cases/expressions/recursion.yaml

@@ -0,0 +1,17 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+cases:
+  - note: deep-parens-over-limit
+    rego: |
+      package test
+      x = ((((((((((((((((((((((((((((((((1))))))))))))))))))))))))))))))))
+    error: "expression nesting too deep"
+    skip: true
+
+  - note: compr-mixed-deep
+    rego: |
+      package test
+      x = (1 + ( 2 + ( 3 + ( 4 + ((((((( [ [ 1 + (2 + ((((((((((((((((((((1))))))))))))))))))))) | true ] | true ] )))))))))))))
+    error: "expression nesting too deep"
+