feat: Safeguard lookup use

Detect invalid indexes and raise internal errors.

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

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<BindingPlan> {
+    pub(super) fn get_binding_plan_for_expr(&self, expr: &ExprRef) -> Result<Option<BindingPlan>> {
         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<BindingPlan> {
-        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(),
             }

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<Vec<HoistedLoop>> {
-        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<HoistedLoop> {
+    pub(super) fn get_expr_loops(&self, expr: &ExprRef) -> Result<Vec<HoistedLoop>> {
         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,

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() {