Rust-related/RustPython
1
0
Fork 1
mirror of https://github.com/RustPython/RustPython.git synced 2026-06-02 19:39:49 +09:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #2989 from DimitrisJim/repeat_overflow_checks

Browse Source 
Check for MemoryErrors when repeating sequences.
This commit is contained in:
Jeong YunWon
2021-08-30 16:12:37 +09:00
committed by GitHub
parent f4d40bc0c9 97b8d7ca7b
commit 4ddd1d3ee3
13 changed files with 81 additions and 63 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
Lib/test/test_array.py
View File

@@ -1508,7 +1508,6 @@ class DoubleTest(FPTest, unittest.TestCase):
typecode = 'd'
minitemsize = 8
@unittest.skip("TODO: RUSTPYTHON, thread 'main' panicked at 'capacity overflow'")
def test_alloc_overflow(self):
from sys import maxsize
a = array.array('d', [-1]*65536)

1
Lib/test/test_bytes.py
View File

@@ -396,7 +396,6 @@ class BaseBytesTest:
self.assertRaises(TypeError, lambda: b1 + "def")
self.assertRaises(TypeError, lambda: "abc" + b2)
@unittest.skip("TODO: RUSTPYTHON, thread 'main' panicked at 'capacity overflow'")
def test_repeat(self):
for b in b"abc", self.type2test(b"abc"):
self.assertEqual(b * 3, b"abcabcabc")

1
Lib/test/test_list.py
View File

@@ -61,7 +61,6 @@ class ListTest(list_tests.CommonTest):
self.assertEqual(len([0]), 1)
self.assertEqual(len([0, 1, 2]), 3)
@unittest.skip("TODO: RUSTPYTHON, thread 'main' panicked at 'capacity overflow'")
def test_overflow(self):
lst = [4, 5, 6, 7]
n = int((sys.maxsize*2+2) // len(lst))

16
vm/src/builtins/bytearray.rs
View File

@@ -273,7 +273,7 @@ impl PyByteArray {
}
}
fn irepeat(zelf: &PyRef<Self>, n: isize, vm: &VirtualMachine) -> PyResult<()> {
fn irepeat(zelf: &PyRef<Self>, n: usize, vm: &VirtualMachine) -> PyResult<()> {
if n == 1 {
return Ok(());
}
@@ -290,11 +290,9 @@ impl PyByteArray {
};
let elements = &mut w.elements;
if n <= 0 {
if n == 0 {
elements.clear();
} else if n != 1 {
let n = n as usize;
let old = elements.clone();
elements.reserve((n - 1) * old.len());
@@ -602,13 +600,15 @@ impl PyByteArray {
#[pymethod(name = "__rmul__")]
#[pymethod(magic)]
fn mul(&self, n: isize) -> Self {
self.inner().repeat(n).into()
fn mul(&self, value: isize, vm: &VirtualMachine) -> PyResult<Self> {
vm.check_repeat_or_memory_error(self.len(), value)
.map(|value| self.inner().repeat(value).into())
}
#[pymethod(magic)]
fn imul(zelf: PyRef<Self>, n: isize, vm: &VirtualMachine) -> PyResult<PyRef<Self>> {
Self::irepeat(&zelf, n, vm).map(|_| zelf)
fn imul(zelf: PyRef<Self>, value: isize, vm: &VirtualMachine) -> PyResult<PyRef<Self>> {
vm.check_repeat_or_memory_error(zelf.len(), value)
.and_then(|value| Self::irepeat(&zelf, value, vm).map(|_| zelf))
}
#[pymethod(name = "__mod__")]

13
vm/src/builtins/bytes.rs
View File

@@ -440,9 +440,6 @@ impl PyBytes {
#[pymethod(name = "__rmul__")]
#[pymethod(magic)]
fn mul(zelf: PyRef<Self>, value: isize, vm: &VirtualMachine) -> PyResult<PyRef<Self>> {
if value > 0 && zelf.inner.len() as isize > std::isize::MAX / value {
return Err(vm.new_overflow_error("repeated bytes are too long".to_owned()));
}
if value == 1 && zelf.class().is(&vm.ctx.types.bytes_type) {
// Special case: when some `bytes` is multiplied by `1`,
// nothing really happens, we need to return an object itself
@@ -450,8 +447,14 @@ impl PyBytes {
// This only works for `bytes` itself, not its subclasses.
return Ok(zelf);
}
let bytes: PyBytes = zelf.inner.repeat(value).into();
Ok(bytes.into_ref(vm))
// todo: map err to overflow.
vm.check_repeat_or_memory_error(zelf.inner.len(), value)
.map(|value| {
let bytes: PyBytes = zelf.inner.repeat(value).into();
bytes.into_ref(vm)
})
// see issue 45044 on b.p.o.
.map_err(|_| vm.new_overflow_error("repeated bytes are too long".to_owned()))
}
#[pymethod(name = "__mod__")]

18
vm/src/builtins/list.rs
View File

@@ -237,25 +237,21 @@ impl PyList {
}
#[pymethod(magic)]
fn mul(&self, counter: isize, vm: &VirtualMachine) -> PyObjectRef {
let new_elements = sequence::seq_mul(&self.borrow_vec(), counter)
#[pymethod(name = "__rmul__")]
fn mul(&self, value: isize, vm: &VirtualMachine) -> PyResult {
let new_elements = sequence::seq_mul(vm, &self.borrow_vec(), value)?
.cloned()
.collect();
vm.ctx.new_list(new_elements)
Ok(vm.ctx.new_list(new_elements))
}
#[pymethod(magic)]
fn rmul(&self, counter: isize, vm: &VirtualMachine) -> PyObjectRef {
self.mul(counter, vm)
}
#[pymethod(magic)]
fn imul(zelf: PyRef<Self>, counter: isize) -> PyRef<Self> {
fn imul(zelf: PyRef<Self>, value: isize, vm: &VirtualMachine) -> PyResult<PyRef<Self>> {
let mut elements = zelf.borrow_vec_mut();
let mut new_elements: Vec<PyObjectRef> =
sequence::seq_mul(&*elements, counter).cloned().collect();
sequence::seq_mul(vm, &*elements, value)?.cloned().collect();
std::mem::swap(elements.deref_mut(), &mut new_elements);
zelf.clone()
Ok(zelf.clone())
}
#[pymethod]

10
vm/src/builtins/pystr.rs
View File

@@ -342,15 +342,19 @@ impl PyStr {
#[pymethod(name = "__rmul__")]
#[pymethod(magic)]
fn mul(zelf: PyRef<Self>, value: isize, vm: &VirtualMachine) -> PyRef<Self> {
fn mul(zelf: PyRef<Self>, value: isize, vm: &VirtualMachine) -> PyResult<PyRef<Self>> {
if value == 1 && zelf.class().is(&vm.ctx.types.str_type) {
// Special case: when some `str` is multiplied by `1`,
// nothing really happens, we need to return an object itself
// with the same `id()` to be compatible with CPython.
// This only works for `str` itself, not its subclasses.
return zelf;
return Ok(zelf);
}
Self::from(zelf.value.repeat(value.to_usize().unwrap_or(0))).into_ref(vm)
// todo: map err to overflow.
vm.check_repeat_or_memory_error(zelf.len(), value)
.map(|value| Self::from(zelf.value.repeat(value)).into_ref(vm))
// see issue 45044 on b.p.o.
.map_err(|_| vm.new_overflow_error("repeated bytes are too long".to_owned()))
}
#[pymethod(magic)]

10
vm/src/builtins/tuple.rs
View File

@@ -169,22 +169,22 @@ impl PyTuple {
#[pymethod(name = "__rmul__")]
#[pymethod(magic)]
fn mul(zelf: PyRef<Self>, counter: isize, vm: &VirtualMachine) -> PyRef<Self> {
if zelf.elements.is_empty() || counter == 0 {
fn mul(zelf: PyRef<Self>, value: isize, vm: &VirtualMachine) -> PyResult<PyRef<Self>> {
Ok(if zelf.elements.is_empty() || value == 0 {
vm.ctx.empty_tuple.clone()
} else if counter == 1 && zelf.class().is(&vm.ctx.types.tuple_type) {
} else if value == 1 && zelf.class().is(&vm.ctx.types.tuple_type) {
// Special case: when some `tuple` is multiplied by `1`,
// nothing really happens, we need to return an object itself
// with the same `id()` to be compatible with CPython.
// This only works for `tuple` itself, not its subclasses.
zelf
} else {
let elements = sequence::seq_mul(&zelf.elements, counter)
let elements = sequence::seq_mul(vm, &zelf.elements, value)?
.cloned()
.collect::<Vec<_>>()
.into_boxed_slice();
Self { elements }.into_ref(vm)
}
})
}
#[pymethod(magic)]

4
vm/src/bytesinner.rs
View File

@@ -964,8 +964,8 @@ impl PyBytesInner {
.format(vm, values)
}
pub fn repeat(&self, n: isize) -> Vec<u8> {
self.elements.repeat(n.to_usize().unwrap_or(0))
pub fn repeat(&self, n: usize) -> Vec<u8> {
self.elements.repeat(n)
}
}

23
vm/src/sequence.rs
View File

@@ -1,7 +1,6 @@
use crate::slots::PyComparisonOp;
use crate::vm::VirtualMachine;
use crate::{PyObjectRef, PyResult};
use num_traits::cast::ToPrimitive;
pub(super) type DynPyIter<'a> = Box<dyn ExactSizeIterator<Item = &'a PyObjectRef> + 'a>;
@@ -94,15 +93,15 @@ impl<'a> Iterator for SeqMul<'a> {
}
}
pub(crate) fn seq_mul(seq: &impl SimpleSeq, repetitions: isize) -> SeqMul {
let repetitions = if seq.len() > 0 {
repetitions.to_usize().unwrap_or(0)
} else {
0
};
SeqMul {
seq,
repetitions,
iter: None,
}
pub(crate) fn seq_mul<'a>(
vm: &VirtualMachine,
seq: &'a impl SimpleSeq,
repetitions: isize,
) -> PyResult<SeqMul<'a>> {
vm.check_repeat_or_memory_error(seq.len(), repetitions)
.map(|repetitions| SeqMul {
seq,
repetitions,
iter: None,
})
}

22
vm/src/stdlib/array.rs
View File

@@ -301,8 +301,7 @@ macro_rules! def_array_enum {
}
}
fn mul(&self, counter: isize) -> Self {
let counter = if counter < 0 { 0 } else { counter as usize };
fn mul(&self, counter: usize) -> Self {
match self {
$(ArrayContentType::$n(v) => {
let elements = v.repeat(counter);
@@ -317,11 +316,10 @@ macro_rules! def_array_enum {
}
}
fn imul(&mut self, counter: isize) {
if counter <= 0 {
fn imul(&mut self, counter: usize) {
if counter == 0 {
self.clear();
} else if counter != 1 {
let counter = counter as usize;
match self {
$(ArrayContentType::$n(v) => {
let old = v.clone();
@@ -864,14 +862,18 @@ impl PyArray {
#[pymethod(name = "__rmul__")]
#[pymethod(magic)]
fn mul(&self, counter: isize, vm: &VirtualMachine) -> PyRef<Self> {
PyArray::from(self.read().mul(counter)).into_ref(vm)
fn mul(&self, value: isize, vm: &VirtualMachine) -> PyResult<PyRef<Self>> {
vm.check_repeat_or_memory_error(self.len(), value)
.map(|value| PyArray::from(self.read().mul(value)).into_ref(vm))
}
#[pymethod(magic)]
fn imul(zelf: PyRef<Self>, counter: isize, vm: &VirtualMachine) -> PyResult<PyRef<Self>> {
zelf.try_resizable(vm)?.imul(counter);
Ok(zelf)
fn imul(zelf: PyRef<Self>, value: isize, vm: &VirtualMachine) -> PyResult<PyRef<Self>> {
vm.check_repeat_or_memory_error(zelf.len(), value)
.and_then(|value| {
zelf.try_resizable(vm)?.imul(value);
Ok(zelf)
})
}
#[pymethod(magic)]

8
vm/src/stdlib/collections.rs
View File

@@ -454,20 +454,20 @@ mod _collections {
#[pymethod(magic)]
#[pymethod(name = "__rmul__")]
fn mul(&self, n: isize) -> Self {
fn mul(&self, value: isize, vm: &VirtualMachine) -> PyResult<Self> {
let deque: SimpleSeqDeque = self.borrow_deque().into();
let mul = sequence::seq_mul(&deque, n);
let mul = sequence::seq_mul(vm, &deque, value)?;
let skipped = self
.maxlen
.and_then(|maxlen| mul.len().checked_sub(maxlen))
.unwrap_or(0);
let deque = mul.skip(skipped).cloned().collect();
PyDeque {
Ok(PyDeque {
deque: PyRwLock::new(deque),
maxlen: self.maxlen,
state: AtomicCell::new(0),
}
})
}
#[pymethod(magic)]

17
vm/src/vm.rs
View File

@@ -784,6 +784,11 @@ impl VirtualMachine {
self.new_exception_msg(runtime_error, msg)
}
pub fn new_memory_error(&self, msg: String) -> PyBaseExceptionRef {
let memory_error_type = self.ctx.exceptions.memory_error.clone();
self.new_exception_msg(memory_error_type, msg)
}
pub fn new_stop_iteration(&self) -> PyBaseExceptionRef {
let stop_iteration_type = self.ctx.exceptions.stop_iteration.clone();
self.new_exception_empty(stop_iteration_type)
@@ -1883,6 +1888,18 @@ impl VirtualMachine {
})
}
/// Checks that the multiplication is able to be performed. On Ok returns the
/// index as a usize for sequences to be able to use immediately.
pub fn check_repeat_or_memory_error(&self, length: usize, n: isize) -> PyResult<usize> {
let n = n.to_usize().unwrap_or(0);
if n > 0 && length > isize::MAX as usize / n {
// Empty message is currently used in CPython.
Err(self.new_memory_error("".to_owned()))
} else {
Ok(n)
}
}
// https://docs.python.org/3/reference/expressions.html#membership-test-operations
fn _membership_iter_search(&self, haystack: PyObjectRef, needle: PyObjectRef) -> PyResult {
let iter = iterator::get_iter(self, haystack)?;