diff --git a/src/util-lua-sandbox.c b/src/util-lua-sandbox.c index a1faec9921..c81e5d5187 100644 --- a/src/util-lua-sandbox.c +++ b/src/util-lua-sandbox.c @@ -73,7 +73,7 @@ static void *LuaAlloc(void *ud, void *ptr, size_t osize, size_t nsize) /* Resizing existing data. */ ssize_t diff = nsize - osize; - if (ctx->alloc_bytes + diff > ctx->alloc_limit) { + if (ctx->alloc_limit != 0 && ctx->alloc_bytes + diff > ctx->alloc_limit) { /* This request will exceed the allocation limit. Act as * though allocation failed. */ ctx->memory_limit_error = true; @@ -382,6 +382,33 @@ static void HookFunc(lua_State *L, lua_Debug *ar) } } +uint64_t SCLuaSbResetBytesLimit(lua_State *L) +{ + uint64_t cfg_limit = 0; + SCLuaSbState *sb = SCLuaSbGetContext(L); + if (sb != NULL) { + cfg_limit = sb->alloc_limit; + sb->alloc_limit = 0; + } + return cfg_limit; +} + +void SCLuaSbUpdateBytesLimit(lua_State *L) +{ + SCLuaSbState *sb = SCLuaSbGetContext(L); + if (sb != NULL) { + sb->alloc_limit = sb->alloc_bytes + sb->alloc_limit; + } +} + +void SCLuaSbRestoreBytesLimit(lua_State *L, const uint64_t cfg_limit) +{ + SCLuaSbState *sb = SCLuaSbGetContext(L); + if (sb != NULL) { + sb->alloc_limit = cfg_limit; + } +} + /** * Reset the instruction counter for the provided state. */ diff --git a/src/util-lua-sandbox.h b/src/util-lua-sandbox.h index c7b15ba08b..37b2087741 100644 --- a/src/util-lua-sandbox.h +++ b/src/util-lua-sandbox.h @@ -88,4 +88,47 @@ void SCLuaSbResetInstructionCounter(lua_State *sb); */ void SCLuaSbLoadLibs(lua_State *L); +/* alloc_limit (max-bytes) handling. To give the script the full budget + * during it's execution, temporary disable the limit and then restore + * it with the existing use as baseline + * + * ``` + * // get configured limit and reset active limit to 0 + * const uint64_t cfg_limit = SCLuaSbResetBytesLimit(tlua->luastate); + * // push some data to the state + * LuaPushStringBuffer(tlua->luastate, input, (size_t)input_len); + * // restore the limit, which may now be lower than `alloc_bytes`. + * SCLuaSbRestoreBytesLimit(tlua->luastate, cfg_limit); + * // update allowance to take `alloc_bytes` as the baseline, so + * // effectively: sb->alloc_bytes + sb->alloc_limit + * SCLuaSbUpdateBytesLimit(tlua->luastate); + * + * ... run script, lua_pcall... + * + * while (lua_gettop(tlua->luastate) > 0) { + * lua_pop(tlua->luastate, 1); + * } + * // restore the original alloc_limit + * SCLuaSbRestoreBytesLimit(tlua->luastate, cfg_limit); + * ``` + */ + +/* + * Resets the byte limit and returns the existing limit. + * Meant to be used to temporarily disable the limit during preparation + * time for a script to run. */ +uint64_t SCLuaSbResetBytesLimit(lua_State *L); + +/* + * Update the alloc_limit to take the pre-script initialization bytes + * as the base line. + */ +void SCLuaSbUpdateBytesLimit(lua_State *L); + +/* + * Set the bytes limit. Meant to be used with the value returned from + * `SCLuaSbResetBytesLimit` + */ +void SCLuaSbRestoreBytesLimit(lua_State *L, const uint64_t cfg_limit); + #endif /* SURICATA_UTIL_LUA_SANDBOX_H */