1 : // Copyright 2012 Google Inc. All Rights Reserved.
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 : // Unittests for iteration primitives.
16 :
17 : #include "syzygy/instrument/transforms/entry_thunk_transform.h"
18 :
19 : #include <vector>
20 :
21 : #include "gmock/gmock.h"
22 : #include "gtest/gtest.h"
23 : #include "syzygy/block_graph/typed_block.h"
24 : #include "syzygy/common/defs.h"
25 : #include "syzygy/pe/pe_utils.h"
26 :
27 : namespace instrument {
28 : namespace transforms {
29 :
30 : namespace {
31 :
32 : using block_graph::BlockGraph;
33 : using block_graph::ConstBlockVector;
34 : using block_graph::Immediate;
35 : using block_graph::TypedBlock;
36 : using core::AbsoluteAddress;
37 : using testing::_;
38 : using testing::Return;
39 :
40 : // This defines the memory layout for the thunks that are created by the
41 : // transform.
42 : #pragma pack(push)
43 : #pragma pack(1)
44 : struct Thunk {
45 : BYTE push;
46 : DWORD func_addr; // The real function to invoke.
47 : WORD indirect_jmp;
48 : DWORD hook_addr; // The instrumentation hook that gets called indirectly.
49 : };
50 : struct ParamThunk {
51 : BYTE push1;
52 : DWORD param; // The parameter for the instrumentation hook.
53 : BYTE push2;
54 : DWORD func_addr; // The real function to invoke.
55 : WORD indirect_jmp;
56 : DWORD hook_addr; // The instrumentation hook that gets called indirectly.
57 : };
58 : #pragma pack(pop)
59 :
60 : class EntryThunkTransformTest : public testing::Test {
61 : public:
62 : EntryThunkTransformTest()
63 : : dos_header_block_(NULL),
64 : nt_headers_block_(NULL),
65 : foo_(NULL),
66 : bar_(NULL),
67 E : array_(NULL) {
68 E : }
69 :
70 E : virtual void SetUp() {
71 : // TODO(siggi): We have a lot of code that does this sort of thing, maybe
72 : // it should be concentrated in a test fixture in pe someplace.
73 : // Create the DOS/NT headers.
74 : dos_header_block_ = bg_.AddBlock(BlockGraph::DATA_BLOCK,
75 : sizeof(IMAGE_DOS_HEADER),
76 E : "DOS Header");
77 E : ASSERT_TRUE(
78 : dos_header_block_->AllocateData(dos_header_block_->size()) != NULL);
79 :
80 : nt_headers_block_ = bg_.AddBlock(BlockGraph::DATA_BLOCK,
81 : sizeof(IMAGE_NT_HEADERS),
82 E : "NT Headers");
83 :
84 E : ASSERT_TRUE(
85 : nt_headers_block_->AllocateData(nt_headers_block_->size()) != NULL);
86 E : TypedBlock<IMAGE_NT_HEADERS> nt_headers;
87 E : ASSERT_TRUE(nt_headers.Init(0, nt_headers_block_));
88 E : nt_headers->Signature = IMAGE_NT_SIGNATURE;
89 E : nt_headers->FileHeader.SizeOfOptionalHeader = sizeof(IMAGE_OPTIONAL_HEADER);
90 E : nt_headers->OptionalHeader.Magic = IMAGE_NT_OPTIONAL_HDR_MAGIC;
91 :
92 E : TypedBlock<IMAGE_DOS_HEADER> dos_header;
93 E : ASSERT_TRUE(dos_header.Init(0, dos_header_block_));
94 E : ASSERT_TRUE(dos_header.SetReference(BlockGraph::RELATIVE_REF,
95 : dos_header->e_lfanew,
96 : nt_headers));
97 :
98 : // Make the DOS header valid just for giggles.
99 E : ASSERT_TRUE(pe::UpdateDosHeader(dos_header_block_));
100 :
101 : // Get the .text section.
102 : BlockGraph::Section* text =
103 E : bg_.FindOrAddSection(pe::kCodeSectionName, pe::kCodeCharacteristics);
104 :
105 : // Create a couple of code blocks for "functions".
106 E : foo_ = bg_.AddBlock(BlockGraph::CODE_BLOCK, 20, "foo");
107 E : foo_->set_section(text->id());
108 : foo_->source_ranges().Push(BlockGraph::Block::DataRange(0, 20),
109 E : BlockGraph::Block::SourceRange(core::RelativeAddress(0x1000), 20));
110 :
111 E : bar_ = bg_.AddBlock(BlockGraph::CODE_BLOCK, 20, "bar");
112 E : bar_->set_section(text->id());
113 : bar_->source_ranges().Push(BlockGraph::Block::DataRange(0, 20),
114 E : BlockGraph::Block::SourceRange(core::RelativeAddress(0x1020), 20));
115 :
116 : // Get the .rdata section.
117 : BlockGraph::Section* rdata =
118 : bg_.FindOrAddSection(pe::kReadOnlyDataSectionName,
119 E : pe::kReadOnlyDataCharacteristics);
120 :
121 : // Create a data array block.
122 : array_ = bg_.AddBlock(BlockGraph::DATA_BLOCK,
123 : 30 * sizeof(AbsoluteAddress),
124 E : "array");
125 E : array_->set_section(rdata->id());
126 :
127 : // foo() refers to the start of bar() with a PC-relative reference.
128 : foo_->SetReference(5, BlockGraph::Reference(BlockGraph::PC_RELATIVE_REF,
129 : sizeof(AbsoluteAddress),
130 : bar_,
131 E : 0, 0));
132 : // foo() is self-referential.
133 : foo_->SetReference(10, BlockGraph::Reference(BlockGraph::PC_RELATIVE_REF,
134 : sizeof(AbsoluteAddress),
135 : foo_,
136 E : 0, 0));
137 :
138 : // bar() refers to foo() five bytes in.
139 : bar_->SetReference(5, BlockGraph::Reference(BlockGraph::PC_RELATIVE_REF,
140 : sizeof(AbsoluteAddress),
141 : foo_,
142 E : 5, 5));
143 :
144 : // The array refers to the start of both foo() and bar().
145 : array_->SetReference(0, BlockGraph::Reference(BlockGraph::ABSOLUTE_REF,
146 : sizeof(AbsoluteAddress),
147 : foo_,
148 E : 0, 0));
149 :
150 : array_->SetReference(4, BlockGraph::Reference(BlockGraph::ABSOLUTE_REF,
151 : sizeof(AbsoluteAddress),
152 : bar_,
153 E : 0, 0));
154 :
155 : // And the array refers 5 bytes into foo().
156 : array_->SetReference(8, BlockGraph::Reference(BlockGraph::ABSOLUTE_REF,
157 : sizeof(AbsoluteAddress),
158 : foo_,
159 E : 5, 5));
160 :
161 E : num_sections_pre_transform_ = bg_.sections().size();
162 :
163 : // No thunks so far.
164 E : ASSERT_NO_FATAL_FAILURE(VerifyThunks(0, 0, 0, 0));
165 E : }
166 :
167 : // Retrieves the thunks.
168 E : void FindThunks(ConstBlockVector* ret, int* param_thunks) {
169 E : ASSERT_TRUE(ret != NULL);
170 E : EXPECT_TRUE(ret->empty());
171 E : ASSERT_TRUE(param_thunks != NULL);
172 :
173 E : *param_thunks = 0;
174 :
175 : BlockGraph::Section* thunk_section =
176 E : bg_.FindSection(common::kThunkSectionName);
177 E : if (thunk_section == NULL)
178 E : return;
179 :
180 E : BlockGraph::BlockMap::const_iterator it = bg_.blocks().begin();
181 E : for (; it != bg_.blocks().end(); ++it) {
182 E : const BlockGraph::Block& block = it->second;
183 E : if (block.section() == thunk_section->id()) {
184 E : EXPECT_EQ(BlockGraph::CODE_BLOCK, block.type());
185 E : EXPECT_EQ(2, block.references().size());
186 E : EXPECT_TRUE(block.size() == sizeof(Thunk) ||
187 : block.size() == sizeof(ParamThunk));
188 :
189 E : if (block.size() == sizeof(ParamThunk))
190 E : ++(*param_thunks);
191 :
192 : // It's a thunk.
193 E : ret->push_back(&block);
194 : }
195 E : }
196 E : }
197 :
198 E : size_t CountDestinations(const ConstBlockVector& blocks) {
199 : typedef std::set<std::pair<BlockGraph::Block*, BlockGraph::Offset>>
200 : ReferenceMap;
201 E : ReferenceMap destinations;
202 E : for (size_t i = 0; i < blocks.size(); ++i) {
203 E : size_t func_addr_offset = offsetof(Thunk, func_addr);
204 E : if (blocks[i]->size() == sizeof(ParamThunk))
205 E : func_addr_offset = offsetof(ParamThunk, func_addr);
206 :
207 : // Lookup and record the destination.
208 E : BlockGraph::Reference ref;
209 E : EXPECT_TRUE(blocks[i]->GetReference(func_addr_offset, &ref));
210 E : EXPECT_EQ(BlockGraph::ABSOLUTE_REF, ref.type());
211 E : destinations.insert(std::make_pair(ref.referenced(), ref.offset()));
212 E : }
213 E : return destinations.size();
214 E : }
215 :
216 E : size_t CountEntryPoints(const ConstBlockVector& blocks) {
217 : typedef std::set<std::pair<BlockGraph::Block*, BlockGraph::Offset>>
218 : ReferenceMap;
219 E : ReferenceMap entrypoints;
220 E : for (size_t i = 0; i < blocks.size(); ++i) {
221 E : size_t hook_addr_offset = offsetof(Thunk, hook_addr);
222 E : if (blocks[i]->size() == sizeof(ParamThunk))
223 E : hook_addr_offset = offsetof(ParamThunk, hook_addr);
224 :
225 : // Lookup and record the entrypoint.
226 E : BlockGraph::Reference ref;
227 E : EXPECT_TRUE(blocks[i]->GetReference(
228 : hook_addr_offset, &ref));
229 E : EXPECT_EQ(BlockGraph::ABSOLUTE_REF, ref.type());
230 E : entrypoints.insert(std::make_pair(ref.referenced(), ref.offset()));
231 E : }
232 E : return entrypoints.size();
233 E : }
234 :
235 E : void VerifySourceRanges(const ConstBlockVector& thunks) {
236 E : for (size_t i = 0; i < thunks.size(); ++i) {
237 : // Test the source ranges on the thunk.
238 E : ASSERT_EQ(1, thunks[i]->source_ranges().size());
239 : BlockGraph::Block::SourceRanges::RangePair r =
240 E : thunks[i]->source_ranges().range_pairs()[0];
241 E : ASSERT_EQ(0, r.first.start());
242 :
243 E : ASSERT_TRUE(r.first.size() == sizeof(Thunk) ||
244 : r.first.size() == sizeof(ParamThunk));
245 :
246 E : BlockGraph::Reference ref;
247 E : EXPECT_TRUE(thunks[i]->GetReference(
248 : offsetof(Thunk, func_addr), &ref));
249 :
250 : // Retrieve the referenced block's source ranges to calculate
251 : // the destination start address.
252 E : EXPECT_EQ(1, ref.referenced()->source_ranges().size());
253 : BlockGraph::Block::SourceRanges::RangePair o =
254 E : ref.referenced()->source_ranges().range_pairs()[0];
255 :
256 : // The thunk's destination should be the block's start, plus the
257 : // reference offset.
258 E : EXPECT_EQ(o.second.start() + ref.offset(), r.second.start());
259 E : EXPECT_TRUE(r.second.size() == sizeof(Thunk) ||
260 : r.second.size() == sizeof(ParamThunk));
261 E : }
262 E : }
263 :
264 : // Verifies that there are num_thunks thunks in the image, and that they
265 : // have the expected properties.
266 : void VerifyThunks(size_t expected_total_thunks,
267 : size_t expected_param_thunks,
268 : size_t expected_destinations,
269 E : size_t expected_entrypoints) {
270 E : ConstBlockVector thunks;
271 E : int param_thunks = 0;
272 E : ASSERT_NO_FATAL_FAILURE(FindThunks(&thunks, ¶m_thunks));
273 :
274 E : EXPECT_EQ(expected_total_thunks, thunks.size());
275 E : EXPECT_EQ(expected_param_thunks, param_thunks);
276 E : EXPECT_EQ(expected_destinations, CountDestinations(thunks));
277 E : EXPECT_EQ(expected_entrypoints, CountEntryPoints(thunks));
278 E : }
279 :
280 : enum ImageType {
281 : DLL_IMAGE,
282 : EXE_IMAGE,
283 : };
284 :
285 E : void SetEmptyDllEntryPoint() {
286 E : TypedBlock<IMAGE_NT_HEADERS> nt_headers;
287 E : ASSERT_TRUE(nt_headers.Init(0, nt_headers_block_));
288 E : nt_headers->FileHeader.Characteristics |= IMAGE_FILE_DLL;
289 E : }
290 :
291 : // Sets the image entrypoint and sets or clears the DLL flag
292 : // in the NT headers.
293 E : void SetEntryPoint(BlockGraph::Block* entrypoint, ImageType image_type) {
294 : // Set the image entrypoint.
295 E : TypedBlock<IMAGE_NT_HEADERS> nt_headers;
296 E : ASSERT_TRUE(nt_headers.Init(0, nt_headers_block_));
297 E : ASSERT_TRUE(
298 : nt_headers.SetReference(
299 : BlockGraph::RELATIVE_REF,
300 : nt_headers->OptionalHeader.AddressOfEntryPoint,
301 : entrypoint,
302 : 0, 0));
303 :
304 : // Set or clear the DLL flag.
305 E : if (image_type == DLL_IMAGE)
306 E : nt_headers->FileHeader.Characteristics |= IMAGE_FILE_DLL;
307 E : else
308 E : nt_headers->FileHeader.Characteristics &= ~IMAGE_FILE_DLL;
309 E : }
310 :
311 : // Creates a TLS directory with the given block for entrypoint, and sets or
312 : // clears the DLL flag in the NT headers.
313 E : void SetTLSEntryPoint(BlockGraph::Block* entrypoint, ImageType image_type) {
314 : // Set the image entrypoint.
315 E : TypedBlock<IMAGE_NT_HEADERS> nt_headers;
316 E : ASSERT_TRUE(nt_headers.Init(0, nt_headers_block_));
317 : IMAGE_DATA_DIRECTORY& data_dir =
318 E : nt_headers->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_TLS];
319 E : ASSERT_EQ(0, data_dir.Size);
320 :
321 : // Create the TLS directory block.
322 : BlockGraph::Block* tls_dir_block = bg_.AddBlock(BlockGraph::DATA_BLOCK,
323 : sizeof(IMAGE_TLS_DIRECTORY),
324 E : "TLS Directory");
325 E : ASSERT_TRUE(tls_dir_block != NULL);
326 E : ASSERT_TRUE(tls_dir_block->AllocateData(tls_dir_block->size()));
327 :
328 : // Hook the TLS dir up to the NT headers.
329 E : ASSERT_TRUE(nt_headers.SetReference(BlockGraph::ABSOLUTE_REF,
330 : data_dir.VirtualAddress,
331 : tls_dir_block,
332 : 0, 0));
333 E : data_dir.Size = tls_dir_block->size();
334 :
335 E : TypedBlock<IMAGE_TLS_DIRECTORY> tls_dir;
336 E : ASSERT_TRUE(tls_dir.Init(0, tls_dir_block));
337 :
338 : BlockGraph::Block* tls_callbacks = bg_.AddBlock(BlockGraph::DATA_BLOCK,
339 : 2 * sizeof(AbsoluteAddress),
340 E : "TLS Callbacks");
341 E : ASSERT_TRUE(tls_callbacks != NULL);
342 E : ASSERT_TRUE(tls_callbacks->AllocateData(tls_callbacks->size()) != NULL);
343 E : ASSERT_TRUE(tls_dir.SetReference(BlockGraph::ABSOLUTE_REF,
344 : tls_dir->AddressOfCallBacks,
345 : tls_callbacks,
346 : 0, 0));
347 :
348 E : ASSERT_TRUE(tls_callbacks->SetReference(0,
349 : BlockGraph::Reference(BlockGraph::ABSOLUTE_REF,
350 : sizeof(AbsoluteAddress),
351 : entrypoint,
352 : 0, 0)));
353 :
354 : // Set or clear the DLL flag.
355 E : if (image_type == DLL_IMAGE)
356 E : nt_headers->FileHeader.Characteristics |= IMAGE_FILE_DLL;
357 E : else
358 E : nt_headers->FileHeader.Characteristics &= ~IMAGE_FILE_DLL;
359 E : }
360 :
361 : protected:
362 : size_t num_sections_pre_transform_;
363 :
364 : BlockGraph bg_;
365 : BlockGraph::Block* dos_header_block_;
366 : BlockGraph::Block* nt_headers_block_;
367 :
368 : BlockGraph::Block* foo_;
369 : BlockGraph::Block* bar_;
370 : BlockGraph::Block* array_;
371 : };
372 :
373 : } // namespace
374 :
375 E : TEST_F(EntryThunkTransformTest, AccessorsAndMutators) {
376 E : EntryThunkTransform tx;
377 :
378 E : EXPECT_TRUE(tx.instrument_unsafe_references());
379 E : EXPECT_FALSE(tx.src_ranges_for_thunks());
380 E : EXPECT_FALSE(tx.only_instrument_module_entry());
381 :
382 E : tx.set_instrument_unsafe_references(false);
383 E : tx.set_src_ranges_for_thunks(true);
384 E : tx.set_only_instrument_module_entry(true);
385 :
386 E : EXPECT_FALSE(tx.instrument_unsafe_references());
387 E : EXPECT_TRUE(tx.src_ranges_for_thunks());
388 E : EXPECT_TRUE(tx.only_instrument_module_entry());
389 E : }
390 :
391 E : TEST_F(EntryThunkTransformTest, ParameterizedThunks) {
392 E : EntryThunkTransform tx;
393 :
394 E : EXPECT_FALSE(tx.EntryThunkIsParameterized());
395 E : EXPECT_FALSE(tx.FunctionThunkIsParameterized());
396 E : EXPECT_EQ(core::kSizeNone, tx.entry_thunk_parameter().size());
397 E : EXPECT_EQ(core::kSizeNone, tx.function_thunk_parameter().size());
398 :
399 : // We shouldn't be allowed to set an 8-bit parameter.
400 E : Immediate imm8(43, core::kSize8Bit);
401 E : EXPECT_FALSE(tx.SetEntryThunkParameter(imm8));
402 E : EXPECT_FALSE(tx.SetFunctionThunkParameter(imm8));
403 :
404 E : EXPECT_FALSE(tx.EntryThunkIsParameterized());
405 E : EXPECT_FALSE(tx.FunctionThunkIsParameterized());
406 E : EXPECT_EQ(core::kSizeNone, tx.entry_thunk_parameter().size());
407 E : EXPECT_EQ(core::kSizeNone, tx.function_thunk_parameter().size());
408 :
409 : // A 32-bit parameter should be accepted just fine.
410 E : Immediate imm32(static_cast<int32>(0x11223344));
411 E : EXPECT_TRUE(tx.SetEntryThunkParameter(imm32));
412 E : EXPECT_TRUE(tx.SetFunctionThunkParameter(imm32));
413 :
414 E : EXPECT_TRUE(tx.EntryThunkIsParameterized());
415 E : EXPECT_TRUE(tx.FunctionThunkIsParameterized());
416 E : EXPECT_EQ(imm32, tx.entry_thunk_parameter());
417 E : EXPECT_EQ(imm32, tx.function_thunk_parameter());
418 :
419 : // A default contructured (with no size) parameter should be accepted.
420 E : EXPECT_TRUE(tx.SetEntryThunkParameter(Immediate()));
421 E : EXPECT_TRUE(tx.SetFunctionThunkParameter(Immediate()));
422 :
423 E : EXPECT_FALSE(tx.EntryThunkIsParameterized());
424 E : EXPECT_FALSE(tx.FunctionThunkIsParameterized());
425 E : EXPECT_EQ(core::kSizeNone, tx.entry_thunk_parameter().size());
426 E : EXPECT_EQ(core::kSizeNone, tx.function_thunk_parameter().size());
427 E : }
428 :
429 E : TEST_F(EntryThunkTransformTest, InstrumentAll) {
430 E : EntryThunkTransform transform;
431 E : ASSERT_NO_FATAL_FAILURE(SetEmptyDllEntryPoint());
432 :
433 E : ASSERT_TRUE(ApplyBlockGraphTransform(&transform, &bg_, dos_header_block_));
434 :
435 : // We should have three thunks - one each for the start of foo() and bar(),
436 : // and one for the middle of foo().
437 E : ASSERT_NO_FATAL_FAILURE(VerifyThunks(3, 0, 3, 1));
438 :
439 : // The .thunks section should have been added.
440 E : EXPECT_EQ(num_sections_pre_transform_ + 1, bg_.sections().size());
441 E : }
442 :
443 E : TEST_F(EntryThunkTransformTest, InstrumentAllWithParam) {
444 E : EntryThunkTransform transform;
445 E : ASSERT_NO_FATAL_FAILURE(SetEmptyDllEntryPoint());
446 E : transform.SetEntryThunkParameter(Immediate(0x11223344));
447 E : transform.SetFunctionThunkParameter(Immediate(0x11223344));
448 :
449 E : ASSERT_TRUE(ApplyBlockGraphTransform(&transform, &bg_, dos_header_block_));
450 :
451 : // We should have three thunks - one each for the start of foo() and bar(),
452 : // and one for the middle of foo().
453 E : ASSERT_NO_FATAL_FAILURE(VerifyThunks(3, 3, 3, 1));
454 :
455 : // The .thunks section should have been added.
456 E : EXPECT_EQ(num_sections_pre_transform_ + 1, bg_.sections().size());
457 E : }
458 :
459 E : TEST_F(EntryThunkTransformTest, InstrumentModuleEntriesOnlyNone) {
460 E : EntryThunkTransform transform;
461 E : ASSERT_NO_FATAL_FAILURE(SetEmptyDllEntryPoint());
462 E : transform.set_only_instrument_module_entry(true);
463 :
464 E : ASSERT_TRUE(ApplyBlockGraphTransform(&transform, &bg_, dos_header_block_));
465 :
466 : // We should have no thunks.
467 E : ASSERT_NO_FATAL_FAILURE(VerifyThunks(0, 0, 0, 0));
468 :
469 : // The .thunks section should not have been added, as there are no hooks
470 : // added.
471 E : EXPECT_EQ(num_sections_pre_transform_, bg_.sections().size());
472 E : }
473 :
474 E : TEST_F(EntryThunkTransformTest, InstrumentModuleEntriesOnlyDllMainOnly) {
475 E : EntryThunkTransform transform;
476 E : ASSERT_NO_FATAL_FAILURE(SetEntryPoint(foo_, DLL_IMAGE));
477 E : transform.set_only_instrument_module_entry(true);
478 :
479 E : ASSERT_TRUE(ApplyBlockGraphTransform(&transform, &bg_, dos_header_block_));
480 :
481 : // We should have one thunk, for the DLL main entry point to the start of
482 : // foo_.
483 E : ASSERT_NO_FATAL_FAILURE(VerifyThunks(1, 0, 1, 1));
484 :
485 : // The .thunks section should have been added.
486 E : EXPECT_EQ(num_sections_pre_transform_ + 1, bg_.sections().size());
487 E : }
488 :
489 E : TEST_F(EntryThunkTransformTest, InstrumentOnlyDllMainWithParamThunk) {
490 E : EntryThunkTransform transform;
491 E : ASSERT_NO_FATAL_FAILURE(SetEntryPoint(foo_, DLL_IMAGE));
492 E : transform.set_only_instrument_module_entry(true);
493 E : transform.SetEntryThunkParameter(Immediate(0x11223344));
494 :
495 E : ASSERT_TRUE(ApplyBlockGraphTransform(&transform, &bg_, dos_header_block_));
496 :
497 : // We should have one thunk, for the DLL main entry point to the start of
498 : // foo_ and it should be parameterized.
499 E : ASSERT_NO_FATAL_FAILURE(VerifyThunks(1, 1, 1, 1));
500 :
501 : // The .thunks section should have been added.
502 E : EXPECT_EQ(num_sections_pre_transform_ + 1, bg_.sections().size());
503 E : }
504 :
505 E : TEST_F(EntryThunkTransformTest, InstrumentModuleEntriesOnlyDllMainAndTls) {
506 E : EntryThunkTransform transform;
507 E : ASSERT_NO_FATAL_FAILURE(SetEntryPoint(foo_, DLL_IMAGE));
508 E : ASSERT_NO_FATAL_FAILURE(SetTLSEntryPoint(bar_, DLL_IMAGE));
509 E : transform.set_only_instrument_module_entry(true);
510 :
511 E : ASSERT_TRUE(ApplyBlockGraphTransform(&transform, &bg_, dos_header_block_));
512 :
513 : // We should have two thunk, for the DLL main entry point and another for the
514 : // TLS. One is to foo_ and one is to bar_.
515 E : ASSERT_NO_FATAL_FAILURE(VerifyThunks(2, 0, 2, 1));
516 :
517 : // The .thunks section should have been added.
518 E : EXPECT_EQ(num_sections_pre_transform_ + 1, bg_.sections().size());
519 E : }
520 :
521 E : TEST_F(EntryThunkTransformTest, InstrumentModuleEntriesOnlyExeMainAndTls) {
522 E : EntryThunkTransform transform;
523 E : ASSERT_NO_FATAL_FAILURE(SetEntryPoint(foo_, EXE_IMAGE));
524 E : ASSERT_NO_FATAL_FAILURE(SetTLSEntryPoint(bar_, EXE_IMAGE));
525 E : transform.set_only_instrument_module_entry(true);
526 :
527 E : ASSERT_TRUE(ApplyBlockGraphTransform(&transform, &bg_, dos_header_block_));
528 :
529 : // We should have one TLS thunk and an EXE entry thunk.
530 E : ASSERT_NO_FATAL_FAILURE(VerifyThunks(2, 0, 2, 2));
531 :
532 : // The .thunks section should have been added.
533 E : EXPECT_EQ(num_sections_pre_transform_ + 1, bg_.sections().size());
534 E : }
535 :
536 E : TEST_F(EntryThunkTransformTest, InstrumentAllDebugFriendly) {
537 E : EntryThunkTransform transform;
538 E : ASSERT_NO_FATAL_FAILURE(SetEntryPoint(foo_, EXE_IMAGE));
539 E : transform.set_src_ranges_for_thunks(true);
540 :
541 E : ASSERT_TRUE(ApplyBlockGraphTransform(&transform, &bg_, dos_header_block_));
542 :
543 : // Verify the source ranges on the thunks.
544 E : ConstBlockVector thunks;
545 E : int param_thunks = 0;
546 E : ASSERT_NO_FATAL_FAILURE(FindThunks(&thunks, ¶m_thunks));
547 E : EXPECT_EQ(0u, param_thunks);
548 E : ASSERT_NO_FATAL_FAILURE(VerifySourceRanges(thunks));
549 E : }
550 :
551 E : TEST_F(EntryThunkTransformTest, InstrumentNoUnsafe) {
552 E : EntryThunkTransform transform;
553 E : ASSERT_NO_FATAL_FAILURE(SetEmptyDllEntryPoint());
554 :
555 : // No unsafe reference instrumentation.
556 E : transform.set_instrument_unsafe_references(false);
557 :
558 : // Tag both foo and bar with unsafe attributes.
559 E : foo_->set_attribute(BlockGraph::HAS_INLINE_ASSEMBLY);
560 E : bar_->set_attribute(BlockGraph::BUILT_BY_UNSUPPORTED_COMPILER);
561 :
562 E : ASSERT_TRUE(ApplyBlockGraphTransform(&transform, &bg_, dos_header_block_));
563 :
564 : // We should have two thunks - one each for the start of foo() and bar().
565 E : ASSERT_NO_FATAL_FAILURE(VerifyThunks(2, 0, 2, 1));
566 :
567 : // The foo->bar reference should not have been thunked.
568 E : BlockGraph::Reference ref;
569 E : ASSERT_TRUE(foo_->GetReference(5, &ref));
570 E : ASSERT_EQ(bar_, ref.referenced());
571 :
572 : // The .thunks section should have been added.
573 E : EXPECT_EQ(num_sections_pre_transform_ + 1, bg_.sections().size());
574 E : }
575 :
576 E : TEST_F(EntryThunkTransformTest, InstrumentDllEntrypoint) {
577 E : EntryThunkTransform transform;
578 E : ASSERT_NO_FATAL_FAILURE(SetEntryPoint(foo_, DLL_IMAGE));
579 :
580 E : ASSERT_TRUE(ApplyBlockGraphTransform(&transform, &bg_, dos_header_block_));
581 :
582 : // We should have three thunks - one each for the start of foo() and bar().
583 : // One of the thunks should use the DllMain entrypoint.
584 E : ASSERT_NO_FATAL_FAILURE(VerifyThunks(3, 0, 3, 2));
585 :
586 : // The .thunks section should have been added.
587 E : EXPECT_EQ(num_sections_pre_transform_ + 1, bg_.sections().size());
588 E : }
589 :
590 E : TEST_F(EntryThunkTransformTest, InstrumentExeEntrypoint) {
591 E : EntryThunkTransform transform;
592 E : ASSERT_NO_FATAL_FAILURE(SetEntryPoint(foo_, EXE_IMAGE));
593 :
594 E : ASSERT_TRUE(ApplyBlockGraphTransform(&transform, &bg_, dos_header_block_));
595 :
596 : // We should have three thunks - one each for the start of foo() and bar().
597 E : ASSERT_NO_FATAL_FAILURE(VerifyThunks(3, 0, 3, 2));
598 :
599 : // The .thunks section should have been added.
600 E : EXPECT_EQ(num_sections_pre_transform_ + 1, bg_.sections().size());
601 E : }
602 :
603 E : TEST_F(EntryThunkTransformTest, InstrumentDllTLSEntrypoint) {
604 E : EntryThunkTransform transform;
605 E : ASSERT_NO_FATAL_FAILURE(SetEntryPoint(foo_, DLL_IMAGE));
606 E : ASSERT_NO_FATAL_FAILURE(SetTLSEntryPoint(bar_, DLL_IMAGE));
607 :
608 E : ASSERT_TRUE(ApplyBlockGraphTransform(&transform, &bg_, dos_header_block_));
609 :
610 : // We should have three thunks - one each for the start of foo() and bar().
611 : // One of the thunks should use the DllMain entrypoint.
612 E : ASSERT_NO_FATAL_FAILURE(VerifyThunks(3, 0, 3, 2));
613 :
614 : // The .thunks section should have been added.
615 E : EXPECT_EQ(num_sections_pre_transform_ + 1, bg_.sections().size());
616 E : }
617 :
618 E : TEST_F(EntryThunkTransformTest, InstrumentExeTLSEntrypoint) {
619 E : EntryThunkTransform transform;
620 E : ASSERT_NO_FATAL_FAILURE(SetEntryPoint(foo_, EXE_IMAGE));
621 E : ASSERT_NO_FATAL_FAILURE(SetTLSEntryPoint(bar_, EXE_IMAGE));
622 :
623 E : ASSERT_TRUE(ApplyBlockGraphTransform(&transform, &bg_, dos_header_block_));
624 :
625 : // We should have three thunks - one each for the start of foo() and bar().
626 E : ASSERT_NO_FATAL_FAILURE(VerifyThunks(3, 0, 3, 3));
627 :
628 : // The .thunks section should have been added.
629 E : EXPECT_EQ(num_sections_pre_transform_ + 1, bg_.sections().size());
630 E : }
631 :
632 : } // namespace transforms
633 : } // namespace instrument
|