Add max memory usage restrictions (#147)

Author: Trey2k

doc_classes/LuaAPI.xml

@@ -104,11 +104,21 @@
 			<description>
 				Controls the garbage collector. The option can be one of the following: [code]GC_STOP[/code], [code]GC_RESTART[/code], [code]GC_COLLECT[/code], [code]GC_COUNT[/code], [code]GC_STEP[/code], [code]GC_SETPAUSE[/code], [code]GC_SETSTEPMUL[/code]. The data is the argument for the option. Returns the result of the option.
 			</description>
+		</method>
+		<method name="get_memory_usage">
+			<return type="int" />
+			<description>
+				Returns the current memory usage of the state in bytes.
+			</description>
+		</method>
 	</methods>
 	<members>
 		<member name="permissive" type="bool" setter="set_permissive" getter="get_permissive" default="true">
 			When set to true all methods will be allowed on Objects be default and lua_fields is treated as a blacklist. When set to false, lua_fields is treated as a whitelist.
 		</member>
+		<member name="memory_limit" type="int" setter="set_memory_limit" getter="get_memory_limit" default="0">
+			Sets the memory limit for the state in bytes. If the limit is 0, there is no limit.
+		</member>
 	</members>
 	<constants>
 		<constant name="GC_STOP" value="1" enum="HookMask">

src/classes/luaAPI.cpp

@@ -9,7 +9,8 @@
 #endif
 
 LuaAPI::LuaAPI() {
-	lState = luaL_newstate();
+	lState = lua_newstate(&LuaAPI::luaAlloc, (void *)&luaAllocData);
+
 	// Creating lua state instance
 	state.setState(lState, this, true);
 }
@@ -25,7 +26,8 @@ void LuaAPI::_bind_methods() {
 
 	ClassDB::bind_method(D_METHOD("bind_libraries", "Array"), &LuaAPI::bindLibraries);
 	ClassDB::bind_method(D_METHOD("set_hook", "Hook", "HookMask", "Count"), &LuaAPI::setHook);
-	ClassDB::bind_method(D_METHOD("configure_gc", "What", "Data"), &LuaAPI::configure_gc);
+	ClassDB::bind_method(D_METHOD("configure_gc", "What", "Data"), &LuaAPI::configureGC);
+	ClassDB::bind_method(D_METHOD("get_memory_usage"), &LuaAPI::getMemoryUsage);
 	ClassDB::bind_method(D_METHOD("push_variant", "Name", "var"), &LuaAPI::pushGlobalVariant);
 	ClassDB::bind_method(D_METHOD("pull_variant", "Name"), &LuaAPI::pullVariant);
 	ClassDB::bind_method(D_METHOD("expose_constructor", "LuaConstructorName", "Object"), &LuaAPI::exposeObjectConstructor);
@@ -39,7 +41,11 @@ void LuaAPI::_bind_methods() {
 	ClassDB::bind_method(D_METHOD("set_permissive", "value"), &LuaAPI::setPermissive);
 	ClassDB::bind_method(D_METHOD("get_permissive"), &LuaAPI::getPermissive);
 
+	ClassDB::bind_method(D_METHOD("set_memory_limit", "limit"), &LuaAPI::setMemoryLimit);
+	ClassDB::bind_method(D_METHOD("get_memory_limit"), &LuaAPI::getMemoryLimit);
+
 	ADD_PROPERTY(PropertyInfo(Variant::INT, "permissive"), "set_permissive", "get_permissive");
+	ADD_PROPERTY(PropertyInfo(Variant::INT, "memory_limit"), "set_memory_limit", "get_memory_limit");
 
 	BIND_ENUM_CONSTANT(HOOK_MASK_CALL);
 	BIND_ENUM_CONSTANT(HOOK_MASK_RETURN);
@@ -65,6 +71,30 @@ void LuaAPI::setHook(Callable hook, int mask, int count) {
 	return state.setHook(hook, mask, count);
 }
 
+int LuaAPI::configureGC(int what, int data) {
+	return lua_gc(lState, what, data);
+}
+
+void LuaAPI::setPermissive(bool value) {
+	permissive = value;
+}
+
+bool LuaAPI::getPermissive() const {
+	return permissive;
+}
+
+void LuaAPI::setMemoryLimit(int limit) {
+	luaAllocData.memoryLimit = limit;
+}
+
+int LuaAPI::getMemoryLimit() const {
+	return luaAllocData.memoryLimit;
+}
+
+int LuaAPI::getMemoryUsage() const {
+	return luaAllocData.memoryUsed;
+}
+
 // Calls LuaState::luaFunctionExists()
 bool LuaAPI::luaFunctionExists(String functionName) {
 	return state.luaFunctionExists(functionName);
@@ -204,3 +234,31 @@ lua_State *LuaAPI::newThreadState() {
 lua_State *LuaAPI::getState() {
 	return lState;
 }
+
+void *LuaAPI::luaAlloc(void *ud, void *ptr, size_t osize, size_t nsize) {
+	LuaAllocData *data = (LuaAllocData *)ud;
+	if (nsize == 0) {
+		if (ptr != nullptr) {
+			data->memoryUsed -= osize;
+			memfree(ptr);
+		}
+		return nullptr;
+	}
+
+	if (ptr == nullptr) {
+		if (data->memoryLimit != 0 && data->memoryUsed + nsize > data->memoryLimit) {
+			return nullptr;
+		}
+
+		data->memoryUsed += nsize;
+		return memalloc(nsize);
+	}
+
+	if (data->memoryLimit != 0 && data->memoryUsed - osize + nsize > data->memoryLimit) {
+		return nullptr;
+	}
+
+	data->memoryUsed -= osize;
+	data->memoryUsed += nsize;
+	return memrealloc(ptr, nsize);
+}

src/classes/luaAPI.h

@@ -32,17 +32,14 @@ class LuaAPI : public RefCounted {
 	void bindLibraries(Array libs);
 	void setHook(Callable hook, int mask, int count);
 
-	inline int configure_gc(int what, int data) {
-		return lua_gc(lState, what, data);
-	}
+	void setPermissive(bool value);
+	bool getPermissive() const;
 
-	inline void setPermissive(bool value) {
-		permissive = value;
-	}
+	void setMemoryLimit(int limit);
+	int getMemoryLimit() const;
 
-	inline bool getPermissive() const {
-		return permissive;
-	}
+	int configureGC(int what, int data);
+	int getMemoryUsage() const;
 
 	bool luaFunctionExists(String functionName);
 
@@ -83,6 +80,15 @@ class LuaAPI : public RefCounted {
 	LuaState state;
 	lua_State *lState = nullptr;
 
+	static void *luaAlloc(void *ud, void *ptr, size_t osize, size_t nsize);
+
+	struct LuaAllocData {
+		int memoryUsed = 0;
+		int memoryLimit = 0;
+	};
+
+	LuaAllocData luaAllocData;
+
 	bool permissive = true;
 
 	LuaError *execute(int handlerIndex);