From 15745eebe1ae919cc6a41dbbc7cdaed5911a773d Mon Sep 17 00:00:00 2001 From: Quentin Date: Tue, 11 Jul 2023 09:24:44 +0200 Subject: [PATCH] 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 https://github.com/YimMenu/YimMenu/commit/309c37460208d23402b17e8b252e379434e2fe8d due to fiber pool being potentially exhausted, we want commands to have their cleanup code ran in priority. --- src/lua/bindings/script.hpp | 55 +++++++++++++++------- src/lua/lua_module.cpp | 12 ++++- src/main.cpp | 6 +-- src/script_mgr.cpp | 9 +++- src/script_mgr.hpp | 8 ++++ src/services/tunables/tunables_service.cpp | 2 +- src/views/core/view_heading.cpp | 22 ++++++--- 7 files changed, 84 insertions(+), 30 deletions(-) diff --git a/src/lua/bindings/script.hpp b/src/lua/bindings/script.hpp index 8b913d8a..389c0d7f 100644 --- a/src/lua/bindings/script.hpp +++ b/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(); module->m_registered_scripts.push_back(big::g_script_mgr.add_script(std::make_unique( - [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", - "yield", &script_util::yield, - "sleep", &script_util::sleep); + "yield", lua_coroutine_yield, + "sleep", lua_coroutine_yield); //clang-format on } } \ No newline at end of file diff --git a/src/lua/lua_module.cpp b/src/lua/lua_module.cpp index 274e6997..039ba259 100644 --- a/src/lua/lua_module.cpp +++ b/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(); diff --git a/src/main.cpp b/src/main.cpp index 72e73cff..ddf533b2 100644 --- a/src/main.cpp +++ b/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(); diff --git a/src/script_mgr.cpp b/src/script_mgr.cpp index 8005ee16..3ac0ef91 100644 --- a/src/script_mgr.cpp +++ b/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); diff --git a/src/script_mgr.hpp b/src/script_mgr.hpp index 9568375e..6fa37441 100644 --- a/src/script_mgr.hpp +++ b/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; diff --git a/src/services/tunables/tunables_service.cpp b/src/services/tunables/tunables_service.cpp index 03b61366..5d9361e7 100644 --- a/src/services/tunables/tunables_service.cpp +++ b/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(); diff --git a/src/views/core/view_heading.cpp b/src/views/core/view_heading.cpp index b6e9b192..d9674545 100644 --- a/src/views/core/view_heading.cpp +++ b/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(); }