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.
This commit is contained in:
Peace-Maker 2016-12-13 17:45:22 -07:00
parent 6e96e0fb84
commit b18e3284e1
8 changed files with 82 additions and 33 deletions

View File

@ -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

View File

@ -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();

View File

@ -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();

View File

@ -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();

View File

@ -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);

View File

@ -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

View File

@ -213,7 +213,7 @@ ICallingConvention *ConstructCallingConvention(HookSetup *setup)
return pCallConv;
}
bool HandleDetour(HookType_t hookType, CHook* pDetour)
ReturnAction_t HandleDetour(HookType_t hookType, CHook* pDetour)
{
DetourMap *map;
if (hookType == HOOKTYPE_PRE)
@ -224,7 +224,7 @@ bool HandleDetour(HookType_t hookType, CHook* pDetour)
// Find the callback list for this detour.
DetourMap::Result r = map->find(pDetour);
if (!r.found())
return false;
return ReturnAction_Ignored;
// List of all callbacks.
PluginCallbackList *wrappers = r->value;
@ -236,7 +236,7 @@ bool HandleDetour(HookType_t hookType, CHook* pDetour)
Handle_t pHndl = BAD_HANDLE;
int argNum = pDetour->m_pCallingConvention->m_vecArgTypes.length();
MRESReturn finalRet = MRES_Ignored;
ReturnAction_t finalRet = ReturnAction_Ignored;
ke::AutoPtr<void> finalRetBuf(new uint8_t[pDetour->m_pCallingConvention->m_returnType.size]);
// Call all the plugin functions..
@ -244,7 +244,7 @@ bool HandleDetour(HookType_t hookType, CHook* pDetour)
{
CDynamicHooksSourcePawn *pWrapper = wrappers->at(i);
IPluginFunction *pCallback = pWrapper->plugin_callback;
MRESReturn tempRet = MRES_Ignored;
ReturnAction_t tempRet = ReturnAction_Ignored;
ke::AutoPtr<void> tempRetBuf(new uint8_t[pDetour->m_pCallingConvention->m_returnType.size]);
// Find the this pointer.
@ -307,10 +307,10 @@ bool HandleDetour(HookType_t hookType, CHook* pDetour)
switch ((MRESReturn)result)
{
case MRES_Handled:
tempRet = MRES_Handled;
tempRet = ReturnAction_Handled;
break;
case MRES_ChangedHandled:
tempRet = MRES_Handled;
tempRet = ReturnAction_Handled;
pWrapper->UpdateParamsFromStruct(paramStruct);
break;
case MRES_ChangedOverride:
@ -333,12 +333,12 @@ bool HandleDetour(HookType_t hookType, CHook* pDetour)
}
else //Throw an error if no override was set
{
tempRet = MRES_Ignored;
tempRet = ReturnAction_Ignored;
pCallback->GetParentRuntime()->GetDefaultContext()->BlamePluginError(pCallback, "Tried to override return value without return value being set");
break;
}
}
tempRet = MRES_Override;
tempRet = ReturnAction_Override;
pWrapper->UpdateParamsFromStruct(paramStruct);
break;
case MRES_Override:
@ -346,7 +346,7 @@ bool HandleDetour(HookType_t hookType, CHook* pDetour)
{
if (returnStruct->isChanged)
{
tempRet = MRES_Override;
tempRet = ReturnAction_Override;
if (pWrapper->returnType == ReturnType_String || pWrapper->returnType == ReturnType_Int || pWrapper->returnType == ReturnType_Bool)
{
tempRetBuf = *(void **)returnStruct->newResult;
@ -362,7 +362,7 @@ bool HandleDetour(HookType_t hookType, CHook* pDetour)
}
else //Throw an error if no override was set
{
tempRet = MRES_Ignored;
tempRet = ReturnAction_Ignored;
pCallback->GetParentRuntime()->GetDefaultContext()->BlamePluginError(pCallback, "Tried to override return value without return value being set");
}
}
@ -372,7 +372,7 @@ bool HandleDetour(HookType_t hookType, CHook* pDetour)
{
if (returnStruct->isChanged)
{
tempRet = MRES_Supercede;
tempRet = ReturnAction_Supercede;
if (pWrapper->returnType == ReturnType_String || pWrapper->returnType == ReturnType_Int || pWrapper->returnType == ReturnType_Bool)
{
tempRetBuf = *(void **)returnStruct->newResult;
@ -388,17 +388,17 @@ bool HandleDetour(HookType_t hookType, CHook* pDetour)
}
else //Throw an error if no override was set
{
tempRet = MRES_Ignored;
tempRet = ReturnAction_Ignored;
pCallback->GetParentRuntime()->GetDefaultContext()->BlamePluginError(pCallback, "Tried to override return value without return value being set");
}
}
else
{
tempRet = MRES_Supercede;
tempRet = ReturnAction_Supercede;
}
break;
default:
tempRet = MRES_Ignored;
tempRet = ReturnAction_Ignored;
break;
}
@ -426,14 +426,14 @@ bool HandleDetour(HookType_t hookType, CHook* pDetour)
}
// If we want to use our own return value, write it back.
if (finalRet >= MRES_Override)
if (finalRet >= ReturnAction_Override)
{
void* pPtr = pDetour->m_pCallingConvention->GetReturnPtr(pDetour->m_pRegisters);
memcpy(pPtr, *finalRetBuf, pDetour->m_pCallingConvention->m_returnType.size);
pDetour->m_pCallingConvention->ReturnPtrChanged(pDetour->m_pRegisters, pPtr);
}
return finalRet == MRES_Supercede;
return finalRet;
}
CDynamicHooksSourcePawn::CDynamicHooksSourcePawn(HookSetup *setup, CHook *pDetour, IPluginFunction *pCallback, bool post)

View File

@ -33,7 +33,7 @@ public:
};
ICallingConvention *ConstructCallingConvention(HookSetup *setup);
bool HandleDetour(HookType_t hookType, CHook* pDetour);
ReturnAction_t HandleDetour(HookType_t hookType, CHook* pDetour);
bool AddDetourPluginHook(HookType_t hookType, CHook *pDetour, HookSetup *setup, IPluginFunction *pCallback);
bool RemoveDetourPluginHook(HookType_t hookType, CHook *pDetour, IPluginFunction *pCallback);
void RemoveAllCallbacksForContext(IPluginContext *pContext);