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 : event_record.Header.TimeStamp =
202 E : reinterpret_cast<LARGE_INTEGER&>(time.ToFileTime());
203 E : event_record.Header.ProcessId = kProcessId;
204 E : event_record.Header.ThreadId = kThreadId;
205 E : event_record.Header.Guid = kCallTraceEventClass;
206 E : event_record.Header.Class.Type = TRACE_PROCESS_ATTACH_EVENT;
207 E : event_record.MofData = &dll_data;
208 E : event_record.MofLength = sizeof(dll_data);
209 E : if (!DispatchEvent(&event_record))
210 i : return false;
211 :
212 E : event_record.Header.Class.Type = TRACE_THREAD_ATTACH_EVENT;
213 E : if (!DispatchEvent(&event_record))
214 i : return false;
215 :
216 : // Get all of the non-padding code blocks in the original image. (Padding
217 : // blocks don't make it to the instrumented DLL, so any events with addresses
218 : // that refer to padding blocks will fail to resolve via the OMAP info.)
219 : BlockGraph::AddressSpace::RangeMapConstIter block_it =
220 E : reorderer_->playback()->image()->blocks.begin();
221 E : for (; block_it != reorderer_->playback()->image()->blocks.end();
222 E : ++block_it) {
223 : if (block_it->second->type() == BlockGraph::CODE_BLOCK &&
224 E : (block_it->second->attributes() & BlockGraph::PADDING_BLOCK) == 0) {
225 E : blocks.push_back(block_it->second);
226 : }
227 E : }
228 :
229 : // Shuffle the code blocks.
230 E : std::random_shuffle(blocks.begin(), blocks.end());
231 :
232 : // Simulate half of the blocks using batch events.
233 : static const size_t kBatchCallCount = 5;
234 E : size_t i = 0;
235 E : for (; i < blocks.size() / 2; i += kBatchCallCount) {
236 : uint8 raw_data[sizeof(TraceBatchEnterData) +
237 E : kBatchCallCount * sizeof(TraceEnterEventData)] = {};
238 : TraceBatchEnterData& event_data =
239 E : *reinterpret_cast<TraceBatchEnterData*>(&raw_data);
240 E : event_data.thread_id = kThreadId;
241 E : event_data.num_calls = kBatchCallCount;
242 :
243 E : for (size_t j = 0; j < kBatchCallCount; ++j) {
244 : // Get the address of this block as an RVA in the instrumented module.
245 E : RelativeAddress rva = blocks[i + j]->addr();
246 : rva = pdb::TranslateAddressViaOmap(reorderer_->playback()->omap_from(),
247 E : rva);
248 :
249 : // Convert this to an absolute address using the base address from above.
250 E : uint64 abs = dll_info.base_address.value() + rva.value();
251 E : void* block_pointer = reinterpret_cast<void*>(abs);
252 :
253 E : event_data.calls[j].function = block_pointer;
254 E : }
255 :
256 E : event_record.Header.Class.Type = TRACE_BATCH_ENTER;
257 E : event_record.MofData = &raw_data;
258 E : event_record.MofLength = sizeof(raw_data);
259 E : if (!DispatchEvent(&event_record))
260 i : return false;
261 E : }
262 :
263 : // Simulate entry/exit pairs with the remaining blocks.
264 E : for (; i < blocks.size(); ++i) {
265 : // Get the address of this block as an RVA in the instrumented module.
266 E : RelativeAddress rva = blocks[i]->addr();
267 : rva = pdb::TranslateAddressViaOmap(reorderer_->playback()->omap_from(),
268 E : rva);
269 :
270 : // Convert this to an absolute address using the base address from above.
271 E : uint64 abs = dll_info.base_address.value() + rva.value();
272 E : void* block_pointer = reinterpret_cast<void*>(abs);
273 :
274 E : TraceEnterEventData event_data = {};
275 E : event_data.function = block_pointer;
276 :
277 : // Simulate an entry event.
278 E : event_record.Header.Class.Type = TRACE_ENTER_EVENT;
279 E : event_record.MofData = &event_data;
280 E : event_record.MofLength = sizeof(event_data);
281 E : if (!DispatchEvent(&event_record))
282 i : return false;
283 :
284 : // Simulate a corresponding exit event.
285 E : event_record.Header.Class.Type = TRACE_EXIT_EVENT;
286 E : if (!DispatchEvent(&event_record))
287 i : return false;
288 E : }
289 :
290 : // Simulate the thread and process detaching from the DLL.
291 E : event_record.Header.Class.Type = TRACE_THREAD_DETACH_EVENT;
292 E : event_record.MofData = &dll_data;
293 E : event_record.MofLength = sizeof(dll_data);
294 E : if (!DispatchEvent(&event_record))
295 i : return false;
296 :
297 E : event_record.Header.Class.Type = TRACE_PROCESS_DETACH_EVENT;
298 E : if (!DispatchEvent(&event_record))
299 i : return false;
300 :
301 : // Simulate the process ending.
302 E : event_handler_->OnProcessEnded(time, kProcessId);
303 :
304 E : return true;
305 E : }
306 :
307 : class ReordererTest : public testing::PELibUnitTest {
308 : public:
309 : typedef testing::PELibUnitTest Super;
310 :
311 : typedef block_graph::BlockGraph BlockGraph;
312 : typedef core::RelativeAddress RelativeAddress;
313 : typedef Reorderer::ImageLayout ImageLayout;
314 : typedef Reorderer::PEFile PEFile;
315 :
316 E : ReordererTest() : test_parse_engine_(NULL) {
317 E : }
318 :
319 E : void SetUp() OVERRIDE {
320 E : Super::SetUp();
321 :
322 : // Create the dummy trace file list.
323 E : Reorderer::TraceFileList trace_file_list;
324 E : trace_file_list.push_back(base::FilePath(L"foo"));
325 :
326 : // Set up the reorderer. These tests rely on
327 : // call_trace_instrumented_test_dll.dll, as generated by the test_data
328 : // project.
329 : const Reorderer::Flags kFlags = Reorderer::kFlagReorderCode |
330 E : Reorderer::kFlagReorderData;
331 : test_reorderer_.reset(new TestReorderer(
332 : testing::GetExeTestDataRelativePath(testing::kTestDllName),
333 : testing::GetExeTestDataRelativePath(
334 : testing::kCallTraceInstrumentedTestDllName),
335 : trace_file_list,
336 E : kFlags));
337 :
338 : // Setup the test parse engine and register it with the parser used
339 : // by the test reorderer. Note that ownership of the pointer is also
340 : // being passed.
341 E : ASSERT_TRUE(test_parse_engine_ == NULL);
342 E : test_parse_engine_ = new TestParseEngine(test_reorderer_.get());
343 E : ASSERT_TRUE(test_parse_engine_ != NULL);
344 E : test_reorderer_->parser().AddParseEngine(test_parse_engine_);
345 E : }
346 :
347 : // A reorderer will be initialized, in SetUp(), for each test run.
348 : scoped_ptr<TestReorderer> test_reorderer_;
349 :
350 : // The reorderer needs to be set up to use a custom parse engine before a
351 : // call to Reorder. This must be heap allocated and the responsibility for
352 : // deleting it rests with the parser.
353 : TestParseEngine* test_parse_engine_;
354 : };
355 :
356 : } // namespace
357 :
358 E : TEST_F(ReordererTest, ValidateCallbacks) {
359 E : MockOrderGenerator mock_order_generator;
360 :
361 : // Setup the expected calls.
362 E : InSequence s;
363 : EXPECT_CALL(mock_order_generator, OnProcessStarted(_, _))
364 E : .WillOnce(Return(true));
365 : EXPECT_CALL(mock_order_generator, OnCodeBlockEntry(_, _, _, _, _))
366 E : .WillRepeatedly(Return(true));
367 : EXPECT_CALL(mock_order_generator, OnProcessEnded(_, _))
368 E : .WillOnce(Return(true));
369 : EXPECT_CALL(mock_order_generator, CalculateReordering(_, _, _, _, _))
370 E : .WillOnce(Return(true));
371 :
372 : // Run the reorderer.
373 E : Reorderer::Order order;
374 E : PEFile pe_file;
375 E : BlockGraph block_graph;
376 E : ImageLayout image_layout(&block_graph);
377 : EXPECT_TRUE(test_reorderer_->Reorder(&mock_order_generator,
378 : &order,
379 : &pe_file,
380 E : &image_layout));
381 E : }
382 :
383 E : TEST_F(ReordererTest, Reorder) {
384 E : TestOrderGenerator test_order_generator;
385 :
386 : // Run the reorderer.
387 E : Reorderer::Order order;
388 E : PEFile pe_file;
389 E : BlockGraph block_graph;
390 E : ImageLayout image_layout(&block_graph);
391 : EXPECT_TRUE(test_reorderer_->Reorder(&test_order_generator,
392 : &order,
393 : &pe_file,
394 E : &image_layout));
395 :
396 : // We expect the order generator to have come up with the same list of
397 : // blocks that the parse engine used for generating dummy trace events.
398 : EXPECT_EQ(test_parse_engine_->blocks,
399 E : test_order_generator.blocks);
400 E : }
401 :
402 E : TEST(OrderTest, OrderConstructor) {
403 E : Reorderer::Order order;
404 E : EXPECT_TRUE(order.sections.empty());
405 E : EXPECT_TRUE(order.comment.empty());
406 E : }
407 :
408 E : TEST(OrderTest, SectionSpecConstructorsAndCopies) {
409 : // Create default constructed section spec.
410 E : Reorderer::Order::SectionSpec default_section_spec;
411 : EXPECT_EQ(Reorderer::Order::SectionSpec::kNewSectionId,
412 E : default_section_spec.id);
413 E : EXPECT_TRUE(default_section_spec.name.empty());
414 E : EXPECT_EQ(0U, default_section_spec.characteristics);
415 E : EXPECT_TRUE(default_section_spec.blocks.empty());
416 :
417 : // Create a customized section spec.
418 E : Reorderer::Order::SectionSpec other_section_spec;
419 E : EXPECT_TRUE(SectionSpecsAreEqual(default_section_spec, other_section_spec));
420 E : other_section_spec.id = 7;
421 E : other_section_spec.characteristics = 19;
422 E : other_section_spec.name = "other";
423 E : EXPECT_FALSE(SectionSpecsAreEqual(default_section_spec, other_section_spec));
424 :
425 : // Copy construct a section spec.
426 E : Reorderer::Order::SectionSpec copied_section_spec(other_section_spec);
427 E : EXPECT_TRUE(SectionSpecsAreEqual(other_section_spec, copied_section_spec));
428 E : EXPECT_FALSE(SectionSpecsAreEqual(copied_section_spec, default_section_spec));
429 :
430 : // Assign to the default section spec.
431 E : default_section_spec = other_section_spec;
432 E : EXPECT_TRUE(SectionSpecsAreEqual(copied_section_spec, default_section_spec));
433 E : }
434 :
435 E : TEST(OrderTest, BlockSpecConstructorsAndCopier) {
436 : static const BlockGraph::Block* kFauxBlockPtr =
437 : reinterpret_cast<BlockGraph::Block*>(0xCCCCCCCC);
438 :
439 : // Default construct a block spec.
440 E : Reorderer::Order::BlockSpec default_block_spec;
441 E : EXPECT_EQ(NULL, default_block_spec.block);
442 E : EXPECT_TRUE(default_block_spec.basic_block_offsets.empty());
443 :
444 : // Explicitly construct a block spec.
445 E : Reorderer::Order::BlockSpec explicit_block_spec(kFauxBlockPtr);
446 E : EXPECT_EQ(kFauxBlockPtr, explicit_block_spec.block);
447 E : EXPECT_TRUE(explicit_block_spec.basic_block_offsets.empty());
448 :
449 : // Default construct another block spec.
450 E : Reorderer::Order::BlockSpec block_spec;
451 E : EXPECT_TRUE(BlockSpecsAreEqual(default_block_spec, block_spec));
452 :
453 : // Modify the block spec.
454 E : block_spec.block = kFauxBlockPtr;
455 E : block_spec.basic_block_offsets.push_back(0);
456 E : block_spec.basic_block_offsets.push_back(8);
457 E : EXPECT_FALSE(BlockSpecsAreEqual(default_block_spec, block_spec));
458 E : EXPECT_FALSE(BlockSpecsAreEqual(explicit_block_spec, block_spec));
459 :
460 : // Copy construct a block spec.
461 E : Reorderer::Order::BlockSpec copied_block_spec(block_spec);
462 E : EXPECT_TRUE(BlockSpecsAreEqual(block_spec, copied_block_spec));
463 E : EXPECT_FALSE(BlockSpecsAreEqual(default_block_spec, block_spec));
464 E : EXPECT_FALSE(BlockSpecsAreEqual(explicit_block_spec, block_spec));
465 :
466 : // Assign to the explicit block spec.
467 E : explicit_block_spec = block_spec;
468 E : EXPECT_TRUE(BlockSpecsAreEqual(copied_block_spec, explicit_block_spec));
469 E : }
470 :
471 E : TEST(OrderTest, SerializeToJsonRoundTrip) {
472 : // Build a dummy block graph.
473 E : BlockGraph block_graph;
474 E : BlockGraph::Section* section1 = block_graph.AddSection(".text", 0);
475 E : BlockGraph::Section* section2 = block_graph.AddSection(".rdata", 0);
476 : BlockGraph::Block* block1 = block_graph.AddBlock(BlockGraph::CODE_BLOCK, 10,
477 E : "block1");
478 : BlockGraph::Block* block2 = block_graph.AddBlock(BlockGraph::DATA_BLOCK, 10,
479 E : "block2");
480 : BlockGraph::Block* block3 = block_graph.AddBlock(BlockGraph::DATA_BLOCK, 10,
481 E : "block3");
482 E : block1->set_section(section1->id());
483 E : block2->set_section(section2->id());
484 E : block3->set_section(section2->id());
485 :
486 : // Build a dummy image layout.
487 E : pe::ImageLayout layout(&block_graph);
488 E : pe::ImageLayout::SectionInfo section_info1 = {};
489 E : section_info1.name = section1->name();
490 E : section_info1.addr = core::RelativeAddress(0x1000);
491 E : section_info1.size = 0x1000;
492 E : section_info1.data_size = 0x1000;
493 E : layout.sections.push_back(section_info1);
494 :
495 E : pe::ImageLayout::SectionInfo section_info2 = {};
496 E : section_info2.name = section2->name();
497 E : section_info2.addr = core::RelativeAddress(0x2000);
498 E : section_info2.size = 0x1000;
499 E : section_info2.data_size = 0x1000;
500 E : layout.sections.push_back(section_info2);
501 :
502 : layout.blocks.InsertBlock(section_info1.addr,
503 E : block1);
504 : layout.blocks.InsertBlock(section_info2.addr,
505 E : block2);
506 : layout.blocks.InsertBlock(section_info2.addr + block2->size(),
507 E : block3);
508 :
509 : // Build a dummy order.
510 E : Reorderer::Order order;
511 E : order.comment = "This is a comment.";
512 E : order.sections.resize(2);
513 E : order.sections[0].id = section1->id();
514 E : order.sections[0].name = section1->name();
515 E : order.sections[0].characteristics = section1->characteristics();
516 E : order.sections[0].blocks.push_back(BlockSpec(block1));
517 E : order.sections[0].blocks.back().basic_block_offsets.push_back(0);
518 E : order.sections[0].blocks.back().basic_block_offsets.push_back(8);
519 E : order.sections[1].id = section2->id();
520 E : order.sections[1].name = section2->name();
521 E : order.sections[1].characteristics = section2->characteristics();
522 E : order.sections[1].blocks.push_back(BlockSpec(block2));
523 E : order.sections[1].blocks.push_back(BlockSpec(block3));
524 :
525 : base::FilePath module = testing::GetExeTestDataRelativePath(
526 E : testing::kTestDllName);
527 E : pe::PEFile pe_file;
528 E : ASSERT_TRUE(pe_file.Init(module));
529 :
530 : // Serialize the order.
531 E : base::FilePath temp_file;
532 E : ASSERT_TRUE(base::CreateTemporaryFile(&temp_file));
533 E : EXPECT_TRUE(order.SerializeToJSON(pe_file, temp_file, true));
534 :
535 : // Get the original module from the file.
536 E : base::FilePath orig_module;
537 E : EXPECT_TRUE(Reorderer::Order::GetOriginalModulePath(temp_file, &orig_module));
538 E : EXPECT_EQ(module, orig_module);
539 :
540 : // Deserialize it.
541 E : Reorderer::Order order2;
542 E : EXPECT_FALSE(OrdersAreEqual(order, order2));
543 E : EXPECT_TRUE(order2.LoadFromJSON(pe_file, layout, temp_file));
544 :
545 : // Expect them to be the same.
546 E : EXPECT_TRUE(OrdersAreEqual(order, order2));
547 :
548 E : EXPECT_TRUE(base::DeleteFile(temp_file, false));
549 E : }
550 :
551 : } // namespace reorder
|