2 files changed, 54 insertions, 18 deletions

diff --git a/NorthstarDedicatedTest/squirrel.cpp b/NorthstarDedicatedTest/squirrel.cpp
index 84caa389..b1c5b4a8 100644
--- a/NorthstarDedicatedTest/squirrel.cpp
+++ b/NorthstarDedicatedTest/squirrel.cpp
@@ -6,6 +6,7 @@
 #include "concommand.h"
 #include "modmanager.h"
 #include <iostream>
+#include "gameutils.h"
 
 // hook forward declarations
 typedef SQInteger (*SQPrintType)(void* sqvm, char* fmt, ...);
@@ -34,6 +35,10 @@ CallScriptInitCallbackType ClientCallScriptInitCallback;
 CallScriptInitCallbackType ServerCallScriptInitCallback;
 template <ScriptContext context> char CallScriptInitCallbackHook(void* sqvm, const char* callback);
 
+RegisterSquirrelFuncType ClientRegisterSquirrelFunc;
+RegisterSquirrelFuncType ServerRegisterSquirrelFunc;
+template <ScriptContext context> int64_t RegisterSquirrelFuncHook(void* sqvm, SQFuncRegistration* funcReg, char unknown);
+
 // core sqvm funcs
 sq_compilebufferType ClientSq_compilebuffer;
 sq_compilebufferType ServerSq_compilebuffer;
@@ -44,9 +49,6 @@ sq_pushroottableType ServerSq_pushroottable;
 sq_callType ClientSq_call;
 sq_callType ServerSq_call;
 
-RegisterSquirrelFuncType ClientRegisterSquirrelFunc;
-RegisterSquirrelFuncType ServerRegisterSquirrelFunc;
-
 // sq stack array funcs
 sq_newarrayType ClientSq_newarray;
 sq_newarrayType ServerSq_newarray;
@@ -162,6 +164,11 @@ void InitialiseClientSquirrel(HMODULE baseAddress)
 		(char*)baseAddress + 0x10190,
 		&CallScriptInitCallbackHook<ScriptContext::CLIENT>,
 		reinterpret_cast<LPVOID*>(&ClientCallScriptInitCallback)); // client callscriptinitcallback function
+	ENABLER_CREATEHOOK(
+		hook,
+		(char*)baseAddress + 0x108E0,
+		&RegisterSquirrelFuncHook<ScriptContext::CLIENT>,
+		reinterpret_cast<LPVOID*>(&ClientRegisterSquirrelFunc)); // client registersquirrelfunc function
 }
 
 void InitialiseServerSquirrel(HMODULE baseAddress)
@@ -217,6 +224,12 @@ void InitialiseServerSquirrel(HMODULE baseAddress)
 		&CallScriptInitCallbackHook<ScriptContext::SERVER>,
 		reinterpret_cast<LPVOID*>(&ServerCallScriptInitCallback)); // server callscriptinitcallback function
 
+	ENABLER_CREATEHOOK(
+		hook,
+		(char*)baseAddress + 0x1DD10,
+		&RegisterSquirrelFuncHook<ScriptContext::SERVER>,
+		reinterpret_cast<LPVOID*>(&ServerRegisterSquirrelFunc)); // server registersquirrelfunc function
+
 	// cheat and clientcmd_can_execute allows clients to execute this, but since it's unsafe we only allow it when cheats are enabled
 	// for script_client and script_ui, we don't use cheats, so clients can execute them on themselves all they want
 	RegisterConCommand(
@@ -447,4 +460,29 @@ template <ScriptContext context> void ExecuteCodeCommand(const CCommand& args)
 		g_UISquirrelManager->ExecuteCode(args.ArgS());
 	else if (context == ScriptContext::SERVER)
 		g_ServerSquirrelManager->ExecuteCode(args.ArgS());
+}
+
+SQRESULT SQ_DevFuncStub(void* sqvm)
+{
+	spdlog::warn("Blocked execution of squirrel developer function for security reasons. To re-enable them use start parameter "
+				 "-allowSquirrelDevFunctions.");
+	return SQRESULT_NULL;
+}
+
+template <ScriptContext context> int64_t RegisterSquirrelFuncHook(void* sqvm, SQFuncRegistration* funcReg, char unknown)
+{
+	static std::set<std::string> allowedDevFunctions = {
+		"Dev_CommandLineHasParm",
+		"Dev_CommandLineParmValue",
+		"Dev_CommandLineRemoveParm",
+	};
+
+	if ((funcReg->devLevel == 1) && (!CommandLine()->CheckParm("-allowSquirrelDevFunctions")) &&
+		(!allowedDevFunctions.count(funcReg->squirrelFuncName)))
+		funcReg->funcPtr = SQ_DevFuncStub;
+
+	if (context == ScriptContext::SERVER)
+		return ServerRegisterSquirrelFunc(sqvm, funcReg, unknown);
+	else
+		return ClientRegisterSquirrelFunc(sqvm, funcReg, unknown);
 }
\ No newline at end of file
diff --git a/NorthstarDedicatedTest/squirrel.h b/NorthstarDedicatedTest/squirrel.h
index cc39cc2a..d6dfcc57 100644
--- a/NorthstarDedicatedTest/squirrel.h
+++ b/NorthstarDedicatedTest/squirrel.h
@@ -37,25 +37,23 @@ struct SQFuncRegistration
 	const char* squirrelFuncName;
 	const char* cppFuncName;
 	const char* helpText;
-	const char* returnValueType;
+	const char* returnTypeString;
 	const char* argTypes;
-	int16_t somethingThatsZero;
-	int16_t padding1;
-	int32_t unknown1;
-	int64_t unknown2;
-	int32_t unknown3;
-	int32_t padding2;
-	int64_t unknown4;
-	int64_t unknown5;
-	int64_t unknown6;
-	int32_t unknown7;
-	int32_t padding3;
+	__int32 unknown1;
+	__int32 devLevel;
+	const char* shortNameMaybe;
+	__int32 unknown2;
+	__int32 returnTypeEnum;
+	__int32* externalBufferPointer;
+	__int64 externalBufferSize;
+	__int64 unknown3;
+	__int64 unknown4;
 	void* funcPtr;
 
 	SQFuncRegistration()
 	{
 		memset(this, 0, sizeof(SQFuncRegistration));
-		this->padding2 = 32;
+		this->returnTypeEnum = 32;
 	}
 };
 
@@ -277,8 +275,8 @@ template <ScriptContext context> class SquirrelManager
 		reg->helpText = new char[helpText.size() + 1];
 		strcpy((char*)reg->helpText, helpText.c_str());
 
-		reg->returnValueType = new char[returnType.size() + 1];
-		strcpy((char*)reg->returnValueType, returnType.c_str());
+		reg->returnTypeString = new char[returnType.size() + 1];
+		strcpy((char*)reg->returnTypeString, returnType.c_str());
 
 		reg->argTypes = new char[argTypes.size() + 1];
 		strcpy((char*)reg->argTypes, argTypes.c_str());