-
Notifications
You must be signed in to change notification settings - Fork 0
/
VMT_Manager.h
123 lines (99 loc) · 2.34 KB
/
VMT_Manager.h
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
#pragma once
#include "DLL_MAIN.h"
#include <assert.h>
class ProtectGuard
{
public:
ProtectGuard(void *base, uint32_t len, uint32_t protect)
{
this->base = base;
this->len = len;
if (!VirtualProtect(base, len, protect, (PDWORD)&old_protect))
throw std::runtime_error(str("Failed to protect region!"));
}
~ProtectGuard()
{
VirtualProtect(base, len, old_protect, (PDWORD)&old_protect);
}
private:
void *base;
uint32_t len;
uint32_t old_protect;
};
class VMTHook
{
public:
bool hooked;
uintptr_t* shadow_vtable;
VMTHook() : class_base(nullptr), method_count(0), shadow_vtable(nullptr), original_vtable(nullptr) {}
VMTHook(void *base) : class_base(base), method_count(0), shadow_vtable(nullptr), original_vtable(nullptr) {}
~VMTHook()
{
RestoreTable();
delete[] shadow_vtable;
}
inline void Setup(void *base = nullptr)
{
if (base != nullptr)
class_base = base;
if (class_base == nullptr)
return;
original_vtable = *(uintptr_t**)class_base;
method_count = GetMethodCount(original_vtable);
if (method_count == 0)
return;
shadow_vtable = new uintptr_t[method_count + 1]();
shadow_vtable[0] = original_vtable[-1];
std::memcpy(&shadow_vtable[1], original_vtable, method_count * sizeof(uintptr_t));
try
{
auto guard = ProtectGuard{ class_base, sizeof(uintptr_t), PAGE_READWRITE };
*(uintptr_t**)class_base = &shadow_vtable[1];
}
catch (...)
{
delete[] shadow_vtable;
}
}
template<typename T>
inline void Hook(uint32_t index, T method)
{
//assert(index < method_count);
shadow_vtable[index + 1] = reinterpret_cast<uintptr_t>(method);
hooked = true;
}
inline void Unhook(uint32_t index)
{
//assert(index < method_count);
shadow_vtable[index + 1] = original_vtable[index];
hooked = false;
}
template<typename T>
inline T GetOriginal(uint32_t index)
{
return (T)original_vtable[index];
}
inline void RestoreTable()
{
try
{
if (original_vtable != nullptr)
{
auto guard = ProtectGuard{ class_base, sizeof(uintptr_t), PAGE_READWRITE };
*(uintptr_t**)class_base = original_vtable;
original_vtable = nullptr;
}
}
catch (...) {}
}
private:
inline uint32_t GetMethodCount(uintptr_t *vtable_start)
{
uint32_t len = -1;
do ++len; while (vtable_start[len]);
return len;
}
void *class_base;
uint32_t method_count;
uintptr_t *original_vtable;
};