1 : // Copyright 2012 Google Inc.
2 : //
3 : // Licensed under the Apache License, Version 2.0 (the "License");
4 : // you may not use this file except in compliance with the License.
5 : // You may obtain a copy of the License at
6 : //
7 : // http://www.apache.org/licenses/LICENSE-2.0
8 : //
9 : // Unless required by applicable law or agreed to in writing, software
10 : // distributed under the License is distributed on an "AS IS" BASIS,
11 : // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 : // See the License for the specific language governing permissions and
13 : // limitations under the License.
14 :
15 : #include "syzygy/instrument/transforms/thunk_import_references_transform.h"
16 :
17 : #include "base/logging.h"
18 : #include "base/stringprintf.h"
19 : #include "syzygy/block_graph/typed_block.h"
20 : #include "syzygy/pe/pe_utils.h"
21 : #include "syzygy/pe/transforms/add_imports_transform.h"
22 :
23 : namespace instrument {
24 : namespace transforms {
25 :
26 : namespace {
27 :
28 : using block_graph::BlockGraph;
29 : using block_graph::TypedBlock;
30 :
31 : // We add this suffix to the destination
32 : const char kThunkSuffix[] = "_thunk";
33 :
34 : } // namespace
35 :
36 : using pe::transforms::AddImportsTransform;
37 :
38 : const char ThunkImportReferencesTransform::kTransformName[] =
39 : "ThunkImportReferencesTransform";
40 :
41 : const char ThunkImportReferencesTransform::kEntryHookName[] =
42 : "_indirect_penter";
43 : const char ThunkImportReferencesTransform::kDefaultInstrumentDll[] =
44 : "call_trace.dll";
45 :
46 : // We look up the absolute address of the function to be called on the
47 : // stack, and then we invoke the instrumentation function indirectly
48 : // through the import table.
49 : // Ff6844332211 push dword ptr [(11223344)]
50 : // FF2588776655 jmp dword ptr [(55667788)]
51 : const ThunkImportReferencesTransform::Thunk
52 : ThunkImportReferencesTransform::kThunkTemplate = {
53 : 0x35FF, NULL, // push DWORD PTR[immediate]
54 : 0x25FF, NULL // jmp DWORD PTR[immediate]
55 : };
56 :
57 : ThunkImportReferencesTransform::ThunkImportReferencesTransform()
58 : : thunk_section_(NULL),
59 E : instrument_dll_name_(kDefaultInstrumentDll) {
60 E : }
61 :
62 : bool ThunkImportReferencesTransform::TransformBlockGraph(
63 : BlockGraph* block_graph,
64 E : BlockGraph::Block* header_block) {
65 E : DCHECK(thunk_section_ == NULL);
66 :
67 : AddImportsTransform::ImportedModule import_module(
68 E : instrument_dll_name_.c_str());
69 E : size_t hook_index = import_module.AddSymbol(kEntryHookName);
70 :
71 E : add_imports_transform_.AddModule(&import_module);
72 :
73 E : if (!add_imports_transform_.TransformBlockGraph(block_graph, header_block)) {
74 i : LOG(ERROR) << "Unable to add imports for instrumentation DLL.";
75 i : return false;
76 : }
77 :
78 E : if (!import_module.GetSymbolReference(hook_index, &hook_ref_)) {
79 i : LOG(ERROR) << "Unable to get import reference for hook.";
80 i : return false;
81 : }
82 :
83 : // Now grab the block containing the IAT so that we can instrument references
84 : // to it. We also get the image import descriptor table block so that we
85 : // can exclude that - we don't want to instrument that.
86 : BlockGraph::Block* iat_block =
87 E : add_imports_transform_.import_address_table_block();
88 E : DCHECK(iat_block != NULL);
89 : BlockGraph::Block* iidt_block =
90 E : add_imports_transform_.image_import_descriptor_block();
91 E : DCHECK(iat_block != NULL);
92 :
93 E : if (!InstrumentIATReferences(block_graph, iat_block, iidt_block)) {
94 i : LOG(ERROR) << "Unable to instrument references to the IAT.";
95 i : return false;
96 : }
97 :
98 E : return true;
99 E : }
100 :
101 : void ThunkImportReferencesTransform::ExcludeModule(
102 i : const base::StringPiece& module_name) {
103 i : modules_to_exclude_.insert(module_name.as_string());
104 i : }
105 :
106 : // This method builds up a set of thunk blocks as well as a thunk table
107 : // containing pointers to these blocks. Existing import references are then
108 : // replaced by references to the thunk table. Since imports are invoked via
109 : // an indirect call or jump instruction, changing the address of the call
110 : // statement from an address into the IAT to an address into the thunk table
111 : // gets the thunk called properly.
112 : bool ThunkImportReferencesTransform::InstrumentIATReferences(
113 : BlockGraph* block_graph,
114 : BlockGraph::Block* iat_block,
115 E : BlockGraph::Block* iidt_block) {
116 :
117 : // Find or create the section we put our thunks in.
118 : thunk_section_ = block_graph->FindOrAddSection(".thunks",
119 E : pe::kCodeCharacteristics);
120 E : if (thunk_section_ == NULL) {
121 i : NOTREACHED();
122 i : return false;
123 : }
124 :
125 : // Typedef for the thunk block map. The key is the offset into the IAT block
126 : // (since all callers can use the same thunk) and the value is the offset into
127 : // the thunk table that points to the thunk block for that IAT entry.
128 : typedef std::map<BlockGraph::Offset, BlockGraph::Offset> ThunkBlockMap;
129 E : ThunkBlockMap thunk_block_map;
130 :
131 : // Create the thunk table. Make it the same size as the IAT, assuming that
132 : // we will need a thunk for each import.
133 : // TODO(robertshield): Resize the block afterwards if not all imports are
134 : // thunked.
135 : BlockGraph::Block* thunk_table_block =
136 : block_graph->AddBlock(BlockGraph::DATA_BLOCK,
137 : iat_block->size(),
138 E : "ImportsThunkTable");
139 E : thunk_table_block->AllocateData(iat_block->size());
140 E : thunk_table_block->set_section(thunk_section_->id());
141 E : BlockGraph::Offset thunk_table_offset = 0;
142 :
143 : // Next, list all Referrers to get all References into the IAT. For each
144 : // Reference, create a thunk (in its own block) and add a pointer to it to
145 : // the thunk table.
146 E : BlockGraph::Block::ReferrerSet iat_referrers(iat_block->referrers());
147 : BlockGraph::Block::ReferrerSet::const_iterator iat_referrer_iter(
148 E : iat_referrers.begin());
149 E : for (; iat_referrer_iter != iat_referrers.end(); ++iat_referrer_iter) {
150 E : const BlockGraph::Block::Referrer& referrer = *iat_referrer_iter;
151 :
152 E : if (referrer.first == iat_block) {
153 i : LOG(WARNING) << "Unexpected self-reference in IAT.";
154 i : continue;
155 : }
156 :
157 E : if (referrer.first->type() != BlockGraph::CODE_BLOCK) {
158 E : LOG(INFO) << "Skipping non-code block reference.";
159 E : continue;
160 : }
161 :
162 : // Now that we know the referring block, we need to find out where in the
163 : // IAT it refers to.
164 E : BlockGraph::Reference ref;
165 E : if (!referrer.first->GetReference(referrer.second, &ref)) {
166 i : LOG(ERROR) << "Unable to get reference from referrer.";
167 i : return false;
168 : }
169 :
170 : // Now we need to figure out if the IAT entry being referred to points
171 : // to one of the excluded modules.
172 : // To do that, figure out the ranges for each of our excluded modules.
173 : // TODO(robertshield): ^ this.
174 :
175 : // Look for the reference in the thunk block map, and only create a new one
176 : // if it does not already exist.
177 E : BlockGraph::Block* thunk_block = NULL;
178 E : BlockGraph::Offset new_ref_offset = 0;
179 :
180 E : ThunkBlockMap::const_iterator thunk_it = thunk_block_map.find(ref.offset());
181 E : if (thunk_it == thunk_block_map.end()) {
182 : // Create the thunk block for this offset into the IAT.
183 E : thunk_block = CreateOneThunk(block_graph, ref);
184 E : if (thunk_block == NULL) {
185 i : LOG(DFATAL) << "Unable to create thunk block.";
186 i : return false;
187 : }
188 :
189 : // Now add a reference to the thunk in the thunk table.
190 : BlockGraph::Reference thunk_ref(BlockGraph::ABSOLUTE_REF,
191 : sizeof(core::AbsoluteAddress),
192 E : thunk_block, 0, 0);
193 E : thunk_table_block->SetReference(thunk_table_offset, thunk_ref);
194 :
195 : // Remember this thunk in case we need to use it again.
196 E : thunk_block_map[ref.offset()] = thunk_table_offset;
197 :
198 E : new_ref_offset = thunk_table_offset;
199 :
200 : // Move to the next empty entry in the thunk table.
201 E : thunk_table_offset += sizeof(core::AbsoluteAddress);
202 : DCHECK_LT(static_cast<BlockGraph::Size>(thunk_table_offset),
203 E : thunk_table_block->size());
204 E : } else {
205 E : new_ref_offset = thunk_it->second;
206 : }
207 :
208 : // Update the referrer to point to the new location in the thunk table.
209 : BlockGraph::Reference new_ref(ref.type(),
210 : ref.size(),
211 : thunk_table_block,
212 : new_ref_offset,
213 E : 0);
214 E : referrer.first->SetReference(referrer.second, new_ref);
215 E : }
216 :
217 E : return true;
218 E : }
219 :
220 : BlockGraph::Block* ThunkImportReferencesTransform::CreateOneThunk(
221 : BlockGraph* block_graph,
222 E : const BlockGraph::Reference& destination) {
223 E : std::string name;
224 : // TODO(robertshield): Name the thunks according to the import they are
225 : // thunking.
226 E : if (destination.offset() == 0) {
227 : name = base::StringPrintf("%s%s",
228 : destination.referenced()->name().c_str(),
229 E : kThunkSuffix);
230 E : } else {
231 : name = base::StringPrintf("%s%s+%d",
232 : destination.referenced()->name().c_str(),
233 : kThunkSuffix,
234 E : destination.offset());
235 : }
236 :
237 : // Create and initialize the new thunk.
238 : BlockGraph::Block* thunk = block_graph->AddBlock(BlockGraph::CODE_BLOCK,
239 : sizeof(kThunkTemplate),
240 E : name.c_str());
241 E : if (thunk == NULL)
242 i : return NULL;
243 :
244 E : thunk->set_section(thunk_section_->id());
245 : thunk->SetData(reinterpret_cast<const uint8*>(&kThunkTemplate),
246 E : sizeof(kThunkTemplate));
247 :
248 E : if (!InitializeThunk(thunk, destination, hook_ref_)) {
249 i : bool removed = block_graph->RemoveBlock(thunk);
250 i : DCHECK(removed);
251 :
252 i : thunk = NULL;
253 : }
254 :
255 E : return thunk;
256 E : }
257 :
258 : bool ThunkImportReferencesTransform::InitializeThunk(
259 : BlockGraph::Block* thunk_block,
260 : const BlockGraph::Reference& destination,
261 E : const BlockGraph::Reference& import_entry) {
262 E : TypedBlock<Thunk> thunk;
263 E : if (!thunk.Init(0, thunk_block))
264 i : return false;
265 :
266 : if (!thunk.SetReference(BlockGraph::ABSOLUTE_REF,
267 : thunk->func_addr,
268 : destination.referenced(),
269 : destination.offset(),
270 E : destination.offset())) {
271 i : return false;
272 : }
273 :
274 : if (!thunk.SetReference(BlockGraph::ABSOLUTE_REF,
275 : thunk->hook_addr,
276 : import_entry.referenced(),
277 : import_entry.offset(),
278 E : import_entry.offset())) {
279 i : return false;
280 : }
281 :
282 E : return true;
283 E : }
284 :
285 : } // namespace transforms
286 : } // namespace instrument
|