Skip to content

Commit 79ca9c1

Browse files
youknowoneyouknowone
committed
Address review: bounds checks, UB fix, version overflow, error handling
- Add bounds checks to read_cache_u16/u32/u64 - Fix quicken() aliasing UB by using &mut directly - Add JumpBackwardJit/JumpBackwardNoJit to deoptimize() - Guard can_specialize_call with NEWLOCALS flag check - Use compare_exchange_weak for version tag to prevent wraparound - Propagate dict lookup errors in LoadAttrMethodWithValues - Apply adaptive backoff on version tag assignment failure - Remove duplicate imports in frame.rs
1 parent a831e4f commit 79ca9c1

File tree

6 files changed

+63
-23
lines changed

6 files changed

+63
-23
lines changed

‎crates/compiler-core/src/bytecode.rs‎

Lines changed: 12 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -466,8 +466,12 @@ impl CodeUnits {
466466
}
467467

468468
/// Read a u16 value from a CACHE code unit at `index`.
469+
///
470+
/// # Panics
471+
/// Panics if `index` is out of bounds.
469472
pub fn read_cache_u16(&self, index: usize) -> u16 {
470473
let units = unsafe { &*self.0.get() };
474+
assert!(index < units.len(), "read_cache_u16: index out of bounds");
471475
let ptr = units.as_ptr().wrapping_add(index) as *const u8;
472476
unsafe { core::ptr::read_unaligned(ptr as *const u16) }
473477
}
@@ -484,6 +488,9 @@ impl CodeUnits {
484488
}
485489

486490
/// Read a u32 value from two consecutive CACHE code units starting at `index`.
491+
///
492+
/// # Panics
493+
/// Panics if `index + 1` is out of bounds.
487494
pub fn read_cache_u32(&self, index: usize) -> u32 {
488495
let lo = self.read_cache_u16(index) as u32;
489496
let hi = self.read_cache_u16(index + 1) as u32;
@@ -502,6 +509,9 @@ impl CodeUnits {
502509
}
503510

504511
/// Read a u64 value from four consecutive CACHE code units starting at `index`.
512+
///
513+
/// # Panics
514+
/// Panics if `index + 3` is out of bounds.
505515
pub fn read_cache_u64(&self, index: usize) -> u64 {
506516
let lo = self.read_cache_u32(index) as u64;
507517
let hi = self.read_cache_u32(index + 2) as u64;
@@ -553,7 +563,7 @@ impl CodeUnits {
553563
/// Called lazily at RESUME (first execution of a code object).
554564
/// Uses the `arg` byte of the first CACHE entry, preserving `op = Instruction::Cache`.
555565
pub fn quicken(&self) {
556-
let units = unsafe { &*self.0.get() };
566+
let units = unsafe { &mut *self.0.get() };
557567
let len = units.len();
558568
let mut i = 0;
559569
while i < len {
@@ -565,9 +575,7 @@ impl CodeUnits {
565575
if !op.is_instrumented() {
566576
let cache_base = i + 1;
567577
if cache_base < len {
568-
unsafe {
569-
self.write_adaptive_counter(cache_base, ADAPTIVE_WARMUP_VALUE);
570-
}
578+
units[cache_base].arg = OpArgByte::from(ADAPTIVE_WARMUP_VALUE);
571579
}
572580
}
573581
i += 1 + caches;

‎crates/compiler-core/src/bytecode/instruction.rs‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -620,6 +620,10 @@ impl Instruction {
620620
}
621621
// RESUME specializations
622622
Self::ResumeCheck => Self::Resume { arg: Arg::marker() },
623+
// JUMP_BACKWARD specializations
624+
Self::JumpBackwardJit | Self::JumpBackwardNoJit => Self::JumpBackward {
625+
target: Arg::marker(),
626+
},
623627
// Instrumented opcodes map back to their base
624628
_ => match self.to_base() {
625629
Some(base) => base,

‎crates/vm/src/builtins/function.rs‎

Lines changed: 8 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -611,12 +611,14 @@ impl Py<PyFunction> {
611611
pub(crate) fn can_specialize_call(&self, effective_nargs: u32) -> bool {
612612
let code = self.code.lock();
613613
let flags = code.flags;
614-
!flags.intersects(
615-
bytecode::CodeFlags::VARARGS
616-
| bytecode::CodeFlags::VARKEYWORDS
617-
| bytecode::CodeFlags::GENERATOR
618-
| bytecode::CodeFlags::COROUTINE,
619-
) && code.kwonlyarg_count == 0
614+
flags.contains(bytecode::CodeFlags::NEWLOCALS)
615+
&& !flags.intersects(
616+
bytecode::CodeFlags::VARARGS
617+
| bytecode::CodeFlags::VARKEYWORDS
618+
| bytecode::CodeFlags::GENERATOR
619+
| bytecode::CodeFlags::COROUTINE,
620+
)
621+
&& code.kwonlyarg_count == 0
620622
&& code.arg_count == effective_nargs
621623
}
622624

‎crates/vm/src/builtins/type.rs‎

Lines changed: 12 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -201,12 +201,19 @@ fn is_subtype_with_mro(a_mro: &[PyTypeRef], a: &Py<PyType>, b: &Py<PyType>) -> b
201201
impl PyType {
202202
/// Assign a fresh version tag. Returns 0 on overflow (all caches invalidated).
203203
pub fn assign_version_tag(&self) -> u32 {
204-
let v = NEXT_TYPE_VERSION.fetch_add(1, Ordering::Relaxed);
205-
if v == 0 {
206-
return 0;
204+
loop {
205+
let current = NEXT_TYPE_VERSION.load(Ordering::Relaxed);
206+
let Some(next) = current.checked_add(1) else {
207+
return 0; // Overflow: version space exhausted
208+
};
209+
if NEXT_TYPE_VERSION
210+
.compare_exchange_weak(current, next, Ordering::Relaxed, Ordering::Relaxed)
211+
.is_ok()
212+
{
213+
self.tp_version_tag.store(current, Ordering::Release);
214+
return current;
215+
}
207216
}
208-
self.tp_version_tag.store(v, Ordering::Release);
209-
v
210217
}
211218

212219
/// Invalidate this type's version tag and cascade to all subclasses.

‎crates/vm/src/frame.rs‎

Lines changed: 23 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -8,10 +8,8 @@ use crate::{
88
PyFloat, PyGenerator, PyInt, PyInterpolation, PyList, PySet, PySlice, PyStr, PyStrInterned,
99
PyTemplate, PyTraceback, PyType, PyUtf8Str,
1010
asyncgenerator::PyAsyncGenWrappedValue,
11-
float::PyFloat,
1211
frame::stack_analysis,
1312
function::{PyCell, PyCellRef, PyFunction},
14-
int::PyInt,
1513
range::PyRangeIterator,
1614
tuple::{PyTuple, PyTupleRef},
1715
},
@@ -2722,7 +2720,23 @@ impl ExecutingFrame<'_> {
27222720
if type_version != 0 && owner.class().tp_version_tag.load(Acquire) == type_version {
27232721
// Check instance dict doesn't shadow the method
27242722
let shadowed = if let Some(dict) = owner.dict() {
2725-
dict.get_item_opt(attr_name, vm).ok().flatten().is_some()
2723+
match dict.get_item_opt(attr_name, vm) {
2724+
Ok(Some(_)) => true,
2725+
Ok(None) => false,
2726+
Err(_) => {
2727+
// Dict lookup error → deoptimize to safe path
2728+
unsafe {
2729+
self.code.instructions.replace_op(
2730+
instr_idx,
2731+
Instruction::LoadAttr { idx: Arg::marker() },
2732+
);
2733+
self.code
2734+
.instructions
2735+
.write_adaptive_counter(cache_base, ADAPTIVE_BACKOFF_VALUE);
2736+
}
2737+
return self.load_attr_slow(vm, oparg);
2738+
}
2739+
}
27262740
} else {
27272741
false
27282742
};
@@ -4471,7 +4485,12 @@ impl ExecutingFrame<'_> {
44714485
type_version = cls.assign_version_tag();
44724486
}
44734487
if type_version == 0 {
4474-
// Version counter overflow
4488+
// Version counter overflow — backoff to avoid re-attempting every execution
4489+
unsafe {
4490+
self.code
4491+
.instructions
4492+
.write_adaptive_counter(cache_base, ADAPTIVE_BACKOFF_VALUE);
4493+
}
44754494
return;
44764495
}
44774496

‎crates/vm/src/stdlib/sys/monitoring.rs‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -276,13 +276,13 @@ pub fn instrument_code(code: &PyCode, events: u32) {
276276
let mut i = 0;
277277
while i < len {
278278
let op = code.code.instructions[i].op;
279-
let de_opt = op.deoptimize();
280-
if u8::from(de_opt) != u8::from(op) {
279+
let base_op = op.deoptimize();
280+
if u8::from(base_op) != u8::from(op) {
281281
unsafe {
282-
code.code.instructions.replace_op(i, de_opt);
282+
code.code.instructions.replace_op(i, base_op);
283283
}
284284
}
285-
let caches = de_opt.cache_entries();
285+
let caches = base_op.cache_entries();
286286
// Zero all CACHE entries (the op+arg bytes may have been overwritten
287287
// by specialization with arbitrary data like pointers).
288288
for c in 1..=caches {

0 commit comments

Comments
 (0)