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