Skip to content

Commit 4eb4d2e

Browse files
youknowdotyouknowdot
committed
Align bytearray mutation and export validation
Assisted-by: OpenAI Codex:model version not exposed
1 parent 4ec53b3 commit 4eb4d2e

1 file changed

Lines changed: 119 additions & 61 deletions

File tree

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

Lines changed: 119 additions & 61 deletions
Original file line numberDiff line numberDiff line change
@@ -29,8 +29,8 @@ use crate::{
2929
PyMappingMethods, PyNumber, PyNumberMethods, PySequenceMethods,
3030
},
3131
sliceable::{
32-
NameBytearray, SaturatedSlice, SaturatedSliceIter, SequenceIndex, SliceableSequenceMutOp,
33-
SliceableSequenceOp,
32+
NameBytearray, SaturatedSlice, SaturatedSliceIter, SequenceIndex, SequenceIndexOp,
33+
SliceableSequenceMutOp, SliceableSequenceOp,
3434
},
3535
types::{
3636
AsBuffer, AsMapping, AsNumber, AsSequence, Callable, Comparable, Constructor,
@@ -98,6 +98,15 @@ enum ResizeAccess<'a> {
9898
Exported { len: usize },
9999
}
100100

101+
impl ResizeAccess<'_> {
102+
fn len(&self) -> usize {
103+
match self {
104+
Self::Exclusive(inner) => inner.elements.len(),
105+
Self::Exported { len } => *len,
106+
}
107+
}
108+
}
109+
101110
pub(crate) type PyByteArrayRef = PyRef<PyByteArray>;
102111

103112
impl From<PyBytesInner> for PyByteArray {
@@ -201,6 +210,11 @@ impl PyByteArray {
201210
}
202211
}
203212

213+
fn len_snapshot(&self) -> usize {
214+
let exported_len = { self.exports.lock().snapshot() };
215+
exported_len.unwrap_or_else(|| self.inner().elements.len())
216+
}
217+
204218
fn with_export<R>(&self, f: impl FnOnce(usize) -> R) -> R {
205219
let export = self.begin_export();
206220
f(export.len)
@@ -438,44 +452,48 @@ impl PyByteArray {
438452
}
439453
}
440454

455+
fn remove_by_index(&self, index: isize, pop: bool, vm: &VirtualMachine) -> PyResult<u8> {
456+
let access = self.resize_access();
457+
let len = access.len();
458+
let index = index.wrapped_at(len).ok_or_else(|| {
459+
let message = if !pop {
460+
"bytearray index out of range"
461+
} else if len == 0 {
462+
"pop from empty bytearray"
463+
} else {
464+
"pop index out of range"
465+
};
466+
vm.new_index_error(message)
467+
})?;
468+
// The exported length permits exact bounds checking without waiting
469+
// on a payload lock already held by a codec callback.
470+
let ResizeAccess::Exclusive(mut inner) = access else {
471+
return Err(vm.new_buffer_error("Existing exports of data: object cannot be re-sized"));
472+
};
473+
Ok(inner.elements.remove(index))
474+
}
475+
441476
fn delitem_by_index(&self, index: isize, vm: &VirtualMachine) -> PyResult<()> {
442-
// Validate before export rejection when the bytes are available without
443-
// waiting. A codec callback may already hold a read guard with a writer
444-
// queued, so a blocking read here could deadlock before that rejection.
445-
let valid = self
446-
.inner
447-
.try_read()
448-
.map(|inner| inner.elements.wrap_index(index).is_some());
449-
if valid == Some(false) {
450-
return Err(vm.new_index_error("bytearray index out of range"));
451-
}
452-
let mut inner = self.try_resizable(vm)?;
453-
// Another thread may have resized the bytearray between the locks.
454-
let index = inner
455-
.elements
456-
.wrap_index(index)
457-
.ok_or_else(|| vm.new_index_error("bytearray index out of range"))?;
458-
inner.elements.remove(index);
459-
Ok(())
477+
self.remove_by_index(index, false, vm).map(drop)
460478
}
461479

462480
fn irepeat(zelf: &Py<Self>, n: isize, vm: &VirtualMachine) -> PyResult<()> {
481+
if n <= 0 {
482+
return zelf.clear(vm);
483+
}
463484
if n == 1 {
464485
return Ok(());
465486
}
466-
let mut w = match zelf.try_resizable(vm) {
467-
Ok(w) => w,
468-
Err(err) => {
469-
return if zelf.borrow_buf().is_empty() {
470-
// We can multiple an empty vector by any integer
471-
Ok(())
472-
} else {
473-
Err(err)
474-
};
475-
}
487+
let access = zelf.resize_access();
488+
let len = access.len();
489+
vm.check_repeat_or_overflow_error(len, n)?;
490+
if len == 0 {
491+
return Ok(());
492+
}
493+
let ResizeAccess::Exclusive(mut inner) = access else {
494+
return Err(vm.new_buffer_error("Existing exports of data: object cannot be re-sized"));
476495
};
477-
478-
w.imul(n, vm)
496+
inner.imul(n, vm)
479497
}
480498
}
481499

@@ -516,12 +534,30 @@ impl PyByteArray {
516534

517535
fn __iadd__(
518536
zelf: PyRef<Self>,
519-
other: ArgBytesLike,
537+
other: PyObjectRef,
520538
vm: &VirtualMachine,
521539
) -> PyResult<PyRef<Self>> {
522-
zelf.try_resizable(vm)?
523-
.elements
524-
.extend(&*other.borrow_buf());
540+
let buffer = ArgBytesLike::try_from_object(vm, other.clone()).map_err(|_| {
541+
vm.new_type_error(format!(
542+
"can't concat {} to {}",
543+
other.class().name(),
544+
zelf.class().name()
545+
))
546+
})?;
547+
if !buffer.is_empty() {
548+
// Reject aliases before reading their payload: a codec callback
549+
// may already hold that read lock with a writer queued behind it.
550+
drop(zelf.try_resizable(vm)?);
551+
// Snapshot before locking the destination, but retain the export
552+
// until the write finishes, then recheck exports in assign_slice.
553+
let buffer = PyBuffer::from(buffer);
554+
let items = Self::buffer_snapshot(&buffer, vm)?;
555+
zelf.assign_slice(
556+
SaturatedSlice::from_parts(isize::MAX, isize::MAX, 1),
557+
&items,
558+
vm,
559+
)?;
560+
}
525561
Ok(zelf)
526562
}
527563

@@ -624,33 +660,50 @@ impl Py<PyByteArray> {
624660

625661
#[pymethod]
626662
fn pop(&self, index: PopArgs, vm: &VirtualMachine) -> PyResult<u8> {
627-
let elements = &mut self.try_resizable(vm)?.elements;
628-
let index = elements
629-
.wrap_index(index.index)
630-
.ok_or_else(|| vm.new_index_error("index out of range"))?;
631-
Ok(elements.remove(index))
663+
self.remove_by_index(index.index, true, vm)
632664
}
633665

634666
#[pymethod]
635667
fn insert(&self, index: isize, item: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> {
636668
let value = value_from_object(vm, &item)?;
637669
let elements = &mut self.try_resizable(vm)?.elements;
638670
let index = elements.saturate_index(index);
671+
elements.try_reserve(1).map_err(|_| vm.no_memory_error())?;
639672
elements.insert(index, value);
640673
Ok(())
641674
}
642675

643676
#[pymethod]
644677
fn append(&self, item: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> {
645678
let value = value_from_object(vm, &item)?;
646-
self.try_resizable(vm)?.elements.push(value);
679+
let elements = &mut self.try_resizable(vm)?.elements;
680+
elements.try_reserve(1).map_err(|_| vm.no_memory_error())?;
681+
elements.push(value);
647682
Ok(())
648683
}
649684

650685
#[pymethod]
651686
fn remove(&self, value: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> {
652687
let byte = value_from_object(vm, &value)?;
653-
let elements = &mut self.try_resizable(vm)?.elements;
688+
let mut inner = match self.resize_access() {
689+
ResizeAccess::Exclusive(inner) => inner,
690+
ResizeAccess::Exported { len } => {
691+
// Missing values precede export rejection. Do not wait for a
692+
// read lock: a reentrant codec can have a writer queued behind
693+
// its existing read guard.
694+
let missing = len == 0
695+
|| self
696+
.inner
697+
.try_read()
698+
.is_some_and(|inner| inner.elements.find_byte(byte).is_none());
699+
return Err(if missing {
700+
vm.new_value_error("value not found in bytearray")
701+
} else {
702+
vm.new_buffer_error("Existing exports of data: object cannot be re-sized")
703+
});
704+
}
705+
};
706+
let elements = &mut inner.elements;
654707
let index = elements
655708
.find_byte(byte)
656709
.ok_or_else(|| vm.new_value_error("value not found in bytearray"))?;
@@ -663,6 +716,13 @@ impl Py<PyByteArray> {
663716
if self.is(&iterable_of_ints) {
664717
return PyByteArray::irepeat(self, 2, vm);
665718
}
719+
// Buffer acquisition can mutate self. The buffer path inserts at the
720+
// old end, whereas the iterable path appends after all conversions.
721+
let start = if iterable_of_ints.check_buffer() {
722+
self.len_snapshot() as isize
723+
} else {
724+
isize::MAX
725+
};
666726
// bytearray_setslice keeps the export alive across the resize, so a value
667727
// looking at this bytearray is what stops it from growing.
668728
let buffer = iterable_of_ints
@@ -679,16 +739,21 @@ impl Py<PyByteArray> {
679739
})
680740
.transpose()?;
681741
let items = match &buffer {
682-
Some(buffer) => buffer
683-
.as_contiguous()
684-
.ok_or_else(|| {
685-
vm.new_buffer_error("non-contiguous buffer is not a bytes-like object")
686-
})?
687-
.to_vec(),
742+
Some(buffer) => {
743+
if !buffer.desc.is_contiguous() {
744+
return Err(
745+
vm.new_buffer_error("non-contiguous buffer is not a bytes-like object")
746+
);
747+
}
748+
if buffer.desc.len == 0 {
749+
return Ok(());
750+
}
751+
drop(self.try_resizable(vm)?);
752+
PyByteArray::buffer_snapshot(buffer, vm)?
753+
}
688754
None => bytearray_extend_from_object(vm, &iterable_of_ints)?,
689755
};
690-
self.try_resizable(vm)?.elements.extend(items);
691-
Ok(())
756+
self.assign_slice(SaturatedSlice::from_parts(start, start, 1), &items, vm)
692757
}
693758

694759
#[pymethod]
@@ -734,10 +799,7 @@ impl Py<PyByteArray> {
734799
}
735800
};
736801
let access = self.resize_access();
737-
let len = match &access {
738-
ResizeAccess::Exclusive(inner) => inner.elements.len(),
739-
ResizeAccess::Exported { len } => *len,
740-
};
802+
let len = access.len();
741803
// Vec<u8> lengths fit isize. Adding a nonnegative length to a negative
742804
// index cannot overflow; preserve the adjusted index in the error.
743805
let to_take = match n {
@@ -1195,10 +1257,7 @@ impl Py<PyByteArray> {
11951257
let size = size as usize;
11961258
// An export's stable length avoids re-entering the payload lock.
11971259
let access = self.resize_access();
1198-
let len = match &access {
1199-
ResizeAccess::Exclusive(inner) => inner.elements.len(),
1200-
ResizeAccess::Exported { len } => *len,
1201-
};
1260+
let len = access.len();
12021261
if size == len {
12031262
return Ok(());
12041263
}
@@ -1458,9 +1517,8 @@ impl AsSequence for PyByteArray {
14581517
PyByteArray::sequence_downcast(seq).__contains__(other.to_owned(), vm)
14591518
}),
14601519
inplace_concat: atomic_func!(|seq, other, vm| {
1461-
let other = ArgBytesLike::try_from_object(vm, other.to_owned())?;
14621520
let zelf = PyByteArray::sequence_downcast(seq).to_owned();
1463-
PyByteArray::__iadd__(zelf, other, vm).map(|x| x.into())
1521+
PyByteArray::__iadd__(zelf, other.to_owned(), vm).map(|x| x.into())
14641522
}),
14651523
inplace_repeat: atomic_func!(|seq, n, vm| {
14661524
let zelf = PyByteArray::sequence_downcast(seq).to_owned();

0 commit comments

Comments
 (0)