From b18e3284e117898f37327b5006187c61fea3d631 Mon Sep 17 00:00:00 2001 From: Peace-Maker Date: Tue, 13 Dec 2016 17:45:22 -0700 Subject: [PATCH] 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 | 23 +++++++++++++++ DynamicHooks/conventions/x86MsCdecl.h | 2 +- DynamicHooks/conventions/x86MsStdcall.h | 2 +- DynamicHooks/conventions/x86MsThiscall.h | 2 +- DynamicHooks/hook.cpp | 37 +++++++++++++++++------- DynamicHooks/hook.h | 15 ++++++++-- dynhooks_sourcepawn.cpp | 32 ++++++++++---------- dynhooks_sourcepawn.h | 2 +- 8 files changed, 82 insertions(+), 33 deletions(-) diff --git a/DynamicHooks/convention.h b/DynamicHooks/convention.h index 461d338..6157a1c 100644 --- a/DynamicHooks/convention.h +++ b/DynamicHooks/convention.h @@ -35,6 +35,7 @@ // >> INCLUDES // ============================================================================ #include "registers.h" +#include #include // ============================================================================ @@ -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 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 \ No newline at end of file diff --git a/DynamicHooks/conventions/x86MsCdecl.h b/DynamicHooks/conventions/x86MsCdecl.h index c3b829d..e965468 100644 --- a/DynamicHooks/conventions/x86MsCdecl.h +++ b/DynamicHooks/conventions/x86MsCdecl.h @@ -64,7 +64,7 @@ class x86MsCdecl: public ICallingConvention { public: x86MsCdecl(ke::Vector &vecArgTypes, DataTypeSized_t returnType, int iAlignment=4); - ~x86MsCdecl(); + virtual ~x86MsCdecl(); virtual ke::Vector GetRegisters(); virtual int GetPopSize(); diff --git a/DynamicHooks/conventions/x86MsStdcall.h b/DynamicHooks/conventions/x86MsStdcall.h index 2d1c46f..a1c5ad0 100644 --- a/DynamicHooks/conventions/x86MsStdcall.h +++ b/DynamicHooks/conventions/x86MsStdcall.h @@ -64,7 +64,7 @@ class x86MsStdcall: public ICallingConvention { public: x86MsStdcall(ke::Vector &vecArgTypes, DataTypeSized_t returnType, int iAlignment=4); - ~x86MsStdcall(); + virtual ~x86MsStdcall(); virtual ke::Vector GetRegisters(); virtual int GetPopSize(); diff --git a/DynamicHooks/conventions/x86MsThiscall.h b/DynamicHooks/conventions/x86MsThiscall.h index fc0e99c..e809c92 100644 --- a/DynamicHooks/conventions/x86MsThiscall.h +++ b/DynamicHooks/conventions/x86MsThiscall.h @@ -65,7 +65,7 @@ class x86MsThiscall: public ICallingConvention { public: x86MsThiscall(ke::Vector &vecArgTypes, DataTypeSized_t returnType, int iAlignment=4); - ~x86MsThiscall(); + virtual ~x86MsThiscall(); virtual ke::Vector GetRegisters(); virtual int GetPopSize(); diff --git a/DynamicHooks/hook.cpp b/DynamicHooks/hook.cpp index a0706d4..af4cefe 100644 --- a/DynamicHooks/hook.cpp +++ b/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); diff --git a/DynamicHooks/hook.h b/DynamicHooks/hook.h index fba0c5f..8ee23ae 100644 --- a/DynamicHooks/hook.h +++ b/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 \ No newline at end of file diff --git a/dynhooks_sourcepawn.cpp b/dynhooks_sourcepawn.cpp index f06a537..e4e4c99 100644 --- a/dynhooks_sourcepawn.cpp +++ b/dynhooks_sourcepawn.cpp @@ -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 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 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) diff --git a/dynhooks_sourcepawn.h b/dynhooks_sourcepawn.h index 6c86820..e0cd1cb 100644 --- a/dynhooks_sourcepawn.h +++ b/dynhooks_sourcepawn.h @@ -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);