fix: Various places where panic might happen in case of malicious data

src/bits.rs

@@ -2,7 +2,7 @@
 pub fn roundup(x: u64, round: u64) -> u64 {
     debug_assert!(round.is_power_of_two());
     // x + (((!x) + 1) & (round - 1))
-    x + ((!x).wrapping_add(1) & (round.wrapping_sub(1)))
+    x.wrapping_add((!x).wrapping_add(1) & (round.wrapping_sub(1)))
 }
 
 #[inline(always)]
@@ -25,6 +25,7 @@ mod tests {
         assert_eq!(16, roundup(16, 16));
         assert_eq!(32, roundup(17, 16));
         assert_eq!(u64::max_value() - 15, roundup(u64::max_value() - 15, 16));
+        assert_eq!(0, roundup(u64::max_value(), 16));
     }
 
     #[test]

src/instructions/rvc.rs

@@ -308,15 +308,17 @@ pub fn factory<R: Register>(instruction_bits: u32) -> Option<Instruction> {
                         // Invalid instruction
                         (0b_00_000_00000_00, 0) => None,
                         // SRLI
-                        (0b_00_000_00000_00, uimm) => {
-                            Some(Itype::new(insts::OP_RVC_SRLI, rd, rd, uimm).0)
-                        }
+                        (0b_00_000_00000_00, uimm) => Some(
+                            Itype::new(insts::OP_RVC_SRLI, rd, rd, uimm & u32::from(R::SHIFT_MASK))
+                                .0,
+                        ),
                         // Invalid instruction
                         (0b_01_000_00000_00, 0) => None,
                         // SRAI
-                        (0b_01_000_00000_00, uimm) => {
-                            Some(Itype::new(insts::OP_RVC_SRAI, rd, rd, uimm).0)
-                        }
+                        (0b_01_000_00000_00, uimm) => Some(
+                            Itype::new(insts::OP_RVC_SRAI, rd, rd, uimm & u32::from(R::SHIFT_MASK))
+                                .0,
+                        ),
                         // ANDI
                         (0b_10_000_00000_00, _) => Some(
                             Itype::new_s(insts::OP_RVC_ANDI, rd, rd, immediate(instruction_bits)).0,
@@ -355,7 +357,7 @@ pub fn factory<R: Register>(instruction_bits: u32) -> Option<Instruction> {
                 // Reserved
                 None
             } else if uimm != 0 {
-                Some(Itype::new(insts::OP_RVC_SLLI, rd, rd, uimm).0)
+                Some(Itype::new(insts::OP_RVC_SLLI, rd, rd, uimm & u32::from(R::SHIFT_MASK)).0)
             } else {
                 Some(blank_instruction(insts::OP_RVC_SLLI64))
             }

src/machine/mod.rs

@@ -101,10 +101,12 @@ pub trait SupportMachine: CoreMachine {
         for program_header in &elf.program_headers {
             if program_header.p_type == PT_LOAD {
                 let aligned_start = round_page_down(program_header.p_vaddr);
-                let padding_start = program_header.p_vaddr - aligned_start;
-                let size = round_page_up(program_header.p_memsz + padding_start);
+                let padding_start = program_header.p_vaddr.wrapping_sub(aligned_start);
+                let size = round_page_up(program_header.p_memsz.wrapping_add(padding_start));
                 let slice_start = program_header.p_offset;
-                let slice_end = program_header.p_offset + program_header.p_filesz;
+                let slice_end = program_header
+                    .p_offset
+                    .wrapping_add(program_header.p_filesz);
                 if slice_start > slice_end || slice_end > program.len() as u64 {
                     return Err(Error::OutOfBound);
                 }

src/memory/mod.rs

@@ -88,7 +88,12 @@ pub fn check_permission<R: Register>(
     size: u64,
     flag: u8,
 ) -> Result<(), Error> {
-    let e = addr + size;
+    // fetch_flag below will check if requested memory is within bound. Here
+    // we only need to test for overflow first
+    let (e, overflowed) = addr.overflowing_add(size);
+    if overflowed {
+        return Err(Error::OutOfBound);
+    }
     let mut current_addr = round_page_down(addr);
     while current_addr < e {
         let page = current_addr / RISCV_PAGESIZE as u64;

tests/test_aot.rs

@@ -233,3 +233,37 @@ pub fn test_aot_invalid_read64() {
     assert!(result.is_err());
     assert_eq!(result.err(), Some(Error::OutOfBound));
 }
+
+#[test]
+pub fn test_aot_load_elf_crash_64() {
+    let mut file = File::open("tests/programs/load_elf_crash_64").unwrap();
+    let mut buffer = Vec::new();
+    file.read_to_end(&mut buffer).unwrap();
+    let buffer: Bytes = buffer.into();
+
+    let mut aot_machine = AotCompilingMachine::load(&buffer.clone(), None).unwrap();
+    let code = aot_machine.compile().unwrap();
+    let mut machine = AsmMachine::default_with_aot_code(&code);
+    machine
+        .load_program(&buffer, &vec!["load_elf_crash_64".into()])
+        .unwrap();
+    let result = machine.run();
+    assert_eq!(result.err(), Some(Error::InvalidPermission));
+}
+
+#[test]
+pub fn test_aot_wxorx_crash_64() {
+    let mut file = File::open("tests/programs/wxorx_crash_64").unwrap();
+    let mut buffer = Vec::new();
+    file.read_to_end(&mut buffer).unwrap();
+    let buffer: Bytes = buffer.into();
+
+    let mut aot_machine = AotCompilingMachine::load(&buffer.clone(), None).unwrap();
+    let code = aot_machine.compile().unwrap();
+    let mut machine = AsmMachine::default_with_aot_code(&code);
+    machine
+        .load_program(&buffer, &vec!["wxorx_crash_64".into()])
+        .unwrap();
+    let result = machine.run();
+    assert_eq!(result.err(), Some(Error::OutOfBound));
+}

tests/test_asm.rs

@@ -208,3 +208,33 @@ pub fn test_invalid_read64() {
     assert!(result.is_err());
     assert_eq!(result.err(), Some(Error::OutOfBound));
 }
+
+#[test]
+pub fn test_asm_load_elf_crash_64() {
+    let mut file = File::open("tests/programs/load_elf_crash_64").unwrap();
+    let mut buffer = Vec::new();
+    file.read_to_end(&mut buffer).unwrap();
+    let buffer: Bytes = buffer.into();
+
+    let mut machine = AsmMachine::default();
+    machine
+        .load_program(&buffer, &vec!["load_elf_crash_64".into()])
+        .unwrap();
+    let result = machine.run();
+    assert_eq!(result.err(), Some(Error::InvalidPermission));
+}
+
+#[test]
+pub fn test_asm_wxorx_crash_64() {
+    let mut file = File::open("tests/programs/wxorx_crash_64").unwrap();
+    let mut buffer = Vec::new();
+    file.read_to_end(&mut buffer).unwrap();
+    let buffer: Bytes = buffer.into();
+
+    let mut machine = AsmMachine::default();
+    machine
+        .load_program(&buffer, &vec!["wxorx_crash_64".into()])
+        .unwrap();
+    let result = machine.run();
+    assert_eq!(result.err(), Some(Error::OutOfBound));
+}

tests/test_misc.rs

@@ -131,3 +131,58 @@ pub fn test_invalid_file_offset64() {
     let result = run::<u64, SparseMemory<u64>>(&buffer, &vec!["invalid_file_offset64".into()]);
     assert_eq!(result.err(), Some(Error::OutOfBound));
 }
+
+#[test]
+pub fn test_op_rvc_srli_crash_32() {
+    let mut file = File::open("tests/programs/op_rvc_srli_crash_32").unwrap();
+    let mut buffer = Vec::new();
+    file.read_to_end(&mut buffer).unwrap();
+    let buffer: Bytes = buffer.into();
+
+    let result = run::<u32, SparseMemory<u32>>(&buffer, &vec!["op_rvc_srli_crash_32".into()]);
+    assert_eq!(result.err(), Some(Error::InvalidPermission));
+}
+
+#[test]
+pub fn test_op_rvc_srai_crash_32() {
+    let mut file = File::open("tests/programs/op_rvc_srai_crash_32").unwrap();
+    let mut buffer = Vec::new();
+    file.read_to_end(&mut buffer).unwrap();
+    let buffer: Bytes = buffer.into();
+
+    let result = run::<u32, SparseMemory<u32>>(&buffer, &vec!["op_rvc_srai_crash_32".into()]);
+    assert!(result.is_ok());
+}
+
+#[test]
+pub fn test_op_rvc_slli_crash_32() {
+    let mut file = File::open("tests/programs/op_rvc_slli_crash_32").unwrap();
+    let mut buffer = Vec::new();
+    file.read_to_end(&mut buffer).unwrap();
+    let buffer: Bytes = buffer.into();
+
+    let result = run::<u32, SparseMemory<u32>>(&buffer, &vec!["op_rvc_slli_crash_32".into()]);
+    assert!(result.is_ok());
+}
+
+#[test]
+pub fn test_load_elf_crash_64() {
+    let mut file = File::open("tests/programs/load_elf_crash_64").unwrap();
+    let mut buffer = Vec::new();
+    file.read_to_end(&mut buffer).unwrap();
+    let buffer: Bytes = buffer.into();
+
+    let result = run::<u64, SparseMemory<u64>>(&buffer, &vec!["load_elf_crash_64".into()]);
+    assert_eq!(result.err(), Some(Error::InvalidPermission));
+}
+
+#[test]
+pub fn test_wxorx_crash_64() {
+    let mut file = File::open("tests/programs/wxorx_crash_64").unwrap();
+    let mut buffer = Vec::new();
+    file.read_to_end(&mut buffer).unwrap();
+    let buffer: Bytes = buffer.into();
+
+    let result = run::<u64, SparseMemory<u64>>(&buffer, &vec!["wxorx_crash_64".into()]);
+    assert_eq!(result.err(), Some(Error::OutOfBound));
+}