Add support for custom calling convention passing arguments in registers

If the compiler decided to pass an argument in a register on an internal function instead of pushing it on the stack to save time, allow us to specify the register the parameter is going to be in.

DHookAddParam received another parameter to set the register.

DynamicHooks/conventions/x86MsCdecl.cpp

@@ -78,6 +78,16 @@ ke::Vector<Register_t> x86MsCdecl::GetRegisters()
 		}
 	}
 
+	// Save all the custom calling convention registers as well.
+	for (unsigned int i = 0; i < m_vecArgTypes.length(); i++)
+	{
+		if (m_vecArgTypes[i].custom_register == None)
+			continue;
+
+		// TODO: Make sure the list is unique? Set?
+		registers.append(m_vecArgTypes[i].custom_register);
+	}
+
 	return registers;
 }
 
@@ -92,7 +102,8 @@ int x86MsCdecl::GetArgStackSize()
 
 	for (unsigned int i = 0; i < m_vecArgTypes.length(); i++)
 	{
-		iArgStackSize += m_vecArgTypes[i].size;
+		if (m_vecArgTypes[i].custom_register == None)
+			iArgStackSize += m_vecArgTypes[i].size;
 	}
 
 	return iArgStackSize;
@@ -103,12 +114,36 @@ void** x86MsCdecl::GetStackArgumentPtr(CRegisters* pRegisters)
 	return (void **)(pRegisters->m_esp->GetValue<unsigned long>() + 4);
 }
 
+int x86MsCdecl::GetArgRegisterSize()
+{
+	int iArgRegisterSize = 0;
+
+	for (unsigned int i = 0; i < m_vecArgTypes.length(); i++)
+	{
+		if (m_vecArgTypes[i].custom_register != None)
+			iArgRegisterSize += m_vecArgTypes[i].size;
+	}
+
+	return iArgRegisterSize;
+}
+
 void* x86MsCdecl::GetArgumentPtr(int iIndex, CRegisters* pRegisters)
 {
+	// Check if this argument was passed in a register.
+	if (m_vecArgTypes[iIndex].custom_register != None)
+	{
+		CRegister *pRegister = pRegisters->GetRegister(m_vecArgTypes[iIndex].custom_register);
+		if (!pRegister)
+			return NULL;
+
+		return pRegister->m_pAddress;
+	}
+
 	int iOffset = 4;
 	for(int i=0; i < iIndex; i++)
 	{
-		iOffset += m_vecArgTypes[i].size;
+		if (m_vecArgTypes[i].custom_register == None)
+			iOffset += m_vecArgTypes[i].size;
 	}
 
 	return (void *) (pRegisters->m_esp->GetValue<unsigned long>() + iOffset);

DynamicHooks/conventions/x86MsCdecl.h

@@ -70,6 +70,7 @@ public:
 	virtual int GetPopSize();
 	virtual int GetArgStackSize();
 	virtual void** GetStackArgumentPtr(CRegisters* pRegisters);
+	virtual int GetArgRegisterSize();
 	
 	virtual void* GetArgumentPtr(int iIndex, CRegisters* pRegisters);
 	virtual void ArgumentPtrChanged(int iIndex, CRegisters* pRegisters, void* pArgumentPtr);

DynamicHooks/conventions/x86MsStdcall.cpp

@@ -78,6 +78,16 @@ ke::Vector<Register_t> x86MsStdcall::GetRegisters()
 		}
 	}
 
+	// Save all the custom calling convention registers as well.
+	for (unsigned int i = 0; i < m_vecArgTypes.length(); i++)
+	{
+		if (m_vecArgTypes[i].custom_register == None)
+			continue;
+
+		// TODO: Make sure the list is unique? Set?
+		registers.append(m_vecArgTypes[i].custom_register);
+	}
+
 	return registers;
 }
 
@@ -87,7 +97,8 @@ int x86MsStdcall::GetPopSize()
 
 	for(unsigned int i=0; i < m_vecArgTypes.length(); i++)
 	{
-		iPopSize += m_vecArgTypes[i].size;
+		if (m_vecArgTypes[i].custom_register == None)
+			iPopSize += m_vecArgTypes[i].size;
 	}
 
 	return iPopSize;
@@ -99,7 +110,8 @@ int x86MsStdcall::GetArgStackSize()
 
 	for (unsigned int i = 0; i < m_vecArgTypes.length(); i++)
 	{
-		iArgStackSize += m_vecArgTypes[i].size;
+		if (m_vecArgTypes[i].custom_register == None)
+			iArgStackSize += m_vecArgTypes[i].size;
 	}
 
 	return iArgStackSize;
@@ -110,15 +122,39 @@ void** x86MsStdcall::GetStackArgumentPtr(CRegisters* pRegisters)
 	return (void **)(pRegisters->m_esp->GetValue<unsigned long>() + 4);
 }
 
-void* x86MsStdcall::GetArgumentPtr(int iIndex, CRegisters* pRegisters)
+int x86MsStdcall::GetArgRegisterSize()
 {
-	int iOffset = 4;
-	for(int i=0; i < iIndex; i++)
+	int iArgRegisterSize = 0;
+
+	for (unsigned int i = 0; i < m_vecArgTypes.length(); i++)
 	{
-		iOffset += m_vecArgTypes[i].size;
+		if (m_vecArgTypes[i].custom_register != None)
+			iArgRegisterSize += m_vecArgTypes[i].size;
 	}
 
-	return (void *) (pRegisters->m_esp->GetValue<unsigned long>() + iOffset);
+	return iArgRegisterSize;
+}
+
+void* x86MsStdcall::GetArgumentPtr(int iIndex, CRegisters* pRegisters)
+{
+	// Check if this argument was passed in a register.
+	if (m_vecArgTypes[iIndex].custom_register != None)
+	{
+		CRegister *pRegister = pRegisters->GetRegister(m_vecArgTypes[iIndex].custom_register);
+		if (!pRegister)
+			return NULL;
+
+		return pRegister->m_pAddress;
+	}
+
+	int iOffset = 4;
+	for (int i = 0; i < iIndex; i++)
+	{
+		if (m_vecArgTypes[i].custom_register == None)
+			iOffset += m_vecArgTypes[i].size;
+	}
+
+	return (void *)(pRegisters->m_esp->GetValue<unsigned long>() + iOffset);
 }
 
 void x86MsStdcall::ArgumentPtrChanged(int iIndex, CRegisters* pRegisters, void* pArgumentPtr)

DynamicHooks/conventions/x86MsStdcall.h

@@ -70,6 +70,7 @@ public:
 	virtual int GetPopSize();
 	virtual int GetArgStackSize();
 	virtual void** GetStackArgumentPtr(CRegisters* pRegisters);
+	virtual int GetArgRegisterSize();
 	
 	virtual void* GetArgumentPtr(int iIndex, CRegisters* pRegisters);
 	virtual void ArgumentPtrChanged(int iIndex, CRegisters* pRegisters, void* pArgumentPtr);

DynamicHooks/conventions/x86MsThiscall.cpp

@@ -66,6 +66,7 @@ ke::Vector<Register_t> x86MsThiscall::GetRegisters()
 	ke::Vector<Register_t> registers;
 	
 	registers.append(ESP);
+	// TODO: Allow custom this register.
 	registers.append(ECX);
 
 	if (m_returnType.type == DATA_TYPE_FLOAT || m_returnType.type == DATA_TYPE_DOUBLE)
@@ -81,6 +82,16 @@ ke::Vector<Register_t> x86MsThiscall::GetRegisters()
 		}
 	}
 
+	// Save all the custom calling convention registers as well.
+	for (unsigned int i = 0; i < m_vecArgTypes.length(); i++)
+	{
+		if (m_vecArgTypes[i].custom_register == None)
+			continue;
+
+		// TODO: Make sure the list is unique? Set?
+		registers.append(m_vecArgTypes[i].custom_register);
+	}
+
 	return registers;
 }
 
@@ -93,7 +104,9 @@ int x86MsThiscall::GetPopSize()
 
 	for(unsigned int i=0; i < m_vecArgTypes.length(); i++)
 	{
-		iPopSize += m_vecArgTypes[i].size;
+		// Only pop arguments that are actually on the stack.
+		if (m_vecArgTypes[i].custom_register == None)
+			iPopSize += m_vecArgTypes[i].size;
 	}
 
 	return iPopSize;
@@ -105,7 +118,8 @@ int x86MsThiscall::GetArgStackSize()
 
 	for (unsigned int i = 0; i < m_vecArgTypes.length(); i++)
 	{
-		iArgStackSize += m_vecArgTypes[i].size;
+		if (m_vecArgTypes[i].custom_register == None)
+			iArgStackSize += m_vecArgTypes[i].size;
 	}
 
 	return iArgStackSize;
@@ -116,17 +130,42 @@ void** x86MsThiscall::GetStackArgumentPtr(CRegisters* pRegisters)
 	return (void **)(pRegisters->m_esp->GetValue<unsigned long>() + 4);
 }
 
+int x86MsThiscall::GetArgRegisterSize()
+{
+	int iArgRegisterSize = 0;
+
+	for (unsigned int i = 0; i < m_vecArgTypes.length(); i++)
+	{
+		if (m_vecArgTypes[i].custom_register != None)
+			iArgRegisterSize += m_vecArgTypes[i].size;
+	}
+
+	return iArgRegisterSize;
+}
+
 void* x86MsThiscall::GetArgumentPtr(int iIndex, CRegisters* pRegisters)
 {
 	if (iIndex == 0)
 	{
+		// TODO: Allow custom this register.
 		return pRegisters->m_ecx->m_pAddress;
 	}
 
+	// Check if this argument was passed in a register.
+	if (m_vecArgTypes[iIndex-1].custom_register != None)
+	{
+		CRegister *pRegister = pRegisters->GetRegister(m_vecArgTypes[iIndex-1].custom_register);
+		if (!pRegister)
+			return NULL;
+
+		return pRegister->m_pAddress;
+	}
+
 	int iOffset = 4;
 	for(int i=0; i < iIndex-1; i++)
 	{
-		iOffset += m_vecArgTypes[i].size;
+		if (m_vecArgTypes[i].custom_register == None)
+			iOffset += m_vecArgTypes[i].size;
 	}
 
 	return (void *) (pRegisters->m_esp->GetValue<unsigned long>() + iOffset);

DynamicHooks/conventions/x86MsThiscall.h

@@ -69,8 +69,9 @@ public:
 
 	virtual ke::Vector<Register_t> GetRegisters();
 	virtual int GetPopSize();
-	virtual int x86MsThiscall::GetArgStackSize();
+	virtual int GetArgStackSize();
 	virtual void** GetStackArgumentPtr(CRegisters* pRegisters);
+	virtual int GetArgRegisterSize();
 	
 	virtual void* GetArgumentPtr(int iIndex, CRegisters* pRegisters);
 	virtual void ArgumentPtrChanged(int iIndex, CRegisters* pRegisters, void* pArgumentPtr);