diff --git a/src/memory.rs b/src/memory.rs index 982f10a1d..dfe7d8636 100644 --- a/src/memory.rs +++ b/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,11 +56,6 @@ 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; @@ -54,6 +63,19 @@ impl MemoryState { 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); + } + } +} diff --git a/src/state.rs b/src/state.rs index 6b2e7ea9d..6f54cab52 100644 --- a/src/state.rs +++ b/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 { + 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), diff --git a/src/state/extra.rs b/src/state/extra.rs index 038855d3a..80265dc14 100644 --- a/src/state/extra.rs +++ b/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, diff --git a/src/state/raw.rs b/src/state/raw.rs index 8060414a6..6fbd5efda 100644 --- a/src/state/raw.rs +++ b/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) } } diff --git a/tests/memory.rs b/tests/memory.rs index 392a6bd56..f2d73d0bf 100644 --- a/tests/memory.rs +++ b/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::(()) { + 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::()?, 2); + + assert_eq!(lua.set_allocation_limit(0)?, LIMIT); + assert_eq!(grow.call::(())?, 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();