1// Copyright (c) 2010 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 "preamble_patcher.h"
6#include "memory_hook.h"
7#include "mini_disassembler.h"
8
9// compatibility shims
10#include "base/logging.h"
11
12// Definitions of assembly statements we need
13#define ASM_JMP32REL 0xE9
14#define ASM_INT3 0xCC
15
16namespace sidestep {
17
18SideStepError PreamblePatcher::RawPatchWithStubAndProtections(
19    void* target_function, void *replacement_function,
20    unsigned char* preamble_stub, unsigned long stub_size,
21    unsigned long* bytes_needed) {
22  // We need to be able to write to a process-local copy of the first
23  // MAX_PREAMBLE_STUB_SIZE bytes of target_function. We may be giving execute
24  // privilege to something that doesn't have it, but that's the price to pay
25  // for tools.
26  DWORD old_target_function_protect = 0;
27  BOOL succeeded = ::VirtualProtect(reinterpret_cast<void*>(target_function),
28                                    MAX_PREAMBLE_STUB_SIZE,
29                                    PAGE_EXECUTE_READWRITE,
30                                    &old_target_function_protect);
31  if (!succeeded) {
32    ASSERT(false, "Failed to make page containing target function "
33                   "copy-on-write.");
34    return SIDESTEP_ACCESS_DENIED;
35  }
36
37  SideStepError error_code = RawPatchWithStub(target_function,
38                                              replacement_function,
39                                              preamble_stub,
40                                              stub_size,
41                                              bytes_needed);
42  if (SIDESTEP_SUCCESS != error_code) {
43    ASSERT1(false);
44    return error_code;
45  }
46
47  // Restore the protection of the first MAX_PREAMBLE_STUB_SIZE bytes of
48  // pTargetFunction to what they were before we started goofing around.
49  succeeded = ::VirtualProtect(reinterpret_cast<void*>(target_function),
50                               MAX_PREAMBLE_STUB_SIZE,
51                               old_target_function_protect,
52                               &old_target_function_protect);
53  if (!succeeded) {
54    ASSERT(false, "Failed to restore protection to target function.");
55    // We must not return an error here because the function has actually
56    // been patched, and returning an error would likely cause our client
57    // code not to unpatch it.  So we just keep going.
58  }
59
60  // Flush the instruction cache to make sure the processor doesn't execute the
61  // old version of the instructions (before our patch).
62  //
63  // FlushInstructionCache is actually a no-op at least on single-processor
64  // XP machines.  I'm not sure why this is so, but it is, yet I want to keep
65  // the call to the API here for correctness in case there is a difference in
66  // some variants of Windows/hardware.
67  succeeded = ::FlushInstructionCache(::GetCurrentProcess(),
68                                      target_function,
69                                      MAX_PREAMBLE_STUB_SIZE);
70  if (!succeeded) {
71    ASSERT(false, "Failed to flush instruction cache.");
72    // We must not return an error here because the function has actually
73    // been patched, and returning an error would likely cause our client
74    // code not to unpatch it.  So we just keep going.
75  }
76
77  return SIDESTEP_SUCCESS;
78}
79
80SideStepError PreamblePatcher::RawPatch(void* target_function,
81                                        void* replacement_function,
82                                        void** original_function_stub) {
83  if (!target_function || !replacement_function || !original_function_stub ||
84      (*original_function_stub) || target_function == replacement_function) {
85    ASSERT(false, "Preconditions not met");
86    return SIDESTEP_INVALID_PARAMETER;
87  }
88
89  // @see MAX_PREAMBLE_STUB_SIZE for an explanation of how we arrives at
90  // this size
91  unsigned char* preamble_stub =
92    reinterpret_cast<unsigned char*>(
93      MemoryHook::Alloc(sizeof(unsigned char) * MAX_PREAMBLE_STUB_SIZE));
94  if (!preamble_stub) {
95    ASSERT(false, "Unable to allocate preamble-stub.");
96    return SIDESTEP_INSUFFICIENT_BUFFER;
97  }
98
99  // Change the protection of the newly allocated preamble stub to
100  // PAGE_EXECUTE_READWRITE. This is required to work with DEP (Data
101  // Execution Prevention) which will cause an exception if code is executed
102  // from a page on which you do not have read access.
103  DWORD old_stub_protect = 0;
104  BOOL succeeded = VirtualProtect(preamble_stub, MAX_PREAMBLE_STUB_SIZE,
105                             PAGE_EXECUTE_READWRITE, &old_stub_protect);
106  if (!succeeded) {
107    ASSERT(false, "Failed to make page preamble stub read-write-execute.");
108    delete[] preamble_stub;
109    return SIDESTEP_ACCESS_DENIED;
110  }
111
112  SideStepError error_code = RawPatchWithStubAndProtections(target_function,
113                                              replacement_function,
114                                              preamble_stub,
115                                              MAX_PREAMBLE_STUB_SIZE,
116                                              NULL);
117  if (SIDESTEP_SUCCESS != error_code) {
118    ASSERT1(false);
119    delete[] preamble_stub;
120    return error_code;
121  }
122
123  *original_function_stub = reinterpret_cast<void*>(preamble_stub);
124
125  // NOTE: For hooking malloc/free, we don't want to use streams which
126  // allocate.  Basically, we've hooked malloc, but not necessarily
127  // hooked free yet.  To do anything which uses the heap could crash
128  // with a mismatched malloc/free!
129  //VLOG(1) << "PreamblePatcher::RawPatch successfully patched 0x"
130  //        << target_function;
131
132  return SIDESTEP_SUCCESS;
133}
134
135SideStepError PreamblePatcher::Unpatch(void* target_function,
136                                       void* replacement_function,
137                                       void* original_function_stub) {
138  ASSERT1(target_function && original_function_stub);
139  if (!target_function || !original_function_stub) {
140    return SIDESTEP_INVALID_PARAMETER;
141  }
142
143  // We disassemble the preamble of the _stub_ to see how many bytes we
144  // originally copied to the stub.
145  MiniDisassembler disassembler;
146  unsigned int preamble_bytes = 0;
147  while (preamble_bytes < 5) {
148    InstructionType instruction_type = disassembler.Disassemble(
149        reinterpret_cast<unsigned char*>(original_function_stub) +
150        preamble_bytes, preamble_bytes);
151    if (IT_GENERIC != instruction_type) {
152      ASSERT(false, "Should only have generic instructions in stub!!");
153      return SIDESTEP_UNSUPPORTED_INSTRUCTION;
154    }
155  }
156
157  // Before unpatching, target_function should be a JMP to
158  // replacement_function.  If it's not, then either it's an error, or
159  // we're falling into the case where the original instruction was a
160  // JMP, and we patched the jumped_to address rather than the JMP
161  // itself.  (For instance, if malloc() is just a JMP to __malloc(),
162  // we patched __malloc() and not malloc().)
163  unsigned char* target = reinterpret_cast<unsigned char*>(target_function);
164  while (1) {    // we stop when target is a JMP to replacement_function
165    if (target[0] != ASM_JMP32REL) {
166      ASSERT(false, "target_function does not look like it was patched.");
167      return SIDESTEP_INVALID_PARAMETER;
168    }
169    int relative_offset;   // Windows guarantees int is 4 bytes
170    ASSERT1(sizeof(relative_offset) == 4);
171    memcpy(reinterpret_cast<void*>(&relative_offset),
172           reinterpret_cast<void*>(target + 1), 4);
173    unsigned char* jump_to = target + 5 + relative_offset;
174    if (jump_to == replacement_function)
175      break;
176    target = jump_to;      // follow the jmp
177  }
178
179  // We need to be able to write to a process-local copy of the first
180  // MAX_PREAMBLE_STUB_SIZE bytes of target_function. We may be giving execute
181  // privilege to something that doesn't have it, but that's the price to pay
182  // for tools.
183  DWORD old_target_function_protect = 0;
184  BOOL succeeded = ::VirtualProtect(reinterpret_cast<void*>(target),
185                                    MAX_PREAMBLE_STUB_SIZE,
186                                    PAGE_EXECUTE_READWRITE,
187                                    &old_target_function_protect);
188  if (!succeeded) {
189    ASSERT(false, "Failed to make page containing target function "
190                   "copy-on-write.");
191    return SIDESTEP_ACCESS_DENIED;
192  }
193
194  // Replace the first few bytes of the original function with the bytes we
195  // previously moved to the preamble stub.
196  memcpy(reinterpret_cast<void*>(target),
197         original_function_stub, preamble_bytes);
198
199  // Stub is now useless so delete it.
200  // [csilvers: Commented out for perftools because it causes big problems
201  //  when we're unpatching malloc.  We just let this live on as a leak.]
202  //delete original_function_stub;
203
204  // Restore the protection of the first MAX_PREAMBLE_STUB_SIZE bytes of
205  // target to what they were before we started goofing around.
206  succeeded = ::VirtualProtect(reinterpret_cast<void*>(target),
207                               MAX_PREAMBLE_STUB_SIZE,
208                               old_target_function_protect,
209                               &old_target_function_protect);
210
211  // Flush the instruction cache to make sure the processor doesn't execute the
212  // old version of the instructions (before our patch).
213  //
214  // See comment on FlushInstructionCache elsewhere in this file.
215  succeeded = ::FlushInstructionCache(::GetCurrentProcess(),
216                                      target,
217                                      MAX_PREAMBLE_STUB_SIZE);
218  if (!succeeded) {
219    ASSERT(false, "Failed to flush instruction cache.");
220    return SIDESTEP_UNEXPECTED;
221  }
222
223  VLOG(1) << "PreamblePatcher::Unpatch successfully unpatched 0x"
224          << target_function;
225  return SIDESTEP_SUCCESS;
226}
227
228};  // namespace sidestep
229