From 5055626968d5fc19f81e1391216b411edd3348f3 Mon Sep 17 00:00:00 2001
From: Erin <erin@erindesu.cz>
Date: Wed, 26 Jul 2023 00:01:25 +0200
Subject: [PATCH] Added runtime bound checking

---
 hbbytecode/src/gen_valider.rs |  6 ++++--
 hbvm/src/vm/mod.rs            | 21 +++++++++++++++++++++
 2 files changed, 25 insertions(+), 2 deletions(-)

diff --git a/hbbytecode/src/gen_valider.rs b/hbbytecode/src/gen_valider.rs
index a1223c1..f82a017 100644
--- a/hbbytecode/src/gen_valider.rs
+++ b/hbbytecode/src/gen_valider.rs
@@ -59,14 +59,16 @@ macro_rules! gen_valider {
                 let start = program;
                 loop {
                     use crate::opcode::*;
+                    extern crate std;
                     program = match program {
                         // End of program
                         [] => return Ok(()),
 
                         // Memory load/store cannot go out-of-bounds register array
-                        [LD..=ST, reg, _, _, _, _, _, _, _, _, count_0, count_1, ..]
+                        //         B   B  D1 D2 D3 D4 D5 D6 D7 D8    H1      H2
+                        [LD..=ST, reg, _, _, _, _, _, _, _, _, _, count_0, count_1, ..]
                             if usize::from(*reg) * 8
-                                + usize::from(u16::from_le_bytes([*count_1, *count_0]))
+                                + usize::from(u16::from_le_bytes([*count_0, *count_1]))
                                 > 2048 =>
                         {
                             return Err(Error {
diff --git a/hbvm/src/vm/mod.rs b/hbvm/src/vm/mod.rs
index b1808f8..fcbfacb 100644
--- a/hbvm/src/vm/mod.rs
+++ b/hbvm/src/vm/mod.rs
@@ -215,6 +215,8 @@ impl<'a, PfHandler: HandlePageFault, const TIMER_QUOTIENT: usize>
                     LD => {
                         // Load. If loading more than register size, continue on adjecent registers
                         let ParamBBDH(dst, base, off, count) = self.decode();
+                        ldst_bound_check(dst, count)?;
+                        
                         let n: usize = match dst {
                             0 => 1,
                             _ => 0,
@@ -230,6 +232,8 @@ impl<'a, PfHandler: HandlePageFault, const TIMER_QUOTIENT: usize>
                     ST => {
                         // Store. Same rules apply as to LD
                         let ParamBBDH(dst, base, off, count) = self.decode();
+                        ldst_bound_check(dst, count)?;
+
                         self.memory.store(
                             self.read_reg(base).cast::<u64>() + off,
                             self.registers.as_ptr().add(usize::from(dst)).cast(),
@@ -250,6 +254,10 @@ impl<'a, PfHandler: HandlePageFault, const TIMER_QUOTIENT: usize>
                     BRC => {
                         // Block register copy
                         let ParamBBB(src, dst, count) = self.decode();
+                        if src.checked_add(count).is_none() || dst.checked_add(count).is_none() {
+                            return Err(VmRunError::RegOutOfBounds);
+                        }
+
                         core::ptr::copy(
                             self.registers.get_unchecked(usize::from(src)),
                             self.registers.get_unchecked_mut(usize::from(dst)),
@@ -394,6 +402,16 @@ impl<'a, PfHandler: HandlePageFault, const TIMER_QUOTIENT: usize>
     }
 }
 
+/// Load/Store target/source register range bound checking
+#[inline]
+fn ldst_bound_check(reg: u8, size: u16) -> Result<(), VmRunError> {
+    if usize::from(reg) * 8 + usize::from(size) > 2048 {
+        Err(VmRunError::RegOutOfBounds)
+    } else {
+        Ok(())
+    }
+}
+
 /// Virtual machine halt error
 #[derive(Copy, Clone, Debug, PartialEq, Eq)]
 #[repr(u8)]
@@ -407,6 +425,9 @@ pub enum VmRunError {
     /// Unhandled store access exception
     StoreAccessEx(u64),
 
+    /// Register out-of-bounds access
+    RegOutOfBounds,
+
     /// Reached unreachable code
     Unreachable,
 }