chore: Harden RVM implementation (#537)

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

src/rvm/vm/comprehension.rs

@@ -1,13 +1,5 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT License.
-#![allow(
-    clippy::indexing_slicing,
-    clippy::arithmetic_side_effects,
-    clippy::option_if_let_else,
-    clippy::useless_let_if_seq,
-    clippy::as_conversions,
-    clippy::pattern_type_mismatch
-)]
 
 use crate::rvm::instructions::{ComprehensionBeginParams, ComprehensionMode};
 use crate::value::Value;
@@ -44,99 +36,11 @@ impl RegoVM {
             ComprehensionMode::Array => Value::new_array(),
             ComprehensionMode::Object => Value::Object(Rc::new(BTreeMap::new())),
         };
-        self.registers[params.result_reg as usize] = initial_result.clone();
+        self.set_register(params.result_reg, initial_result.clone())?;
 
         let auto_iterate = params.collection_reg != params.result_reg;
         let iteration_state = if auto_iterate {
-            let source_value = self.registers[params.collection_reg as usize].clone();
-            match source_value {
-                Value::Array(items) => {
-                    if items.is_empty() {
-                        None
-                    } else {
-                        Some(IterationState::Array { items, index: 0 })
-                    }
-                }
-                Value::Object(obj) => {
-                    if obj.is_empty() {
-                        None
-                    } else {
-                        Some(IterationState::Object {
-                            obj,
-                            current_key: None,
-                            first_iteration: true,
-                        })
-                    }
-                }
-                Value::Set(set) => {
-                    if set.is_empty() {
-                        None
-                    } else {
-                        Some(IterationState::Set {
-                            items: set,
-                            current_item: None,
-                            first_iteration: true,
-                        })
-                    }
-                }
-                Value::Undefined => None,
-                Value::Null => None,
-                _ => None,
-            }
-        } else {
-            None
-        };
-
-        let mut has_iteration = false;
-        if let Some(state) = iteration_state.as_ref() {
-            has_iteration = self.setup_next_iteration(state, params.key_reg, params.value_reg)?;
-        }
-
-        let resume_pc = if auto_iterate {
-            params.comprehension_end as usize
-        } else {
-            params.comprehension_end.saturating_sub(1) as usize
-        };
-
-        let mut comprehension_context = ComprehensionContext {
-            mode: params.mode.clone(),
-            result_reg: params.result_reg,
-            key_reg: params.key_reg,
-            value_reg: params.value_reg,
-            body_start: params.body_start,
-            comprehension_end: params.comprehension_end,
-            iteration_state,
-            resume_pc,
-        };
-
-        if auto_iterate {
-            if has_iteration {
-                self.pc = params.body_start as usize - 1;
-            } else {
-                comprehension_context.iteration_state = None;
-                self.pc = params.comprehension_end as usize - 1;
-            }
-        }
-
-        self.comprehension_stack.push(comprehension_context);
-
-        Ok(())
-    }
-
-    fn execute_comprehension_begin_suspendable(
-        &mut self,
-        params: &ComprehensionBeginParams,
-    ) -> Result<()> {
-        let initial_result = match params.mode {
-            ComprehensionMode::Set => Value::new_set(),
-            ComprehensionMode::Array => Value::new_array(),
-            ComprehensionMode::Object => Value::Object(Rc::new(BTreeMap::new())),
-        };
-        self.registers[params.result_reg as usize] = initial_result.clone();
-
-        let auto_iterate = params.collection_reg != params.result_reg;
-        let iteration_state = if auto_iterate {
-            let source_value = self.registers[params.collection_reg as usize].clone();
+            let source_value = self.get_register(params.collection_reg)?.clone();
             match source_value {
                 Value::Array(items) => {
                     if items.is_empty() {
@@ -182,9 +86,98 @@ impl RegoVM {
         };
 
         let resume_pc = if auto_iterate {
-            params.comprehension_end as usize
+            usize::from(params.comprehension_end)
         } else {
-            params.comprehension_end.saturating_sub(1) as usize
+            usize::from(params.comprehension_end.saturating_sub(1))
+        };
+
+        let mut comprehension_context = ComprehensionContext {
+            mode: params.mode.clone(),
+            result_reg: params.result_reg,
+            key_reg: params.key_reg,
+            value_reg: params.value_reg,
+            body_start: params.body_start,
+            comprehension_end: params.comprehension_end,
+            iteration_state,
+            resume_pc,
+        };
+
+        if auto_iterate {
+            if has_iteration {
+                self.pc = usize::from(params.body_start).saturating_sub(1);
+            } else {
+                comprehension_context.iteration_state = None;
+                self.pc = usize::from(params.comprehension_end).saturating_sub(1);
+            }
+        }
+
+        self.comprehension_stack.push(comprehension_context);
+
+        Ok(())
+    }
+
+    fn execute_comprehension_begin_suspendable(
+        &mut self,
+        params: &ComprehensionBeginParams,
+    ) -> Result<()> {
+        let initial_result = match params.mode {
+            ComprehensionMode::Set => Value::new_set(),
+            ComprehensionMode::Array => Value::new_array(),
+            ComprehensionMode::Object => Value::Object(Rc::new(BTreeMap::new())),
+        };
+        self.set_register(params.result_reg, initial_result.clone())?;
+
+        let auto_iterate = params.collection_reg != params.result_reg;
+        let iteration_state = if auto_iterate {
+            let source_value = self.get_register(params.collection_reg)?.clone();
+            match source_value {
+                Value::Array(items) => {
+                    if items.is_empty() {
+                        None
+                    } else {
+                        Some(IterationState::Array { items, index: 0 })
+                    }
+                }
+                Value::Object(obj) => {
+                    if obj.is_empty() {
+                        None
+                    } else {
+                        Some(IterationState::Object {
+                            obj,
+                            current_key: None,
+                            first_iteration: true,
+                        })
+                    }
+                }
+                Value::Set(set) => {
+                    if set.is_empty() {
+                        None
+                    } else {
+                        Some(IterationState::Set {
+                            items: set,
+                            current_item: None,
+                            first_iteration: true,
+                        })
+                    }
+                }
+                Value::Undefined => None,
+                Value::Null => None,
+                _ => None,
+            }
+        } else {
+            None
+        };
+
+        let has_iteration = if let Some(state) = iteration_state.as_ref() {
+            self.setup_next_iteration(state, params.key_reg, params.value_reg)?
+        } else {
+            false
+        };
+
+        let resume_pc = if auto_iterate {
+            usize::from(params.comprehension_end)
+        } else {
+            usize::from(params.comprehension_end.saturating_sub(1))
         };
 
         let mut comprehension_context = ComprehensionContext {
@@ -200,13 +193,13 @@ impl RegoVM {
 
         let next_pc = if auto_iterate {
             if has_iteration {
-                params.body_start as usize
+                usize::from(params.body_start)
             } else {
                 comprehension_context.iteration_state = None;
-                params.comprehension_end as usize
+                usize::from(params.comprehension_end)
             }
         } else {
-            self.pc + 1
+            self.pc.saturating_add(1)
         };
 
         let return_pc = comprehension_context.resume_pc;
@@ -248,20 +241,21 @@ impl RegoVM {
         } else {
             return Err(VmError::InvalidIteration {
                 value: Value::String(Arc::from("No active comprehension")),
+                pc: self.pc,
             });
         };
 
-        let value_to_add = self.registers[value_reg as usize].clone();
+        let value_to_add = self.get_register(value_reg)?.clone();
         let key_value = if let Some(key_reg) = key_reg {
-            Some(self.registers[key_reg as usize].clone())
+            Some(self.get_register(key_reg)?.clone())
         } else if matches!(comprehension_context.mode, ComprehensionMode::Object) {
-            Some(self.registers[comprehension_context.key_reg as usize].clone())
+            Some(self.get_register(comprehension_context.key_reg)?.clone())
         } else {
             None
         };
 
-        let result_reg = comprehension_context.result_reg as usize;
-        let current_result = self.registers[result_reg].clone();
+        let result_reg = comprehension_context.result_reg;
+        let current_result = self.get_register(result_reg)?.clone();
         let mode = comprehension_context.mode.clone();
 
         let updated_result = match (mode, current_result) {
@@ -284,31 +278,41 @@ impl RegoVM {
                     self.comprehension_stack.push(comprehension_context);
                     return Err(VmError::InvalidIteration {
                         value: Value::String(Arc::from("Object comprehension requires key")),
+                        pc: self.pc,
                     });
                 }
             }
             (_mode, other) => {
                 self.comprehension_stack.push(comprehension_context);
-                return Err(VmError::InvalidIteration { value: other });
+                return Err(VmError::InvalidIteration {
+                    value: other,
+                    pc: self.pc,
+                });
             }
         };
 
-        self.registers[result_reg] = updated_result;
+        self.set_register(result_reg, updated_result)?;
 
         if let Some(iter_state) = comprehension_context.iteration_state.as_mut() {
-            match iter_state {
-                IterationState::Object { current_key, .. } => {
+            match *iter_state {
+                IterationState::Object {
+                    ref mut current_key,
+                    ..
+                } => {
                     let tracked_key =
                         if comprehension_context.key_reg != comprehension_context.value_reg {
-                            self.registers[comprehension_context.key_reg as usize].clone()
+                            self.get_register(comprehension_context.key_reg)?.clone()
                         } else {
-                            self.registers[comprehension_context.value_reg as usize].clone()
+                            self.get_register(comprehension_context.value_reg)?.clone()
                         };
                     *current_key = Some(tracked_key);
                 }
-                IterationState::Set { current_item, .. } => {
+                IterationState::Set {
+                    ref mut current_item,
+                    ..
+                } => {
                     *current_item =
-                        Some(self.registers[comprehension_context.value_reg as usize].clone());
+                        Some(self.get_register(comprehension_context.value_reg)?.clone());
                 }
                 IterationState::Array { .. } => {}
             }
@@ -321,10 +325,10 @@ impl RegoVM {
             )?;
 
             if has_next {
-                self.pc = comprehension_context.body_start as usize - 1;
+                self.pc = usize::from(comprehension_context.body_start).saturating_sub(1);
             } else {
                 comprehension_context.iteration_state = None;
-                self.pc = comprehension_context.comprehension_end as usize - 1;
+                self.pc = usize::from(comprehension_context.comprehension_end).saturating_sub(1);
             }
         }
 
@@ -347,110 +351,162 @@ impl RegoVM {
             })
             .ok_or(VmError::InvalidIteration {
                 value: Value::String(Arc::from("No active comprehension")),
+                pc: self.pc,
             })?;
 
-        let (iteration_state_snapshot, key_reg_idx, value_reg_idx, body_start, comprehension_end) = {
-            let frame = self.execution_stack.get_mut(comprehension_index).ok_or(
-                VmError::InvalidIteration {
-                    value: Value::String(Arc::from("No active comprehension")),
-                },
-            )?;
-
-            match &mut frame.kind {
-                FrameKind::Comprehension { context, .. } => {
-                    let value_to_add = self.registers[value_reg as usize].clone();
-                    let key_value = if let Some(key_reg) = key_reg {
-                        Some(self.registers[key_reg as usize].clone())
-                    } else if matches!(context.mode, ComprehensionMode::Object) {
-                        Some(self.registers[context.key_reg as usize].clone())
-                    } else {
-                        None
-                    };
-
-                    let result_reg_idx = context.result_reg as usize;
-                    let current_result = self.registers[result_reg_idx].clone();
-                    let mode = context.mode.clone();
-
-                    let updated_result = match (mode, current_result) {
-                        (ComprehensionMode::Set, Value::Set(set)) => {
-                            let mut new_set = set.as_ref().clone();
-                            new_set.insert(value_to_add);
-                            Value::Set(crate::Rc::new(new_set))
-                        }
-                        (ComprehensionMode::Array, Value::Array(arr)) => {
-                            let mut new_arr = arr.as_ref().to_vec();
-                            new_arr.push(value_to_add);
-                            Value::Array(crate::Rc::new(new_arr))
-                        }
-                        (ComprehensionMode::Object, Value::Object(obj)) => {
-                            if let Some(key) = key_value {
-                                let mut new_obj = obj.as_ref().clone();
-                                new_obj.insert(key, value_to_add);
-                                Value::Object(crate::Rc::new(new_obj))
-                            } else {
-                                return Err(VmError::InvalidIteration {
-                                    value: Value::String(Arc::from(
-                                        "Object comprehension requires key",
-                                    )),
-                                });
-                            }
-                        }
-                        (_mode, other) => {
-                            return Err(VmError::InvalidIteration { value: other });
-                        }
-                    };
-
-                    self.registers[result_reg_idx] = updated_result;
-
-                    if let Some(iter_state) = context.iteration_state.as_mut() {
-                        match iter_state {
-                            IterationState::Object { current_key, .. } => {
-                                let tracked_key = if context.key_reg != context.value_reg {
-                                    self.registers[context.key_reg as usize].clone()
-                                } else {
-                                    self.registers[context.value_reg as usize].clone()
-                                };
-                                *current_key = Some(tracked_key);
-                            }
-                            IterationState::Set { current_item, .. } => {
-                                *current_item =
-                                    Some(self.registers[context.value_reg as usize].clone());
-                            }
-                            IterationState::Array { .. } => {}
-                        }
-
-                        iter_state.advance();
-                    }
-
-                    (
-                        context.iteration_state.clone(),
-                        context.key_reg,
-                        context.value_reg,
-                        context.body_start,
-                        context.comprehension_end,
-                    )
-                }
-                _ => {
-                    return Err(VmError::InvalidIteration {
+        let (
+            value_to_add,
+            key_value,
+            current_result,
+            mode,
+            result_reg_idx,
+            key_reg_idx,
+            value_reg_idx,
+            iteration_key,
+            iteration_value,
+        ) = {
+            let frame =
+                self.execution_stack
+                    .get(comprehension_index)
+                    .ok_or(VmError::InvalidIteration {
                         value: Value::String(Arc::from("No active comprehension")),
+                        pc: self.pc,
+                    })?;
+
+            if let FrameKind::Comprehension { ref context, .. } = frame.kind {
+                let value_to_add = self.get_register(value_reg)?.clone();
+                let key_value = if let Some(key_reg) = key_reg {
+                    Some(self.get_register(key_reg)?.clone())
+                } else if matches!(context.mode, ComprehensionMode::Object) {
+                    Some(self.get_register(context.key_reg)?.clone())
+                } else {
+                    None
+                };
+
+                let result_reg_idx = context.result_reg;
+                let current_result = self.get_register(result_reg_idx)?.clone();
+                let mode = context.mode.clone();
+                let iteration_key = self.get_register(context.key_reg)?.clone();
+                let iteration_value = self.get_register(context.value_reg)?.clone();
+
+                (
+                    value_to_add,
+                    key_value,
+                    current_result,
+                    mode,
+                    result_reg_idx,
+                    context.key_reg,
+                    context.value_reg,
+                    iteration_key,
+                    iteration_value,
+                )
+            } else {
+                return Err(VmError::InvalidIteration {
+                    value: Value::String(Arc::from("No active comprehension")),
+                    pc: self.pc,
+                });
+            }
+        };
+
+        let updated_result = match (mode, current_result) {
+            (ComprehensionMode::Set, Value::Set(set)) => {
+                let mut new_set = set.as_ref().clone();
+                new_set.insert(value_to_add);
+                Value::Set(crate::Rc::new(new_set))
+            }
+            (ComprehensionMode::Array, Value::Array(arr)) => {
+                let mut new_arr = arr.as_ref().to_vec();
+                new_arr.push(value_to_add);
+                Value::Array(crate::Rc::new(new_arr))
+            }
+            (ComprehensionMode::Object, Value::Object(obj)) => {
+                if let Some(key) = key_value {
+                    let mut new_obj = obj.as_ref().clone();
+                    new_obj.insert(key, value_to_add);
+                    Value::Object(crate::Rc::new(new_obj))
+                } else {
+                    return Err(VmError::InvalidIteration {
+                        value: Value::String(Arc::from("Object comprehension requires key")),
+                        pc: self.pc,
                     });
                 }
             }
+            (_mode, other) => {
+                return Err(VmError::InvalidIteration {
+                    value: other,
+                    pc: self.pc,
+                });
+            }
         };
 
+        let (iteration_state_snapshot, body_start, comprehension_end) = {
+            let frame = self.execution_stack.get_mut(comprehension_index).ok_or(
+                VmError::InvalidIteration {
+                    value: Value::String(Arc::from("No active comprehension")),
+                    pc: self.pc,
+                },
+            )?;
+
+            if let &mut FrameKind::Comprehension {
+                ref mut context, ..
+            } = &mut frame.kind
+            {
+                if let Some(iter_state) = context.iteration_state.as_mut() {
+                    match *iter_state {
+                        IterationState::Object {
+                            ref mut current_key,
+                            ..
+                        } => {
+                            let tracked_key = if context.key_reg != context.value_reg {
+                                iteration_key.clone()
+                            } else {
+                                iteration_value.clone()
+                            };
+                            *current_key = Some(tracked_key);
+                        }
+                        IterationState::Set {
+                            ref mut current_item,
+                            ..
+                        } => {
+                            *current_item = Some(iteration_value.clone());
+                        }
+                        IterationState::Array { .. } => {}
+                    }
+
+                    iter_state.advance();
+                }
+
+                (
+                    context.iteration_state.clone(),
+                    context.body_start,
+                    context.comprehension_end,
+                )
+            } else {
+                return Err(VmError::InvalidIteration {
+                    value: Value::String(Arc::from("No active comprehension")),
+                    pc: self.pc,
+                });
+            }
+        };
+
+        self.set_register(result_reg_idx, updated_result)?;
+
         if let Some(state) = iteration_state_snapshot.as_ref() {
             let has_next = self.setup_next_iteration(state, key_reg_idx, value_reg_idx)?;
 
             if has_next {
                 if let Some(frame) = self.execution_stack.get_mut(comprehension_index) {
-                    frame.pc = body_start as usize;
+                    frame.pc = usize::from(body_start);
                     self.frame_pc_overridden = true;
                 }
             } else if let Some(frame) = self.execution_stack.get_mut(comprehension_index) {
-                if let FrameKind::Comprehension { context, .. } = &mut frame.kind {
+                if let &mut FrameKind::Comprehension {
+                    ref mut context, ..
+                } = &mut frame.kind
+                {
                     context.iteration_state = None;
                 }
-                frame.pc = comprehension_end as usize;
+                frame.pc = usize::from(comprehension_end);
                 self.frame_pc_overridden = true;
             }
         }
@@ -479,7 +535,10 @@ impl RegoVM {
 
     pub(super) fn handle_comprehension_condition_failure_suspendable(&mut self) -> Result<bool> {
         if let Some(mut frame) = self.execution_stack.pop() {
-            let handled = if let FrameKind::Comprehension { context, .. } = &mut frame.kind {
+            let handled = if let &mut FrameKind::Comprehension {
+                ref mut context, ..
+            } = &mut frame.kind
+            {
                 self.advance_comprehension_after_failure(context)?;
                 true
             } else {
@@ -503,18 +562,18 @@ impl RegoVM {
                 iter_state,
                 context.key_reg,
                 context.value_reg,
-            );
+            )?;
             iter_state.advance();
             let has_next =
                 self.setup_next_iteration(iter_state, context.key_reg, context.value_reg)?;
             if has_next {
-                self.pc = context.body_start.saturating_sub(1) as usize;
+                self.pc = usize::from(context.body_start.saturating_sub(1));
             } else {
                 context.iteration_state = None;
-                self.pc = context.comprehension_end.saturating_sub(1) as usize;
+                self.pc = usize::from(context.comprehension_end.saturating_sub(1));
             }
         } else {
-            self.pc = context.comprehension_end.saturating_sub(1) as usize;
+            self.pc = usize::from(context.comprehension_end.saturating_sub(1));
         }
 
         Ok(())
@@ -525,31 +584,41 @@ impl RegoVM {
         iter_state: &mut IterationState,
         key_reg: u8,
         value_reg: u8,
-    ) {
-        match iter_state {
-            IterationState::Object { current_key, .. } => {
+    ) -> Result<()> {
+        match *iter_state {
+            IterationState::Object {
+                ref mut current_key,
+                ..
+            } => {
                 let tracked_key = if key_reg != value_reg {
-                    self.registers[key_reg as usize].clone()
+                    self.get_register(key_reg)?.clone()
                 } else {
-                    self.registers[value_reg as usize].clone()
+                    self.get_register(value_reg)?.clone()
                 };
                 *current_key = Some(tracked_key);
             }
-            IterationState::Set { current_item, .. } => {
-                *current_item = Some(self.registers[value_reg as usize].clone());
+            IterationState::Set {
+                ref mut current_item,
+                ..
+            } => {
+                *current_item = Some(self.get_register(value_reg)?.clone());
             }
             IterationState::Array { .. } => {}
         }
+
+        Ok(())
     }
 
     fn execute_comprehension_end_run_to_completion(&mut self) -> Result<()> {
-        if let Some(_context) = self.comprehension_stack.pop() {
-            Ok(())
-        } else {
-            Err(VmError::InvalidIteration {
-                value: Value::String(Arc::from("No active comprehension context")),
-            })
-        }
+        self.comprehension_stack.pop().map_or_else(
+            || {
+                Err(VmError::InvalidIteration {
+                    value: Value::String(Arc::from("No active comprehension context")),
+                    pc: self.pc,
+                })
+            },
+            |_context| Ok(()),
+        )
     }
 
     fn execute_comprehension_end_suspendable(&mut self) -> Result<()> {
@@ -565,6 +634,7 @@ impl RegoVM {
                     }
                     return Err(VmError::InvalidIteration {
                         value: Value::String(Arc::from("No active comprehension context")),
+                        pc: self.pc,
                     });
                 }
             };
@@ -620,6 +690,7 @@ impl RegoVM {
                     }
                     return Err(VmError::InvalidIteration {
                         value: Value::String(Arc::from(message.into_boxed_str())),
+                        pc: self.pc,
                     });
                 }
             }