← Back to branch summary

~ubuntu-branches/ubuntu/karmic/gears/karmic

« back to all changes in this revision

Viewing changes to gears/base/ie/iat_patch.cc

  • Committer: Bazaar Package Importer
  • Author(s): Stefan Lesicnik
  • Date: 2009-04-30 19:15:25 UTC
  • Revision ID: james.westby@ubuntu.com-20090430191525-0790sb5wzg8ou0xb
Tags: upstream-0.5.21.0~svn3334+dfsg
Import upstream version 0.5.21.0~svn3334+dfsg

Show diffs side-by-side

added added

removed removed

Lines of Context:
gears/base/ie/iat_patch.cc
 
1
// Copyright 2006-2009, Google Inc.
 
2
//
 
3
// Redistribution and use in source and binary forms, with or without
 
4
// modification, are permitted provided that the following conditions are met:
 
5
//
 
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.
 
14
//
 
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.
 
25
 
 
26
#include "gears/base/ie/iat_patch.h"
 
27
 
 
28
#include <atlbase.h>  // for _pAtlModule
 
29
#include <assert.h>
 
30
 
 
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)
 
37
 
 
38
namespace iat_patch {
 
39
 
 
40
struct InterceptFunctionInformation {
 
41
  bool finished_operation;
 
42
  const char* imported_from_module;
 
43
  const char* function_name;
 
44
  void* new_function;
 
45
  void** old_function;
 
46
  IMAGE_THUNK_DATA** iat_thunk;
 
47
  DWORD return_code;
 
48
};
 
49
 
 
50
static void* GetIATFunction(IMAGE_THUNK_DATA* iat_thunk) {
 
51
  if (NULL == iat_thunk) {
 
52
    NOTREACHED();
 
53
    return NULL;
 
54
  }
 
55
 
 
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.
 
60
  union FunctionThunk {
 
61
    IMAGE_THUNK_DATA thunk;
 
62
    void* pointer;
 
63
  } iat_function;
 
64
 
 
65
  iat_function.thunk = *iat_thunk;
 
66
  return iat_function.pointer;
 
67
}
 
68
 
 
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);
 
73
}
 
74
 
 
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);
 
80
 
 
81
  if (NULL == intercept_information) {
 
82
    NOTREACHED();
 
83
    return false;
 
84
  }
 
85
 
 
86
  DCHECK(module);
 
87
 
 
88
  if ((0 == lstrcmpiA(module, intercept_information->imported_from_module)) &&
 
89
     (NULL != name) &&
 
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)
 
102
                                      : NULL;
 
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;
 
107
          DCHECK(false);
 
108
          return false;
 
109
        }
 
110
        old_function = exported_function;
 
111
      }
 
112
      *(intercept_information->old_function) = old_function;
 
113
    }
 
114
 
 
115
    if (NULL != intercept_information->iat_thunk) {
 
116
      *(intercept_information->iat_thunk) = iat;
 
117
    }
 
118
 
 
119
    // portability check
 
120
    COMPILE_ASSERT(sizeof(iat->u1.Function) ==
 
121
      sizeof(intercept_information->new_function), unknown_IAT_thunk_format);
 
122
 
 
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));
 
128
 
 
129
    // Terminate further enumeration.
 
130
    intercept_information->finished_operation = true;
 
131
    return false;
 
132
  }
 
133
 
 
134
  return true;
 
135
}
 
136
 
 
137
DWORD InterceptImportedFunction(HMODULE module_handle,
 
138
                                const char* imported_from_module,
 
139
                                const char* function_name, void* new_function,
 
140
                                void** old_function,
 
141
                                IMAGE_THUNK_DATA** iat_thunk) {
 
142
  if ((NULL == module_handle) || (NULL == imported_from_module) ||
 
143
     (NULL == function_name) || (NULL == new_function)) {
 
144
    NOTREACHED();
 
145
    return ERROR_INVALID_PARAMETER;
 
146
  }
 
147
 
 
148
  PEImage target_image(module_handle);
 
149
  if (!target_image.VerifyMagic()) {
 
150
    NOTREACHED();
 
151
    return ERROR_INVALID_PARAMETER;
 
152
  }
 
153
 
 
154
  InterceptFunctionInformation intercept_information = {
 
155
    false,
 
156
    imported_from_module,
 
157
    function_name,
 
158
    new_function,
 
159
    old_function,
 
160
    iat_thunk,
 
161
    ERROR_GEN_FAILURE};
 
162
 
 
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);
 
169
  }
 
170
 
 
171
  return intercept_information.return_code;
 
172
}
 
173
 
 
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)) {
 
179
    NOTREACHED();
 
180
    return ERROR_INVALID_PARAMETER;
 
181
  }
 
182
 
 
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.
 
186
    NOTREACHED();
 
187
    return ERROR_INVALID_FUNCTION;
 
188
  }
 
189
 
 
190
  return ModifyCode(&(iat_thunk->u1.Function),
 
191
                    &original_function,
 
192
                    sizeof(original_function));
 
193
}
 
194
 
 
195
DWORD ModifyCode(void* old_code, void* new_code, int length) {
 
196
  if ((NULL == old_code) || (NULL == new_code) || (0 == length)) {
 
197
    NOTREACHED();
 
198
    return ERROR_INVALID_PARAMETER;
 
199
  }
 
200
 
 
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,
 
205
                     length,
 
206
                     PAGE_READWRITE,
 
207
                     &old_page_protection)) {
 
208
 
 
209
    // Write the data.
 
210
    CopyMemory(old_code, new_code, length);
 
211
 
 
212
    // Restore the old page protection.
 
213
    error = ERROR_SUCCESS;
 
214
    VirtualProtect(old_code,
 
215
                  length,
 
216
                  old_page_protection,
 
217
                  &old_page_protection);
 
218
  } else {
 
219
    error = GetLastError();
 
220
    NOTREACHED();
 
221
  }
 
222
 
 
223
  return error;
 
224
}
 
225
 
 
226
IATPatchFunction::IATPatchFunction()
 
227
    : original_function_(NULL),
 
228
      iat_thunk_(NULL),
 
229
      intercept_function_(NULL) {
 
230
}
 
231
 
 
232
IATPatchFunction::~IATPatchFunction() {
 
233
  if (NULL != intercept_function_) {
 
234
    DWORD error = Unpatch();
 
235
    DCHECK_EQ(NO_ERROR, error);
 
236
  }
 
237
}
 
238
 
 
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_);
 
246
 
 
247
  DWORD error = InterceptImportedFunction(module_handle,
 
248
                                          imported_from_module,
 
249
                                          function_name,
 
250
                                          new_function,
 
251
                                          &original_function_,
 
252
                                          &iat_thunk_);
 
253
 
 
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());
 
259
 
 
260
    // Since we've now patched this process, we must lock the module
 
261
    // to prevent from being unloaded.
 
262
    _pAtlModule->Lock();
 
263
  }
 
264
 
 
265
  return error;
 
266
}
 
267
 
 
268
DWORD IATPatchFunction::Unpatch() {
 
269
  DWORD error = RestoreImportedFunction(intercept_function_,
 
270
                                        original_function_,
 
271
                                        iat_thunk_);
 
272
 
 
273
  if (NO_ERROR == error) {
 
274
    intercept_function_ = NULL;
 
275
    original_function_ = NULL;
 
276
    iat_thunk_ = NULL;
 
277
 
 
278
    // Give back the reference we held while the patch was in place.
 
279
    _pAtlModule->Unlock();
 
280
  }
 
281
 
 
282
  return error;
 
283
}
 
284
 
 
285
bool IATPatchFunction::check_patch() const {
 
286
  DCHECK(is_patched());
 
287
  return GetIATFunction(iat_thunk_) == intercept_function_;
 
288
}
 
289
}  // namespace iat_patch