rgw/lua: add configurable memory limit for the lua state

Signed-off-by: Yuval Lifshitz <ylifshit@redhat.com>

1 files changed, 93 insertions, 7 deletions

diff --git a/src/rgw/rgw_lua_utils.cc b/src/rgw/rgw_lua_utils.cc
index 1dd214b7bab..4118bed42e0 100644
--- a/src/rgw/rgw_lua_utils.cc
+++ b/src/rgw/rgw_lua_utils.cc
@@ -1,7 +1,7 @@
 #include <string>
 #include <lua.hpp>
 #include "common/ceph_context.h"
-#include "common/dout.h"
+#include "common/debug.h"
 #include "rgw_lua_utils.h"
 #include "rgw_lua_version.h"
 
@@ -64,8 +64,7 @@ void set_package_path(lua_State* L, const std::string& install_dir) {
   if (install_dir.empty()) {
     return;
   }
-  lua_getglobal(L, "package");
-  if (!lua_istable(L, -1)) {
+  if (lua_getglobal(L, "package") != LUA_TTABLE) {
     return;
   }
   const auto path = install_dir+"/share/lua/"+CEPH_LUA_VERSION+"/?.lua";
@@ -85,10 +84,97 @@ void open_standard_libs(lua_State* L) {
   unsetglobal(L, "dofile");
   unsetglobal(L, "debug");
   // remove os.exit()
-  lua_getglobal(L, "os");
-  lua_pushstring(L, "exit");
-  lua_pushnil(L);
-  lua_settable(L, -3);
+  if (lua_getglobal(L, "os") == LUA_TTABLE) {
+    lua_pushstring(L, "exit");
+    lua_pushnil(L);
+    lua_settable(L, -3);
+  }
+}
+
+// allocator function that verifies against maximum allowed memory value
+void* allocator(void* ud, void* ptr, std::size_t osize, std::size_t nsize) {
+  auto mem = reinterpret_cast<std::size_t*>(ud); // remaining memory
+  // free memory
+  if (nsize == 0) {
+    if (mem && ptr) {
+      *mem += osize;
+    }
+    free(ptr);
+    return nullptr;
+  }
+  // re/alloc memory
+  if (mem) {
+    const std::size_t realloc_size = ptr ? osize : 0;
+    if (nsize > realloc_size && (nsize - realloc_size) > *mem) {
+      return nullptr;
+    }
+    *mem += realloc_size;
+    *mem -= nsize;
+  }
+  return realloc(ptr, nsize);
+}
+
+// create new lua state together with its memory counter
+lua_State* newstate(int max_memory) {
+  std::size_t* remaining_memory = nullptr;
+  if (max_memory > 0) {
+    remaining_memory = new std::size_t(max_memory);
+  }
+  lua_State* L = lua_newstate(allocator, remaining_memory);
+  if (!L) {
+    delete remaining_memory;
+    remaining_memory = nullptr;
+  }
+  if (L) {
+    lua_atpanic(L, [](lua_State* L) -> int {
+      const char* msg = lua_tostring(L, -1);
+      if (msg == nullptr) msg = "error object is not a string";
+      throw std::runtime_error(msg);
+    });
+  }
+  return L;
+}
+
+// lua_state_guard ctor
+lua_state_guard::lua_state_guard(std::size_t  _max_memory, const DoutPrefixProvider* _dpp) :
+  max_memory(_max_memory),
+  dpp(_dpp),
+  state(newstate(_max_memory)) {
+  if (state &&  perfcounter) {
+    perfcounter->inc(l_rgw_lua_current_vms, 1);
+  }
+}
+
+// lua_state_guard dtor
+lua_state_guard::~lua_state_guard() {
+  lua_State* L = state;
+  if (!L) {
+    return;
+  }
+  void* ud = nullptr;
+  lua_getallocf(L, &ud);
+  auto remaining_memory = static_cast<std::size_t*>(ud);
+
+  if (remaining_memory) {
+    const auto used_memory = max_memory - *remaining_memory;
+    ldpp_dout(dpp, 20) << "Lua is using: " << used_memory << 
+      " bytes (" << 100.0*used_memory/max_memory << "%)" << dendl;
+    // dont limit memory during cleanup
+    *remaining_memory = 0;
+  }
+  try {
+    lua_close(L);
+  } catch (const std::runtime_error& e) {
+    ldpp_dout(dpp, 20) << "Lua cleanup failed with: " << e.what() << dendl;
+  }
+
+  // TODO: use max_memory and remaining memory to check for leaks
+  // this could be done only if we don't zero the remianing memory during clanup
+
+  delete remaining_memory;
+  if (perfcounter) {
+    perfcounter->dec(l_rgw_lua_current_vms, 1);
+  }
 }
 
 } // namespace rgw::lua