8000 Add MAX_MEMORY_SIZE, new_memory_error, and size checks to objlist / · RustPython/RustPython@822badd · GitHub
[go: up one dir, main page]
Skip to content

Commit 822badd

Browse files
philippeitisphilippeitis
philippeitis
authored and
philippeitis
committed
Add MAX_MEMORY_SIZE, new_memory_error, and size checks to objlist /
objstr
1 parent bd6a6ce commit 822badd

File tree

3 files changed

+40
-23
lines changed

3 files changed

+40
-23
lines changed

‎vm/src/obj/objlist.rs

Lines changed: 23 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@ use crate::pyobject::{
1919
PyObjectRef, PyRef, PyResult, PyValue, TryFromObject, TypeProtocol,
2020
};
2121
use crate::sequence::{self, SimpleSeq};
22-
use crate::vm::{ReprGuard, VirtualMachine};
22+
use crate::vm::{ReprGuard, VirtualMachine, MAX_MEMORY_SIZE};
2323

2424
/// Built-in mutable sequence.
2525
///
@@ -119,14 +119,16 @@ impl PyList {
119119
Ok(result) => match result.as_bigint().to_u8() {
120120
Some(result) => elements.push(result),
121121
None => {
122-
return Err(vm.new_value_error("bytes must be in range (0, 256)".to_owned()))
122+
return Err(
123+
vm.new_value_error("bytes must be in range (0, 256)".to_owned())
124+
);
123125
}
124126
},
125127
_ => {
126128
return Err(vm.new_type_error(format!(
127129
"'{}' object cannot be interpreted as an integer",
128130
elem.class().name
129-
)))
131+
)));
130132
}
131133
}
132134
}
@@ -468,25 +470,31 @@ impl PyList {
468470
}
469471

470472
#[pymethod(name = "__mul__")]
471-
fn mul(&self, counter: isize, vm: &VirtualMachine) -> PyObjectRef {
472-
let new_elements = sequence::seq_mul(&self.borrow_sequence(), counter)
473-
.cloned()
474-
.collect();
475-
vm.ctx.new_list(new_elements)
473+
fn mul(&self, counter: isize, vm: &VirtualMachine) -> PyResult<PyObjectRef> {
474+
if counter < 0 || self.len() * (counter as usize) < MAX_MEMORY_SIZE {
475+
let new_elements = sequence::seq_mul(&self.borrow_sequence(), counter)
476+
.cloned()
477+
.collect();
478+
return Ok(vm.ctx.new_list(new_elements));
479+
}
480+
return Err(vm.new_memory_error("".to_owned()));
476481
}
477482

478483
#[pymethod(name = "__rmul__")]
479-
fn rmul(&self, counter: isize, vm: &VirtualMachine) -> PyObjectRef {
484+
fn rmul(&self, counter: isize, vm: &VirtualMachine) -> PyResult<PyObjectRef> {
480485
self.mul(counter, &vm)
481486
}
482487

483488
#[pymethod(name = "__imul__")]
484-
fn imul(zelf: PyRef<Self>, counter: isize) -> PyRef<Self> {
485-
let new_elements = sequence::seq_mul(&zelf.borrow_sequence(), counter)
486-
.cloned()
487-
.collect();
488-
zelf.elements.replace(new_elements);
489-
zelf
489+
fn imul(zelf: PyRef<Self>, counter: isize) -> PyResult<Self> {
490+
if counter < 0 || zelf.len() * (counter as usize) < MAX_MEMORY_SIZE {
491+
let new_elements = sequence::seq_mul(&zelf.borrow_sequence(), counter)
492+
.cloned()
493+
.collect();
494+
zelf.elements.replace(new_elements);
495+
return Ok(*zelf);
496+
}
497+
return Err(vm.new_memory_error("".to_owned()));
490498
}
491499

492500
#[pymethod]

‎vm/src/obj/objstr.rs

Lines changed: 11 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -33,7 +33,7 @@ use crate::pyobject::{
3333
Either, IdProtocol, IntoPyObject, ItemProtocol, PyClassImpl, PyContext, PyIterable,
3434
PyObjectRef, PyRef, PyResult, PyValue, TryIntoRef, TypeProtocol,
3535
};
36-
use crate::vm::VirtualMachine;
36+
use crate::vm::{VirtualMachine, MAX_MEMORY_SIZE};
3737

3838
/// str(object='') -> str
3939
/// str(bytes_or_buffer[, encoding[, errors]]) -> str
@@ -328,13 +328,16 @@ impl PyString {
328328

329329
#[pymethod(name = "__mul__")]
330330
fn mul(&self, multiplier: isize, vm: &VirtualMachine) -> PyResult<String> {
331-
multiplier
332-
.max(0)
333-
.to_usize()
334-
.map(|multiplier| self.value.repeat(multiplier))
335-
.ok_or_else(|| {
336-
vm.new_overflow_error("cannot fit 'int' into an index-sized integer".to_owned())
337-
})
331+
if multiplier < 0 || self.len() * (multiplier as usize) < MAX_MEMORY_SIZE {
332+
return multiplier
333+
.max(0)
334+
.to_usize()
335+
.map(|multiplier| self.value.repeat(multiplier))
336+
.ok_or_else(|| {
337+
vm.new_overflow_error("cannot fit 'int' into an index-sized integer".to_owned())
338+
});
339+
}
340+
return Err(vm.new_memory_error("".to_owned()));
338341
}
339342

340343
#[pymethod(name = "__rmul__")]

‎vm/src/vm.rs

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -73,6 +73,7 @@ pub struct VirtualMachine {
7373
}
7474

7575
pub const NSIG: usize = 64;
76+
pub const MAX_MEMORY_SIZE: usize = 0x2000_0000_0000_0000;
7677

7778
#[derive(Copy, Clone)]
7879
pub enum InitParameter {
@@ -383,6 +384,11 @@ impl VirtualMachine {
383384
self.new_exception_msg(type_error, msg)
384385
}
385386

387+
pub fn new_memory_error(&self, msg: String) -> PyBaseExceptionRef {
388+
let memory_error = self.ctx.exceptions.memory_error.clone();
389+
self.new_exception_msg(memory_error, msg)
390+
}
391+
386392
pub fn new_name_error(&self, msg: String) -> PyBaseExceptionRef {
387393
let name_error = self.ctx.exceptions.name_error.clone();
388394
self.new_exception_msg(name_error, msg)

0 commit comments

Comments
 (0)
0