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