superacli 1.1.4 → 1.1.6

package/ref-monty/crates/monty/src/resource.rs

@@ -0,0 +1,559 @@
+use std::{
+    fmt,
+    sync::atomic::{AtomicU16, Ordering},
+    time::{Duration, Instant},
+};
+
+use crate::{
+    ExcType, MontyException,
+    exception_private::{ExceptionRaise, RawStackFrame, RunError, SimpleException},
+};
+
+/// Threshold in bytes above which `check_large_result` is called.
+///
+/// Operations that may produce results larger than this threshold (100KB) should call
+/// `check_large_result` before performing the operation. This prevents DoS attacks
+/// where operations like `2 ** 10_000_000` allocate huge amounts of memory before
+/// the allocation check can catch them.
+pub const LARGE_RESULT_THRESHOLD: usize = 100_000;
+
+/// Pre-checks that an operation producing `item_len * count` bytes won't exceed resource limits.
+///
+/// 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)
+}
+
+/// Pre-checks that `base ** exponent` won't exceed resource limits before computing.
+///
+/// The result of `base ** exp` has approximately `base_bits * exp` bits.
+/// For bases with 0 or 1 significant bits (0, 1, -1), the result is always
+/// small regardless of exponent, so the check is skipped.
+///
+/// The estimate includes a 4× safety multiplier because `BigInt::pow` uses repeated squaring,
+/// which allocates intermediate values on the Rust heap (not tracked by the resource tracker).
+/// At peak, old/new base and old/new accumulator coexist simultaneously during each
+/// multiplication step, requiring roughly 4× the final result size in memory.
+pub fn check_pow_size(base_bits: u64, exponent: u64, tracker: &impl ResourceTracker) -> Result<(), ResourceError> {
+    // 0**n = 0, 1**n = 1, (-1)**n = ±1 — always small
+    if base_bits <= 1 {
+        return Ok(());
+    }
+    let result_bytes = estimate_bits_to_bytes(base_bits.saturating_mul(exponent));
+    // Repeated squaring needs ~4× result size in peak memory (old/new base + old/new accumulator
+    // coexist during each multiplication step), and these are Rust heap allocations not tracked
+    // by the resource tracker.
+    check_estimated_size(result_bytes.saturating_mul(4), tracker)
+}
+
+/// Pre-checks that an integer multiplication won't exceed resource limits.
+///
+/// The result of multiplying two numbers has at most `a_bits + b_bits` bits.
+pub fn check_mult_size(a_bits: u64, b_bits: u64, tracker: &impl ResourceTracker) -> Result<(), ResourceError> {
+    check_estimated_size(estimate_bits_to_bytes(a_bits.saturating_add(b_bits)), tracker)
+}
+
+/// Pre-checks that a left shift won't exceed resource limits.
+///
+/// The result of `value << shift` has approximately `value_bits + shift` bits.
+/// For zero values the result is always zero, so the check is skipped.
+pub fn check_lshift_size(
+    value_bits: u64,
+    shift_amount: u64,
+    tracker: &impl ResourceTracker,
+) -> Result<(), ResourceError> {
+    if value_bits == 0 {
+        return Ok(());
+    }
+    check_estimated_size(estimate_bits_to_bytes(value_bits.saturating_add(shift_amount)), tracker)
+}
+
+/// Pre-checks that an integer division overflow promotion won't exceed resource limits.
+///
+/// Division results are bounded by the dividend size, but we still check for consistency
+/// with other BigInt promotion paths.
+pub fn check_div_size(dividend_bits: u64, tracker: &impl ResourceTracker) -> Result<(), ResourceError> {
+    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`
+/// to avoid overhead on small operations.
+pub(crate) fn check_estimated_size(
+    estimated_bytes: usize,
+    tracker: &impl ResourceTracker,
+) -> Result<(), ResourceError> {
+    if estimated_bytes > LARGE_RESULT_THRESHOLD {
+        tracker.check_large_result(estimated_bytes)?;
+    }
+    Ok(())
+}
+
+/// Converts an estimated bit count to bytes, saturating to `usize::MAX` on overflow.
+///
+/// Overflow means the result is astronomically large, so saturating ensures
+/// the resource limit check always triggers rather than being silently skipped.
+fn estimate_bits_to_bytes(bits: u64) -> usize {
+    usize::try_from(bits.saturating_add(7) / 8).unwrap_or(usize::MAX)
+}
+
+/// Error returned when a resource limit is exceeded during execution.
+///
+/// This allows the sandbox to enforce strict limits on allocation count,
+/// execution time, and memory usage.
+#[derive(Debug, Clone)]
+pub enum ResourceError {
+    /// Maximum number of allocations exceeded.
+    Allocation { limit: usize, count: usize },
+    /// Maximum execution time exceeded.
+    Time { limit: Duration, elapsed: Duration },
+    /// Maximum memory usage exceeded.
+    Memory { limit: usize, used: usize },
+    /// Maximum recursion depth exceeded.
+    Recursion { limit: usize, depth: usize },
+    /// Any other error, e.g. when propagating a python exception
+    Exception(MontyException),
+}
+
+impl fmt::Display for ResourceError {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        match self {
+            Self::Allocation { limit, count } => {
+                write!(f, "allocation limit exceeded: {count} > {limit}")
+            }
+            Self::Time { limit, elapsed } => {
+                write!(f, "time limit exceeded: {elapsed:?} > {limit:?}")
+            }
+            Self::Memory { limit, used } => {
+                write!(f, "memory limit exceeded: {used} bytes > {limit} bytes")
+            }
+            Self::Recursion { .. } => {
+                write!(f, "maximum recursion depth exceeded")
+            }
+            Self::Exception(exc) => {
+                write!(f, "{exc}")
+            }
+        }
+    }
+}
+
+impl std::error::Error for ResourceError {}
+
+impl ResourceError {
+    /// Converts this resource error to a Python exception with optional stack frame.
+    ///
+    /// Maps resource error types to Python exception types:
+    /// - `Allocation` → `MemoryError`
+    /// - `Memory` → `MemoryError`
+    /// - `Time` → `TimeoutError`
+    /// - `Recursion` → `RecursionError`
+    #[must_use]
+    pub(crate) fn into_exception(self, frame: Option<RawStackFrame>) -> ExceptionRaise {
+        let (exc_type, msg) = match self {
+            Self::Allocation { limit, count } => (
+                ExcType::MemoryError,
+                Some(format!("allocation limit exceeded: {count} > {limit}")),
+            ),
+            Self::Memory { limit, used } => (
+                ExcType::MemoryError,
+                Some(format!("memory limit exceeded: {used} bytes > {limit} bytes")),
+            ),
+            Self::Time { limit, elapsed } => (
+                ExcType::TimeoutError,
+                Some(format!("time limit exceeded: {elapsed:?} > {limit:?}")),
+            ),
+            Self::Recursion { .. } => (
+                ExcType::RecursionError,
+                Some("maximum recursion depth exceeded".to_string()),
+            ),
+            Self::Exception(exc) => (exc.exc_type(), exc.into_message()),
+        };
+        let exc = SimpleException::new(exc_type, msg);
+        match frame {
+            Some(f) => exc.with_frame(f),
+            None => exc.into(),
+        }
+    }
+}
+
+impl From<ResourceError> for RunError {
+    fn from(err: ResourceError) -> Self {
+        // RecursionError is catchable in CPython, so it must be catchable here too.
+        // Other resource errors (memory, time, allocation) remain uncatchable to prevent
+        // untrusted code from suppressing resource limit violations.
+        if matches!(err, ResourceError::Recursion { .. }) {
+            Self::Exc(err.into_exception(None))
+        } else {
+            Self::UncatchableExc(err.into_exception(None))
+        }
+    }
+}
+
+/// Trait for tracking resource usage and scheduling garbage collection.
+///
+/// Implementations can enforce limits on allocations, time, and memory,
+/// as well as schedule periodic garbage collection.
+///
+/// All implementations should eventually trigger garbage collection to handle
+/// reference cycles. The `should_gc` method controls *frequency*, not whether
+/// GC runs at all.
+pub trait ResourceTracker: fmt::Debug {
+    /// Called before each heap allocation.
+    ///
+    /// Returns `Ok(())` if the allocation should proceed, or `Err(ResourceError)`
+    /// if a limit would be exceeded.
+    ///
+    /// # Arguments
+    /// * `size` - Approximate size in bytes of the allocation
+    fn on_allocate(&mut self, get_size: impl FnOnce() -> usize) -> Result<(), ResourceError>;
+
+    /// Called when memory is freed (during dec_ref or garbage collection).
+    ///
+    /// # Arguments
+    /// * `size` - Size in bytes of the freed allocation
+    fn on_free(&mut self, get_size: impl FnOnce() -> usize);
+
+    /// Called periodically (at statement boundaries) to check time limits.
+    ///
+    /// Returns `Ok(())` if within time limit, or `Err(ResourceError::Time)`
+    /// if the limit is exceeded.
+    ///
+    /// Takes `&self` rather than `&mut self` because checking elapsed time is a
+    /// read-only operation. This allows time checks in contexts that only have
+    /// an immutable heap reference, such as `py_repr_fmt`.
+    fn check_time(&self) -> Result<(), ResourceError>;
+
+    /// Called before pushing a new call frame to check recursion depth.
+    ///
+    /// Returns `Ok(())` if within recursion limit, or `Err(ResourceError::Recursion)`
+    /// if the limit would be exceeded.
+    ///
+    /// # Arguments
+    /// * `current_depth` - Current call stack depth (before the new frame is pushed)
+    fn check_recursion_depth(&self, current_depth: usize) -> Result<(), ResourceError>;
+
+    /// Called before operations that may produce large results (>100KB).
+    ///
+    /// This allows pre-emptive rejection of operations like `2 ** 10_000_000`
+    /// before the memory is actually allocated. The check only happens for
+    /// estimated result sizes above `LARGE_RESULT_THRESHOLD` to avoid overhead
+    /// on small operations.
+    ///
+    /// # Arguments
+    /// * `estimated_bytes` - Approximate size of the result in bytes
+    ///
+    /// Returns `Ok(())` to allow the operation, or `Err(ResourceError)` to reject.
+    fn check_large_result(&self, estimated_bytes: usize) -> Result<(), ResourceError>;
+}
+
+/// A resource tracker that imposes no limits except default recursion limit.
+///
+/// Recursion limit is set to the cpython default of 1000.
+#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
+pub struct NoLimitTracker;
+
+impl ResourceTracker for NoLimitTracker {
+    #[inline]
+    fn on_allocate(&mut self, _: impl FnOnce() -> usize) -> Result<(), ResourceError> {
+        Ok(())
+    }
+
+    #[inline]
+    fn on_free(&mut self, _: impl FnOnce() -> usize) {}
+
+    #[inline]
+    fn check_time(&self) -> Result<(), ResourceError> {
+        Ok(())
+    }
+
+    /// Set the recursion limit to 1000.
+    ///
+    /// The high limit here may cause stack overflow errors in debug mode, but do not those errors should
+    /// not occur with release builds.
+    #[inline]
+    fn check_recursion_depth(&self, current_depth: usize) -> Result<(), ResourceError> {
+        const DEFAULT_RECURSION_LIMIT: usize = 1000;
+        if current_depth >= DEFAULT_RECURSION_LIMIT {
+            Err(ResourceError::Recursion {
+                limit: DEFAULT_RECURSION_LIMIT,
+                depth: current_depth + 1,
+            })
+        } else {
+            Ok(())
+        }
+    }
+
+    #[inline]
+    fn check_large_result(&self, _estimated_bytes: usize) -> Result<(), ResourceError> {
+        // No limit - always allow operations regardless of result size
+        Ok(())
+    }
+}
+
+/// Configuration for resource limits.
+///
+/// All limits are optional - set to `None` to disable a specific limit.
+/// Use `ResourceLimits::default()` for no limits, or build custom limits
+/// with the builder pattern.
+#[derive(Debug, Clone, Default, serde::Serialize, serde::Deserialize)]
+pub struct ResourceLimits {
+    /// Maximum number of heap allocations allowed.
+    pub max_allocations: Option<usize>,
+    /// Maximum execution time.
+    pub max_duration: Option<Duration>,
+    /// Maximum heap memory in bytes (approximate).
+    pub max_memory: Option<usize>,
+    /// Run garbage collection every N allocations.
+    pub gc_interval: Option<usize>,
+    /// Maximum recursion depth (function call stack depth).
+    pub max_recursion_depth: Option<usize>,
+}
+
+/// Recommended maximum recursion depth if not otherwise specified.
+pub const DEFAULT_MAX_RECURSION_DEPTH: usize = 1000;
+
+impl ResourceLimits {
+    /// Creates a new ResourceLimits with all limits disabled, except max recursion which is set to 1000.
+    #[must_use]
+    pub fn new() -> Self {
+        Self {
+            max_recursion_depth: Some(1000),
+            ..Default::default()
+        }
+    }
+
+    /// Sets the maximum number of allocations.
+    #[must_use]
+    pub fn max_allocations(mut self, limit: usize) -> Self {
+        self.max_allocations = Some(limit);
+        self
+    }
+
+    /// Sets the maximum execution duration.
+    #[must_use]
+    pub fn max_duration(mut self, limit: Duration) -> Self {
+        self.max_duration = Some(limit);
+        self
+    }
+
+    /// Sets the maximum memory usage in bytes.
+    #[must_use]
+    pub fn max_memory(mut self, limit: usize) -> Self {
+        self.max_memory = Some(limit);
+        self
+    }
+
+    /// Sets the garbage collection interval (run GC every N allocations).
+    #[must_use]
+    pub fn gc_interval(mut self, interval: usize) -> Self {
+        self.gc_interval = Some(interval);
+        self
+    }
+
+    /// Sets the maximum recursion depth (function call stack depth).
+    #[must_use]
+    pub fn max_recursion_depth(mut self, limit: Option<usize>) -> Self {
+        self.max_recursion_depth = limit;
+        self
+    }
+}
+
+/// How often to actually check `Instant::elapsed()` in `check_time`.
+///
+/// Calling `Instant::elapsed()` on every `check_time` invocation adds measurable
+/// overhead in tight loops (the VM calls `check_time` on every instruction).
+/// By only checking every N calls, we reduce this overhead while still catching
+/// timeouts promptly.
+const TIME_CHECK_INTERVAL: u16 = 10;
+
+/// A resource tracker that enforces configurable limits.
+///
+/// Tracks allocation count, memory usage, and execution time, returning
+/// errors when limits are exceeded. Also schedules garbage collection
+/// at configurable intervals.
+///
+/// When serialized/deserialized, the `start_time` is reset to `Instant::now()`.
+/// This means time limits restart from zero after deserialization.
+#[derive(Debug, serde::Serialize, serde::Deserialize)]
+pub struct LimitedTracker {
+    limits: ResourceLimits,
+    /// When execution started (for time limit checking).
+    /// Reset to `Instant::now()` on deserialization.
+    #[serde(skip, default = "Instant::now")]
+    start_time: Instant,
+    /// Total number of allocations made.
+    allocation_count: usize,
+    /// Current approximate memory usage in bytes.
+    current_memory: usize,
+    /// Counter for rate-limiting `Instant::elapsed()` calls in `check_time`.
+    ///
+    /// Uses `AtomicU16` for interior mutability since `check_time` takes `&self`
+    /// and `LimitedTracker` must be `Sync` (it ends up inside PyO3 pyclass types).
+    check_counter: AtomicU16,
+}
+
+impl LimitedTracker {
+    /// Creates a new LimitedTracker with the given limits.
+    ///
+    /// The start time is recorded when the tracker is created, so create
+    /// it immediately before starting execution.
+    #[must_use]
+    pub fn new(limits: ResourceLimits) -> Self {
+        Self {
+            limits,
+            start_time: Instant::now(),
+            allocation_count: 0,
+            current_memory: 0,
+            check_counter: AtomicU16::new(0),
+        }
+    }
+
+    /// Returns the current allocation count.
+    #[must_use]
+    pub fn allocation_count(&self) -> usize {
+        self.allocation_count
+    }
+
+    /// Returns the current approximate memory usage.
+    #[must_use]
+    pub fn current_memory(&self) -> usize {
+        self.current_memory
+    }
+
+    /// Returns the elapsed time since tracker creation.
+    #[must_use]
+    pub fn elapsed(&self) -> Duration {
+        self.start_time.elapsed()
+    }
+
+    /// Sets the maximum execution duration and resets the start time to now.
+    ///
+    /// This is useful when resuming execution after an external function call
+    /// where you want to enforce a different (typically shorter) time limit
+    /// for the resumed phase without counting the time spent in the host.
+    pub fn set_max_duration(&mut self, duration: Duration) {
+        self.limits.max_duration = Some(duration);
+        self.start_time = Instant::now();
+    }
+}
+
+impl ResourceTracker for LimitedTracker {
+    fn on_allocate(&mut self, get_size: impl FnOnce() -> usize) -> Result<(), ResourceError> {
+        // Check allocation count limit
+        if let Some(max) = self.limits.max_allocations
+            && self.allocation_count >= max
+        {
+            return Err(ResourceError::Allocation {
+                limit: max,
+                count: self.allocation_count + 1,
+            });
+        }
+
+        let size = get_size();
+        // Check memory limit
+        if let Some(max) = self.limits.max_memory {
+            let new_memory = self.current_memory + size;
+            if new_memory > max {
+                return Err(ResourceError::Memory {
+                    limit: max,
+                    used: new_memory,
+                });
+            }
+        }
+
+        // Update tracking state
+        self.allocation_count += 1;
+        self.current_memory += size;
+
+        Ok(())
+    }
+
+    fn on_free(&mut self, get_size: impl FnOnce() -> usize) {
+        self.current_memory = self.current_memory.saturating_sub(get_size());
+    }
+
+    fn check_time(&self) -> Result<(), ResourceError> {
+        if let Some(max) = self.limits.max_duration {
+            let count = self.check_counter.fetch_add(1, Ordering::Relaxed).wrapping_add(1);
+            if count.is_multiple_of(TIME_CHECK_INTERVAL) {
+                // Only call Instant::elapsed() every TIME_CHECK_INTERVAL calls
+                let elapsed = self.start_time.elapsed();
+                if elapsed > max {
+                    // Reset counter so the very next check_time call also triggers
+                    // an elapsed check. This is important because some callers
+                    // (e.g. repr_sequence_fmt) catch the error and return normally,
+                    // and we need the VM loop's next check_time to re-detect timeout.
+                    self.check_counter
+                        .store(TIME_CHECK_INTERVAL.wrapping_sub(1), Ordering::Relaxed);
+                    return Err(ResourceError::Time { limit: max, elapsed });
+                }
+            }
+        }
+        Ok(())
+    }
+
+    fn check_recursion_depth(&self, current_depth: usize) -> Result<(), ResourceError> {
+        if let Some(max) = self.limits.max_recursion_depth {
+            // current_depth is before push, so new depth would be current_depth + 1
+            if current_depth >= max {
+                return Err(ResourceError::Recursion {
+                    limit: max,
+                    depth: current_depth + 1,
+                });
+            }
+        }
+        Ok(())
+    }
+
+    fn check_large_result(&self, estimated_bytes: usize) -> Result<(), ResourceError> {
+        // Check if this would exceed memory limit
+        if let Some(max) = self.limits.max_memory {
+            let new_memory = self.current_memory.saturating_add(estimated_bytes);
+            if new_memory > max {
+                return Err(ResourceError::Memory {
+                    limit: max,
+                    used: new_memory,
+                });
+            }
+        }
+        Ok(())
+    }
+}