Fix overriding return value

The custom return value was lost when calling the original function.
Save and restore our own return value, if we're about to call the original function.

DynamicHooks/convention.h

@@ -35,6 +35,7 @@
 // >> INCLUDES
 // ============================================================================
 #include "registers.h"
+#include <string.h>
 #include <am-vector.h>
 
 // ============================================================================
@@ -161,6 +162,12 @@ public:
 		if (!m_returnType.size)
 			m_returnType.size = GetDataTypeSize(m_returnType);
 		m_iAlignment = iAlignment;
+		m_pSavedReturnBuffer = malloc(m_returnType.size);
+	}
+
+	virtual ~ICallingConvention()
+	{
+		free(m_pSavedReturnBuffer);
 	}
 
 	/*
@@ -204,10 +211,26 @@ public:
 	*/
 	virtual void ReturnPtrChanged(CRegisters* pRegisters, void* pReturnPtr) = 0;
 
+	/*
+	Save the return value in a seperate buffer, so we can restore it after calling the original function.
+	*/
+	virtual void SaveReturnValue(CRegisters* pRegisters)
+	{
+		memcpy(m_pSavedReturnBuffer, GetReturnPtr(pRegisters), m_returnType.size);
+	}
+
+	virtual void RestoreReturnValue(CRegisters* pRegisters)
+	{
+		memcpy(GetReturnPtr(pRegisters), m_pSavedReturnBuffer, m_returnType.size);
+		ReturnPtrChanged(pRegisters, m_pSavedReturnBuffer);
+	}
+
 public:
 	ke::Vector<DataTypeSized_t> m_vecArgTypes;
 	DataTypeSized_t m_returnType;
 	int m_iAlignment;
+	// Save the return in case we call the original function and want to override the return again.
+	void* m_pSavedReturnBuffer;
 };
 
 #endif // _CONVENTION_H

DynamicHooks/conventions/x86MsCdecl.h

@@ -64,7 +64,7 @@ class x86MsCdecl: public ICallingConvention
 {	
 public:
 	x86MsCdecl(ke::Vector<DataTypeSized_t> &vecArgTypes, DataTypeSized_t returnType, int iAlignment=4);
-	~x86MsCdecl();
+	virtual ~x86MsCdecl();
 
 	virtual ke::Vector<Register_t> GetRegisters();
 	virtual int GetPopSize();

DynamicHooks/conventions/x86MsStdcall.h

@@ -64,7 +64,7 @@ class x86MsStdcall: public ICallingConvention
 {	
 public:
 	x86MsStdcall(ke::Vector<DataTypeSized_t> &vecArgTypes, DataTypeSized_t returnType, int iAlignment=4);
-	~x86MsStdcall();
+	virtual ~x86MsStdcall();
 
 	virtual ke::Vector<Register_t> GetRegisters();
 	virtual int GetPopSize();

DynamicHooks/conventions/x86MsThiscall.h

@@ -65,7 +65,7 @@ class x86MsThiscall: public ICallingConvention
 {	
 public:
 	x86MsThiscall(ke::Vector<DataTypeSized_t> &vecArgTypes, DataTypeSized_t returnType, int iAlignment=4);
-	~x86MsThiscall();
+	virtual ~x86MsThiscall();
 
 	virtual ke::Vector<Register_t> GetRegisters();
 	virtual int GetPopSize();

DynamicHooks/hook.cpp

@@ -52,6 +52,7 @@ CHook::CHook(void* pFunc, ICallingConvention* pConvention)
 	m_pFunc = pFunc;
 	m_pRegisters = new CRegisters(pConvention->GetRegisters());
 	m_pCallingConvention = pConvention;
+	m_LastPreReturnAction = ReturnAction_Ignored;
 
 	if (!m_hookHandler.init())
 		return;
@@ -150,21 +151,35 @@ bool CHook::AreCallbacksRegistered()
 	return false;
 }
 
-bool CHook::HookHandler(HookType_t eHookType)
+ReturnAction_t CHook::HookHandler(HookType_t eHookType)
 {
-	bool bOverride = false;
+	if (eHookType == HOOKTYPE_POST)
+	{
+		if (m_LastPreReturnAction == ReturnAction_Override)
+			m_pCallingConvention->RestoreReturnValue(m_pRegisters);
+	}
+
+	ReturnAction_t returnAction = ReturnAction_Ignored;
 	HookTypeMap::Result r = m_hookHandler.find(eHookType);
 	if (!r.found())
-		return bOverride;
+		return returnAction;
 
 	HookHandlerSet &callbacks = r->value;
 	for(HookHandlerSet::iterator it=callbacks.iter(); !it.empty(); it.next())
 	{
-		bool result = ((HookHandlerFn) *it)(eHookType, this);
-		if (result)
-			bOverride = true;
+		ReturnAction_t result = ((HookHandlerFn) *it)(eHookType, this);
+		if (result > returnAction)
+			returnAction = result;
 	}
-	return bOverride;
+
+	if (eHookType == HOOKTYPE_PRE)
+	{
+		m_LastPreReturnAction = returnAction;
+		if (returnAction == ReturnAction_Override)
+			m_pCallingConvention->SaveReturnValue(m_pRegisters);
+	}
+
+	return returnAction;
 }
 
 void* __cdecl CHook::GetReturnAddress(void* pESP)
@@ -190,9 +205,9 @@ void* CHook::CreateBridge()
 	// Write a redirect to the post-hook code
 	Write_ModifyReturnAddress(masm);
 
-	// Call the pre-hook handler and jump to label_supercede if true was returned
+	// Call the pre-hook handler and jump to label_supercede if ReturnAction_Supercede was returned
 	Write_CallHandler(masm, HOOKTYPE_PRE);
-	masm.cmpb(r8_al, true);
+	masm.cmpb(r8_al, ReturnAction_Supercede);
 	
 	// Restore the previously saved registers, so any changes will be applied
 	Write_RestoreRegisters(masm);
@@ -202,7 +217,7 @@ void* CHook::CreateBridge()
 	// Jump to the trampoline
 	masm.jmp(ExternalAddress(m_pTrampoline));
 
-	// This code will be executed if a pre-hook returns true
+	// This code will be executed if a pre-hook returns ReturnAction_Supercede
 	masm.bind(&label_supercede);
 
 	// Finally, return to the caller
@@ -302,7 +317,7 @@ void* CHook::CreatePostCallback()
 
 void CHook::Write_CallHandler(sp::MacroAssembler& masm, HookType_t type)
 {
-	bool (__cdecl CHook::*HookHandler)(HookType_t) = &CHook::HookHandler;
+	ReturnAction_t (__cdecl CHook::*HookHandler)(HookType_t) = &CHook::HookHandler;
 
 	// Save the registers so that we can access them in our handlers
 	Write_SaveRegisters(masm);

DynamicHooks/hook.h

@@ -51,12 +51,20 @@ enum HookType_t
 	HOOKTYPE_POST
 };
 
+enum ReturnAction_t
+{
+	ReturnAction_Ignored,				// handler didn't take any action
+	ReturnAction_Handled,				// plugin did something, but real function should still be called
+	ReturnAction_Override,				// call real function, but use my return value
+	ReturnAction_Supercede				// skip real function; use my return value
+};
+
 
 // ============================================================================
 // >> TYPEDEFS
 // ============================================================================
 class CHook;
-typedef bool (*HookHandlerFn)(HookType_t, CHook*);
+typedef ReturnAction_t (*HookHandlerFn)(HookType_t, CHook*);
 
 #ifdef __linux__
 #define __cdecl
@@ -171,7 +179,7 @@ private:
 
 	void* CreatePostCallback();
 
-	bool __cdecl HookHandler(HookType_t type);
+	ReturnAction_t __cdecl HookHandler(HookType_t type);
 
 	void* __cdecl GetReturnAddress(void* pESP);
 	void __cdecl SetReturnAddress(void* pRetAddr, void* pESP);
@@ -198,6 +206,9 @@ public:
 	void* m_pNewRetAddr;
 
 	ReturnAddressMap m_RetAddr;
+
+	// Save the last return action of the pre HookHandler for use in the post handler.
+	ReturnAction_t m_LastPreReturnAction;
 };
 
 #endif // _HOOK_H