Add memory pre-checks to padding, replace and f-string operations (#238)

Author: MRoci

crates/monty/src/bytecode/vm/format.rs

@@ -5,7 +5,7 @@ use crate::{
     defer_drop,
     exception_private::{ExcType, RunError, SimpleException},
     fstring::{ParsedFormatSpec, ascii_escape, decode_format_spec, format_string, format_with_spec},
-    resource::ResourceTracker,
+    resource::{ResourceTracker, check_repeat_size},
     types::{PyTrait, str::allocate_string},
     value::Value,
 };
@@ -61,6 +61,10 @@ impl<T: ResourceTracker> VM<'_, '_, T> {
 
             let spec = this.get_format_spec(spec_value, value)?;
 
+            // Pre-check: reject format specs with huge width before pad_string
+            // allocates an untracked Rust String.
+            check_repeat_size(spec.width, spec.fill.len_utf8(), this.heap.tracker())?;
+
             match conversion {
                 // No conversion - format original value
                 0 => format_with_spec(value, &spec, this.heap, this.interns)?,

crates/monty/src/resource.rs

@@ -17,10 +17,11 @@ use crate::{
 /// the allocation check can catch them.
 pub const LARGE_RESULT_THRESHOLD: usize = 100_000;
 
-/// Pre-checks that a sequence repeat won't exceed resource limits before allocating.
+/// Pre-checks that an operation producing `item_len * count` bytes won't exceed resource limits.
 ///
-/// This prevents DoS via expressions like `'x' * 999_999_999` or `b'ab' * huge_int`
-/// by estimating the result size and checking against the resource tracker.
+/// Used for sequence repeats (`'x' * 999_999_999`), padding operations
+/// (`str.ljust`, `str.center`, `str.zfill`, etc.), and any other operation
+/// where the result size is a simple product of two known values.
 pub fn check_repeat_size(item_len: usize, count: usize, tracker: &impl ResourceTracker) -> Result<(), ResourceError> {
     check_estimated_size(item_len.saturating_mul(count), tracker)
 }
@@ -77,6 +78,41 @@ pub fn check_div_size(dividend_bits: u64, tracker: &impl ResourceTracker) -> Res
     check_estimated_size(estimate_bits_to_bytes(dividend_bits), tracker)
 }
 
+/// Pre-checks that a string/bytes replace won't exceed resource limits before allocating.
+///
+/// This prevents DoS via expressions like `('a' * 1000).replace('a', 'b' * 10_000_000)`
+/// where a small tracked input is amplified into a huge untracked Rust `String`/`Vec`
+/// by `String::replace()` before `allocate_string()` can check the result.
+///
+/// The upper bound on result size is: if `old` is non-empty, at most `input_len / old_len`
+/// replacements can occur, each producing `new_len` bytes instead of `old_len`. When `count`
+/// is specified, replacements are capped to that value.
+pub fn check_replace_size(
+    input_len: usize,
+    old_len: usize,
+    new_len: usize,
+    count: i64,
+    tracker: &impl ResourceTracker,
+) -> Result<(), ResourceError> {
+    // Empty pattern (old_len == 0): inserts before each element + after the last = input_len + 1
+    let max_replacements = input_len
+        .checked_div(old_len)
+        .unwrap_or_else(|| input_len.saturating_add(1));
+
+    let replacements = if count < 0 {
+        max_replacements
+    } else {
+        max_replacements.min(usize::try_from(count).unwrap_or(usize::MAX))
+    };
+
+    // Result = input_len - (replacements * old_len) + (replacements * new_len)
+    let removed = replacements.saturating_mul(old_len);
+    let added = replacements.saturating_mul(new_len);
+    let estimated = input_len.saturating_sub(removed).saturating_add(added);
+
+    check_estimated_size(estimated, tracker)
+}
+
 /// Checks an estimated result size against the resource tracker.
 ///
 /// Only calls the tracker when the estimate exceeds `LARGE_RESULT_THRESHOLD`

crates/monty/src/types/bytes.rs

@@ -79,7 +79,7 @@ use crate::{
     exception_private::{ExcType, RunResult, SimpleException},
     heap::{DropWithHeap, Heap, HeapData, HeapGuard, HeapId},
     intern::{Interns, StaticStrings, StringId},
-    resource::{ResourceError, ResourceTracker},
+    resource::{ResourceError, ResourceTracker, check_repeat_size, check_replace_size},
     types::List,
     value::{EitherStr, Value},
 };
@@ -1774,6 +1774,8 @@ fn bytes_replace(
 ) -> RunResult<Value> {
     let (old, new, count) = parse_bytes_replace_args("bytes.replace", args, heap, interns)?;
 
+    check_replace_size(bytes.len(), old.len(), new.len(), count, heap.tracker())?;
+
     let result = if count < 0 {
         bytes_replace_all(bytes, &old, &new, heap)?
     } else {
@@ -1949,6 +1951,7 @@ fn bytes_center(
     let result = if width <= len {
         bytes.to_vec()
     } else {
+        check_repeat_size(width, 1, heap.tracker())?;
         let total_pad = width - len;
         let left_pad = total_pad / 2;
         let right_pad = total_pad - left_pad;
@@ -1981,6 +1984,7 @@ fn bytes_ljust(
     let result = if width <= len {
         bytes.to_vec()
     } else {
+        check_repeat_size(width, 1, heap.tracker())?;
         let pad = width - len;
         let mut result = Vec::with_capacity(width);
         result.extend_from_slice(bytes);
@@ -2008,6 +2012,7 @@ fn bytes_rjust(
     let result = if width <= len {
         bytes.to_vec()
     } else {
+        check_repeat_size(width, 1, heap.tracker())?;
         let pad = width - len;
         let mut result = Vec::with_capacity(width);
         for _ in 0..pad {
@@ -2076,6 +2081,7 @@ fn bytes_zfill(bytes: &[u8], args: ArgValues, heap: &mut Heap<impl ResourceTrack
     let result = if width <= len {
         bytes.to_vec()
     } else {
+        check_repeat_size(width, 1, heap.tracker())?;
         let pad = width - len;
         let mut result = Vec::with_capacity(width);
 

crates/monty/src/types/str.rs

@@ -16,7 +16,7 @@ use crate::{
     exception_private::{ExcType, RunResult},
     heap::{DropWithHeap, Heap, HeapData, HeapGuard, HeapId},
     intern::{Interns, StaticStrings, StringId},
-    resource::{ResourceError, ResourceTracker},
+    resource::{ResourceError, ResourceTracker, check_repeat_size, check_replace_size},
     types::Type,
     value::{EitherStr, Value},
 };
@@ -1728,6 +1728,8 @@ fn str_rpartition(
 fn str_replace(s: &str, args: ArgValues, heap: &mut Heap<impl ResourceTracker>, interns: &Interns) -> RunResult<Value> {
     let (old, new, count) = parse_replace_args("str.replace", args, heap, interns)?;
 
+    check_replace_size(s.len(), old.len(), new.len(), count, heap.tracker())?;
+
     let result = if count < 0 {
         s.replace(&old, &new)
     } else {
@@ -1820,6 +1822,7 @@ fn str_center(s: &str, args: ArgValues, heap: &mut Heap<impl ResourceTracker>, i
     let result = if width <= len {
         s.to_owned()
     } else {
+        check_repeat_size(width, fillchar.len_utf8(), heap.tracker())?;
         let total_pad = width - len;
         let left_pad = total_pad / 2;
         let right_pad = total_pad - left_pad;
@@ -1847,6 +1850,7 @@ fn str_ljust(s: &str, args: ArgValues, heap: &mut Heap<impl ResourceTracker>, in
     let result = if width <= len {
         s.to_owned()
     } else {
+        check_repeat_size(width, fillchar.len_utf8(), heap.tracker())?;
         let pad = width - len;
         let mut result = String::with_capacity(width);
         result.push_str(s);
@@ -1869,6 +1873,7 @@ fn str_rjust(s: &str, args: ArgValues, heap: &mut Heap<impl ResourceTracker>, in
     let result = if width <= len {
         s.to_owned()
     } else {
+        check_repeat_size(width, fillchar.len_utf8(), heap.tracker())?;
         let pad = width - len;
         let mut result = String::with_capacity(width);
         for _ in 0..pad {
@@ -1936,6 +1941,8 @@ fn str_zfill(s: &str, args: ArgValues, heap: &mut Heap<impl ResourceTracker>) ->
     let result = if width <= len {
         s.to_owned()
     } else {
+        // zfill always pads with ASCII '0' (1 byte)
+        check_repeat_size(width, 1, heap.tracker())?;
         let pad = width - len;
         let mut chars = s.chars();
         let first = chars.next();