From 93a633750c45942cb607ce446954b9c8314ef010 Mon Sep 17 00:00:00 2001 From: Anand Krishnamoorthi Date: Thu, 18 Dec 2025 10:31:00 -0600 Subject: [PATCH 1/9] feat: Guard against runtime panics in lexer Add guardrails for operations to ensure that they won't panic at runtime. Signed-off-by: Anand Krishnamoorthi --- src/lexer.rs | 41 +++++++++++++++++++++++++++++++++++++---- tests/lexer/mod.rs | 28 ++++++++++++++++++++++++++++ 2 files changed, 65 insertions(+), 4 deletions(-) diff --git a/src/lexer.rs b/src/lexer.rs index e01eb77..6b37e78 100644 --- a/src/lexer.rs +++ b/src/lexer.rs @@ -98,7 +98,13 @@ impl SourceStr { } pub fn text(&self) -> &str { - &self.source.contents()[self.start as usize..self.end as usize] + let start = self.start as usize; + let end = self.end as usize; + // Use safe slicing to avoid panics on malformed spans + self.source + .contents() + .get(start..end) + .unwrap_or("") } pub fn clone_empty(&self) -> SourceStr { @@ -241,7 +247,13 @@ pub struct Span { impl Span { pub fn text(&self) -> &str { - &self.source.contents()[self.start as usize..self.end as usize] + let start = self.start as usize; + let end = self.end as usize; + // Use safe slicing to avoid panics on malformed spans + self.source + .contents() + .get(start..end) + .unwrap_or("") } pub fn source_str(&self) -> SourceStr { @@ -445,7 +457,13 @@ impl<'source> Lexer<'source> { } // Ensure that the number is parsable in Rust. - match serde_json::from_str::(&self.source.contents()[start..end]) { + let num_slice = self + .source + .contents() + .get(start..end) + .ok_or_else(|| self.source.error(self.line, col, "invalid number span"))?; + + match serde_json::from_str::(num_slice) { Ok(_) => (), Err(e) => { let serde_msg = &e.to_string(); @@ -506,6 +524,10 @@ impl<'source> Lexer<'source> { } } let end = self.peek().0; + if end <= start { + // Guard against invalid span that would underflow end - 1 + return Err(self.source.error(line, col, "invalid raw string span")); + } Ok(Token( TokenKind::RawString, Span { @@ -574,8 +596,19 @@ impl<'source> Lexer<'source> { let end = self.peek().0; self.col += (end - start) as u32; + if start == 0 || end <= start { + // Reject invalid spans before slicing/serde to avoid panic + return Err(self.source.error(line, col, "invalid string span")); + } + + let str_slice = self + .source + .contents() + .get(start - 1..end) + .ok_or_else(|| self.source.error(line, col, "invalid string span"))?; + // Ensure that the string is parsable in Rust. - match serde_json::from_str::(&self.source.contents()[start - 1..end]) { + match serde_json::from_str::(str_slice) { Ok(_) => (), Err(e) => { let serde_msg = &e.to_string(); diff --git a/tests/lexer/mod.rs b/tests/lexer/mod.rs index ec8f8c6..d21c946 100644 --- a/tests/lexer/mod.rs +++ b/tests/lexer/mod.rs @@ -224,6 +224,34 @@ fn invalid_line() -> Result<()> { Ok(()) } +#[test] +fn invalid_span_text_fallbacks() -> Result<()> { + let rego = "abc"; + let source = Source::from_contents("case.rego".to_string(), rego.to_string())?; + + let ss = SourceStr::new(source.clone(), 100, 200); + assert_eq!( + ss.text(), + "", + "SourceStr should return fallback for out-of-bounds span" + ); + + let span = Span { + source: source.clone(), + line: 1, + col: 1, + start: 5, + end: 2, + }; + assert_eq!( + span.text(), + "", + "Span should return fallback for malformed span" + ); + + Ok(()) +} + #[test] #[cfg(feature = "std")] fn file_more_than_64_kb_size() -> Result<()> { From fd4bb3081f5d29cc31e09c0d1003f3aa524c704a Mon Sep 17 00:00:00 2001 From: Anand Krishnamoorthi Date: Thu, 18 Dec 2025 11:01:12 -0600 Subject: [PATCH 2/9] feat: Safeguard against panics in parser Signed-off-by: Anand Krishnamoorthi --- src/parser.rs | 47 +++++++++++++++++++------- tests/parser/cases/import/future.yaml | 6 ++++ tests/parser/cases/some/some.vars.yaml | 18 ++++++++++ 3 files changed, 58 insertions(+), 13 deletions(-) diff --git a/src/parser.rs b/src/parser.rs index 61f2fb2..62e96ea 100644 --- a/src/parser.rs +++ b/src/parser.rs @@ -180,19 +180,27 @@ impl<'source> Parser<'source> { fn handle_import_future_keywords(&mut self, comps: &[Span]) -> Result { if comps.len() >= 2 && comps[0].text() == "future" && comps[1].text() == "keywords" { - match comps.len() - 2 { - 1 => self.set_future_keyword(comps[2].text(), &Some(comps[2].clone()))?, + match comps.len().saturating_sub(2) { + 1 if comps.len() >= 3 => { + self.set_future_keyword(comps[2].text(), &Some(comps[2].clone()))? + } 0 => { let span = &comps[1]; for kw in FUTURE_KEYWORDS.iter() { self.set_future_keyword(kw, &Some(span.clone()))?; } } - _ => { + _ if comps.len() >= 4 => { let s = &comps[3]; - return Err(self - .source - .error(s.line, s.col - 1, "invalid future keyword")); + return Err(self.source.error( + s.line, + s.col.saturating_sub(1), + "invalid future keyword", + )); + } + _ => { + let s = &comps[1]; + return Err(self.source.error(s.line, s.col, "invalid future keyword")); } } Ok(true) @@ -1067,16 +1075,14 @@ impl<'source> Parser<'source> { } span.end = self.end; - // Since exprs are discarded, adjust the expression index counter. - self.eidx -= vars.len() as u32; + // Since exprs are discarded, adjust the expression index counter (saturating to avoid underflow). + self.eidx = self.eidx.saturating_sub(vars.len() as u32); return Ok(Literal::SomeVars { span, vars }); } - let (key, value) = match refs.len() { - 2 => (Some(refs[0].clone()), refs[1].clone()), - 1 => (None, refs[0].clone()), - _ => { - let span = &vars[2]; + if refs.len() >= 3 { + // Too many identifiers before `in`. + if let Some(span) = vars.get(2).or_else(|| vars.last()) { return Err(anyhow!( "{}:{}:{} error: encountered `{}` while expecting `in`", span.source.file(), @@ -1085,6 +1091,21 @@ impl<'source> Parser<'source> { span.text() )); } + return Err(anyhow!( + "invalid some-decl: expected `in` after variable names" + )); + } + + let (key, value) = match refs.len() { + 2 => (Some(refs[0].clone()), refs[1].clone()), + 1 => (None, refs[0].clone()), + _ => { + // We always parse at least one identifier before `in`; guard defensively. + // parse_ident rejects `in` when no vars are present, so this is effectively unreachable. + return Err(anyhow!( + "invalid some-decl: expected variable names before `in`" + )); + } }; self.parse_future_keyword("in", false, "while parsing some-decl")?; diff --git a/tests/parser/cases/import/future.yaml b/tests/parser/cases/import/future.yaml index f538d99..4b7a281 100644 --- a/tests/parser/cases/import/future.yaml +++ b/tests/parser/cases/import/future.yaml @@ -163,6 +163,12 @@ cases: num_statements: 0 want_result: {} + - note: invalid-future-component + rego: | + package test + import future.keywords.foo.bar + error: "invalid future keyword" + - note: shadow/1 rego: | package test diff --git a/tests/parser/cases/some/some.vars.yaml b/tests/parser/cases/some/some.vars.yaml index e3a294d..4a9cf9d 100644 --- a/tests/parser/cases/some/some.vars.yaml +++ b/tests/parser/cases/some/some.vars.yaml @@ -78,3 +78,21 @@ cases: some a, 5 } error: encountered `5` while expecting identifier + + - note: no-vars-before-in + rego: | + package test + import future.keywords.in + x = y { + some in xs + } + error: "unexpected keyword `in`" + + - note: too-many-before-in + rego: | + package test + import future.keywords.in + x = y { + some a, b, c in xs + } + error: "encountered `c` while expecting `in`" From 5d0cf95332905eb8b856828a2d4872aa51a4e5bc Mon Sep 17 00:00:00 2001 From: Anand Krishnamoorthi Date: Thu, 18 Dec 2025 12:26:19 -0600 Subject: [PATCH 3/9] feat: add recursion limit to parser Signed-off-by: Anand Krishnamoorthi --- src/parser.rs | 85 +++++++++++++++++-- tests/parser/cases/expressions/recursion.yaml | 17 ++++ 2 files changed, 93 insertions(+), 9 deletions(-) create mode 100644 tests/parser/cases/expressions/recursion.yaml diff --git a/src/parser.rs b/src/parser.rs index 62e96ea..a898e0c 100644 --- a/src/parser.rs +++ b/src/parser.rs @@ -22,6 +22,11 @@ pub struct Parser<'source> { future_keywords: BTreeMap>, 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 { @@ -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(&mut self, f: impl FnOnce(&mut Self) -> Result) -> Result { + // 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 { - #[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> { - // 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()?; diff --git a/tests/parser/cases/expressions/recursion.yaml b/tests/parser/cases/expressions/recursion.yaml new file mode 100644 index 0000000..949b427 --- /dev/null +++ b/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" + From 6bc1249dc807e08bc1fe171c89f0a8048de2a032 Mon Sep 17 00:00:00 2001 From: Anand Krishnamoorthi Date: Fri, 19 Dec 2025 07:53:02 -0600 Subject: [PATCH 4/9] feat: Safeguard lookup use Detect invalid indexes and raise internal errors. Signed-off-by: Anand Krishnamoorthi --- src/compiler/hoist.rs | 130 ++++++++++++++++++------- src/engine.rs | 4 + src/interpreter.rs | 99 +++++++++++-------- src/languages/rego/compiler/core.rs | 12 ++- src/languages/rego/compiler/loops.rs | 51 +++++++--- src/languages/rego/compiler/queries.rs | 15 ++- src/lookup.rs | 104 ++++++++++++++++---- src/scheduler.rs | 15 ++- src/tests/scheduler/analyzer/mod.rs | 8 +- 9 files changed, 319 insertions(+), 119 deletions(-) diff --git a/src/compiler/hoist.rs b/src/compiler/hoist.rs index bba31e0..8f6e392 100644 --- a/src/compiler/hoist.rs +++ b/src/compiler/hoist.rs @@ -13,6 +13,7 @@ use super::destructuring_planner::{ use crate::ast::{Expr, ExprRef, Literal, LiteralStmt, Module, Query, Ref, Rule, RuleHead}; use crate::compiler::context::{ContextType, ScopeContext}; use crate::lookup::Lookup; +use crate::lookup::LookupResult; use crate::scheduler::compute_module_globals; use crate::*; use anyhow::{anyhow, Result}; @@ -147,43 +148,87 @@ impl HoistedLoopsLookup { } /// Store hoisted loops for a statement - pub fn set_statement_loops(&mut self, module_idx: u32, stmt_idx: u32, loops: Vec) { - self.statement_loops.set(module_idx, stmt_idx, loops); + pub fn set_statement_loops( + &mut self, + module_idx: u32, + stmt_idx: u32, + loops: Vec, + ) -> Result<()> { + self.statement_loops + .set_checked(module_idx, stmt_idx, loops) + .map_err(|err| anyhow!("statement_loops out of bounds: {err}")) } /// Get hoisted loops for a statement - pub fn get_statement_loops(&self, module_idx: u32, stmt_idx: u32) -> Option<&Vec> { + pub fn get_statement_loops( + &self, + module_idx: u32, + stmt_idx: u32, + ) -> LookupResult>> { self.statement_loops.get_checked(module_idx, stmt_idx) } /// Store hoisted loops for an expression (output expressions) - pub fn set_expr_loops(&mut self, module_idx: u32, expr_idx: u32, loops: Vec) { - self.expr_loops.set(module_idx, expr_idx, loops); + pub fn set_expr_loops( + &mut self, + module_idx: u32, + expr_idx: u32, + loops: Vec, + ) -> Result<()> { + self.expr_loops + .set_checked(module_idx, expr_idx, loops) + .map_err(|err| anyhow!("expr_loops out of bounds: {err}")) } /// Get hoisted loops for an expression - pub fn get_expr_loops(&self, module_idx: u32, expr_idx: u32) -> Option<&Vec> { + pub fn get_expr_loops( + &self, + module_idx: u32, + expr_idx: u32, + ) -> LookupResult>> { self.expr_loops.get_checked(module_idx, expr_idx) } /// Store the compilation context for a query - pub fn set_query_context(&mut self, module_idx: u32, query_idx: u32, context: ScopeContext) { - self.query_contexts.set(module_idx, query_idx, context); + pub fn set_query_context( + &mut self, + module_idx: u32, + query_idx: u32, + context: ScopeContext, + ) -> Result<()> { + self.query_contexts + .set_checked(module_idx, query_idx, context) + .map_err(|err| anyhow!("query_contexts out of bounds: {err}")) } /// Store a binding plan for an expression - pub fn set_expr_binding_plan(&mut self, module_idx: u32, expr_idx: u32, plan: BindingPlan) { - self.expr_binding_plans.set(module_idx, expr_idx, plan); + pub fn set_expr_binding_plan( + &mut self, + module_idx: u32, + expr_idx: u32, + plan: BindingPlan, + ) -> Result<()> { + self.expr_binding_plans + .set_checked(module_idx, expr_idx, plan) + .map_err(|err| anyhow!("expr_binding_plans out of bounds: {err}")) } /// Get the compilation context for a query #[allow(dead_code)] - pub fn get_query_context(&self, module_idx: u32, query_idx: u32) -> Option<&ScopeContext> { + pub fn get_query_context( + &self, + module_idx: u32, + query_idx: u32, + ) -> LookupResult> { self.query_contexts.get_checked(module_idx, query_idx) } /// Get the binding plan for an expression - pub fn get_expr_binding_plan(&self, module_idx: u32, expr_idx: u32) -> Option<&BindingPlan> { + pub fn get_expr_binding_plan( + &self, + module_idx: u32, + expr_idx: u32, + ) -> LookupResult> { self.expr_binding_plans.get_checked(module_idx, expr_idx) } @@ -275,14 +320,18 @@ impl LoopHoister { Ok(self.lookup) } - fn create_scope_context(&self, module_idx: u32) -> ScopeContext { + fn create_scope_context(&self, module_idx: u32) -> Result { let mut context = ScopeContext::new(); - if let Some(globals) = self.module_globals.get_checked(module_idx, 0) { + if let Some(globals) = self + .module_globals + .get_checked(module_idx, 0) + .map_err(|err| anyhow!("module_globals out of bounds: {err}"))? + { context.module_globals = Some(globals.clone()); } - context + Ok(context) } /// Populate loop hoisting information for all modules, with extra capacity @@ -309,7 +358,8 @@ impl LoopHoister { self.lookup.ensure_expr_capacity(last_module_idx + i, 0); self.module_globals.ensure_capacity(last_module_idx + i, 0); self.module_globals - .set(last_module_idx + i, 0, crate::Rc::new(BTreeSet::new())); + .set_checked(last_module_idx + i, 0, crate::Rc::new(BTreeSet::new())) + .map_err(|err| anyhow!("module_globals out of bounds: {err}"))?; } Ok(self.lookup) } @@ -360,10 +410,11 @@ impl LoopHoister { reserved_globals.insert("data".to_string()); reserved_globals.insert("input".to_string()); self.module_globals - .set(module_idx, 0, crate::Rc::new(reserved_globals)); + .set_checked(module_idx, 0, crate::Rc::new(reserved_globals)) + .map_err(|err| anyhow!("module_globals out of bounds: {err}"))?; // Populate the query with default context - let context = self.create_scope_context(module_idx); + let context = self.create_scope_context(module_idx)?; self.lookup.ensure_query_capacity(module_idx, query.qidx); self.populate_query(module_idx, query, &context)?; Ok(()) @@ -374,7 +425,7 @@ impl LoopHoister { match rule { Rule::Spec { head, bodies, .. } => { // Create a context for this rule - let mut context = self.create_scope_context(module_idx); + let mut context = self.create_scope_context(module_idx)?; // Bind function parameters if this is a function rule if let RuleHead::Func { args, .. } = head { @@ -396,7 +447,7 @@ impl LoopHoister { module_idx, expr_idx, binding_plan, - ); + )?; } Err(err) => return Err(map_binding_error(err)), } @@ -454,7 +505,7 @@ impl LoopHoister { module_idx, body.query.qidx, populated_body_context.clone(), - ); + )?; // Process the key expression if present if let Some(ref key) = key_expr { @@ -497,7 +548,7 @@ impl LoopHoister { } Rule::Default { value, .. } => { // For default rules, just process the value expression - let context = self.create_scope_context(module_idx); + let context = self.create_scope_context(module_idx)?; self.populate_output_expr(module_idx, value, &context)?; } } @@ -518,7 +569,11 @@ impl LoopHoister { // Get the scheduled order if available let stmt_order: Vec = if let Some(ref schedule) = self.schedule { - if let Some(query_schedule) = schedule.queries.get(module_idx, query.qidx) { + if let Some(query_schedule) = schedule + .queries + .get_checked(module_idx, query.qidx) + .map_err(|err| anyhow!("schedule out of bounds: {err}"))? + { query_schedule .order .iter() @@ -566,7 +621,8 @@ impl LoopHoister { } self.lookup.ensure_statement_capacity(module_idx, stmt_idx); - self.lookup.set_statement_loops(module_idx, stmt_idx, loops); + self.lookup + .set_statement_loops(module_idx, stmt_idx, loops)?; Ok(()) } @@ -596,7 +652,7 @@ impl LoopHoister { self.lookup.ensure_expr_capacity(module_idx, expr_idx); Self::bind_vars_from_plan_to_context(&binding_plan, context); self.lookup - .set_expr_binding_plan(module_idx, expr_idx, binding_plan); + .set_expr_binding_plan(module_idx, expr_idx, binding_plan)?; if let Some(key_expr) = key { self.analyze_expr(module_idx, key_expr, context, loops)?; @@ -615,7 +671,7 @@ impl LoopHoister { self.populate_query(module_idx, query.as_ref(), &every_context)?; self.lookup.ensure_query_capacity(module_idx, query.qidx); self.lookup - .set_query_context(module_idx, query.qidx, populated_context); + .set_query_context(module_idx, query.qidx, populated_context)?; } NotExpr { expr, .. } => { self.analyze_expr(module_idx, expr, context, loops)?; @@ -663,7 +719,7 @@ impl LoopHoister { self.populate_query(module_idx, query.as_ref(), &compr_context)?; self.lookup.ensure_query_capacity(module_idx, query.qidx); self.lookup - .set_query_context(module_idx, query.qidx, populated_context.clone()); + .set_query_context(module_idx, query.qidx, populated_context.clone())?; self.populate_output_expr_with_context(module_idx, term, &populated_context)?; } E::ObjectCompr { @@ -678,7 +734,7 @@ impl LoopHoister { self.populate_query(module_idx, query.as_ref(), &compr_context)?; self.lookup.ensure_query_capacity(module_idx, query.qidx); self.lookup - .set_query_context(module_idx, query.qidx, populated_context.clone()); + .set_query_context(module_idx, query.qidx, populated_context.clone())?; self.populate_output_expr_with_context(module_idx, key, &populated_context)?; self.populate_output_expr_with_context(module_idx, value, &populated_context)?; } @@ -721,8 +777,11 @@ impl LoopHoister { // Immediately bind variables from the plan to context Self::bind_vars_from_plan_to_context(&binding_plan, context); - self.lookup - .set_expr_binding_plan(module_idx, expr_idx, binding_plan); + self.lookup.set_expr_binding_plan( + module_idx, + expr_idx, + binding_plan, + )?; } Err(err) => return Err(map_binding_error(err)), } @@ -746,8 +805,11 @@ impl LoopHoister { let expr_idx = index.as_ref().eidx(); self.lookup.ensure_expr_capacity(module_idx, expr_idx); Self::bind_vars_from_plan_to_context(&binding_plan, context); - self.lookup - .set_expr_binding_plan(module_idx, expr_idx, binding_plan); + self.lookup.set_expr_binding_plan( + module_idx, + expr_idx, + binding_plan, + )?; } Err(err) => return Err(map_binding_error(err)), } @@ -780,7 +842,7 @@ impl LoopHoister { self.lookup.ensure_expr_capacity(module_idx, expr_idx); Self::bind_vars_from_plan_to_context(&binding_plan, context); self.lookup - .set_expr_binding_plan(module_idx, expr_idx, binding_plan); + .set_expr_binding_plan(module_idx, expr_idx, binding_plan)?; self.analyze_expr(module_idx, lhs, context, loops)?; self.analyze_expr(module_idx, rhs, context, loops)?; @@ -857,7 +919,7 @@ impl LoopHoister { let expr_idx = expr.as_ref().eidx(); self.lookup.ensure_expr_capacity(module_idx, expr_idx); - self.lookup.set_expr_loops(module_idx, expr_idx, loops); + self.lookup.set_expr_loops(module_idx, expr_idx, loops)?; Ok(()) } diff --git a/src/engine.rs b/src/engine.rs index 3f90f04..c11a212 100644 --- a/src/engine.rs +++ b/src/engine.rs @@ -871,6 +871,8 @@ impl Engine { debug_assert!( query_lookup .get_statement_loops(module_idx, stmt.sidx) + .ok() + .and_then(|entry| entry) .is_some(), "missing hoisted loop entry for query statement index {}", stmt.sidx @@ -895,6 +897,8 @@ impl Engine { debug_assert!( existing_table .get_statement_loops(module_idx, stmt.sidx) + .ok() + .and_then(|entry| entry) .is_some(), "missing hoisted loop entry after merge for module {} stmt {}", module_idx, diff --git a/src/interpreter.rs b/src/interpreter.rs index d0c3bf3..860489b 100644 --- a/src/interpreter.rs +++ b/src/interpreter.rs @@ -359,16 +359,20 @@ impl Interpreter { } // Helper methods for working with ExprLookup - fn set_loop_var_value(&mut self, expr: &ExprRef, value: Value) { + fn set_loop_var_value(&mut self, expr: &ExprRef, value: Value) -> Result<()> { let module_idx = self.current_module_index; let expr_idx = expr.eidx(); - self.loop_var_values.set(module_idx, expr_idx, value); + self.loop_var_values + .set_checked(module_idx, expr_idx, value) + .map_err(|err| anyhow!("internal error: loop var indices out of bounds: {err}")) } - fn get_loop_var_value(&self, expr: &ExprRef) -> Option<&Value> { + fn get_loop_var_value(&self, expr: &ExprRef) -> Result> { let module_idx = self.current_module_index; let expr_idx = expr.eidx(); - self.loop_var_values.get(module_idx, expr_idx) + self.loop_var_values + .get_checked(module_idx, expr_idx) + .map_err(|err| anyhow!("internal error: loop var indices out of bounds: {err}")) } fn remove_loop_var_value(&mut self, expr: &ExprRef) { @@ -406,12 +410,13 @@ impl Interpreter { if let Some(last_param) = params.last() { let module_idx = self.current_module_index; let expr_idx = last_param.as_ref().eidx(); - return match self + let binding_plan = self .compiled_policy .loop_hoisting_table .get_expr_binding_plan(module_idx, expr_idx) - .cloned() - { + .map_err(|err| anyhow!("loop hoisting table out of bounds: {err}"))?; + + return match binding_plan.cloned() { Some(BindingPlan::Parameter { destructuring_plan, .. }) => Ok(Some(destructuring_plan)), @@ -483,7 +488,7 @@ impl Interpreter { // Collect a chaing of '.field' or '["field"]' let mut path = vec![]; loop { - if let Some(v) = self.get_loop_var_value(expr) { + if let Some(v) = self.get_loop_var_value(expr)? { path.reverse(); return Ok(Self::get_value_chained(v.clone(), &path[..])); } @@ -876,15 +881,17 @@ impl Interpreter { let module_idx = self.current_module_index; let expr_idx = collection.as_ref().eidx(); + let binding_plan = self + .compiled_policy + .loop_hoisting_table + .get_expr_binding_plan(module_idx, expr_idx) + .map_err(|err| anyhow!("loop hoisting table out of bounds: {err}"))?; + let Some(BindingPlan::SomeIn { key_plan, value_plan, .. - }) = self - .compiled_policy - .loop_hoisting_table - .get_expr_binding_plan(module_idx, expr_idx) - .cloned() + }) = binding_plan.cloned() else { bail!("internal error: missing binding plan for some..in expression"); }; @@ -1367,7 +1374,7 @@ impl Interpreter { match loop_value { Value::Array(items) => { for item in items.iter() { - self.set_loop_var_value(loop_target_expr, item.clone()); + self.set_loop_var_value(loop_target_expr, item.clone())?; if self.execute_destructuring_plan(&walk_plan, item)? == Value::from(true) @@ -1417,12 +1424,13 @@ impl Interpreter { let index_plan = if let Some(index) = index_expr { let module_idx = self.current_module_index; let expr_idx = index.as_ref().eidx(); - match self + let plan = self .compiled_policy .loop_hoisting_table .get_expr_binding_plan(module_idx, expr_idx) - .cloned() - { + .map_err(|err| anyhow!("loop hoisting table out of bounds: {err}"))?; + + match plan.cloned() { Some(BindingPlan::LoopIndex { destructuring_plan, .. }) => destructuring_plan, @@ -1445,7 +1453,7 @@ impl Interpreter { match loop_value { Value::Array(items) => { for (idx, v) in items.iter().enumerate() { - self.set_loop_var_value(loop_target_expr, v.clone()); + self.set_loop_var_value(loop_target_expr, v.clone())?; if self.execute_destructuring_plan(&index_plan, &Value::from(idx))? == Value::from(true) @@ -1466,7 +1474,7 @@ impl Interpreter { } Value::Set(items) => { for v in items.iter() { - self.set_loop_var_value(loop_target_expr, v.clone()); + self.set_loop_var_value(loop_target_expr, v.clone())?; // For sets, index is also the value. if self.execute_destructuring_plan(&index_plan, v)? == Value::from(true) { @@ -1485,7 +1493,7 @@ impl Interpreter { } Value::Object(obj) => { for (k, v) in obj.iter() { - self.set_loop_var_value(loop_target_expr, v.clone()); + self.set_loop_var_value(loop_target_expr, v.clone())?; // For objects, index is key. if self.execute_destructuring_plan(&index_plan, k)? == Value::from(true) { result = self.eval_stmts_in_loop(stmts, &loops[1..])? || result; @@ -1830,19 +1838,19 @@ impl Interpreter { match self.eval_expr(Self::loop_collection_expr(loop_info))? { Value::Array(items) => { for v in items.iter() { - self.set_loop_var_value(loop_target_expr, v.clone()); + self.set_loop_var_value(loop_target_expr, v.clone())?; result = self.eval_output_expr_in_loop(&loops[1..])? || result; } } Value::Set(items) => { for v in items.iter() { - self.set_loop_var_value(loop_target_expr, v.clone()); + self.set_loop_var_value(loop_target_expr, v.clone())?; result = self.eval_output_expr_in_loop(&loops[1..])? || result; } } Value::Object(obj) => { for (_, v) in obj.iter() { - self.set_loop_var_value(loop_target_expr, v.clone()); + self.set_loop_var_value(loop_target_expr, v.clone())?; result = self.eval_output_expr_in_loop(&loops[1..])? || result; } } @@ -1881,9 +1889,10 @@ impl Interpreter { .compiled_policy .loop_hoisting_table .get_expr_loops(self.current_module_index, ke.as_ref().eidx()) + .map_err(|err| anyhow!("loop hoisting table out of bounds: {err}"))? { Some(hoisted_loops) => { - loops.extend(hoisted_loops.iter().cloned()); + loops.extend(hoisted_loops.clone()); } None => { bail!(ke.span().error("Loop hoisting information not found for key expression. This is likely a bug in the compilation phase.")); @@ -1897,9 +1906,10 @@ impl Interpreter { .compiled_policy .loop_hoisting_table .get_expr_loops(self.current_module_index, oe.as_ref().eidx()) + .map_err(|err| anyhow!("loop hoisting table out of bounds: {err}"))? { Some(hoisted_loops) => { - loops.extend(hoisted_loops.iter().cloned()); + loops.extend(hoisted_loops.clone()); } None => { bail!(oe.span().error("Loop hoisting information not found for output expression. This is likely a bug in the compilation phase.")); @@ -1927,10 +1937,11 @@ impl Interpreter { } // Get pre-computed hoisted loops from compilation phase - let loop_exprs = match self + let loop_exprs: Vec = match self .compiled_policy .loop_hoisting_table .get_statement_loops(self.current_module_index, stmt.sidx) + .map_err(|err| anyhow!("loop hoisting table out of bounds: {err}"))? { Some(hoisted_loops) => { // Use pre-computed loops from compilation phase @@ -1994,11 +2005,14 @@ impl Interpreter { let query_module_index = self.compiled_policy.modules.len() as u32; if self.current_module_index == query_module_index { // Use query schedule for the current module - match self - .query_schedule - .as_ref() - .and_then(|s| s.queries.get(query_module_index, query.qidx)) - { + let schedule = match self.query_schedule.as_ref() { + Some(s) => s + .queries + .get_checked(query_module_index, query.qidx) + .map_err(|err| anyhow!("schedule out of bounds: {err}"))?, + None => None, + }; + match schedule { Some(schedule) => Some(&schedule.order), None => { if self.query_schedule.is_some() { @@ -2011,12 +2025,15 @@ impl Interpreter { } } else { // Use compiled policy schedule for other modules - match self - .compiled_policy - .schedule - .as_ref() - .and_then(|s| s.queries.get(self.current_module_index, query.qidx)) - { + let schedule = match self.compiled_policy.schedule.as_ref() { + Some(s) => s + .queries + .get_checked(self.current_module_index, query.qidx) + .map_err(|err| anyhow!("schedule out of bounds: {err}"))?, + None => None, + }; + + match schedule { Some(schedule) => Some(&schedule.order), None => { if self.compiled_policy.schedule.is_some() { @@ -2349,7 +2366,7 @@ impl Interpreter { params: &[ExprRef], ) -> Result { // Return generated values of walk builtin. - if let Some(v) = self.get_loop_var_value(expr) { + if let Some(v) = self.get_loop_var_value(expr)? { return Ok(v.clone()); } @@ -2497,6 +2514,7 @@ impl Interpreter { .compiled_policy .loop_hoisting_table .get_expr_binding_plan(module_idx, expr_idx) + .map_err(|err| anyhow!("loop hoisting table out of bounds: {err}"))? .cloned() { // Execute the destructuring plan with the parameter value @@ -2642,6 +2660,7 @@ impl Interpreter { .compiled_policy .loop_hoisting_table .get_expr_binding_plan(module_idx, expr_idx) + .map_err(|err| anyhow!("loop hoisting table out of bounds: {err}"))? .cloned() { // Execute the destructuring plan with the return value @@ -2933,6 +2952,7 @@ impl Interpreter { .compiled_policy .loop_hoisting_table .get_expr_binding_plan(module_idx, expr_idx) + .map_err(|err| anyhow!("loop hoisting table out of bounds: {err}"))? .cloned() .ok_or_else(|| { expr.span().error( @@ -3644,7 +3664,8 @@ impl Interpreter { if let Some(ref self_schedule) = &self.query_schedule { if let Some(query_schedule) = self_schedule .queries - .get(current_module_idx, current_query_idx) + .get_checked(current_module_idx, current_query_idx) + .map_err(|err| anyhow!("schedule out of bounds: {err}"))? { for idx in 0..results.result.len() { let e = Expression { diff --git a/src/languages/rego/compiler/core.rs b/src/languages/rego/compiler/core.rs index 5d0d72a..4718d4f 100644 --- a/src/languages/rego/compiler/core.rs +++ b/src/languages/rego/compiler/core.rs @@ -186,14 +186,20 @@ impl<'a> Compiler<'a> { self.lookup_local_var(var_name) } - pub(super) fn get_binding_plan_for_expr(&self, expr: &ExprRef) -> Option { + pub(super) fn get_binding_plan_for_expr(&self, expr: &ExprRef) -> Result> { let module_idx = self.current_module_index; let expr_idx = expr.as_ref().eidx(); self.policy .inner .loop_hoisting_table .get_expr_binding_plan(module_idx, expr_idx) - .cloned() + .map_err(|err| { + CompilerError::General { + message: format!("loop hoisting table out of bounds: {err}"), + } + .at(expr.span()) + }) + .map(|plan: Option<&BindingPlan>| plan.cloned()) } pub(super) fn expect_binding_plan_for_expr( @@ -201,7 +207,7 @@ impl<'a> Compiler<'a> { expr: &ExprRef, context: &str, ) -> Result { - self.get_binding_plan_for_expr(expr).ok_or_else(|| { + self.get_binding_plan_for_expr(expr)?.ok_or_else(|| { CompilerError::MissingBindingPlan { context: context.to_string(), } diff --git a/src/languages/rego/compiler/loops.rs b/src/languages/rego/compiler/loops.rs index f02300c..4809627 100644 --- a/src/languages/rego/compiler/loops.rs +++ b/src/languages/rego/compiler/loops.rs @@ -15,31 +15,54 @@ use alloc::vec::Vec; impl<'a> Compiler<'a> { pub(super) fn get_statement_loops(&self, stmt: &LiteralStmt) -> Result> { - self.policy + let loops = self + .policy .inner .loop_hoisting_table .get_statement_loops(self.current_module_index, stmt.sidx) - .cloned() - .ok_or_else(|| { + .map_err(|err| { CompilerError::General { - message: format!( - "missing loop hoisting data for statement at {}:{}", - stmt.span.line, stmt.span.col - ), + message: format!("loop hoisting table out of bounds: {err}"), } .at(&stmt.span) - }) + })?; + + loops.cloned().ok_or_else(|| { + CompilerError::General { + message: format!( + "missing loop hoisting data for statement at {}:{}", + stmt.span.line, stmt.span.col + ), + } + .at(&stmt.span) + }) } - pub(super) fn get_expr_loops(&self, expr: &ExprRef) -> Vec { + pub(super) fn get_expr_loops(&self, expr: &ExprRef) -> Result> { let module_idx = self.current_module_index; let expr_idx = expr.as_ref().eidx(); - self.policy + let loops = self + .policy .inner .loop_hoisting_table .get_expr_loops(module_idx, expr_idx) - .cloned() - .unwrap_or_default() + .map_err(|err| { + CompilerError::General { + message: format!("loop hoisting table out of bounds: {err}"), + } + .at(expr.span()) + })?; + + loops.cloned().ok_or_else(|| { + CompilerError::General { + message: format!( + "missing loop hoisting data for expression at {}:{}", + expr.span().line, + expr.span().col + ), + } + .at(expr.span()) + }) } pub(super) fn compile_hoisted_loops( @@ -192,7 +215,7 @@ impl<'a> Compiler<'a> { let mut key_binding_plan: Option<(BindingPlan, Span)> = None; if let Some(key_var) = key_var { - if let Some(binding_plan) = self.get_binding_plan_for_expr(key_var) { + if let Some(binding_plan) = self.get_binding_plan_for_expr(key_var)? { if let BindingPlan::LoopIndex { .. } = &binding_plan { key_binding_plan = Some((binding_plan, key_var.span().clone())); } else { @@ -327,7 +350,7 @@ impl<'a> Compiler<'a> { let body_start = self.program.instructions.len() as u16; - if let Some(binding_plan) = self.get_binding_plan_for_expr(collection) { + if let Some(binding_plan) = self.get_binding_plan_for_expr(collection)? { if let BindingPlan::SomeIn { key_plan, value_plan, diff --git a/src/languages/rego/compiler/queries.rs b/src/languages/rego/compiler/queries.rs index a6990a5..c8c3cac 100644 --- a/src/languages/rego/compiler/queries.rs +++ b/src/languages/rego/compiler/queries.rs @@ -5,6 +5,7 @@ use super::{Compiler, CompilerError, ComprehensionType, ContextType, Result}; use crate::ast::{self, LiteralStmt, Query}; use crate::rvm::program::RuleType; use crate::rvm::Instruction; +use alloc::format; use alloc::vec::Vec; impl<'a> Compiler<'a> { @@ -13,7 +14,15 @@ impl<'a> Compiler<'a> { let result = { let schedule = match &self.policy.inner.schedule { - Some(s) => s.queries.get(self.current_module_index, query.qidx), + Some(s) => s + .queries + .get_checked(self.current_module_index, query.qidx) + .map_err(|err| { + CompilerError::General { + message: format!("schedule out of bounds: {err}"), + } + .at(&query.span) + })?, None => None, }; @@ -94,11 +103,11 @@ impl<'a> Compiler<'a> { let mut key_value_loops = Vec::new(); if let Some(expr) = key_expr.as_ref() { - key_value_loops.extend(self.get_expr_loops(expr)); + key_value_loops.extend(self.get_expr_loops(expr)?); } if let Some(expr) = value_expr.as_ref() { - key_value_loops.extend(self.get_expr_loops(expr)); + key_value_loops.extend(self.get_expr_loops(expr)?); } if !key_value_loops.is_empty() { diff --git a/src/lookup.rs b/src/lookup.rs index 7a8a234..0cdafef 100644 --- a/src/lookup.rs +++ b/src/lookup.rs @@ -7,6 +7,49 @@ //! with expressions, queries, and statements using their respective indices. use crate::*; +use core::fmt; + +/// Error indicating that lookup indices are out of bounds. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum LookupIndexError { + /// The requested module index exceeds the available modules. + ModuleOutOfBounds { module_idx: u32, modules: usize }, + /// The requested node index exceeds the available nodes for the module. + NodeOutOfBounds { + module_idx: u32, + node_idx: u32, + nodes: usize, + }, +} + +impl fmt::Display for LookupIndexError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + LookupIndexError::ModuleOutOfBounds { + module_idx, + modules, + } => { + write!( + f, + "module_idx {module_idx} out of bounds (modules={modules})" + ) + } + LookupIndexError::NodeOutOfBounds { + module_idx, + node_idx, + nodes, + } => write!( + f, + "node_idx {node_idx} out of bounds for module {module_idx} (nodes={nodes})" + ), + } + } +} + +#[cfg(feature = "std")] +impl std::error::Error for LookupIndexError {} + +pub type LookupResult = core::result::Result; /// Generic lookup table that stores data indexed by module and node indices. #[derive(Debug, Clone)] @@ -43,28 +86,19 @@ impl Lookup { } } - /// Set data using direct indices (bounds assumed to be ensured). - pub fn set(&mut self, module_idx: u32, node_idx: u32, value: T) { - *self.get_mut(module_idx, node_idx) = Some(value); - } - - /// Get data using direct indices (bounds assumed to be ensured). - pub fn get(&self, module_idx: u32, node_idx: u32) -> Option<&T> { - self.slots[module_idx as usize][node_idx as usize].as_ref() + /// Set data using direct indices with bounds checking. + /// Returns Ok(()) if written, Err if either index is out of bounds. + pub fn set_checked(&mut self, module_idx: u32, node_idx: u32, value: T) -> LookupResult<()> { + let (m, n) = self.validate_indices(module_idx, node_idx)?; + self.slots[m][n] = Some(value); + Ok(()) } /// Get data using direct indices with bounds checking. - /// Returns None if the module or node index is out of range or unset. - pub fn get_checked(&self, module_idx: u32, node_idx: u32) -> Option<&T> { - self.slots - .get(module_idx as usize) - .and_then(|module| module.get(node_idx as usize)) - .and_then(|slot| slot.as_ref()) - } - - /// Get mutable reference to data using direct indices (bounds assumed to be ensured). - pub fn get_mut(&mut self, module_idx: u32, node_idx: u32) -> &mut Option { - &mut self.slots[module_idx as usize][node_idx as usize] + /// Returns Ok(None) if the entry is unset, Err if indices are out of range. + pub fn get_checked(&self, module_idx: u32, node_idx: u32) -> LookupResult> { + let (m, n) = self.validate_indices(module_idx, node_idx)?; + Ok(self.slots[m][n].as_ref()) } /// Clear data at the given indices by setting it to None. @@ -95,4 +129,36 @@ impl Lookup { None } } + + /// Validate indices and return them as usize on success. + fn validate_indices(&self, module_idx: u32, node_idx: u32) -> LookupResult<(usize, usize)> { + let m = module_idx as usize; + if m >= self.slots.len() { + debug_assert!( + m < self.slots.len(), + "module_idx {m} out of bounds (modules={})", + self.slots.len() + ); + return Err(LookupIndexError::ModuleOutOfBounds { + module_idx, + modules: self.slots.len(), + }); + } + + let n = node_idx as usize; + if n >= self.slots[m].len() { + debug_assert!( + n < self.slots[m].len(), + "node_idx {n} out of bounds for module {m} (nodes={})", + self.slots[m].len() + ); + return Err(LookupIndexError::NodeOutOfBounds { + module_idx, + node_idx, + nodes: self.slots[m].len(), + }); + } + + Ok((m, n)) + } } diff --git a/src/scheduler.rs b/src/scheduler.rs index 888228c..c4722df 100644 --- a/src/scheduler.rs +++ b/src/scheduler.rs @@ -17,7 +17,7 @@ use alloc::string::String; use core::cmp; use core::fmt; -use anyhow::{bail, Result}; +use anyhow::{anyhow, bail, Result}; #[derive(Debug)] pub struct Definition { @@ -584,7 +584,8 @@ impl Analyzer { Expr::ArrayCompr { query, term, .. } | Expr::SetCompr { query, term, .. } => { self.analyze_query(None, Some(term.clone()), query, Scope::default())?; self.schedule_table - .get(self.current_module_index, query.qidx) + .get_checked(self.current_module_index, query.qidx) + .map_err(|err| anyhow!("schedule_table out of bounds: {err}"))? .map(|qs| &qs.scope) } Expr::ObjectCompr { @@ -597,7 +598,8 @@ impl Analyzer { Scope::default(), )?; self.schedule_table - .get(self.current_module_index, query.qidx) + .get_checked(self.current_module_index, query.qidx) + .map_err(|err| anyhow!("schedule_table out of bounds: {err}"))? .map(|qs| &qs.scope) } _ => break, @@ -1099,7 +1101,8 @@ impl Analyzer { order, }; self.schedule_table - .set(self.current_module_index, query.qidx, query_schedule); + .set_checked(self.current_module_index, query.qidx, query_schedule) + .map_err(|err| anyhow!("schedule_table out of bounds: {err}"))?; // Propagate input usage to parent scopes if scope.uses_input && !self.scopes.is_empty() { @@ -1174,7 +1177,9 @@ pub fn compute_module_globals( } result.ensure_capacity(module_idx as u32, 0); - result.set(module_idx as u32, 0, module_globals); + result + .set_checked(module_idx as u32, 0, module_globals) + .map_err(|err| anyhow!("module globals out of bounds: {err}"))?; } Ok(result) diff --git a/src/tests/scheduler/analyzer/mod.rs b/src/tests/scheduler/analyzer/mod.rs index 1e24e00..90aa009 100644 --- a/src/tests/scheduler/analyzer/mod.rs +++ b/src/tests/scheduler/analyzer/mod.rs @@ -3,7 +3,7 @@ use crate::*; use crate::{ast::*, lexer::*, parser::*, scheduler::*}; -use anyhow::{bail, Result}; +use anyhow::{anyhow, bail, Result}; use serde::{Deserialize, Serialize}; use test_generator::test_resources; @@ -68,7 +68,11 @@ fn analyze_file(regos: &[String], expected_scopes: &[Scope]) -> Result<()> { let mut scopes = Vec::new(); for (module_idx, qidx, query) in all_queries.iter() { // Find the corresponding query schedule - if let Some(query_schedule) = schedule.queries.get(*module_idx, *qidx) { + if let Some(query_schedule) = schedule + .queries + .get_checked(*module_idx, *qidx) + .map_err(|err| anyhow!("schedule out of bounds: {err}"))? + { scopes.push((query.clone(), &query_schedule.scope)); } } From 70f63a09829a165d28418fcffc799d07f8910b91 Mon Sep 17 00:00:00 2001 From: Anand Krishnamoorthi Date: Fri, 19 Dec 2025 11:01:02 -0600 Subject: [PATCH 5/9] fix: Avoid unrap/expect in context management Signed-off-by: Anand Krishnamoorthi --- src/interpreter.rs | 33 ++++++++++++++++++++++++--------- 1 file changed, 24 insertions(+), 9 deletions(-) diff --git a/src/interpreter.rs b/src/interpreter.rs index 860489b..8da660a 100644 --- a/src/interpreter.rs +++ b/src/interpreter.rs @@ -1700,7 +1700,7 @@ impl Interpreter { }; let comps_defined = comps.iter().all(|v| v != &Value::Undefined); - let ctx = self.contexts.last_mut().expect("no current context"); + let ctx = self.get_current_context_mut()?; if is_const_rule { ctx.early_return = true; @@ -1753,7 +1753,7 @@ impl Interpreter { let key = self.eval_expr(&ke)?; let value = self.eval_expr(&oe)?; - let ctx = self.contexts.last_mut().unwrap(); + let ctx = self.get_current_context_mut()?; if key != Value::Undefined && value != Value::Undefined { let map = ctx.value.as_object_mut()?; match map.get(&key) { @@ -1782,7 +1782,7 @@ impl Interpreter { } (None, Some(oe)) => { let output = self.eval_expr(&oe)?; - let ctx = self.contexts.last_mut().unwrap(); + let ctx = self.get_current_context_mut()?; if output != Value::Undefined { match &mut ctx.value { Value::Array(a) => { @@ -1806,9 +1806,12 @@ impl Interpreter { } // If a query snippet is being run, gather results. - let ctx = self.contexts.last_mut().expect("no current context"); - if let Some(result) = &ctx.result { - let mut result = result.clone(); + let result_opt = { + let ctx = self.get_current_context_mut()?; + ctx.result.clone() + }; + + if let Some(mut result) = result_opt { if let Some(scope) = self.scopes.last() { for (name, value) in scope.iter() { result @@ -1824,6 +1827,7 @@ impl Interpreter { .all(|v| v.value != Value::Undefined && v.value != Value::Bool(false)) && !result.expressions.is_empty() { + let ctx = self.get_current_context_mut()?; ctx.results.result.push(result); } } @@ -1872,6 +1876,13 @@ impl Interpreter { } } + fn get_current_context_mut(&mut self) -> Result<&mut Context> { + match self.contexts.last_mut() { + Some(ctx) => Ok(ctx), + _ => bail!("internal error: no active context found"), + } + } + fn get_exprs_from_context(&self) -> Result { let ctx = self.get_current_context()?; Ok((ctx.key_expr.clone(), ctx.output_expr.clone())) @@ -1971,9 +1982,12 @@ impl Interpreter { result = self.eval_output_expr()?; } else { // If a query snippet is being run, gather results. - let ctx = self.contexts.last_mut().expect("no current context"); - if let Some(result) = &ctx.result { - let mut result = result.clone(); + let result_opt = { + let ctx = self.get_current_context_mut()?; + ctx.result.clone() + }; + + if let Some(mut result) = result_opt { if let Some(scope) = self.scopes.last() { for (name, value) in scope.iter() { result @@ -1990,6 +2004,7 @@ impl Interpreter { .all(|v| v.value != Value::Undefined && v.value != Value::Bool(false)) && !result.expressions.is_empty() { + let ctx = self.get_current_context_mut()?; ctx.results.result.push(result); } } From 889a02ddd64276a062d5d15b819621acc465efa5 Mon Sep 17 00:00:00 2001 From: Anand Krishnamoorthi Date: Fri, 19 Dec 2025 11:12:31 -0600 Subject: [PATCH 6/9] fix: Avoid unwrap when accesssing current module Signed-off-by: Anand Krishnamoorthi --- src/interpreter.rs | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/src/interpreter.rs b/src/interpreter.rs index 8da660a..fe1974a 100644 --- a/src/interpreter.rs +++ b/src/interpreter.rs @@ -2858,7 +2858,8 @@ impl Interpreter { Ok(Self::get_value_chained(self.data.clone(), fields)) } else if !self.compiled_policy.modules.is_empty() { - let path = Parser::get_path_ref_components(&self.module.clone().unwrap().package.refr)?; + let module = self.current_module()?; + let path = Parser::get_path_ref_components(&module.package.refr)?; let mut path: Vec<&str> = path.iter().map(|s| s.text()).collect(); path.push(name.text()); @@ -3036,7 +3037,8 @@ impl Interpreter { } fn make_rule_context(&self, head: &RuleHead) -> Result<(Context, Vec)> { - let mut path = Parser::get_path_ref_components(&self.module.clone().unwrap().package.refr)?; + let module = self.current_module()?; + let mut path = Parser::get_path_ref_components(&module.package.refr)?; match head { RuleHead::Compr { refr, assign, .. } => { let output_expr = assign.as_ref().map(|assign| assign.value.clone()); @@ -3356,8 +3358,8 @@ impl Interpreter { let scopes = core::mem::take(&mut self.scopes); - let mut path = - Parser::get_path_ref_components(&self.module.clone().unwrap().package.refr)?; + let module = self.current_module()?; + let mut path = Parser::get_path_ref_components(&module.package.refr)?; let (refr, index) = match refr.as_ref() { Expr::RefBrack { refr, index, .. } => (refr, Some(index.clone())), From 3f29eb2fa6298878674e8014d378615362e1dfdb Mon Sep 17 00:00:00 2001 From: Anand Krishnamoorthi Date: Fri, 19 Dec 2025 11:19:31 -0600 Subject: [PATCH 7/9] fix: Create ordered statements in a safe way Signed-off-by: Anand Krishnamoorthi --- src/interpreter.rs | 25 ++++++++++++++++++++++++- 1 file changed, 24 insertions(+), 1 deletion(-) diff --git a/src/interpreter.rs b/src/interpreter.rs index fe1974a..2e551f1 100644 --- a/src/interpreter.rs +++ b/src/interpreter.rs @@ -2063,7 +2063,30 @@ impl Interpreter { }; let ordered_stmts: Vec<&LiteralStmt> = match order_indices { - Some(order) => order.iter().map(|i| &query.stmts[*i as usize]).collect(), + Some(order) => { + let stmts_len = query.stmts.len(); + if order.len() != stmts_len { + let msg = format!( + "invalid schedule: expected {stmts_len} statement indices, found {}", + order.len() + ); + bail!(query.span.error(msg.as_str())); + } + + let mut ordered = Vec::with_capacity(stmts_len); + for idx in order { + let stmt_idx = *idx as usize; + if stmt_idx >= stmts_len { + let msg = format!( + "invalid schedule index {stmt_idx} for {} statements", + stmts_len + ); + bail!(query.span.error(msg.as_str())); + } + ordered.push(&query.stmts[stmt_idx]); + } + ordered + } None => query.stmts.iter().collect(), }; From 273a80571eb5975b4a7b1f856c323e01dcbcd64f Mon Sep 17 00:00:00 2001 From: Anand Krishnamoorthi Date: Fri, 19 Dec 2025 11:21:39 -0600 Subject: [PATCH 8/9] fix: apply expression ordering to schedule in a safe way Signed-off-by: Anand Krishnamoorthi --- src/interpreter.rs | 18 +++++++++++++++++- 1 file changed, 17 insertions(+), 1 deletion(-) diff --git a/src/interpreter.rs b/src/interpreter.rs index 2e551f1..af9faf0 100644 --- a/src/interpreter.rs +++ b/src/interpreter.rs @@ -3708,14 +3708,30 @@ impl Interpreter { .map_err(|err| anyhow!("schedule out of bounds: {err}"))? { for idx in 0..results.result.len() { + let exprs_len = results.result[idx].expressions.len(); + if query_schedule.order.len() != exprs_len { + let msg = format!( + "invalid schedule: expected {exprs_len} expression indices, found {}", + query_schedule.order.len() + ); + bail!(query.span.error(msg.as_str())); + } + let e = Expression { value: Value::Undefined, text: "".into(), location: Location { row: 0, col: 0 }, }; - let mut ordered_expressions = vec![e; results.result[idx].expressions.len()]; + let mut ordered_expressions = vec![e; exprs_len]; for (expr_idx, value) in results.result[idx].expressions.iter().enumerate() { let orig_idx = query_schedule.order[expr_idx] as usize; + if orig_idx >= exprs_len { + let msg = format!( + "invalid schedule expression index {orig_idx} for {} expressions", + exprs_len + ); + bail!(query.span.error(msg.as_str())); + } ordered_expressions[orig_idx] = value.clone(); } if !ordered_expressions From dbfb8e38a86aad6bf7b335b211a7c84a545310ee Mon Sep 17 00:00:00 2001 From: Anand Krishnamoorthi Date: Fri, 19 Dec 2025 12:08:09 -0600 Subject: [PATCH 9/9] fix: Skip reordering in truncated queries. In case all the statements of a query don't execute, skip reordering the result expressions to match the source order. Doing so requires maintaining additional data structures not worth the complexity for now. Additionally we want to discourage queries and encourage evaluating rules. Queries are inherently less performant than rules which can be precompiled. Signed-off-by: Anand Krishnamoorthi --- src/interpreter.rs | 36 ++++++++++++++++++++---------------- 1 file changed, 20 insertions(+), 16 deletions(-) diff --git a/src/interpreter.rs b/src/interpreter.rs index af9faf0..c800415 100644 --- a/src/interpreter.rs +++ b/src/interpreter.rs @@ -3698,7 +3698,14 @@ impl Interpreter { _ => bail!("internal error: no context"), }; - // Apply expression ordering from the schedule if needed + // Apply expression ordering from the schedule when it is safe to do so. + // If the schedule references statements that did not produce expressions, the lengths + // will differ; in that case we keep the collected order to avoid spurious errors. + // Example: a schedule for `1 == 2; 1 == 1` may list both statements. The first produces + // an expression (false), but evaluation stops and the second never runs, so only one + // expression is collected. In that case `order.len()` can be 2 while one expression is + // available. In these cases, we avoid reordering - doing so requires maintaining additional + // data during evaluation and is not worth the complexity for now. let current_module_idx = compiled_modules_len; let current_query_idx = query.qidx; if let Some(ref self_schedule) = &self.query_schedule { @@ -3709,34 +3716,31 @@ impl Interpreter { { for idx in 0..results.result.len() { let exprs_len = results.result[idx].expressions.len(); + // Skip reordering when the schedule length does not match produced expressions. if query_schedule.order.len() != exprs_len { - let msg = format!( - "invalid schedule: expected {exprs_len} expression indices, found {}", - query_schedule.order.len() - ); - bail!(query.span.error(msg.as_str())); + continue; } - let e = Expression { + let placeholder = Expression { value: Value::Undefined, text: "".into(), location: Location { row: 0, col: 0 }, }; - let mut ordered_expressions = vec![e; exprs_len]; + let mut ordered_expressions = vec![placeholder; exprs_len]; + let mut invalid = false; for (expr_idx, value) in results.result[idx].expressions.iter().enumerate() { let orig_idx = query_schedule.order[expr_idx] as usize; if orig_idx >= exprs_len { - let msg = format!( - "invalid schedule expression index {orig_idx} for {} expressions", - exprs_len - ); - bail!(query.span.error(msg.as_str())); + invalid = true; + break; } ordered_expressions[orig_idx] = value.clone(); } - if !ordered_expressions - .iter() - .any(|v| v.value == Value::Undefined) + + if !invalid + && !ordered_expressions + .iter() + .any(|v| v.value == Value::Undefined) { results.result[idx].expressions = ordered_expressions; }