superacli 1.1.6 → 1.1.7

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

@@ -1,559 +0,0 @@
-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(())
-    }
-}