Skip to content

Commit 09a5a5e

Browse files
youknowoneyouknowone
committed
Use checked arithmetic in LocalsPlus and DataStack allocators
1 parent f7907ce commit 09a5a5e

File tree

3 files changed

+71
-36
lines changed

3 files changed

+71
-36
lines changed

‎crates/common/src/lock.rs‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,10 @@ cfg_if::cfg_if! {
2929
pub use std::sync::LazyLock;
3030
} else {
3131
pub struct LazyLock<T, F = fn() -> T>(core::cell::LazyCell<T, F>);
32-
// SAFETY: Without std, there can be no threads.
32+
// SAFETY: This branch is only active when both "std" and "threading"
33+
// features are absent — i.e., truly single-threaded no_std environments
34+
// (e.g., embedded or bare-metal WASM). Without std, the Rust runtime
35+
// cannot spawn threads, so Sync is trivially satisfied.
3336
unsafe impl<T, F> Sync for LazyLock<T, F> {}
3437

3538
impl<T, F: FnOnce() -> T> LazyLock<T, F> {

‎crates/vm/src/datastack.rs‎

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -102,10 +102,15 @@ impl DataStack {
102102
#[inline(never)]
103103
fn push_slow(&mut self, aligned_size: usize) -> *mut u8 {
104104
let mut chunk_size = MIN_CHUNK_SIZE;
105-
let needed =
106-
aligned_size + MINIMUM_OVERHEAD + core::mem::size_of::<DataStackChunk>() + ALIGN;
105+
let needed = aligned_size
106+
.checked_add(MINIMUM_OVERHEAD)
107+
.and_then(|v| v.checked_add(core::mem::size_of::<DataStackChunk>()))
108+
.and_then(|v| v.checked_add(ALIGN))
109+
.expect("DataStack chunk size overflow");
107110
while chunk_size < needed {
108-
chunk_size *= 2;
111+
chunk_size = chunk_size
112+
.checked_mul(2)
113+
.expect("DataStack chunk size overflow");
109114
}
110115
// Save current position in old chunk.
111116
unsafe {

‎crates/vm/src/frame.rs‎

Lines changed: 59 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -163,10 +163,13 @@ const _: () = {
163163
impl LocalsPlus {
164164
/// Create a new heap-backed LocalsPlus. All slots start as None (0).
165165
fn new(nlocalsplus: usize, stacksize: usize) -> Self {
166-
let capacity = nlocalsplus + stacksize;
166+
let capacity = nlocalsplus
167+
.checked_add(stacksize)
168+
.expect("LocalsPlus capacity overflow");
169+
let nlocalsplus_u32 = u32::try_from(nlocalsplus).expect("nlocalsplus exceeds u32");
167170
Self {
168171
data: LocalsPlusData::Heap(vec![0usize; capacity].into_boxed_slice()),
169-
nlocalsplus: nlocalsplus as u32,
172+
nlocalsplus: nlocalsplus_u32,
170173
stack_top: 0,
171174
}
172175
}
@@ -177,14 +180,19 @@ impl LocalsPlus {
177180
/// The caller must call `materialize_localsplus()` when the frame finishes
178181
/// to migrate data to the heap, then `datastack_pop()` to free the memory.
179182
fn new_on_datastack(nlocalsplus: usize, stacksize: usize, vm: &VirtualMachine) -> Self {
180-
let capacity = nlocalsplus + stacksize;
181-
let byte_size = capacity * core::mem::size_of::<usize>();
183+
let capacity = nlocalsplus
184+
.checked_add(stacksize)
185+
.expect("LocalsPlus capacity overflow");
186+
let byte_size = capacity
187+
.checked_mul(core::mem::size_of::<usize>())
188+
.expect("LocalsPlus byte size overflow");
189+
let nlocalsplus_u32 = u32::try_from(nlocalsplus).expect("nlocalsplus exceeds u32");
182190
let ptr = vm.datastack_push(byte_size) as *mut usize;
183191
// Zero-initialize all slots (0 = None for both PyObjectRef and PyStackRef).
184192
unsafe { core::ptr::write_bytes(ptr, 0, capacity) };
185193
Self {
186194
data: LocalsPlusData::DataStack { ptr, capacity },
187-
nlocalsplus: nlocalsplus as u32,
195+
nlocalsplus: nlocalsplus_u32,
188196
stack_top: 0,
189197
}
190198
}
@@ -2362,18 +2370,16 @@ impl ExecutingFrame<'_> {
23622370
// Same as LoadFast but explicitly checks for unbound locals
23632371
// (LoadFast in RustPython already does this check)
23642372
let idx = idx.get(arg) as usize;
2365-
let x = self.localsplus.fastlocals()[idx]
2366-
.clone()
2367-
.ok_or_else(|| {
2368-
vm.new_exception_msg(
2369-
vm.ctx.exceptions.unbound_local_error.to_owned(),
2370-
format!(
2371-
"local variable '{}' referenced before assignment",
2372-
self.code.varnames[idx]
2373-
)
2374-
.into(),
2373+
let x = self.localsplus.fastlocals()[idx].clone().ok_or_else(|| {
2374+
vm.new_exception_msg(
2375+
vm.ctx.exceptions.unbound_local_error.to_owned(),
2376+
format!(
2377+
"local variable '{}' referenced before assignment",
2378+
self.code.varnames[idx]
23752379
)
2376-
})?;
2380+
.into(),
2381+
)
2382+
})?;
23772383
self.push_value(x);
23782384
Ok(None)
23792385
}
@@ -2413,18 +2419,16 @@ impl ExecutingFrame<'_> {
24132419
// lifetime issues at yield/exception points are resolved.
24142420
Instruction::LoadFastBorrow { var_num: idx } => {
24152421
let idx = idx.get(arg) as usize;
2416-
let x = self.localsplus.fastlocals()[idx]
2417-
.clone()
2418-
.ok_or_else(|| {
2419-
vm.new_exception_msg(
2420-
vm.ctx.exceptions.unbound_local_error.to_owned(),
2421-
format!(
2422-
"local variable '{}' referenced before assignment",
2423-
self.code.varnames[idx]
2424-
)
2425-
.into(),
2422+
let x = self.localsplus.fastlocals()[idx].clone().ok_or_else(|| {
2423+
vm.new_exception_msg(
2424+
vm.ctx.exceptions.unbound_local_error.to_owned(),
2425+
format!(
2426+
"local variable '{}' referenced before assignment",
2427+
self.code.varnames[idx]
24262428
)
2427-
})?;
2429+
.into(),
2430+
)
2431+
})?;
24282432
self.push_value(x);
24292433
Ok(None)
24302434
}
@@ -4329,7 +4333,10 @@ impl ExecutingFrame<'_> {
43294333
let nargs: u32 = arg.into();
43304334
let callable = self.nth_value(nargs + 1);
43314335
let stack_len = self.localsplus.stack_len();
4332-
let self_or_null_is_some = self.localsplus.stack_index(stack_len - nargs as usize - 1).is_some();
4336+
let self_or_null_is_some = self
4337+
.localsplus
4338+
.stack_index(stack_len - nargs as usize - 1)
4339+
.is_some();
43334340
if !self_or_null_is_some
43344341
&& cached_version != 0
43354342
&& let Some(cls) = callable.downcast_ref::<PyType>()
@@ -6029,11 +6036,21 @@ impl ExecutingFrame<'_> {
60296036
args_vec.push(self_val);
60306037
}
60316038
for stack_idx in args_start..stack_len {
6032-
let val = self.localsplus.stack_index_mut(stack_idx).take().unwrap().to_pyobj();
6039+
let val = self
6040+
.localsplus
6041+
.stack_index_mut(stack_idx)
6042+
.take()
6043+
.unwrap()
6044+
.to_pyobj();
60336045
args_vec.push(val);
60346046
}
60356047

6036-
let callable_obj = self.localsplus.stack_index_mut(callable_idx).take().unwrap().to_pyobj();
6048+
let callable_obj = self
6049+
.localsplus
6050+
.stack_index_mut(callable_idx)
6051+
.take()
6052+
.unwrap()
6053+
.to_pyobj();
60376054
self.localsplus.stack_truncate(callable_idx);
60386055

60396056
let result = callable_obj.vectorcall(args_vec, effective_nargs, None, vm)?;
@@ -6109,11 +6126,21 @@ impl ExecutingFrame<'_> {
61096126
args_vec.push(self_val);
61106127
}
61116128
for stack_idx in args_start..stack_len {
6112-
let val = self.localsplus.stack_index_mut(stack_idx).take().unwrap().to_pyobj();
6129+
let val = self
6130+
.localsplus
6131+
.stack_index_mut(stack_idx)
6132+
.take()
6133+
.unwrap()
6134+
.to_pyobj();
61136135
args_vec.push(val);
61146136
}
61156137

6116-
let callable_obj = self.localsplus.stack_index_mut(callable_idx).take().unwrap().to_pyobj();
6138+
let callable_obj = self
6139+
.localsplus
6140+
.stack_index_mut(callable_idx)
6141+
.take()
6142+
.unwrap()
6143+
.to_pyobj();
61176144
self.localsplus.stack_truncate(callable_idx);
61186145

61196146
let kwnames = kwarg_names_tuple.as_slice();

0 commit comments

Comments
 (0)