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