📝 CodeRabbit Chat: Add Security Controls and Tests for Krates Operations

rust/src/modules/krates/mod.rs

@@ -8,7 +8,6 @@
  * http://www.apache.org/licenses/LICENSE-2.0
  */
 
-// TODO: Krates
 pub(crate) mod gas_monitor;
 pub(crate) mod has_nightly_opcodes;
 pub(crate) mod validate_bytecode;

rust/src/modules/krates/validate_bytecode.rs

@@ -12,7 +12,6 @@ use crate::modules::vmerror::VMError;
 use crate::types::instructions::Instructions;
 use crate::types::value::FuncMetadata;
 use ahash::AHashMap;
-use smol_str::SmolStr;
 #[cold]
 pub fn validate_bytecode(
   bytecode: &[Instructions],
@@ -33,13 +32,38 @@ pub fn validate_bytecode(
       _ => {}
     }
   }
-  for (name, meta) in functions {
+  for meta in functions.values() {
     if meta.start >= len {
-      return Err(VMError::SystemError(SmolStr::from(format!(
-        "Function '{}' start address {} is out of bounds (len: {})",
-        name, meta.start, len
-      ))));
+      return Err(VMError::OutOfBounds {
+        ip: meta.start,
+        index: meta.start,
+        len,
+      });
     }
   }
   Ok(())
 }
+
+#[cfg(test)]
+mod tests {
+  use super::*;
+
+  #[test]
+  fn accepts_valid_bytecode() {
+    let bytecode = vec![Instructions::Jump(0)];
+    assert!(validate_bytecode(&bytecode, &AHashMap::new()).is_ok());
+  }
+
+  #[test]
+  fn rejects_invalid_jump_with_structured_error() {
+    let result = validate_bytecode(&[Instructions::Jump(1)], &AHashMap::new());
+    assert!(matches!(
+      result,
+      Err(VMError::OutOfBounds {
+        ip: 0,
+        index: 1,
+        len: 1
+      })
+    ));
+  }
+}

rust/src/modules/krates/validate_security.rs

@@ -76,8 +76,42 @@ pub fn validate_security(
       _ => {}
     }
   }
-  if total_instr > 10 && (nop_count * 10) > total_instr {
+  if total_instr > 10 && nop_count > total_instr / 10 {
     return Err(VMError::ExcessiveNopPadding);
   }
   Ok(())
 }
+
+#[cfg(test)]
+mod tests {
+  use super::*;
+  use smol_str::SmolStr;
+
+  #[test]
+  fn accepts_allowed_import() {
+    let bytecode = vec![Instructions::Import(SmolStr::new("math"), 0)];
+    assert!(validate_security(&bytecode, &SecurityConfig::default()).is_ok());
+  }
+
+  #[test]
+  fn rejects_untrusted_import() {
+    let bytecode = vec![Instructions::Import(SmolStr::new("private"), 0)];
+    assert!(matches!(
+      validate_security(&bytecode, &SecurityConfig::default()),
+      Err(VMError::UnauthorizedModule { ip: 0, .. })
+    ));
+  }
+
+  #[test]
+  fn enforces_io_limit() {
+    let config = SecurityConfig {
+      max_io: 1,
+      ..Default::default()
+    };
+    let bytecode = vec![Instructions::Print, Instructions::Println];
+    assert!(matches!(
+      validate_security(&bytecode, &config),
+      Err(VMError::IoFlood { ip: 1 })
+    ));
+  }
+}

rust/src/vm/execute.rs

@@ -86,6 +86,35 @@ pub fn execute(
       }
       tick += 1;
       let instr = &bytecode[ip];
+      if security_config.max_stack_size > 0
+        && stack.len() >= security_config.max_stack_size
+        && matches!(
+          instr,
+          Instructions::PushInt16(_)
+            | Instructions::PushInt32(_)
+            | Instructions::PushInt64(_)
+            | Instructions::PushInt128(_)
+            | Instructions::PushFloat16(_)
+            | Instructions::PushFloat32(_)
+            | Instructions::PushFloat64(_)
+            | Instructions::PushString(_)
+            | Instructions::PushArray(_)
+            | Instructions::PushBool(_)
+            | Instructions::PushObject(_)
+            | Instructions::PushUndefined
+            | Instructions::PushNull
+            | Instructions::PushNaN
+            | Instructions::Push(_)
+            | Instructions::ValIdx(_)
+            | Instructions::GetIdx(_)
+            | Instructions::Dup
+        )
+      {
+        return Err(VMError::StackOverflow {
+          ip,
+          limit: security_config.max_stack_size,
+        });
+      }
       match instr {
         Instructions::PushInt16(_)
         | Instructions::PushInt32(_)
@@ -384,3 +413,20 @@ fn test_out_of_bounds_jump_is_rejected_before_execution() {
     })
   ));
 }
+
+#[test]
+fn test_configured_stack_limit_is_enforced() {
+  let bytecode = vec![Instructions::PushInt32(1), Instructions::PushInt32(2)];
+  let options = crate::types::value::RunOptions {
+    security_config: crate::types::security_config::SecurityConfig {
+      max_stack_size: 1,
+      ..Default::default()
+    },
+    ..Default::default()
+  };
+  let result = execute(bytecode, &mut Some(options), None);
+  assert!(matches!(
+    result,
+    Err(VMError::StackOverflow { ip: 1, limit: 1 })
+  ));
+}