Add Lua::set_allocation_limit to limit the size of a single allocation

src/memory.rs

@@ -5,10 +5,11 @@ use std::ptr;
 pub(crate) static ALLOCATOR: ffi::lua_Alloc = allocator;
 
 #[repr(C)]
-#[derive(Default)]
 pub(crate) struct MemoryState {
     used_memory: isize,
     memory_limit: isize,
+    // The largest size of a single allocation, `usize::MAX` means no limit.
+    allocation_limit: usize,
     // Can be set to temporary ignore the memory limit.
     // This is used when calling `lua_pushcfunction` for lua5.1/jit/luau.
     ignore_limit: bool,
@@ -17,6 +18,19 @@ pub(crate) struct MemoryState {
     limit_reached: bool,
 }
 
+impl Default for MemoryState {
+    fn default() -> Self {
+        MemoryState {
+            used_memory: 0,
+            memory_limit: 0,
+            allocation_limit: usize::MAX,
+            ignore_limit: false,
+            #[cfg(feature = "luau")]
+            limit_reached: false,
+        }
+    }
+}
+
 impl MemoryState {
     #[cfg(feature = "luau")]
     #[inline]
@@ -42,18 +56,26 @@ impl MemoryState {
         self.used_memory as usize
     }
 
-    #[inline]
-    pub(crate) fn memory_limit(&self) -> usize {
-        self.memory_limit as usize
-    }
-
     #[inline]
     pub(crate) fn set_memory_limit(&mut self, limit: usize) -> usize {
         let prev_limit = self.memory_limit;
         self.memory_limit = limit.min(isize::MAX as usize) as isize;
         prev_limit as usize
     }
 
+    #[inline]
+    pub(crate) fn set_allocation_limit(&mut self, limit: usize) -> usize {
+        let prev_limit = self.allocation_limit;
+        self.allocation_limit = if limit == 0 { usize::MAX } else { limit };
+        if prev_limit == usize::MAX { 0 } else { prev_limit }
+    }
+
+    // Returns `true` if any memory or allocation limit is set
+    #[inline]
+    pub(crate) fn has_limit(&self) -> bool {
+        self.memory_limit > 0 || self.allocation_limit != usize::MAX
+    }
+
     // This function is used primarily for calling `lua_pushcfunction` in lua5.1/jit/luau
     // to bypass the memory limit (if set).
     #[cfg(any(feature = "lua51", feature = "luajit", feature = "luau"))]
@@ -112,6 +134,15 @@ unsafe extern "C" fn allocator(
         return ptr::null_mut();
     }
 
+    // Is a new or growing block within the allocation limit?
+    if nsize > mem_state.allocation_limit && (ptr.is_null() || nsize > osize) && !mem_state.ignore_limit {
+        #[cfg(feature = "luau")]
+        {
+            mem_state.limit_reached = true;
+        }
+        return ptr::null_mut();
+    }
+
     // Are we fit to the memory limits?
     let mut mem_diff = nsize as isize;
     if !ptr.is_null() {
@@ -149,3 +180,59 @@ unsafe extern "C" fn allocator(
     }
     new_ptr
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    unsafe fn realloc(state: &mut MemoryState, ptr: *mut c_void, osize: usize, nsize: usize) -> *mut c_void {
+        allocator(state as *mut MemoryState as *mut c_void, ptr, osize, nsize)
+    }
+
+    #[test]
+    fn test_allocation_limit() {
+        let mut state = MemoryState::default();
+        unsafe {
+            // No limit by default
+            let large = realloc(&mut state, ptr::null_mut(), 0, 4096);
+            assert!(!large.is_null());
+
+            assert_eq!(state.set_allocation_limit(1024), 0);
+            assert!(state.has_limit());
+
+            // Below, at and one past the limit
+            let below = realloc(&mut state, ptr::null_mut(), 0, 1023);
+            assert!(!below.is_null());
+            let at = realloc(&mut state, ptr::null_mut(), 0, 1024);
+            assert!(!at.is_null());
+            assert!(realloc(&mut state, ptr::null_mut(), 0, 1025).is_null());
+            #[cfg(feature = "luau")]
+            assert!(state.limit_reached);
+            assert_eq!(state.used_memory(), 4096 + 1023 + 1024);
+
+            // Growing past the limit fails and keeps the block
+            assert!(realloc(&mut state, at, 1024, 1025).is_null());
+            let at2 = realloc(&mut state, below, 1023, 1024);
+            assert!(!at2.is_null());
+            #[cfg(feature = "luau")]
+            assert!(!state.limit_reached);
+
+            // Shrinking always succeeds, even from above the limit
+            let large = realloc(&mut state, large, 4096, 2048);
+            assert!(!large.is_null());
+            assert!(realloc(&mut state, large, 2048, 2049).is_null());
+            assert_eq!(state.used_memory(), 2048 + 1024 + 1024);
+
+            // Zero removes the limit
+            assert_eq!(state.set_allocation_limit(0), 1024);
+            assert!(!state.has_limit());
+            let large = realloc(&mut state, large, 2048, 8192);
+            assert!(!large.is_null());
+
+            realloc(&mut state, large, 8192, 0);
+            realloc(&mut state, at, 1024, 0);
+            realloc(&mut state, at2, 1024, 0);
+            assert_eq!(state.used_memory(), 0);
+        }
+    }
+}

src/state.rs

@@ -1119,7 +1119,32 @@ impl Lua {
             match MemoryState::get(lua.state()) {
                 mem_state if !mem_state.is_null() => {
                     let prev_limit = (*mem_state).set_memory_limit(limit);
-                    (*lua.extra.get()).unlikely_memory_error = (*mem_state).memory_limit() == 0;
+                    (*lua.extra.get()).unlikely_memory_error = !(*mem_state).has_limit();
+                    Ok(prev_limit)
+                }
+                _ => Err(Error::MemoryControlNotAvailable),
+            }
+        }
+    }
+
+    /// Sets a limit (in bytes) on the size of a single memory allocation in this Lua state.
+    ///
+    /// Once Lua requests a new block, or grows an existing one, larger than this limit, a
+    /// `Error::MemoryError` is generated instead, the same way as for [`Lua::set_memory_limit`].
+    /// Shrinking a block always succeeds. Unlike the memory limit, this bounds how large a single
+    /// string, table or buffer can be, and so how much work a single library call can do on it.
+    ///
+    /// A limit that is too small may prevent Lua from working at all.
+    /// Returns previous limit (zero means no limit).
+    ///
+    /// Does not work in module mode where Lua state is managed externally.
+    pub fn set_allocation_limit(&self, limit: usize) -> Result<usize> {
+        let lua = self.lock();
+        unsafe {
+            match MemoryState::get(lua.state()) {
+                mem_state if !mem_state.is_null() => {
+                    let prev_limit = (*mem_state).set_allocation_limit(limit);
+                    (*lua.extra.get()).unlikely_memory_error = !(*mem_state).has_limit();
                     Ok(prev_limit)
                 }
                 _ => Err(Error::MemoryControlNotAvailable),

src/state/extra.rs

@@ -179,7 +179,7 @@ impl ExtraData {
             libs: StdLib::NONE,
             unlikely_memory_error: owned && {
                 let mem_state = MemoryState::get(state);
-                !mem_state.is_null() && (*mem_state).memory_limit() == 0
+                !mem_state.is_null() && !(*mem_state).has_limit()
             },
             skip_memory_check: false,
             ref_thread,

src/state/raw.rs

@@ -1087,7 +1087,7 @@ impl RawLua {
 
         // Externally managed states may change their allocator.
         match MemoryState::get(self.state()) {
-            mem_state if !mem_state.is_null() => (*mem_state).memory_limit() == 0,
+            mem_state if !mem_state.is_null() => !(*mem_state).has_limit(),
             _ => (*self.extra.get()).skip_memory_check, // Check the special flag (only for module mode)
         }
     }

tests/memory.rs

@@ -111,6 +111,122 @@ fn test_memory_limit_thread() -> Result<()> {
     Ok(())
 }
 
+// Creates a string of `len` bytes that is not interned yet, returns `false` on a memory error
+fn try_create_string(lua: &Lua, len: usize) -> bool {
+    use std::sync::atomic::{AtomicU64, Ordering};
+    static SEQ: AtomicU64 = AtomicU64::new(0);
+
+    let mut bytes = vec![b'x'; len.max(8)];
+    bytes.truncate(len);
+    let seq = SEQ.fetch_add(1, Ordering::Relaxed).to_le_bytes();
+    let n = len.min(8);
+    bytes[..n].copy_from_slice(&seq[..n]);
+    match lua.create_string(&bytes) {
+        Ok(_) => true,
+        Err(Error::MemoryError(_)) => false,
+        Err(err) => panic!("unexpected error: {err:?}"),
+    }
+}
+
+#[test]
+fn test_allocation_limit() -> Result<()> {
+    let lua = Lua::new();
+
+    if cfg!(feature = "luajit") && lua.set_allocation_limit(0).is_err() {
+        // seems this luajit version does not support memory limit
+        return Ok(());
+    }
+
+    // No limit by default
+    assert_eq!(lua.set_allocation_limit(0)?, 0);
+    assert!(try_create_string(&lua, 1 << 24));
+    lua.gc_collect()?;
+
+    const LIMIT: usize = 64 * 1024;
+    assert_eq!(lua.set_allocation_limit(LIMIT)?, 0);
+    assert!(try_create_string(&lua, LIMIT / 2));
+    assert!(!try_create_string(&lua, LIMIT));
+
+    // The largest string that fits grows with the limit byte for byte
+    let largest_string = |lua: &Lua| {
+        let (mut lo, mut hi) = (LIMIT / 2, 2 * LIMIT);
+        while hi - lo > 1 {
+            let mid = (lo + hi) / 2;
+            if try_create_string(lua, mid) {
+                lo = mid;
+            } else {
+                hi = mid;
+            }
+        }
+        assert!(try_create_string(lua, lo));
+        assert!(!try_create_string(lua, lo + 1));
+        lo
+    };
+    let largest = largest_string(&lua);
+    assert!(largest < LIMIT);
+    lua.set_allocation_limit(LIMIT + 64)?;
+    assert_eq!(largest_string(&lua), largest + 64);
+    lua.set_allocation_limit(LIMIT)?;
+
+    // Growing a table past the limit fails
+    let grow = lua
+        .load("local t = {}; for i = 1, 100000 do t[i] = i end; return #t")
+        .into_function()?;
+    match grow.call::<usize>(()) {
+        Err(Error::MemoryError(_)) => {}
+        something_else => panic!("did not trigger memory error: {:?}", something_else),
+    }
+
+    // The error can be caught by `pcall` and Lua keeps working
+    let (ok, len) = lua
+        .load(
+            r#"
+            local ok = pcall(string.rep, "x", 2 * 65536)
+            return ok, #string.rep("x", 1024)
+        "#,
+        )
+        .eval::<(bool, usize)>()?;
+    assert!(!ok);
+    assert_eq!(len, 1024);
+    assert_eq!(lua.load("return 1 + 1").eval::<i32>()?, 2);
+
+    assert_eq!(lua.set_allocation_limit(0)?, LIMIT);
+    assert_eq!(grow.call::<usize>(())?, 100000);
+
+    Ok(())
+}
+
+#[test]
+fn test_allocation_limit_with_memory_limit() -> Result<()> {
+    let lua = Lua::new();
+
+    if cfg!(feature = "luajit") && lua.set_allocation_limit(0).is_err() {
+        // seems this luajit version does not support memory limit
+        return Ok(());
+    }
+
+    const LIMIT: usize = 64 * 1024;
+    lua.set_allocation_limit(LIMIT)?;
+    let memory_limit = lua.used_memory() + (1 << 20);
+    lua.set_memory_limit(memory_limit)?;
+
+    // Each limit applies on its own
+    assert!(try_create_string(&lua, LIMIT / 2));
+    assert!(!try_create_string(&lua, 2 * LIMIT));
+    assert_eq!(lua.set_allocation_limit(0)?, LIMIT);
+    assert!(try_create_string(&lua, 2 * LIMIT));
+    lua.gc_collect()?;
+    assert!(!try_create_string(&lua, 2 << 20));
+
+    // Removing the memory limit keeps the allocation limit
+    lua.set_allocation_limit(LIMIT)?;
+    assert_eq!(lua.set_memory_limit(0)?, memory_limit);
+    assert!(try_create_string(&lua, LIMIT / 2));
+    assert!(!try_create_string(&lua, 2 * LIMIT));
+
+    Ok(())
+}
+
 #[test]
 fn test_gc_control() -> Result<()> {
     let lua = Lua::new();