Handle MemoryError for sequences

crazymerlyn · crazymerlyn · commit ff4df62b2765 · 2024-09-28T03:25:04.000+01:00
diff --git a/vm/src/sequence.rs b/vm/src/sequence.rs
@@ -1,6 +1,10 @@
 use crate::{
-    builtins::PyIntRef, function::OptionalArg, sliceable::SequenceIndexOp, types::PyComparisonOp,
-    vm::VirtualMachine, AsObject, PyObject, PyObjectRef, PyResult,
+    builtins::PyIntRef,
+    function::OptionalArg,
+    sliceable::SequenceIndexOp,
+    types::PyComparisonOp,
+    vm::{VirtualMachine, MAX_MEMORY_SIZE},
+    AsObject, PyObject, PyObjectRef, PyResult,
 };
 use optional::Optioned;
 use std::ops::{Deref, Range};
@@ -95,6 +99,11 @@ where
 {
     fn mul(&self, vm: &VirtualMachine, n: isize) -> PyResult<Vec<T>> {
         let n = vm.check_repeat_or_overflow_error(self.as_ref().len(), n)?;
+
+        if self.as_ref().len() * n >= MAX_MEMORY_SIZE {
+            return Err(vm.new_memory_error("".to_owned()));
+        }
+
         let mut v = Vec::with_capacity(n * self.as_ref().len());
         for _ in 0..n {
             v.extend_from_slice(self.as_ref());
diff --git a/vm/src/vm/mod.rs b/vm/src/vm/mod.rs
@@ -53,6 +53,8 @@ pub use interpreter::Interpreter;
 pub(crate) use method::PyMethod;
 pub use setting::Settings;
 
+pub const MAX_MEMORY_SIZE: usize = (std::usize::MAX >> 3) + 1;
+
 // Objects are live when they are on stack, or referenced by a name (for now)
 
 /// Top level container of a python virtual machine. In theory you could
