Use AMTL instead of the stdlib in DynamicHooks

This commit is contained in:
Peace-Maker 2016-12-13 00:34:28 -07:00
parent 8ab9c14e93
commit c70d3b9ee0
2 changed files with 63 additions and 23 deletions

View File

@ -53,6 +53,12 @@ CHook::CHook(void* pFunc, ICallingConvention* pConvention)
m_pRegisters = new CRegisters(pConvention->GetRegisters());
m_pCallingConvention = pConvention;
if (!m_hookHandler.init())
return;
if (!m_RetAddr.init())
return;
unsigned char* pTarget = (unsigned char *) pFunc;
// Determine the number of bytes we need to copy
@ -103,37 +109,56 @@ void CHook::AddCallback(HookType_t eHookType, HookHandlerFn* pCallback)
if (!pCallback)
return;
if (!IsCallbackRegistered(eHookType, pCallback))
m_hookHandler[eHookType].push_back(pCallback);
HookTypeMap::Insert i = m_hookHandler.findForAdd(eHookType);
if (!i.found()) {
HookHandlerSet set;
set.init();
m_hookHandler.add(i, eHookType, ke::Move(set));
}
i->value.add(pCallback);
}
void CHook::RemoveCallback(HookType_t eHookType, HookHandlerFn* pCallback)
{
if (IsCallbackRegistered(eHookType, pCallback))
m_hookHandler[eHookType].remove(pCallback);
HookTypeMap::Result r = m_hookHandler.find(eHookType);
if (!r.found())
return;
r->value.removeIfExists(pCallback);
}
bool CHook::IsCallbackRegistered(HookType_t eHookType, HookHandlerFn* pCallback)
{
std::list<HookHandlerFn *> callbacks = m_hookHandler[eHookType];
for(std::list<HookHandlerFn *>::iterator it=callbacks.begin(); it != callbacks.end(); it++)
{
if (*it == pCallback)
return true;
}
return false;
HookTypeMap::Result r = m_hookHandler.find(eHookType);
if (!r.found())
return false;
return r->value.has(pCallback);
}
bool CHook::AreCallbacksRegistered()
{
return !m_hookHandler[HOOKTYPE_PRE].empty() || !m_hookHandler[HOOKTYPE_POST].empty();
HookTypeMap::Result r = m_hookHandler.find(HOOKTYPE_PRE);
if (r.found() && r->value.elements() > 0)
return true;
r = m_hookHandler.find(HOOKTYPE_POST);
if (r.found() && r->value.elements() > 0)
return true;
return false;
}
bool CHook::HookHandler(HookType_t eHookType)
{
bool bOverride = false;
std::list<HookHandlerFn *> callbacks = this->m_hookHandler[eHookType];
for(std::list<HookHandlerFn *>::iterator it=callbacks.begin(); it != callbacks.end(); it++)
HookTypeMap::Result r = m_hookHandler.find(eHookType);
if (!r.found())
return bOverride;
HookHandlerSet &callbacks = r->value;
for(HookHandlerSet::iterator it=callbacks.iter(); !it.empty(); it.next())
{
bool result = ((HookHandlerFn) *it)(eHookType, this);
if (result)
@ -144,15 +169,17 @@ bool CHook::HookHandler(HookType_t eHookType)
void* __cdecl CHook::GetReturnAddress(void* pESP)
{
if (m_RetAddr.count(pESP) == 0)
ReturnAddressMap::Result r = m_RetAddr.find(pESP);
if (!r.found())
puts("ESP not present.");
return m_RetAddr[pESP];
return r->value;
}
void __cdecl CHook::SetReturnAddress(void* pRetAddr, void* pESP)
{
m_RetAddr[pESP] = pRetAddr;
ReturnAddressMap::Insert i = m_RetAddr.findForAdd(pESP);
m_RetAddr.add(i, pESP, pRetAddr);
}
void* CHook::CreateBridge()

View File

@ -34,11 +34,10 @@
// ============================================================================
// >> INCLUDES
// ============================================================================
#include <list>
#include <map>
#include "registers.h"
#include "convention.h"
#include <am-hashmap.h>
#include <am-hashset.h>
// ============================================================================
// >> HookType_t
@ -63,6 +62,20 @@ typedef bool (*HookHandlerFn)(HookType_t, CHook*);
#define __cdecl
#endif
struct IntegerPolicy
{
static inline uint32_t hash(size_t i) {
return ke::HashInteger<sizeof(size_t)>(i);
}
static inline bool matches(size_t i1, size_t i2) {
return i1 == i2;
}
};
typedef ke::HashSet<HookHandlerFn*, ke::PointerPolicy<HookHandlerFn>> HookHandlerSet;
typedef ke::HashMap<HookType_t, HookHandlerSet, IntegerPolicy> HookTypeMap;
typedef ke::HashMap<void*, void*, ke::PointerPolicy<void>> ReturnAddressMap;
namespace sp
{
class MacroAssembler;
@ -164,12 +177,12 @@ private:
void __cdecl SetReturnAddress(void* pRetAddr, void* pESP);
public:
std::map<HookType_t, std::list<HookHandlerFn*> > m_hookHandler;
HookTypeMap m_hookHandler;
// Address of the original function
void* m_pFunc;
ICallingConvention* m_pCallingConvention;
// Address of the bridge
@ -184,7 +197,7 @@ public:
// New return address
void* m_pNewRetAddr;
std::map<void*, void*> m_RetAddr;
ReturnAddressMap m_RetAddr;
};
#endif // _HOOK_H