Skip to content

Commit ed6a8f6

Browse files
committed
Add check_initializing_stack to check mem issues
1 parent 68a5331 commit ed6a8f6

File tree

1 file changed

+45
-0
lines changed

1 file changed

+45
-0
lines changed

src/machine/mod.rs

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@ use super::memory::Memory;
1616
use super::syscalls::Syscalls;
1717
use super::{
1818
registers::{A0, A7, REGISTER_ABI_NAMES, SP},
19+
error::OutOfBoundKind,
1920
Error, ISA_MOP, RISCV_GENERAL_REGISTER_NUMBER,
2021
};
2122

@@ -252,6 +253,43 @@ pub trait SupportMachine: CoreMachine {
252253
Ok(stack_start + stack_size - self.registers()[SP].to_u64())
253254
}
254255

256+
// Dry run initialize_stack to see whether it will cause memorh errors.
257+
//
258+
// Returns `Ok(())` when the stack is enough. Returns `Err` when `initialize_stack` will fail because of
259+
// `MemOutOfStack` or `MemOutOfBound`.
260+
fn check_initializing_stack(
261+
&mut self,
262+
args_lens: &[u64],
263+
stack_start: u64,
264+
stack_size: u64,
265+
) -> Result<(), Error> {
266+
if self.version() >= VERSION1 && args_lens.is_empty() {
267+
return Ok(());
268+
}
269+
270+
let mut sp = stack_start + stack_size;
271+
let mut values_len = args_lens.len() as u64 + 1;
272+
for arg_len in args_lens {
273+
sp = wrapping_sub_address_or_mem_out_of_bound(sp, arg_len + 1)?;
274+
}
275+
if self.version() >= VERSION1 {
276+
values_len += 1;
277+
let values_bytes = Self::REG::BITS as u64 / 8 * values_len;
278+
let unaligned_sp_address = wrapping_sub_address_or_mem_out_of_bound(sp, values_bytes)?;
279+
let aligned_sp_address = unaligned_sp_address & (!15);
280+
let aligned_bytes = unaligned_sp_address - aligned_sp_address;
281+
sp = wrapping_sub_address_or_mem_out_of_bound(sp, aligned_bytes)?;
282+
}
283+
284+
sp = wrapping_sub_address_or_mem_out_of_bound(sp, Self::REG::BITS as u64 / 8 * values_len)?;
285+
if sp >= stack_start {
286+
Ok(())
287+
} else {
288+
// args exceed stack size
289+
Err(Error::MemOutOfStack)
290+
}
291+
}
292+
255293
#[cfg(feature = "pprof")]
256294
fn code(&self) -> &Bytes;
257295
}
@@ -766,6 +804,13 @@ impl Pause {
766804
}
767805
}
768806

807+
fn wrapping_sub_address_or_mem_out_of_bound(lhs: u64, rhs: u64) -> Result<u64, Error> {
808+
match lhs.overflowing_sub(rhs) {
809+
(result, false) => Ok(result),
810+
(result, true) => Err(Error::MemOutOfBound(result, OutOfBoundKind::Memory)),
811+
}
812+
}
813+
769814
#[cfg(test)]
770815
mod tests {
771816
use std::sync::atomic::AtomicU8;

0 commit comments

Comments
 (0)