Cleo return arguments count check (#18)

Author: MiranDMC

cleo_sdk/CLEO.h

@@ -370,7 +370,7 @@ void WINAPI CLEO_SetThreadCondResult(CRunningScript* thread, BOOL result);
 void WINAPI CLEO_ThreadJumpAtLabelPtr(CRunningScript* thread, int labelPtr);
 
 eDataType WINAPI CLEO_GetOperandType(const CRunningScript* thread); // peek parameter data type
-DWORD WINAPI CLEO_GetVarArgCount(CRunningScript* thread); // peek var-args count
+DWORD WINAPI CLEO_GetVarArgCount(CRunningScript* thread); // peek remaining var-args count
 
 extern SCRIPT_VAR* opcodeParams;
 extern SCRIPT_VAR* missionLocals;

source/CCustomOpcodeSystem.cpp

@@ -2075,7 +2075,7 @@ namespace CLEO
 		int label = 0;
 
 		char* moduleTxt = nullptr;
-		auto paramType = CLEO_GetOperandType(thread);
+		auto paramType = (eDataType)*thread->GetBytePointer();
 		switch (paramType)
 		{
 			// label of current script
@@ -2104,7 +2104,7 @@ namespace CLEO
 				break;
 
 			default:
-				SHOW_ERROR("Invalid type (%s) of the first argument in opcode [0AB1] in script %s \nScript suspended.", ToKindStr(paramType), ((CCustomScript*)thread)->GetInfoStr().c_str());
+				SHOW_ERROR("Invalid type (%s) of the 'input param count' argument in opcode [0AB1] in script %s \nScript suspended.", ToKindStr(paramType), ((CCustomScript*)thread)->GetInfoStr().c_str());
 				return CCustomOpcodeSystem::ErrorSuspendScript(thread);
 		}
 
@@ -2142,13 +2142,41 @@ namespace CLEO
 			label = scriptRef.offset;
 		}
 
-		DWORD nParams = 0;
-		if(*thread->GetBytePointer()) *thread >> nParams;
-		if(nParams > 32)
+		// "number of input parameters" opcode argument
+		DWORD nParams;
+		paramType = (eDataType)*thread->GetBytePointer();
+		switch (paramType)
 		{
-			SHOW_ERROR("Argument count (%d), out of supported range (32) of opcode [0AB1] in script %s", nParams, ((CCustomScript*)thread)->GetInfoStr().c_str());
+			case DT_END:
+				nParams = 0;
+				break;
 
-			return CCustomOpcodeSystem::ErrorSuspendScript(thread);
+			// literal integers
+			case DT_BYTE:
+			case DT_WORD:
+			case DT_DWORD:
+				*thread >> nParams;
+				break;
+
+			default:
+				SHOW_ERROR("Invalid type of first argument in opcode [0AB1], in script %s", ((CCustomScript*)thread)->GetInfoStr().c_str());
+				return CCustomOpcodeSystem::ErrorSuspendScript(thread);
+		}
+		if (nParams)
+		{
+			auto nVarArg = GetVarArgCount(thread);
+			if (nParams > nVarArg) // if less it means there are return params too
+			{
+				SHOW_ERROR("Opcode [0AB1] declared %d input args, but provided %d in script %s\nScript suspended.", nParams, nVarArg, ((CCustomScript*)thread)->GetInfoStr().c_str());
+				return CCustomOpcodeSystem::ErrorSuspendScript(thread);
+			}
+
+			if (nParams > 32)
+			{
+				SHOW_ERROR("Argument count %d is out of supported range (32) of opcode [0AB1] in script %s", nParams, ((CCustomScript*)thread)->GetInfoStr().c_str());
+
+				return CCustomOpcodeSystem::ErrorSuspendScript(thread);
+			}
 		}
 
 		static SCRIPT_VAR arguments[32];
@@ -2228,18 +2256,43 @@ namespace CLEO
 	{
 		ScmFunction *scmFunc = ScmFunction::Store[reinterpret_cast<CCustomScript*>(thread)->GetScmFunction()];
 
-		DWORD returnParamCount = 0;
-		if (*thread->GetBytePointer()) *thread >> returnParamCount;
+		DWORD returnParamCount = GetVarArgCount(thread);
+		if (returnParamCount)
+		{
+			DWORD declaredParamCount;
+
+			auto paramType = (eDataType)*thread->GetBytePointer();
+			switch (paramType)
+			{
+				// literal integers
+				case DT_BYTE:
+				case DT_WORD:
+				case DT_DWORD:
+					*thread >> declaredParamCount;
+				break;
+
+			default:
+				SHOW_ERROR("Invalid type of first argument in opcode [0AB2], in script %s", ((CCustomScript*)thread)->GetInfoStr().c_str());
+				return CCustomOpcodeSystem::ErrorSuspendScript(thread);
+			}
+
+			if(returnParamCount - 1 != declaredParamCount) // minus 'num args' itself
+			{
+				SHOW_ERROR("Opcode [0AB2] declared %d return args, but provided %d in script %s\nScript suspended.", declaredParamCount, returnParamCount - 1, ((CCustomScript*)thread)->GetInfoStr().c_str());
+				return CCustomOpcodeSystem::ErrorSuspendScript(thread);
+			}
+		}
 		if (returnParamCount) GetScriptParams(thread, returnParamCount);
 
 		scmFunc->Return(thread); // jump back to cleo_call, right after last input param. Return slot var args starts here
 		if (scmFunc->moduleExportRef != nullptr) GetInstance().ModuleSystem.ReleaseModuleRef((char*)scmFunc->moduleExportRef); // export - release module
 		delete scmFunc;
 
 		DWORD returnSlotCount = GetVarArgCount(thread);
-		if (returnSlotCount > returnParamCount)
+		if(returnParamCount) returnParamCount--; // do not count the 'num args' argument itself
+		if (returnSlotCount != returnParamCount)
 		{
-			SHOW_ERROR("Opcode [0AB2] returned fewer params than expected by function caller in script %s\nScript suspended.", ((CCustomScript*)thread)->GetInfoStr().c_str());
+			SHOW_ERROR("Opcode [0AB2] returned %d params, while function caller expected %d in script %s\nScript suspended.", returnParamCount, returnSlotCount, ((CCustomScript*)thread)->GetInfoStr().c_str());
 			return CCustomOpcodeSystem::ErrorSuspendScript(thread);
 		}