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/pe/pe_utils.h"
16 :
17 : #include "base/strings/string_split.h"
18 : #include "gmock/gmock.h"
19 : #include "gtest/gtest.h"
20 : #include "syzygy/block_graph/unittest_util.h"
21 : #include "syzygy/core/unittest_util.h"
22 : #include "syzygy/pe/pe_file.h"
23 : #include "syzygy/pe/unittest_util.h"
24 : #include "syzygy/pe/transforms/pe_add_imports_transform.h"
25 :
26 : namespace pe {
27 :
28 : using block_graph::BlockGraph;
29 : using core::RelativeAddress;
30 : using pe::transforms::PEAddImportsTransform;
31 :
32 : namespace {
33 :
34 : class PEUtilsTest : public testing::Test {
35 : public:
36 : PEUtilsTest()
37 : : nt_headers_block_(NULL),
38 : dos_header_block_(NULL),
39 : main_entry_point_block_(NULL),
40 : tls_initializer_block_(NULL),
41 : nt_headers_(NULL),
42 E : dos_header_(NULL) {
43 E : }
44 :
45 E : virtual void SetUp() {
46 E : block_graph_.set_image_format(BlockGraph::PE_IMAGE);
47 : // Create the NT headers block.
48 E : ASSERT_NO_FATAL_FAILURE(CreateNtHeadersBlock());
49 : // And the DOS header block.
50 E : ASSERT_NO_FATAL_FAILURE(CreateDosHeaderBlock());
51 : // And set-up some entry points.
52 E : ASSERT_NO_FATAL_FAILURE(CreateEntryPoints());
53 E : }
54 :
55 : protected:
56 : void CreateDosHeaderBlock();
57 : void CreateNtHeadersBlock();
58 : void CreateEntryPoints();
59 :
60 : testing::DummyTransformPolicy policy_;
61 : BlockGraph block_graph_;
62 :
63 : BlockGraph::Block* nt_headers_block_;
64 : BlockGraph::Block* dos_header_block_;
65 : BlockGraph::Block* main_entry_point_block_;
66 : BlockGraph::Block* tls_initializer_block_;
67 : IMAGE_DOS_HEADER* dos_header_;
68 : IMAGE_NT_HEADERS* nt_headers_;
69 : };
70 :
71 E : void PEUtilsTest::CreateNtHeadersBlock() {
72 : nt_headers_block_ = block_graph_.AddBlock(BlockGraph::DATA_BLOCK,
73 : sizeof(IMAGE_NT_HEADERS),
74 E : "NT Headers");
75 E : ASSERT_TRUE(nt_headers_block_ != NULL);
76 :
77 : nt_headers_ = reinterpret_cast<IMAGE_NT_HEADERS*>(
78 E : nt_headers_block_->AllocateData(sizeof(IMAGE_NT_HEADERS)));
79 E : ASSERT_TRUE(nt_headers_ != NULL);
80 :
81 E : nt_headers_->Signature = IMAGE_NT_SIGNATURE;
82 : nt_headers_->FileHeader.SizeOfOptionalHeader =
83 E : sizeof(nt_headers_->OptionalHeader);
84 E : nt_headers_->OptionalHeader.Magic = IMAGE_NT_OPTIONAL_HDR_MAGIC;
85 E : }
86 :
87 E : void PEUtilsTest::CreateDosHeaderBlock() {
88 : dos_header_block_ = block_graph_.AddBlock(BlockGraph::DATA_BLOCK,
89 : sizeof(IMAGE_DOS_HEADER),
90 E : "DOS Header");
91 E : ASSERT_TRUE(dos_header_block_ != NULL);
92 :
93 : dos_header_ = reinterpret_cast<IMAGE_DOS_HEADER*>(
94 E : dos_header_block_->AllocateData(dos_header_block_->size()));
95 E : ASSERT_TRUE(dos_header_ != NULL);
96 :
97 : // Set the correct magic constants in the manufactured DOS header.
98 E : dos_header_->e_magic = IMAGE_DOS_SIGNATURE;
99 : // Set the "DOS File Size" headers.
100 E : dos_header_->e_cblp = dos_header_block_->size() % 512;
101 E : dos_header_->e_cp = dos_header_block_->size() / 512;
102 E : if (dos_header_->e_cblp != 0)
103 E : dos_header_->e_cp++;
104 : // Set the header paragraph size.
105 E : dos_header_->e_cparhdr = dos_header_block_->size() / 16;
106 :
107 E : if (nt_headers_block_ != NULL) {
108 : // Set the NT headers reference.
109 : dos_header_block_->SetReference(
110 : offsetof(IMAGE_DOS_HEADER, e_lfanew),
111 : BlockGraph::Reference(BlockGraph::RELATIVE_REF,
112 : sizeof(RelativeAddress),
113 : nt_headers_block_,
114 E : 0, 0));
115 : }
116 E : }
117 :
118 E : void PEUtilsTest::CreateEntryPoints() {
119 : // Setup the main entry-point.
120 : main_entry_point_block_ = block_graph_.AddBlock(
121 E : BlockGraph::CODE_BLOCK, 1, "main_entry_point");
122 E : ASSERT_TRUE(main_entry_point_block_ != NULL);
123 : ASSERT_TRUE(nt_headers_block_->SetReference(
124 : offsetof(IMAGE_NT_HEADERS, OptionalHeader.AddressOfEntryPoint),
125 : BlockGraph::Reference(BlockGraph::RELATIVE_REF,
126 : BlockGraph::Reference::kMaximumSize,
127 E : main_entry_point_block_, 0, 0)));
128 :
129 : // Setup the TLS directory.
130 : static const size_t kTlsDirectoryOffset = offsetof(
131 : IMAGE_NT_HEADERS,
132 : OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_TLS].VirtualAddress);
133 : BlockGraph::Block* tls_directory_block = block_graph_.AddBlock(
134 : BlockGraph::DATA_BLOCK,
135 : sizeof(IMAGE_TLS_DIRECTORY),
136 E : "tls_directory");
137 E : ASSERT_TRUE(tls_directory_block != NULL);
138 E : ASSERT_TRUE(tls_directory_block->AllocateData(tls_directory_block->size()));
139 : nt_headers_->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_TLS].Size =
140 E : tls_directory_block->size();
141 : ASSERT_TRUE(nt_headers_block_->SetReference(
142 : kTlsDirectoryOffset,
143 : BlockGraph::Reference(BlockGraph::RELATIVE_REF,
144 : BlockGraph::Reference::kMaximumSize,
145 E : tls_directory_block, 0, 0)));
146 :
147 : // Setup the TLS callbacks table. Reserving enough space for one callback
148 : // and the trailing NULL sentinel.
149 : BlockGraph::Block* tls_callbacks_block = block_graph_.AddBlock(
150 : BlockGraph::DATA_BLOCK,
151 : 2 * BlockGraph::Reference::kMaximumSize,
152 E : "tls_callbacks");
153 E : ASSERT_TRUE(tls_callbacks_block != NULL);
154 E : ASSERT_TRUE(tls_callbacks_block->AllocateData(tls_callbacks_block->size()));
155 : nt_headers_->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_TLS].Size =
156 E : tls_directory_block->size();
157 : ASSERT_TRUE(tls_directory_block->SetReference(
158 : offsetof(IMAGE_TLS_DIRECTORY, AddressOfCallBacks),
159 : BlockGraph::Reference(BlockGraph::RELATIVE_REF,
160 : BlockGraph::Reference::kMaximumSize,
161 E : tls_callbacks_block, 0, 0)));
162 :
163 : // Add a TLS initializer.
164 : tls_initializer_block_ = block_graph_.AddBlock(
165 E : BlockGraph::CODE_BLOCK, 1, "tls_initializer");
166 E : ASSERT_TRUE(tls_initializer_block_ != NULL);
167 : ASSERT_TRUE(tls_callbacks_block->SetReference(
168 : 0,
169 : BlockGraph::Reference(BlockGraph::RELATIVE_REF,
170 : BlockGraph::Reference::kMaximumSize,
171 E : tls_initializer_block_, 0, 0)));
172 E : }
173 :
174 : } // namespace
175 :
176 E : TEST_F(PEUtilsTest, IsValidDosHeaderBlockSuccess) {
177 : // This DOS header should test valid.
178 E : EXPECT_TRUE(IsValidDosHeaderBlock(dos_header_block_));
179 E : }
180 :
181 E : TEST_F(PEUtilsTest, IsValidDosHeaderBlockNoDataFails) {
182 E : dos_header_block_->SetData(NULL, 0);
183 E : EXPECT_FALSE(IsValidDosHeaderBlock(dos_header_block_));
184 E : }
185 :
186 E : TEST_F(PEUtilsTest, IsValidDosHeaderBlockTooShortFails) {
187 E : dos_header_block_->ResizeData(sizeof(IMAGE_DOS_HEADER) - 1);
188 E : dos_header_block_->set_size(sizeof(IMAGE_DOS_HEADER) - 1);
189 E : EXPECT_FALSE(IsValidDosHeaderBlock(dos_header_block_));
190 E : }
191 :
192 E : TEST_F(PEUtilsTest, IsValidDosHeaderBlockInvalidMagicFails) {
193 E : ++dos_header_->e_magic;
194 E : EXPECT_FALSE(IsValidDosHeaderBlock(dos_header_block_));
195 E : }
196 :
197 E : TEST_F(PEUtilsTest, IsValidDosHeaderBlockInvalidDosFileSizeFails) {
198 E : dos_header_->e_cp = 0;
199 E : dos_header_->e_cblp = 0;
200 E : EXPECT_FALSE(IsValidDosHeaderBlock(dos_header_block_));
201 :
202 : // This is invalid, as there are zero pages, and thus no last page.
203 E : dos_header_->e_cblp = 10;
204 E : EXPECT_FALSE(IsValidDosHeaderBlock(dos_header_block_));
205 E : }
206 :
207 E : TEST_F(PEUtilsTest, IsValidDosHeaderBlockInvalidHeaderSizeFails) {
208 E : --dos_header_->e_cparhdr;
209 E : EXPECT_FALSE(IsValidDosHeaderBlock(dos_header_block_));
210 E : }
211 :
212 E : TEST_F(PEUtilsTest, IsValidDosHeaderBlockInvalidNTHeaderRefFails) {
213 : // Set the NT headers reference to a non-zero offset.
214 : dos_header_block_->SetReference(
215 : offsetof(IMAGE_DOS_HEADER, e_lfanew),
216 : BlockGraph::Reference(BlockGraph::RELATIVE_REF,
217 : sizeof(RelativeAddress),
218 : nt_headers_block_,
219 E : 10, 10));
220 E : EXPECT_FALSE(IsValidDosHeaderBlock(dos_header_block_));
221 E : }
222 :
223 E : TEST_F(PEUtilsTest, IsValidDosHeaderBlockNoNTHeaderRefFails) {
224 : // Clear the NT headers reference.
225 E : dos_header_block_->RemoveReference(offsetof(IMAGE_DOS_HEADER, e_lfanew));
226 E : EXPECT_FALSE(IsValidDosHeaderBlock(dos_header_block_));
227 E : }
228 :
229 E : TEST_F(PEUtilsTest, IsValidNtHeaderBlockSuccess) {
230 : // The NT headers are valid.
231 E : EXPECT_TRUE(IsValidNtHeadersBlock(nt_headers_block_));
232 E : }
233 :
234 E : TEST_F(PEUtilsTest, IsValidNtHeaderBlockInvalidSigFails) {
235 E : ++nt_headers_->Signature;
236 : // Invalid signature.
237 E : EXPECT_FALSE(IsValidNtHeadersBlock(nt_headers_block_));
238 E : }
239 :
240 E : TEST_F(PEUtilsTest, IsValidNtHeaderBlockInvalidOptionalSigFails) {
241 E : ++nt_headers_->OptionalHeader.Magic;
242 : // Invalid signature.
243 E : EXPECT_FALSE(IsValidNtHeadersBlock(nt_headers_block_));
244 E : }
245 :
246 E : TEST_F(PEUtilsTest, IsValidNtHeaderBlockInvalidOptionalSizeFails) {
247 E : ++nt_headers_->FileHeader.SizeOfOptionalHeader;
248 : // Invalid signature.
249 E : EXPECT_FALSE(IsValidNtHeadersBlock(nt_headers_block_));
250 E : }
251 :
252 E : TEST_F(PEUtilsTest, GetNtHeadersBlockFromDosHeaderBlock) {
253 : ASSERT_EQ(nt_headers_block_,
254 E : GetNtHeadersBlockFromDosHeaderBlock(dos_header_block_));
255 E : }
256 :
257 E : TEST_F(PEUtilsTest, GetNtHeadersBlockFromDosHeaderBlockConst) {
258 : ASSERT_EQ(nt_headers_block_,
259 : GetNtHeadersBlockFromDosHeaderBlock(
260 E : const_cast<const BlockGraph::Block*>(dos_header_block_)));
261 E : }
262 :
263 E : TEST_F(PEUtilsTest, GetSectionType) {
264 : // Only the characteristics of the header should matter in determining the
265 : // type of a section.
266 E : IMAGE_SECTION_HEADER code_header = {};
267 : code_header.Characteristics = IMAGE_SCN_CNT_CODE |
268 E : IMAGE_SCN_MEM_READ | IMAGE_SCN_MEM_EXECUTE;
269 :
270 : // Anything that is readable and isn't code will be identified as data.
271 E : IMAGE_SECTION_HEADER data_header = {};
272 : data_header.Characteristics = IMAGE_SCN_CNT_INITIALIZED_DATA |
273 : IMAGE_SCN_MEM_READ | IMAGE_SCN_MEM_WRITE |
274 E : IMAGE_SCN_ALIGN_4096BYTES;
275 E : IMAGE_SECTION_HEADER rdata_header = {};
276 : rdata_header.Characteristics = IMAGE_SCN_CNT_INITIALIZED_DATA |
277 : IMAGE_SCN_MEM_READ |
278 E : IMAGE_SCN_LNK_COMDAT;
279 E : IMAGE_SECTION_HEADER bss_header = {};
280 : bss_header.Characteristics = IMAGE_SCN_CNT_UNINITIALIZED_DATA |
281 : IMAGE_SCN_MEM_READ | IMAGE_SCN_MEM_WRITE |
282 E : IMAGE_SCN_ALIGN_2BYTES;
283 :
284 : // Anything that is neither readable nor code nor initialized data will be
285 : // flagged as unknown.
286 E : IMAGE_SECTION_HEADER strange_header = {};
287 : strange_header.Characteristics = IMAGE_SCN_CNT_UNINITIALIZED_DATA |
288 E : IMAGE_SCN_MEM_WRITE;
289 :
290 E : EXPECT_EQ(kSectionCode, GetSectionType(code_header));
291 E : EXPECT_EQ(kSectionData, GetSectionType(data_header));
292 E : EXPECT_EQ(kSectionData, GetSectionType(rdata_header));
293 E : EXPECT_EQ(kSectionData, GetSectionType(bss_header));
294 E : EXPECT_EQ(kSectionUnknown, GetSectionType(strange_header));
295 E : }
296 :
297 E : TEST_F(PEUtilsTest, AddLabelToBlock) {
298 : BlockGraph::Block* block =
299 E : block_graph_.AddBlock(BlockGraph::CODE_BLOCK, 0x100, "AddLabelToBlock");
300 E : ASSERT_TRUE(block != NULL);
301 :
302 : ASSERT_TRUE(AddLabelToBlock(0x10, "L1", BlockGraph::DEBUG_START_LABEL,
303 E : block));
304 : ASSERT_TRUE(AddLabelToBlock(0x20, "L2", BlockGraph::CALL_SITE_LABEL,
305 E : block));
306 :
307 E : BlockGraph::Label label;
308 E : ASSERT_TRUE(block->GetLabel(0x10, &label));
309 E : EXPECT_TRUE(label.name() == "L1");
310 E : EXPECT_EQ(BlockGraph::DEBUG_START_LABEL, label.attributes());
311 E : ASSERT_TRUE(block->GetLabel(0x20, &label));
312 E : EXPECT_TRUE(label.name() == "L2");
313 E : EXPECT_EQ(BlockGraph::CALL_SITE_LABEL, label.attributes());
314 :
315 E : ASSERT_TRUE(AddLabelToBlock(0x10, "L3", BlockGraph::CODE_LABEL, block));
316 E : ASSERT_TRUE(block->GetLabel(0x10, &label));
317 : EXPECT_EQ(BlockGraph::DEBUG_START_LABEL | BlockGraph::CODE_LABEL,
318 E : label.attributes());
319 :
320 E : std::vector<std::string> names;
321 E : base::SplitStringUsingSubstr(label.name(), kLabelNameSep, &names);
322 E : EXPECT_TRUE(std::find(names.begin(), names.end(), "L1") != names.end());
323 E : EXPECT_TRUE(std::find(names.begin(), names.end(), "L3") != names.end());
324 : EXPECT_EQ(BlockGraph::DEBUG_START_LABEL | BlockGraph::CODE_LABEL,
325 E : label.attributes());
326 :
327 : // Add a data label to L3. It should supercede the code label, and cause it
328 : // to be removed.
329 E : ASSERT_TRUE(AddLabelToBlock(0x10, "L4", BlockGraph::DATA_LABEL, block));
330 E : ASSERT_TRUE(block->GetLabel(0x10, &label));
331 : EXPECT_EQ(BlockGraph::DEBUG_START_LABEL | BlockGraph::DATA_LABEL,
332 E : label.attributes());
333 E : }
334 :
335 E : TEST_F(PEUtilsTest, CopySectionInfoToBlockGraph) {
336 E : pe::PEFile image_file;
337 :
338 : ASSERT_TRUE(image_file.Init(
339 E : testing::GetExeRelativePath(testing::kTestDllName)));
340 E : ASSERT_TRUE(CopySectionInfoToBlockGraph(image_file, &block_graph_));
341 :
342 E : const BlockGraph::SectionMap& sections = block_graph_.sections();
343 E : size_t num_sections = sections.size();
344 E : ASSERT_EQ(image_file.file_header()->NumberOfSections, num_sections);
345 :
346 E : for (size_t i = 0; i < num_sections; ++i) {
347 E : const IMAGE_SECTION_HEADER* header = image_file.section_header(i);
348 E : BlockGraph::SectionMap::const_iterator it = sections.find(i);
349 E : ASSERT_TRUE(it != sections.end());
350 :
351 E : EXPECT_EQ(i, it->second.id());
352 E : EXPECT_EQ(header->Characteristics, it->second.characteristics());
353 :
354 E : char name[IMAGE_SIZEOF_SHORT_NAME + 1] = {};
355 : strncpy(name, reinterpret_cast<const char*>(header->Name),
356 E : IMAGE_SIZEOF_SHORT_NAME);
357 E : EXPECT_EQ(name, it->second.name());
358 E : }
359 E : }
360 :
361 E : TEST_F(PEUtilsTest, GetExeEntryPoint) {
362 E : EntryPoint entry_point;
363 :
364 : // Get the entry point for an EXE.
365 E : nt_headers_->FileHeader.Characteristics &= ~IMAGE_FILE_DLL;
366 E : EXPECT_TRUE(GetExeEntryPoint(dos_header_block_, &entry_point));
367 E : EXPECT_EQ(main_entry_point_block_, entry_point.first);
368 E : EXPECT_EQ(0, entry_point.second);
369 :
370 : // Should return no entry points if the image is a DLL.
371 E : nt_headers_->FileHeader.Characteristics |= IMAGE_FILE_DLL;
372 E : EXPECT_TRUE(GetExeEntryPoint(dos_header_block_, &entry_point));
373 E : EXPECT_EQ(NULL, entry_point.first);
374 :
375 : // Should fail if the image is an EXE with no entry-point.
376 E : nt_headers_->FileHeader.Characteristics &= ~IMAGE_FILE_DLL;
377 : ASSERT_TRUE(nt_headers_block_->RemoveReference(
378 E : offsetof(IMAGE_NT_HEADERS, OptionalHeader.AddressOfEntryPoint)));
379 E : EXPECT_FALSE(GetExeEntryPoint(dos_header_block_, &entry_point));
380 E : }
381 :
382 E : TEST_F(PEUtilsTest, GetDllEntryPoint) {
383 E : EntryPoint entry_point;
384 :
385 : // Get the DLL entry point.
386 E : nt_headers_->FileHeader.Characteristics |= IMAGE_FILE_DLL;
387 E : EXPECT_TRUE(GetDllEntryPoint(dos_header_block_, &entry_point));
388 E : EXPECT_EQ(main_entry_point_block_, entry_point.first);
389 E : EXPECT_EQ(0, entry_point.second);
390 :
391 : // Should return no entry points if the image is an EXE.
392 E : nt_headers_->FileHeader.Characteristics &= ~IMAGE_FILE_DLL;
393 E : EXPECT_TRUE(GetDllEntryPoint(dos_header_block_, &entry_point));
394 E : EXPECT_EQ(NULL, entry_point.first);
395 :
396 : // Should return no entry points if the image is a DLL without an entry-point.
397 E : nt_headers_->FileHeader.Characteristics |= IMAGE_FILE_DLL;
398 : ASSERT_TRUE(nt_headers_block_->RemoveReference(
399 E : offsetof(IMAGE_NT_HEADERS, OptionalHeader.AddressOfEntryPoint)));
400 E : EXPECT_TRUE(GetDllEntryPoint(dos_header_block_, &entry_point));
401 E : EXPECT_EQ(NULL, entry_point.first);
402 E : }
403 :
404 E : TEST_F(PEUtilsTest, GetTlsInitializers) {
405 : // A container to store the entry-points.
406 E : EntryPointSet entry_points;
407 :
408 : // Get the entry points.
409 E : EXPECT_TRUE(GetTlsInitializers(dos_header_block_, &entry_points));
410 E : EXPECT_EQ(1U, entry_points.size());
411 E : EXPECT_EQ(tls_initializer_block_, entry_points.begin()->first);
412 E : }
413 :
414 E : TEST_F(PEUtilsTest, HasImportEntry) {
415 : // Creates an imported module.
416 E : transforms::ImportedModule module("foo.dll");
417 E : const char* kFooFunc = "foo_func";
418 : size_t function_foo = module.AddSymbol(
419 E : kFooFunc, transforms::ImportedModule::kAlwaysImport);
420 E : ASSERT_EQ(kFooFunc, module.GetSymbolName(function_foo));
421 :
422 : // Apply the transform to add this module import to the block-graph.
423 E : PEAddImportsTransform transform;
424 E : transform.AddModule(&module);
425 : ASSERT_TRUE(block_graph::ApplyBlockGraphTransform(
426 E : &transform, &policy_, &block_graph_, dos_header_block_));
427 :
428 : // Ensure that we can find this module, and that we can't find a
429 : // non-imported module.
430 E : bool has_import = false;
431 E : EXPECT_TRUE(HasImportEntry(dos_header_block_, "foo.dll", &has_import));
432 E : EXPECT_TRUE(has_import);
433 E : EXPECT_TRUE(HasImportEntry(dos_header_block_, "bar.dll", &has_import));
434 E : EXPECT_FALSE(has_import);
435 E : }
436 :
437 E : TEST_F(PEUtilsTest, RedirectReferences) {
438 : BlockGraph::Block* b1 = block_graph_.AddBlock(
439 E : BlockGraph::CODE_BLOCK, 12, "b1");
440 : BlockGraph::Block* b2 = block_graph_.AddBlock(
441 E : BlockGraph::CODE_BLOCK, 4, "b2");
442 : BlockGraph::Block* b3 = block_graph_.AddBlock(
443 E : BlockGraph::CODE_BLOCK, 4, "b3");
444 : BlockGraph::Block* b4 = block_graph_.AddBlock(
445 E : BlockGraph::CODE_BLOCK, 4, "b4");
446 :
447 : // Set up a little block-graph.
448 : // - b1 points to b2, b3 and b4. It acts as a fake PE structure.
449 : // - b2 points to block b3.
450 : // - Nothing points to b4.
451 :
452 E : b1->set_attribute(BlockGraph::PE_PARSED);
453 : b1->SetReference(0, BlockGraph::Reference(BlockGraph::ABSOLUTE_REF,
454 E : 4, b2, 0, 0));
455 : b1->SetReference(4, BlockGraph::Reference(BlockGraph::ABSOLUTE_REF,
456 E : 4, b3, 0, 0));
457 : b1->SetReference(8, BlockGraph::Reference(BlockGraph::ABSOLUTE_REF,
458 E : 4, b4, 0, 0));
459 :
460 : b2->SetReference(0, BlockGraph::Reference(BlockGraph::ABSOLUTE_REF,
461 E : 4, b3, 0, 0));
462 :
463 E : ASSERT_EQ(0u, b1->referrers().size());
464 E : ASSERT_EQ(1u, b2->referrers().size());
465 E : ASSERT_EQ(2u, b3->referrers().size());
466 E : ASSERT_EQ(1u, b4->referrers().size());
467 E : ASSERT_EQ(3u, b1->references().size());
468 E : ASSERT_EQ(1u, b2->references().size());
469 E : ASSERT_EQ(0u, b3->references().size());
470 E : ASSERT_EQ(0u, b4->references().size());
471 :
472 : // Set up a redirection from b4:0 to b2:0. Nothing except the PE block b1
473 : // references b4, so this shouldn't change anything.
474 E : ReferenceMap redirects1;
475 E : redirects1[ReferenceDest(b4, 0)] = ReferenceDest(b2, 0);
476 E : RedirectReferences(redirects1);
477 E : EXPECT_EQ(0u, b1->referrers().size());
478 E : EXPECT_EQ(1u, b2->referrers().size());
479 E : EXPECT_EQ(2u, b3->referrers().size());
480 E : EXPECT_EQ(1u, b4->referrers().size());
481 E : EXPECT_EQ(3u, b1->references().size());
482 E : EXPECT_EQ(1u, b2->references().size());
483 E : EXPECT_EQ(0u, b3->references().size());
484 E : EXPECT_EQ(0u, b4->references().size());
485 :
486 : // Set up a redirection from b3:0 to b4:0. This shouldn't affect the
487 : // references in b1, but should modify the reference in b2.
488 E : ReferenceMap redirects2;
489 E : redirects2[ReferenceDest(b3, 0)] = ReferenceDest(b4, 0);
490 E : RedirectReferences(redirects2);
491 E : EXPECT_EQ(0u, b1->referrers().size());
492 E : EXPECT_EQ(1u, b2->referrers().size());
493 E : EXPECT_EQ(1u, b3->referrers().size());
494 E : EXPECT_EQ(2u, b4->referrers().size());
495 E : EXPECT_EQ(3u, b1->references().size());
496 E : EXPECT_EQ(1u, b2->references().size());
497 E : EXPECT_EQ(0u, b3->references().size());
498 E : EXPECT_EQ(0u, b4->references().size());
499 E : }
500 :
501 : } // namespace pe
|