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