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 : #include "syzygy/reorder/reorderer.h"
16 :
17 : #include "gmock/gmock.h"
18 : #include "gtest/gtest.h"
19 : #include "syzygy/core/unittest_util.h"
20 : #include "syzygy/pdb/omap.h"
21 : #include "syzygy/pe/unittest_util.h"
22 : #include "syzygy/reorder/order_generator_test.h"
23 : #include "syzygy/trace/parse/parse_engine.h"
24 : #include "syzygy/trace/parse/parser.h"
25 :
26 : namespace reorder {
27 :
28 : namespace {
29 :
30 : using block_graph::BlockGraph;
31 : using block_graph::ConstBlockVector;
32 : using testing::_;
33 : using testing::BlockSpecsAreEqual;
34 : using testing::DoAll;
35 : using testing::InSequence;
36 : using testing::InvokeWithoutArgs;
37 : using testing::OrdersAreEqual;
38 : using testing::Return;
39 : using testing::SectionSpecsAreEqual;
40 : using trace::parser::ParseEngine;
41 : using trace::parser::Parser;
42 :
43 : typedef Reorderer::Order::BlockSpec BlockSpec;
44 : typedef Reorderer::Order::SectionSpec SectionSpec;
45 : typedef Reorderer::Order::SectionSpecVector SectionSpecVector;
46 :
47 : // A wrapper for Reorderer giving us access to some of its internals.
48 : class TestReorderer : public Reorderer {
49 : public:
50 : TestReorderer(const base::FilePath& module_path,
51 : const base::FilePath& instrumented_path,
52 : const TraceFileList& trace_files,
53 : Flags flags);
54 :
55 E : const PEFile::Signature& instr_signature() const {
56 E : return playback_.instr_signature();
57 E : }
58 :
59 E : Parser& parser() { return parser_; }
60 E : Playback* playback() { return &playback_; }
61 : };
62 :
63 : // A dummy order generator that does nothing but record events fed to it via
64 : // the reorderer.
65 : class TestOrderGenerator : public Reorderer::OrderGenerator {
66 : public:
67 E : TestOrderGenerator() : Reorderer::OrderGenerator("TestOrderGenerator") {
68 E : }
69 :
70 E : virtual ~TestOrderGenerator() {
71 E : }
72 :
73 : virtual bool OnCodeBlockEntry(const BlockGraph::Block* block,
74 : RelativeAddress address,
75 : uint32 process_id,
76 : uint32 thread_id,
77 E : const Reorderer::UniqueTime& time) override {
78 : // Record the visited block.
79 E : blocks.push_back(block);
80 E : return true;
81 E : }
82 :
83 : virtual bool CalculateReordering(const PEFile& pe_file,
84 : const ImageLayout& image,
85 : bool reorder_code,
86 : bool reorder_data,
87 E : Reorderer::Order* order) override {
88 : // We don't actually generate an ordering.
89 E : return true;
90 E : }
91 :
92 : ConstBlockVector blocks;
93 : };
94 :
95 : class MockOrderGenerator : public Reorderer::OrderGenerator {
96 : public:
97 E : MockOrderGenerator() : Reorderer::OrderGenerator("MockOrderGenerator") {
98 E : }
99 :
100 : MOCK_METHOD2(OnProcessStarted,
101 E : bool(uint32 process_id, const Reorderer::UniqueTime& time));
102 :
103 : MOCK_METHOD2(OnProcessEnded,
104 E : bool(uint32 process_id, const Reorderer::UniqueTime& time));
105 :
106 : MOCK_METHOD5(OnCodeBlockEntry,
107 : bool(const BlockGraph::Block* block,
108 : RelativeAddress address,
109 : uint32 process_id,
110 : uint32 thread_id,
111 E : const Reorderer::UniqueTime& time));
112 :
113 : MOCK_METHOD5(CalculateReordering,
114 : bool(const PEFile& pe_file,
115 : const ImageLayout& image,
116 : bool reorder_code,
117 : bool reorder_data,
118 E : Reorderer::Order* order));
119 : };
120 :
121 : // A dummy parse engine. This lets us feed hand-crafted events to any consumer.
122 : class TestParseEngine : public ParseEngine {
123 : public:
124 : typedef block_graph::BlockGraph BlockGraph;
125 : typedef core::RelativeAddress RelativeAddress;
126 : typedef Reorderer::ImageLayout ImageLayout;
127 : typedef Reorderer::PEFile PEFile;
128 :
129 : explicit TestParseEngine(TestReorderer* reorderer)
130 : : ParseEngine("TestParseEngine", true),
131 E : reorderer_(reorderer) {
132 E : }
133 :
134 E : virtual ~TestParseEngine() {
135 E : }
136 :
137 E : virtual bool IsRecognizedTraceFile(const base::FilePath& trace_file_path) {
138 E : return true;
139 E : }
140 :
141 E : virtual bool OpenTraceFile(const base::FilePath& trace_file_path) {
142 E : return true;
143 E : }
144 :
145 : virtual bool ConsumeAllEvents();
146 :
147 E : virtual bool CloseAllTraceFiles() {
148 E : return true;
149 E : }
150 :
151 : using ParseEngine::AddModuleInformation;
152 :
153 : // This will hold the list of blocks that we expect the order generator to
154 : // build.
155 : ConstBlockVector blocks;
156 :
157 : private:
158 : // The parser needs to have a pointer to the reorderer in order to get image
159 : // data from it for producing false events.
160 : TestReorderer* reorderer_;
161 : };
162 :
163 : TestReorderer::TestReorderer(const base::FilePath& module_path,
164 : const base::FilePath& instrumented_path,
165 : const TraceFileList& trace_files,
166 : Flags flags)
167 E : : Reorderer(module_path, instrumented_path, trace_files, flags) {
168 E : }
169 :
170 : const DWORD kProcessId = 0xAAAAAAAA;
171 : const DWORD kThreadId = 0xBBBBBBBB;
172 : const pe::ModuleInformation kExeInfo(
173 : L"file_name.exe", pe::PEFile::AbsoluteAddress(0x11111111), 0x22222222,
174 E : 0x33333333, 0x44444444);
175 :
176 E : bool TestParseEngine::ConsumeAllEvents() {
177 : // Add dummy module information for some running process.
178 E : if (!AddModuleInformation(kProcessId, kExeInfo))
179 i : return false;
180 :
181 : // Simulate a process starting.
182 E : base::Time time = base::Time::Now();
183 E : event_handler_->OnProcessStarted(time, kProcessId, NULL);
184 :
185 E : const pe::ModuleInformation& dll_info = reorderer_->instr_signature();
186 :
187 E : TraceModuleData dll_data = {};
188 : dll_data.module_base_addr =
189 E : reinterpret_cast<ModuleAddr>(dll_info.base_address.value());
190 E : dll_data.module_base_size = dll_info.module_size;
191 E : dll_data.module_exe[0] = 0;
192 E : dll_data.module_checksum = dll_info.module_checksum;
193 E : dll_data.module_time_date_stamp = dll_info.module_time_date_stamp;
194 : wcscpy_s(dll_data.module_name,
195 : arraysize(dll_data.module_name),
196 E : dll_info.path.c_str());
197 :
198 : // Simulate the process and thread attaching to the DLL. This adds the DLL
199 : // to the list of modules.
200 E : EVENT_TRACE event_record = {};
201 E : FILETIME timestamp = time.ToFileTime();
202 E : event_record.Header.TimeStamp.LowPart = timestamp.dwLowDateTime;
203 E : event_record.Header.TimeStamp.HighPart = timestamp.dwHighDateTime;
204 E : event_record.Header.ProcessId = kProcessId;
205 E : event_record.Header.ThreadId = kThreadId;
206 E : event_record.Header.Guid = kCallTraceEventClass;
207 E : event_record.Header.Class.Type = TRACE_PROCESS_ATTACH_EVENT;
208 E : event_record.MofData = &dll_data;
209 E : event_record.MofLength = sizeof(dll_data);
210 E : if (!DispatchEvent(&event_record))
211 i : return false;
212 :
213 E : event_record.Header.Class.Type = TRACE_THREAD_ATTACH_EVENT;
214 E : if (!DispatchEvent(&event_record))
215 i : return false;
216 :
217 : // Get all of the non-padding code blocks in the original image. (Padding
218 : // blocks don't make it to the instrumented DLL, so any events with addresses
219 : // that refer to padding blocks will fail to resolve via the OMAP info.)
220 : BlockGraph::AddressSpace::RangeMapConstIter block_it =
221 E : reorderer_->playback()->image()->blocks.begin();
222 E : for (; block_it != reorderer_->playback()->image()->blocks.end();
223 E : ++block_it) {
224 : if (block_it->second->type() == BlockGraph::CODE_BLOCK &&
225 E : (block_it->second->attributes() & BlockGraph::PADDING_BLOCK) == 0) {
226 E : blocks.push_back(block_it->second);
227 : }
228 E : }
229 :
230 : // Shuffle the code blocks.
231 E : std::random_shuffle(blocks.begin(), blocks.end());
232 :
233 : // Simulate half of the blocks using batch events.
234 : static const size_t kBatchCallCount = 5;
235 E : size_t i = 0;
236 E : for (; i < blocks.size() / 2; i += kBatchCallCount) {
237 : uint8 raw_data[sizeof(TraceBatchEnterData) +
238 E : kBatchCallCount * sizeof(TraceEnterEventData)] = {};
239 : TraceBatchEnterData& event_data =
240 E : *reinterpret_cast<TraceBatchEnterData*>(&raw_data);
241 E : event_data.thread_id = kThreadId;
242 E : event_data.num_calls = kBatchCallCount;
243 :
244 E : for (size_t j = 0; j < kBatchCallCount; ++j) {
245 : // Get the address of this block as an RVA in the instrumented module.
246 E : RelativeAddress rva = blocks[i + j]->addr();
247 : rva = pdb::TranslateAddressViaOmap(reorderer_->playback()->omap_from(),
248 E : rva);
249 :
250 : // Convert this to an absolute address using the base address from above.
251 E : uint64 abs = dll_info.base_address.value() + rva.value();
252 E : void* block_pointer = reinterpret_cast<void*>(abs);
253 :
254 E : event_data.calls[j].function = block_pointer;
255 E : }
256 :
257 E : event_record.Header.Class.Type = TRACE_BATCH_ENTER;
258 E : event_record.MofData = &raw_data;
259 E : event_record.MofLength = sizeof(raw_data);
260 E : if (!DispatchEvent(&event_record))
261 i : return false;
262 E : }
263 :
264 : // Simulate entry/exit pairs with the remaining blocks.
265 E : for (; i < blocks.size(); ++i) {
266 : // Get the address of this block as an RVA in the instrumented module.
267 E : RelativeAddress rva = blocks[i]->addr();
268 : rva = pdb::TranslateAddressViaOmap(reorderer_->playback()->omap_from(),
269 E : rva);
270 :
271 : // Convert this to an absolute address using the base address from above.
272 E : uint64 abs = dll_info.base_address.value() + rva.value();
273 E : void* block_pointer = reinterpret_cast<void*>(abs);
274 :
275 E : TraceEnterEventData event_data = {};
276 E : event_data.function = block_pointer;
277 :
278 : // Simulate an entry event.
279 E : event_record.Header.Class.Type = TRACE_ENTER_EVENT;
280 E : event_record.MofData = &event_data;
281 E : event_record.MofLength = sizeof(event_data);
282 E : if (!DispatchEvent(&event_record))
283 i : return false;
284 :
285 : // Simulate a corresponding exit event.
286 E : event_record.Header.Class.Type = TRACE_EXIT_EVENT;
287 E : if (!DispatchEvent(&event_record))
288 i : return false;
289 E : }
290 :
291 : // Simulate the thread and process detaching from the DLL.
292 E : event_record.Header.Class.Type = TRACE_THREAD_DETACH_EVENT;
293 E : event_record.MofData = &dll_data;
294 E : event_record.MofLength = sizeof(dll_data);
295 E : if (!DispatchEvent(&event_record))
296 i : return false;
297 :
298 E : event_record.Header.Class.Type = TRACE_PROCESS_DETACH_EVENT;
299 E : if (!DispatchEvent(&event_record))
300 i : return false;
301 :
302 : // Simulate the process ending.
303 E : event_handler_->OnProcessEnded(time, kProcessId);
304 :
305 E : return true;
306 E : }
307 :
308 : class ReordererTest : public testing::PELibUnitTest {
309 : public:
310 : typedef testing::PELibUnitTest Super;
311 :
312 : typedef block_graph::BlockGraph BlockGraph;
313 : typedef core::RelativeAddress RelativeAddress;
314 : typedef Reorderer::ImageLayout ImageLayout;
315 : typedef Reorderer::PEFile PEFile;
316 :
317 E : ReordererTest() : test_parse_engine_(NULL) {
318 E : }
319 :
320 E : void SetUp() override {
321 E : Super::SetUp();
322 :
323 : // Create the dummy trace file list.
324 E : Reorderer::TraceFileList trace_file_list;
325 E : trace_file_list.push_back(base::FilePath(L"foo"));
326 :
327 : // Set up the reorderer. These tests rely on
328 : // call_trace_instrumented_test_dll.dll, as generated by the test_data
329 : // project.
330 : const Reorderer::Flags kFlags = Reorderer::kFlagReorderCode |
331 E : Reorderer::kFlagReorderData;
332 : test_reorderer_.reset(new TestReorderer(
333 : testing::GetExeTestDataRelativePath(testing::kTestDllName),
334 : testing::GetExeTestDataRelativePath(
335 : testing::kCallTraceInstrumentedTestDllName),
336 : trace_file_list,
337 E : kFlags));
338 :
339 : // Setup the test parse engine and register it with the parser used
340 : // by the test reorderer. Note that ownership of the pointer is also
341 : // being passed.
342 E : ASSERT_TRUE(test_parse_engine_ == NULL);
343 E : test_parse_engine_ = new TestParseEngine(test_reorderer_.get());
344 E : ASSERT_TRUE(test_parse_engine_ != NULL);
345 E : test_reorderer_->parser().AddParseEngine(test_parse_engine_);
346 E : }
347 :
348 : // A reorderer will be initialized, in SetUp(), for each test run.
349 : scoped_ptr<TestReorderer> test_reorderer_;
350 :
351 : // The reorderer needs to be set up to use a custom parse engine before a
352 : // call to Reorder. This must be heap allocated and the responsibility for
353 : // deleting it rests with the parser.
354 : TestParseEngine* test_parse_engine_;
355 : };
356 :
357 : } // namespace
358 :
359 E : TEST_F(ReordererTest, ValidateCallbacks) {
360 E : MockOrderGenerator mock_order_generator;
361 :
362 : // Setup the expected calls.
363 E : InSequence s;
364 : EXPECT_CALL(mock_order_generator, OnProcessStarted(_, _))
365 E : .WillOnce(Return(true));
366 : EXPECT_CALL(mock_order_generator, OnCodeBlockEntry(_, _, _, _, _))
367 E : .WillRepeatedly(Return(true));
368 : EXPECT_CALL(mock_order_generator, OnProcessEnded(_, _))
369 E : .WillOnce(Return(true));
370 : EXPECT_CALL(mock_order_generator, CalculateReordering(_, _, _, _, _))
371 E : .WillOnce(Return(true));
372 :
373 : // Run the reorderer.
374 E : Reorderer::Order order;
375 E : PEFile pe_file;
376 E : BlockGraph block_graph;
377 E : ImageLayout image_layout(&block_graph);
378 : EXPECT_TRUE(test_reorderer_->Reorder(&mock_order_generator,
379 : &order,
380 : &pe_file,
381 E : &image_layout));
382 E : }
383 :
384 E : TEST_F(ReordererTest, Reorder) {
385 E : TestOrderGenerator test_order_generator;
386 :
387 : // Run the reorderer.
388 E : Reorderer::Order order;
389 E : PEFile pe_file;
390 E : BlockGraph block_graph;
391 E : ImageLayout image_layout(&block_graph);
392 : EXPECT_TRUE(test_reorderer_->Reorder(&test_order_generator,
393 : &order,
394 : &pe_file,
395 E : &image_layout));
396 :
397 : // We expect the order generator to have come up with the same list of
398 : // blocks that the parse engine used for generating dummy trace events.
399 : EXPECT_EQ(test_parse_engine_->blocks,
400 E : test_order_generator.blocks);
401 E : }
402 :
403 E : TEST(OrderTest, OrderConstructor) {
404 E : Reorderer::Order order;
405 E : EXPECT_TRUE(order.sections.empty());
406 E : EXPECT_TRUE(order.comment.empty());
407 E : }
408 :
409 E : TEST(OrderTest, SectionSpecConstructorsAndCopies) {
410 : // Create default constructed section spec.
411 E : Reorderer::Order::SectionSpec default_section_spec;
412 : EXPECT_EQ(Reorderer::Order::SectionSpec::kNewSectionId,
413 E : default_section_spec.id);
414 E : EXPECT_TRUE(default_section_spec.name.empty());
415 E : EXPECT_EQ(0U, default_section_spec.characteristics);
416 E : EXPECT_TRUE(default_section_spec.blocks.empty());
417 :
418 : // Create a customized section spec.
419 E : Reorderer::Order::SectionSpec other_section_spec;
420 E : EXPECT_TRUE(SectionSpecsAreEqual(default_section_spec, other_section_spec));
421 E : other_section_spec.id = 7;
422 E : other_section_spec.characteristics = 19;
423 E : other_section_spec.name = "other";
424 E : EXPECT_FALSE(SectionSpecsAreEqual(default_section_spec, other_section_spec));
425 :
426 : // Copy construct a section spec.
427 E : Reorderer::Order::SectionSpec copied_section_spec(other_section_spec);
428 E : EXPECT_TRUE(SectionSpecsAreEqual(other_section_spec, copied_section_spec));
429 E : EXPECT_FALSE(SectionSpecsAreEqual(copied_section_spec, default_section_spec));
430 :
431 : // Assign to the default section spec.
432 E : default_section_spec = other_section_spec;
433 E : EXPECT_TRUE(SectionSpecsAreEqual(copied_section_spec, default_section_spec));
434 E : }
435 :
436 E : TEST(OrderTest, BlockSpecConstructorsAndCopier) {
437 : static const BlockGraph::Block* kFauxBlockPtr =
438 : reinterpret_cast<BlockGraph::Block*>(0xCCCCCCCC);
439 :
440 : // Default construct a block spec.
441 E : Reorderer::Order::BlockSpec default_block_spec;
442 E : EXPECT_EQ(NULL, default_block_spec.block);
443 E : EXPECT_TRUE(default_block_spec.basic_block_offsets.empty());
444 :
445 : // Explicitly construct a block spec.
446 E : Reorderer::Order::BlockSpec explicit_block_spec(kFauxBlockPtr);
447 E : EXPECT_EQ(kFauxBlockPtr, explicit_block_spec.block);
448 E : EXPECT_TRUE(explicit_block_spec.basic_block_offsets.empty());
449 :
450 : // Default construct another block spec.
451 E : Reorderer::Order::BlockSpec block_spec;
452 E : EXPECT_TRUE(BlockSpecsAreEqual(default_block_spec, block_spec));
453 :
454 : // Modify the block spec.
455 E : block_spec.block = kFauxBlockPtr;
456 E : block_spec.basic_block_offsets.push_back(0);
457 E : block_spec.basic_block_offsets.push_back(8);
458 E : EXPECT_FALSE(BlockSpecsAreEqual(default_block_spec, block_spec));
459 E : EXPECT_FALSE(BlockSpecsAreEqual(explicit_block_spec, block_spec));
460 :
461 : // Copy construct a block spec.
462 E : Reorderer::Order::BlockSpec copied_block_spec(block_spec);
463 E : EXPECT_TRUE(BlockSpecsAreEqual(block_spec, copied_block_spec));
464 E : EXPECT_FALSE(BlockSpecsAreEqual(default_block_spec, block_spec));
465 E : EXPECT_FALSE(BlockSpecsAreEqual(explicit_block_spec, block_spec));
466 :
467 : // Assign to the explicit block spec.
468 E : explicit_block_spec = block_spec;
469 E : EXPECT_TRUE(BlockSpecsAreEqual(copied_block_spec, explicit_block_spec));
470 E : }
471 :
472 E : TEST(OrderTest, SerializeToJsonRoundTrip) {
473 : // Build a dummy block graph.
474 E : BlockGraph block_graph;
475 E : BlockGraph::Section* section1 = block_graph.AddSection(".text", 0);
476 E : BlockGraph::Section* section2 = block_graph.AddSection(".rdata", 0);
477 : BlockGraph::Block* block1 = block_graph.AddBlock(BlockGraph::CODE_BLOCK, 10,
478 E : "block1");
479 : BlockGraph::Block* block2 = block_graph.AddBlock(BlockGraph::DATA_BLOCK, 10,
480 E : "block2");
481 : BlockGraph::Block* block3 = block_graph.AddBlock(BlockGraph::DATA_BLOCK, 10,
482 E : "block3");
483 E : block1->set_section(section1->id());
484 E : block2->set_section(section2->id());
485 E : block3->set_section(section2->id());
486 :
487 : // Build a dummy image layout.
488 E : pe::ImageLayout layout(&block_graph);
489 E : pe::ImageLayout::SectionInfo section_info1 = {};
490 E : section_info1.name = section1->name();
491 E : section_info1.addr = core::RelativeAddress(0x1000);
492 E : section_info1.size = 0x1000;
493 E : section_info1.data_size = 0x1000;
494 E : layout.sections.push_back(section_info1);
495 :
496 E : pe::ImageLayout::SectionInfo section_info2 = {};
497 E : section_info2.name = section2->name();
498 E : section_info2.addr = core::RelativeAddress(0x2000);
499 E : section_info2.size = 0x1000;
500 E : section_info2.data_size = 0x1000;
501 E : layout.sections.push_back(section_info2);
502 :
503 : layout.blocks.InsertBlock(section_info1.addr,
504 E : block1);
505 : layout.blocks.InsertBlock(section_info2.addr,
506 E : block2);
507 : layout.blocks.InsertBlock(section_info2.addr + block2->size(),
508 E : block3);
509 :
510 : // Build a dummy order.
511 E : Reorderer::Order order;
512 E : order.comment = "This is a comment.";
513 E : order.sections.resize(2);
514 E : order.sections[0].id = section1->id();
515 E : order.sections[0].name = section1->name();
516 E : order.sections[0].characteristics = section1->characteristics();
517 E : order.sections[0].blocks.push_back(BlockSpec(block1));
518 E : order.sections[0].blocks.back().basic_block_offsets.push_back(0);
519 E : order.sections[0].blocks.back().basic_block_offsets.push_back(8);
520 E : order.sections[1].id = section2->id();
521 E : order.sections[1].name = section2->name();
522 E : order.sections[1].characteristics = section2->characteristics();
523 E : order.sections[1].blocks.push_back(BlockSpec(block2));
524 E : order.sections[1].blocks.push_back(BlockSpec(block3));
525 :
526 : base::FilePath module = testing::GetExeTestDataRelativePath(
527 E : testing::kTestDllName);
528 E : pe::PEFile pe_file;
529 E : ASSERT_TRUE(pe_file.Init(module));
530 :
531 : // Serialize the order.
532 E : base::FilePath temp_file;
533 E : ASSERT_TRUE(base::CreateTemporaryFile(&temp_file));
534 E : EXPECT_TRUE(order.SerializeToJSON(pe_file, temp_file, true));
535 :
536 : // Get the original module from the file.
537 E : base::FilePath orig_module;
538 E : EXPECT_TRUE(Reorderer::Order::GetOriginalModulePath(temp_file, &orig_module));
539 E : EXPECT_EQ(module, orig_module);
540 :
541 : // Deserialize it.
542 E : Reorderer::Order order2;
543 E : EXPECT_FALSE(OrdersAreEqual(order, order2));
544 E : EXPECT_TRUE(order2.LoadFromJSON(pe_file, layout, temp_file));
545 :
546 : // Expect them to be the same.
547 E : EXPECT_TRUE(OrdersAreEqual(order, order2));
548 :
549 E : EXPECT_TRUE(base::DeleteFile(temp_file, false));
550 E : }
551 :
552 : } // namespace reorder
|