diff --git a/vm/src/obj/objlist.rs b/vm/src/obj/objlist.rs index db4c981ee8..20f05fce9e 100644 --- a/vm/src/obj/objlist.rs +++ b/vm/src/obj/objlist.rs @@ -19,7 +19,7 @@ use crate::pyobject::{ PyObjectRef, PyRef, PyResult, PyValue, TryFromObject, TypeProtocol, }; use crate::sequence::{self, SimpleSeq}; -use crate::vm::{ReprGuard, VirtualMachine}; +use crate::vm::{ReprGuard, VirtualMachine, MAX_MEMORY_SIZE}; /// Built-in mutable sequence. /// @@ -119,14 +119,16 @@ impl PyList { Ok(result) => match result.as_bigint().to_u8() { Some(result) => elements.push(result), None => { - return Err(vm.new_value_error("bytes must be in range (0, 256)".to_owned())) + return Err( + vm.new_value_error("bytes must be in range (0, 256)".to_owned()) + ); } }, _ => { return Err(vm.new_type_error(format!( "'{}' object cannot be interpreted as an integer", elem.class().name - ))) + ))); } } } @@ -468,25 +470,31 @@ impl PyList { } #[pymethod(name = "__mul__")] - fn mul(&self, counter: isize, vm: &VirtualMachine) -> PyObjectRef { - let new_elements = sequence::seq_mul(&self.borrow_sequence(), counter) - .cloned() - .collect(); - vm.ctx.new_list(new_elements) + fn mul(&self, counter: isize, vm: &VirtualMachine) -> PyResult { + if counter < 0 || self.len() * (counter as usize) < MAX_MEMORY_SIZE { + let new_elements = sequence::seq_mul(&self.borrow_sequence(), counter) + .cloned() + .collect(); + return Ok(vm.ctx.new_list(new_elements)); + } + Err(vm.new_memory_error("".to_owned())) } #[pymethod(name = "__rmul__")] - fn rmul(&self, counter: isize, vm: &VirtualMachine) -> PyObjectRef { + fn rmul(&self, counter: isize, vm: &VirtualMachine) -> PyResult { self.mul(counter, &vm) } #[pymethod(name = "__imul__")] - fn imul(zelf: PyRef, counter: isize) -> PyRef { - let new_elements = sequence::seq_mul(&zelf.borrow_sequence(), counter) - .cloned() - .collect(); - zelf.elements.replace(new_elements); - zelf + fn imul(zelf: PyRef, counter: isize, vm: &VirtualMachine) -> PyResult> { + if counter < 0 || zelf.len() * (counter as usize) < MAX_MEMORY_SIZE { + let new_elements = sequence::seq_mul(&zelf.borrow_sequence(), counter) + .cloned() + .collect(); + zelf.elements.replace(new_elements); + return Ok(zelf); + } + Err(vm.new_memory_error("".to_owned())) } #[pymethod] diff --git a/vm/src/obj/objstr.rs b/vm/src/obj/objstr.rs index f4538eb705..02c5c3a8dc 100644 --- a/vm/src/obj/objstr.rs +++ b/vm/src/obj/objstr.rs @@ -33,7 +33,7 @@ use crate::pyobject::{ Either, IdProtocol, IntoPyObject, ItemProtocol, PyClassImpl, PyContext, PyIterable, PyObjectRef, PyRef, PyResult, PyValue, TryIntoRef, TypeProtocol, }; -use crate::vm::VirtualMachine; +use crate::vm::{VirtualMachine, MAX_MEMORY_SIZE}; /// str(object='') -> str /// str(bytes_or_buffer[, encoding[, errors]]) -> str @@ -328,13 +328,16 @@ impl PyString { #[pymethod(name = "__mul__")] fn mul(&self, multiplier: isize, vm: &VirtualMachine) -> PyResult { - multiplier - .max(0) - .to_usize() - .map(|multiplier| self.value.repeat(multiplier)) - .ok_or_else(|| { - vm.new_overflow_error("cannot fit 'int' into an index-sized integer".to_owned()) - }) + if multiplier < 0 || self.len() * (multiplier as usize) < MAX_MEMORY_SIZE { + return multiplier + .max(0) + .to_usize() + .map(|multiplier| self.value.repeat(multiplier)) + .ok_or_else(|| { + vm.new_overflow_error("cannot fit 'int' into an index-sized integer".to_owned()) + }); + } + Err(vm.new_memory_error("".to_owned())) } #[pymethod(name = "__rmul__")] diff --git a/vm/src/vm.rs b/vm/src/vm.rs index 6d3493bf31..cb36114a9e 100644 --- a/vm/src/vm.rs +++ b/vm/src/vm.rs @@ -73,6 +73,7 @@ pub struct VirtualMachine { } pub const NSIG: usize = 64; +pub const MAX_MEMORY_SIZE: usize = (std::usize::MAX >> 3) + 1; #[derive(Copy, Clone)] pub enum InitParameter { @@ -383,6 +384,11 @@ impl VirtualMachine { self.new_exception_msg(type_error, msg) } + pub fn new_memory_error(&self, msg: String) -> PyBaseExceptionRef { + let memory_error = self.ctx.exceptions.memory_error.clone(); + self.new_exception_msg(memory_error, msg) + } + pub fn new_name_error(&self, msg: String) -> PyBaseExceptionRef { let name_error = self.ctx.exceptions.name_error.clone(); self.new_exception_msg(name_error, msg)