From 2896dc0d372270f0f51b0194194b6b7eb29c3c96 Mon Sep 17 00:00:00 2001
From: Kawe Mazidjatari <48657826+Mauler125@users.noreply.github.com>
Date: Mon, 1 Jul 2024 01:30:01 +0200
Subject: [PATCH] Engine: fix code bug when running "uiscript_reset" during
 init or load screens

Make sure this command cannot be executed before the UI VM is initialized or during load screens, doing so will hard crash the engine.
---
 src/common/callback.cpp                       | 26 ++++++++++++++++---
 src/common/callback.h                         |  9 +++++++
 src/vscript/languages/squirrel_re/vsquirrel.h | 10 ++++++-
 3 files changed, 41 insertions(+), 4 deletions(-)

diff --git a/src/common/callback.cpp b/src/common/callback.cpp
index f98f42ed..be604ee0 100644
--- a/src/common/callback.cpp
+++ b/src/common/callback.cpp
@@ -13,8 +13,10 @@
 #include "engine/shared/shared_rcon.h"
 #ifndef CLIENT_DLL
 #include "engine/server/sv_rcon.h"
+#include "engine/server/server.h"
 #endif // !CLIENT_DLL
 #ifndef DEDICATED
+#include "engine/gl_screen.h"
 #include "engine/client/cl_rcon.h"
 #include "engine/client/cdll_engine_int.h"
 #include "engine/client/clientstate.h"
@@ -24,9 +26,6 @@
 #include "engine/host_cmd.h"
 #include "engine/host_state.h"
 #include "engine/enginetrace.h"
-#ifndef CLIENT_DLL
-#include "engine/server/server.h"
-#endif // !CLIENT_DLL
 
 #include "rtech/pak/pakencode.h"
 #include "rtech/pak/pakdecode.h"
@@ -71,6 +70,7 @@
 #include "game/client/cliententitylist.h"
 #include "game/client/viewrender.h"
 #endif // !DEDICATED
+#include "vscript/languages/squirrel_re/vsquirrel.h"
 
 
 /*
@@ -566,8 +566,28 @@ void Cmd_Exec_f(const CCommand& args)
 	v__Cmd_Exec_f(args);
 }
 
+#ifndef DEDICATED
+void UIScript_Reset_f()
+{
+	// NOTE: function 'Script_InitUIVM' hasn't been called yet; nothing to reset,
+	// prevent crashing during init code.
+	if (!*g_bUIScriptInitialized)
+		return;
+
+	// NOTE: don't allow the execution of this command in loading screens, doing
+	// so will hard crash the engine!
+	if (*scr_drawloading)
+		return;
+
+	v__UIScript_Reset_f();
+}
+#endif // !DEDICATED
+
 
 void VCallback::Detour(const bool bAttach) const
 {
 	DetourSetup(&v__Cmd_Exec_f, &Cmd_Exec_f, bAttach);
+#ifndef DEDICATED
+	DetourSetup(&v__UIScript_Reset_f, &UIScript_Reset_f, bAttach);
+#endif // !DEDICATED
 }
diff --git a/src/common/callback.h b/src/common/callback.h
index 47a48858..fba09c28 100644
--- a/src/common/callback.h
+++ b/src/common/callback.h
@@ -4,6 +4,9 @@ inline bool(*v_SetupGamemode)(const char* pszPlayList);
 
 /* ==== CONCOMMANDCALLBACK ============================================================================================================================================== */
 inline void(*v__Cmd_Exec_f)(const CCommand& args);
+#ifndef DEDICATED
+inline void(*v__UIScript_Reset_f)();
+#endif // !DEDICATED
 
 ///////////////////////////////////////////////////////////////////////////////
 void MP_GameMode_Changed_f(IConVar* pConVar, const char* pOldString);
@@ -38,11 +41,17 @@ class VCallback : public IDetour
 	{
 		LogFunAdr("SetupGamemode", v_SetupGamemode);
 		LogFunAdr("Cmd_Exec_f", v__Cmd_Exec_f);
+#ifndef DEDICATED
+		LogFunAdr("UIScript_Reset_f", v__UIScript_Reset_f);
+#endif // !DEDICATED
 	}
 	virtual void GetFun(void) const
 	{
 		g_GameDll.FindPatternSIMD("40 53 48 83 EC 20 48 8B D9 48 C7 C0 ?? ?? ?? ??").GetPtr(v_SetupGamemode);
 		g_GameDll.FindPatternSIMD("40 55 53 48 8D AC 24 ?? ?? ?? ?? B8 ?? ?? ?? ?? E8 ?? ?? ?? ?? 48 2B E0 48 8B D9").GetPtr(v__Cmd_Exec_f);
+#ifndef DEDICATED
+		g_GameDll.FindPatternSIMD("40 55 41 54 48 8D AC 24 ?? ?? ?? ?? 48 81 EC ?? ?? ?? ?? 45 33 E4 48 8D 0D").GetPtr(v__UIScript_Reset_f);
+#endif // !DEDICATED
 	}
 	virtual void GetVar(void) const { }
 	virtual void GetCon(void) const { }
diff --git a/src/vscript/languages/squirrel_re/vsquirrel.h b/src/vscript/languages/squirrel_re/vsquirrel.h
index 212d075b..69f08c09 100644
--- a/src/vscript/languages/squirrel_re/vsquirrel.h
+++ b/src/vscript/languages/squirrel_re/vsquirrel.h
@@ -93,6 +93,8 @@ inline CSquirrelVM* g_pServerScript;
 #ifndef DEDICATED
 inline CSquirrelVM* g_pClientScript;
 inline CSquirrelVM* g_pUIScript;
+
+inline bool* g_bUIScriptInitialized;
 #endif // !DEDICATED
 
 #define DEFINE_SCRIPTENUM_NAMED(s, enumName, startValue, ...) \
@@ -165,7 +167,13 @@ class VSquirrel : public IDetour
 		g_GameDll.FindPatternSIMD("E8 ?? ?? ?? ?? C6 47 1C 01").FollowNearCallSelf().GetPtr(CSquirrelVM__ExecuteCodeCallback);
 		g_GameDll.FindPatternSIMD("E8 ?? ?? ?? ?? BB ?? ?? ?? ?? 8B C3").FollowNearCallSelf().GetPtr(CSquirrelVM__ThrowError);
 	}
-	virtual void GetVar(void) const { }
+	virtual void GetVar(void) const
+	{
+#ifndef DEDICATED
+		CMemory p_Script_InitUIVM = g_GameDll.FindPatternSIMD("48 83 EC ?? 80 3D ?? ?? ?? ?? ?? 74 ?? 32 C0");
+		p_Script_InitUIVM.Offset(0x10).FindPatternSelf("80 3D").ResolveRelativeAddressSelf(2, 7).GetPtr(g_bUIScriptInitialized);
+#endif
+	}
 	virtual void GetCon(void) const { }
 	virtual void Detour(const bool bAttach) const;
 };