1 // Copyright (c) 2006-2008 The Chromium Authors. All rights reserved.
2 // Use of this source code is governed by a BSD-style license that can be
3 // found in the LICENSE file.
4
5 #include "base/iat_patch.h"
6 #include "base/logging.h"
7
8 namespace iat_patch {
9
10 struct InterceptFunctionInformation {
11 bool finished_operation;
12 const char* imported_from_module;
13 const char* function_name;
14 void* new_function;
15 void** old_function;
16 IMAGE_THUNK_DATA** iat_thunk;
17 DWORD return_code;
18 };
19
GetIATFunction(IMAGE_THUNK_DATA * iat_thunk)20 static void* GetIATFunction(IMAGE_THUNK_DATA* iat_thunk) {
21 if (NULL == iat_thunk) {
22 NOTREACHED();
23 return NULL;
24 }
25
26 // Works around the 64 bit portability warning:
27 // The Function member inside IMAGE_THUNK_DATA is really a pointer
28 // to the IAT function. IMAGE_THUNK_DATA correctly maps to IMAGE_THUNK_DATA32
29 // or IMAGE_THUNK_DATA64 for correct pointer size.
30 union FunctionThunk {
31 IMAGE_THUNK_DATA thunk;
32 void* pointer;
33 } iat_function;
34
35 iat_function.thunk = *iat_thunk;
36 return iat_function.pointer;
37 }
38
InterceptEnumCallback(const PEImage & image,const char * module,DWORD ordinal,const char * name,DWORD hint,IMAGE_THUNK_DATA * iat,void * cookie)39 static bool InterceptEnumCallback(const PEImage &image, const char* module,
40 DWORD ordinal, const char* name, DWORD hint,
41 IMAGE_THUNK_DATA* iat, void* cookie) {
42 InterceptFunctionInformation* intercept_information =
43 reinterpret_cast<InterceptFunctionInformation*>(cookie);
44
45 if (NULL == intercept_information) {
46 NOTREACHED();
47 return false;
48 }
49
50 DCHECK(module);
51
52 if ((0 == lstrcmpiA(module, intercept_information->imported_from_module)) &&
53 (NULL != name) &&
54 (0 == lstrcmpiA(name, intercept_information->function_name))) {
55 // Save the old pointer.
56 if (NULL != intercept_information->old_function) {
57 *(intercept_information->old_function) = GetIATFunction(iat);
58 }
59
60 if (NULL != intercept_information->iat_thunk) {
61 *(intercept_information->iat_thunk) = iat;
62 }
63
64 // portability check
65 COMPILE_ASSERT(sizeof(iat->u1.Function) ==
66 sizeof(intercept_information->new_function), unknown_IAT_thunk_format);
67
68 // Patch the function.
69 intercept_information->return_code =
70 ModifyCode(&(iat->u1.Function),
71 &(intercept_information->new_function),
72 sizeof(intercept_information->new_function));
73
74 // Terminate further enumeration.
75 intercept_information->finished_operation = true;
76 return false;
77 }
78
79 return true;
80 }
81
InterceptImportedFunction(HMODULE module_handle,const char * imported_from_module,const char * function_name,void * new_function,void ** old_function,IMAGE_THUNK_DATA ** iat_thunk)82 DWORD InterceptImportedFunction(HMODULE module_handle,
83 const char* imported_from_module,
84 const char* function_name, void* new_function,
85 void** old_function,
86 IMAGE_THUNK_DATA** iat_thunk) {
87 if ((NULL == module_handle) || (NULL == imported_from_module) ||
88 (NULL == function_name) || (NULL == new_function)) {
89 NOTREACHED();
90 return ERROR_INVALID_PARAMETER;
91 }
92
93 PEImage target_image(module_handle);
94 if (!target_image.VerifyMagic()) {
95 NOTREACHED();
96 return ERROR_INVALID_PARAMETER;
97 }
98
99 InterceptFunctionInformation intercept_information = {
100 false,
101 imported_from_module,
102 function_name,
103 new_function,
104 old_function,
105 iat_thunk,
106 ERROR_GEN_FAILURE};
107
108 // First go through the IAT. If we don't find the import we are looking
109 // for in IAT, search delay import table.
110 target_image.EnumAllImports(InterceptEnumCallback, &intercept_information);
111 if (!intercept_information.finished_operation) {
112 target_image.EnumAllDelayImports(InterceptEnumCallback,
113 &intercept_information);
114 }
115
116 return intercept_information.return_code;
117 }
118
RestoreImportedFunction(void * intercept_function,void * original_function,IMAGE_THUNK_DATA * iat_thunk)119 DWORD RestoreImportedFunction(void* intercept_function,
120 void* original_function,
121 IMAGE_THUNK_DATA* iat_thunk) {
122 if ((NULL == intercept_function) || (NULL == original_function) ||
123 (NULL == iat_thunk)) {
124 NOTREACHED();
125 return ERROR_INVALID_PARAMETER;
126 }
127
128 if (GetIATFunction(iat_thunk) != intercept_function) {
129 // Check if someone else has intercepted on top of us.
130 // We cannot unpatch in this case, just raise a red flag.
131 NOTREACHED();
132 return ERROR_INVALID_FUNCTION;
133 }
134
135 return ModifyCode(&(iat_thunk->u1.Function),
136 &original_function,
137 sizeof(original_function));
138 }
139
ModifyCode(void * old_code,void * new_code,int length)140 DWORD ModifyCode(void* old_code, void* new_code, int length) {
141 if ((NULL == old_code) || (NULL == new_code) || (0 == length)) {
142 NOTREACHED();
143 return ERROR_INVALID_PARAMETER;
144 }
145
146 // Change the page protection so that we can write.
147 DWORD error = NO_ERROR;
148 DWORD old_page_protection = 0;
149 if (VirtualProtect(old_code,
150 length,
151 PAGE_READWRITE,
152 &old_page_protection)) {
153
154 // Write the data.
155 CopyMemory(old_code, new_code, length);
156
157 // Restore the old page protection.
158 error = ERROR_SUCCESS;
159 VirtualProtect(old_code,
160 length,
161 old_page_protection,
162 &old_page_protection);
163 } else {
164 error = GetLastError();
165 NOTREACHED();
166 }
167
168 return error;
169 }
170
IATPatchFunction()171 IATPatchFunction::IATPatchFunction()
172 : module_handle_(NULL),
173 original_function_(NULL),
174 iat_thunk_(NULL),
175 intercept_function_(NULL) {
176 }
177
~IATPatchFunction()178 IATPatchFunction::~IATPatchFunction() {
179 if (NULL != intercept_function_) {
180 DWORD error = Unpatch();
181 DCHECK_EQ(NO_ERROR, error);
182 }
183 }
184
Patch(const wchar_t * module,const char * imported_from_module,const char * function_name,void * new_function)185 DWORD IATPatchFunction::Patch(const wchar_t* module,
186 const char* imported_from_module,
187 const char* function_name,
188 void* new_function) {
189 DCHECK_EQ(static_cast<void*>(NULL), original_function_);
190 DCHECK_EQ(static_cast<IMAGE_THUNK_DATA*>(NULL), iat_thunk_);
191 DCHECK_EQ(static_cast<void*>(NULL), intercept_function_);
192
193 HMODULE module_handle = LoadLibraryW(module);
194
195 if (module_handle == NULL) {
196 NOTREACHED();
197 return GetLastError();
198 }
199
200 DWORD error = InterceptImportedFunction(module_handle,
201 imported_from_module,
202 function_name,
203 new_function,
204 &original_function_,
205 &iat_thunk_);
206
207 if (NO_ERROR == error) {
208 DCHECK_NE(original_function_, intercept_function_);
209 module_handle_ = module_handle;
210 intercept_function_ = new_function;
211 } else {
212 FreeLibrary(module_handle);
213 }
214
215 return error;
216 }
217
Unpatch()218 DWORD IATPatchFunction::Unpatch() {
219 DWORD error = RestoreImportedFunction(intercept_function_,
220 original_function_,
221 iat_thunk_);
222 DCHECK(NO_ERROR == error);
223
224 // Hands off the intercept if we fail to unpatch.
225 // If IATPatchFunction::Unpatch fails during RestoreImportedFunction
226 // it means that we cannot safely unpatch the import address table
227 // patch. In this case its better to be hands off the intercept as
228 // trying to unpatch again in the destructor of IATPatchFunction is
229 // not going to be any safer
230 if (module_handle_)
231 FreeLibrary(module_handle_);
232 module_handle_ = NULL;
233 intercept_function_ = NULL;
234 original_function_ = NULL;
235 iat_thunk_ = NULL;
236
237 return error;
238 }
239
240 } // namespace iat_patch
241