Fix(lua / script manager): Potential fix for stack overflow when yielding from lua scripts. Fix io / os lua libs being accessible. (#1681)

* fix(script-mgr / lua): first kill all scripts, then unload lua modules: because the lua scripts depend on lua state (which is stored inside lua module instance), killing the lua module first would not allow proper cleaning because of the lua state getting destroyed while the lua script might still be running.

* fix(unloading): Why even reset the fiber pool here?

* fix(lua): don't allow for io / os lua lib to be accessed for security reasons.

* fix(lua): Potential fix for C stack overflow error by using lua coroutine yielding instead of calling fiber yield directly from lua functions.

* feat(unloading): allow to unload in the main title screen. Also revert 309c374602 due to fiber pool being potentially exhausted, we want commands to have their cleanup code ran in priority.

src/lua/bindings/script.hpp

@@ -12,24 +12,20 @@ namespace lua::script
 	class script_util
 	{
 	public:
+		// We keep the two functions below yield and sleep for backcompat.
+		// The idea of exposing big::script::get_current()->yield directly to lua
+		// on the surface looks like it could works but it doesnt, the lua stack end up exploding.
+
 		// Lua API: Function
 		// Class: script_util
 		// Name: yield
 		// Yield execution.
-		void yield()
-		{
-			big::script::get_current()->yield();
-		}
 
 		// Lua API: Function
 		// Class: script_util
 		// Name: sleep
 		// Param: ms: integer: The amount of time in milliseconds that we will sleep for.
 		// Sleep for the given amount of time, time is in milliseconds.
-		void sleep(int ms)
-		{
-			big::script::get_current()->yield(std::chrono::milliseconds(ms));
-		}
 	};
 	static script_util dummy_script_util;
 
@@ -66,12 +62,12 @@ namespace lua::script
 	//     ENTITY.DELETE_ENTITY(spawnedVehicle)
 	// end)
 	// ```
-	static void register_looped(const std::string& name, sol::function func, sol::this_state state)
+	static void register_looped(const std::string& name, sol::coroutine func, sol::this_state state)
 	{
 		auto module = sol::state_view(state)["!this"].get<big::lua_module*>();
 
 		module->m_registered_scripts.push_back(big::g_script_mgr.add_script(std::make_unique<big::script>(
-		    [func] {
+		    [func] () mutable {
 			    while (big::g_running)
 			    {
 				    auto res = func(dummy_script_util);
@@ -79,7 +75,14 @@ namespace lua::script
 				    if (!res.valid())
 					    big::g_lua_manager->handle_error(res, res.lua_state());
 
-				    big::script::get_current()->yield();
+					if (func.runnable())
+					{
+					    big::script::get_current()->yield(std::chrono::milliseconds(res.return_count() ? res[0] : 0));
+					}
+					else
+					{
+					    big::script::get_current()->yield();
+					}
 			    }
 		    },
 		    name)));
@@ -113,13 +116,27 @@ namespace lua::script
     //     ENTITY.DELETE_ENTITY(spawnedVehicle)
 	// end)
 	// ```
-	static void run_in_fiber(sol::function func)
+	static void run_in_fiber(sol::coroutine func)
 	{
-		big::g_fiber_pool->queue_job([func] {
-			auto res = func(dummy_script_util);
+		big::g_fiber_pool->queue_job([func] () mutable {
+			while (big::g_running)
+			{
+				auto res = func(dummy_script_util);
 
-			if (!res.valid())
-				big::g_lua_manager->handle_error(res, res.lua_state());
+				if (!res.valid())
+					big::g_lua_manager->handle_error(res, res.lua_state());
+
+				// Still runnable, meaning that the user function yielded
+				// and that there is more code to run still.
+				if (func.runnable())
+				{
+					big::script::get_current()->yield(std::chrono::milliseconds(res.return_count() ? res[0] : 0));
+				}
+				else
+				{
+					break;
+				}
+			}
 		});
 	}
 
@@ -129,10 +146,12 @@ namespace lua::script
 		ns["register_looped"] = register_looped;
 		ns["run_in_fiber"]    = run_in_fiber;
 
+		sol::function lua_coroutine_yield = state["coroutine"]["yield"];
+
 		//clang-format off
 		state.new_usertype<script_util>("script_util",
-			"yield", &script_util::yield,
-			"sleep", &script_util::sleep);
+		    "yield", lua_coroutine_yield,
+		    "sleep", lua_coroutine_yield);
 		//clang-format on
 	}
 }

src/lua/lua_module.cpp

@@ -49,7 +49,17 @@ namespace big
 	    m_module_name(module_name),
 	    m_module_id(rage::joaat(module_name))
 	{
-		m_state.open_libraries();
+		m_state.open_libraries(
+			sol::lib::base,
+			sol::lib::package,
+			sol::lib::coroutine,
+		    sol::lib::string,
+		    sol::lib::math,
+			sol::lib::table,
+		    sol::lib::debug,
+			sol::lib::bit32,
+			sol::lib::utf8
+		);
 
 		const auto& scripts_folder = g_lua_manager->get_scripts_folder();
 

src/main.cpp

@@ -149,6 +149,9 @@ BOOL APIENTRY DllMain(HMODULE hmod, DWORD reason, PVOID)
 			    while (g_running)
 				    std::this_thread::sleep_for(500ms);
 
+				g_script_mgr.remove_all_scripts();
+			    LOG(INFO) << "Scripts unregistered.";
+
 			    lua_manager_instance.reset();
 			    LOG(INFO) << "Lua manager uninitialized.";
 
@@ -158,9 +161,6 @@ BOOL APIENTRY DllMain(HMODULE hmod, DWORD reason, PVOID)
 			    native_hooks_instance.reset();
 			    LOG(INFO) << "Dynamic native hooker uninitialized.";
 
-			    g_script_mgr.remove_all_scripts();
-			    LOG(INFO) << "Scripts unregistered.";
-
 			    // cleans up the thread responsible for saving settings
 			    g.destroy();
 

src/script_mgr.cpp

@@ -38,9 +38,16 @@ namespace big
 		gta_util::execute_as_script(RAGE_JOAAT("main_persistent"), std::mem_fn(&script_mgr::tick_internal), this);
 	}
 
+	void script_mgr::ensure_main_fiber()
+	{
+		ConvertThreadToFiber(nullptr);
+
+		m_can_tick = true;
+	}
+
 	void script_mgr::tick_internal()
 	{
-		static bool ensure_main_fiber = (ConvertThreadToFiber(nullptr), true);
+		static bool ensure_it = (ensure_main_fiber(), true);
 
 		std::lock_guard lock(m_mutex);
 

src/script_mgr.hpp

@@ -28,13 +28,21 @@ namespace big
 
 		void tick();
 
+		[[nodiscard]] inline bool can_tick() const
+		{
+			return m_can_tick;
+		}
+
 	private:
+		void ensure_main_fiber();
 		void tick_internal();
 
 	private:
 		std::recursive_mutex m_mutex;
 		script_list m_scripts;
 		script_list m_scripts_to_add;
+
+		bool m_can_tick = false;
 	};
 
 	inline script_mgr g_script_mgr;

src/services/tunables/tunables_service.cpp

@@ -24,7 +24,7 @@ namespace big
 			g_thread_pool->push([this] {
 				while (!g_pointers->m_gta.m_script_globals[1])
 				{
-					if (!m_loading)
+					if (!m_loading || !g_running)
 						return;
 
 					std::this_thread::yield();

src/views/core/view_heading.cpp

@@ -1,5 +1,6 @@
 #include "fiber_pool.hpp"
 #include "views/view.hpp"
+#include "script_mgr.hpp"
 
 namespace big
 {
@@ -24,14 +25,23 @@ namespace big
 			ImGui::PushStyleColor(ImGuiCol_Text, ImVec4(0.69f, 0.29f, 0.29f, 1.00f));
 			if (components::nav_button("UNLOAD"_T))
 			{
-				g_fiber_pool->reset();
-				g_fiber_pool->queue_job([] {
-					for (auto& command : g_looped_commands)
-						if (command->is_enabled())
-							command->on_disable();
+				// allow to unload in the main title screen.
+				if (g_script_mgr.can_tick())
+				{
+					// empty the pool, we want the that job below run no matter what for clean up purposes.
+					g_fiber_pool->reset();
+					g_fiber_pool->queue_job([] {
+						for (auto& command : g_looped_commands)
+							if (command->is_enabled())
+								command->on_disable();
 
+						g_running = false;
+					});
+				}
+				else
+				{
 					g_running = false;
-				});
+				}
 			}
 			ImGui::PopStyleColor();
 		}