1
// Copyright 2006-2009, Google Inc.
3
// Redistribution and use in source and binary forms, with or without
4
// modification, are permitted provided that the following conditions are met:
6
// 1. Redistributions of source code must retain the above copyright notice,
7
// this list of conditions and the following disclaimer.
8
// 2. Redistributions in binary form must reproduce the above copyright notice,
9
// this list of conditions and the following disclaimer in the documentation
10
// and/or other materials provided with the distribution.
11
// 3. Neither the name of Google Inc. nor the names of its contributors may be
12
// used to endorse or promote products derived from this software without
13
// specific prior written permission.
15
// THIS SOFTWARE IS PROVIDED BY THE AUTHOR ``AS IS'' AND ANY EXPRESS OR IMPLIED
16
// WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF
17
// MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO
18
// EVENT SHALL THE AUTHOR BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
19
// SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
20
// PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS;
21
// OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY,
22
// WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR
23
// OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF
24
// ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
26
#include "gears/base/ie/iat_patch.h"
28
#include <atlbase.h> // for _pAtlModule
31
// redifinitions of things from logging.h which is not part of the project
32
#define NOTREACHED() assert(false)
33
#define DCHECK(x) assert(x)
34
#define DCHECK_EQ(val1, val2) assert(val1 == val2)
35
#define DCHECK_NE(val1, val2) assert(val1 != val2)
36
#define COMPILE_ASSERT(x, label) assert(x)
40
struct InterceptFunctionInformation {
41
bool finished_operation;
42
const char* imported_from_module;
43
const char* function_name;
46
IMAGE_THUNK_DATA** iat_thunk;
50
static void* GetIATFunction(IMAGE_THUNK_DATA* iat_thunk) {
51
if (NULL == iat_thunk) {
56
// Works around the 64 bit portability warning:
57
// The Function member inside IMAGE_THUNK_DATA is really a pointer
58
// to the IAT function. IMAGE_THUNK_DATA correctly maps to IMAGE_THUNK_DATA32
59
// or IMAGE_THUNK_DATA64 for correct pointer size.
61
IMAGE_THUNK_DATA thunk;
65
iat_function.thunk = *iat_thunk;
66
return iat_function.pointer;
69
static HMODULE GetModuleHandleFromAddress(void *address) {
70
MEMORY_BASIC_INFORMATION mbi;
71
SIZE_T result = VirtualQuery(address, &mbi, sizeof(mbi));
72
return static_cast<HMODULE>(mbi.AllocationBase);
75
static bool InterceptEnumCallback(const PEImage &image, const char* module,
76
DWORD ordinal, const char* name, DWORD hint,
77
IMAGE_THUNK_DATA* iat, void* cookie) {
78
InterceptFunctionInformation* intercept_information =
79
reinterpret_cast<InterceptFunctionInformation*>(cookie);
81
if (NULL == intercept_information) {
88
if ((0 == lstrcmpiA(module, intercept_information->imported_from_module)) &&
90
(0 == lstrcmpiA(name, intercept_information->function_name))) {
91
// Save the old pointer.
92
if (NULL != intercept_information->old_function) {
93
void* old_function = GetIATFunction(iat);
94
if (image.module() == GetModuleHandleFromAddress(old_function)) {
95
// If the old pointer points back into the importing module, we don't
96
// trust it. It's most likely a deferred loading stub. In this case
97
// we lookup the original function pointer in the exporting module
98
// and return that value instead.
99
HMODULE exporting_module = LoadLibraryA(module);
100
void* exported_function = (exporting_module != NULL)
101
? GetProcAddress(exporting_module, name)
103
if (NULL == exported_function) {
104
// Terminate the enumeration and fail to patch
105
intercept_information->return_code = GetLastError();
106
intercept_information->finished_operation = true;
110
old_function = exported_function;
112
*(intercept_information->old_function) = old_function;
115
if (NULL != intercept_information->iat_thunk) {
116
*(intercept_information->iat_thunk) = iat;
120
COMPILE_ASSERT(sizeof(iat->u1.Function) ==
121
sizeof(intercept_information->new_function), unknown_IAT_thunk_format);
123
// Patch the function.
124
intercept_information->return_code =
125
ModifyCode(&(iat->u1.Function),
126
&(intercept_information->new_function),
127
sizeof(intercept_information->new_function));
129
// Terminate further enumeration.
130
intercept_information->finished_operation = true;
137
DWORD InterceptImportedFunction(HMODULE module_handle,
138
const char* imported_from_module,
139
const char* function_name, void* new_function,
141
IMAGE_THUNK_DATA** iat_thunk) {
142
if ((NULL == module_handle) || (NULL == imported_from_module) ||
143
(NULL == function_name) || (NULL == new_function)) {
145
return ERROR_INVALID_PARAMETER;
148
PEImage target_image(module_handle);
149
if (!target_image.VerifyMagic()) {
151
return ERROR_INVALID_PARAMETER;
154
InterceptFunctionInformation intercept_information = {
156
imported_from_module,
163
// First go through the IAT. If we don't find the import we are looking
164
// for in IAT, search delay import table.
165
target_image.EnumAllImports(InterceptEnumCallback, &intercept_information);
166
if (!intercept_information.finished_operation) {
167
target_image.EnumAllDelayImports(InterceptEnumCallback,
168
&intercept_information);
171
return intercept_information.return_code;
174
DWORD RestoreImportedFunction(void* intercept_function,
175
void* original_function,
176
IMAGE_THUNK_DATA* iat_thunk) {
177
if ((NULL == intercept_function) || (NULL == original_function) ||
178
(NULL == iat_thunk)) {
180
return ERROR_INVALID_PARAMETER;
183
if (GetIATFunction(iat_thunk) != intercept_function) {
184
// Check if someone else has intercepted on top of us.
185
// We cannot unpatch in this case, just raise a red flag.
187
return ERROR_INVALID_FUNCTION;
190
return ModifyCode(&(iat_thunk->u1.Function),
192
sizeof(original_function));
195
DWORD ModifyCode(void* old_code, void* new_code, int length) {
196
if ((NULL == old_code) || (NULL == new_code) || (0 == length)) {
198
return ERROR_INVALID_PARAMETER;
201
// Change the page protection so that we can write.
202
DWORD error = NO_ERROR;
203
DWORD old_page_protection = 0;
204
if (VirtualProtect(old_code,
207
&old_page_protection)) {
210
CopyMemory(old_code, new_code, length);
212
// Restore the old page protection.
213
error = ERROR_SUCCESS;
214
VirtualProtect(old_code,
217
&old_page_protection);
219
error = GetLastError();
226
IATPatchFunction::IATPatchFunction()
227
: original_function_(NULL),
229
intercept_function_(NULL) {
232
IATPatchFunction::~IATPatchFunction() {
233
if (NULL != intercept_function_) {
234
DWORD error = Unpatch();
235
DCHECK_EQ(NO_ERROR, error);
239
DWORD IATPatchFunction::Patch(HMODULE module_handle,
240
const char* imported_from_module,
241
const char* function_name,
242
void* new_function) {
243
DCHECK_EQ(static_cast<void*>(NULL), original_function_);
244
DCHECK_EQ(static_cast<IMAGE_THUNK_DATA*>(NULL), iat_thunk_);
245
DCHECK_EQ(static_cast<void*>(NULL), intercept_function_);
247
DWORD error = InterceptImportedFunction(module_handle,
248
imported_from_module,
254
if (NO_ERROR == error) {
255
DCHECK(original_function_);
256
intercept_function_ = new_function;
257
DCHECK_NE(original_function_, intercept_function_);
258
DCHECK(check_patch());
260
// Since we've now patched this process, we must lock the module
261
// to prevent from being unloaded.
268
DWORD IATPatchFunction::Unpatch() {
269
DWORD error = RestoreImportedFunction(intercept_function_,
273
if (NO_ERROR == error) {
274
intercept_function_ = NULL;
275
original_function_ = NULL;
278
// Give back the reference we held while the patch was in place.
279
_pAtlModule->Unlock();
285
bool IATPatchFunction::check_patch() const {
286
DCHECK(is_patched());
287
return GetIATFunction(iat_thunk_) == intercept_function_;
289
} // namespace iat_patch