1    :  // Copyright 2011 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    :  // Contains unittests for the functionality defined in image_layout.h.
  16    :  
  17    :  #include "syzygy/pe/image_layout.h"
  18    :  
  19    :  #include "gmock/gmock.h"
  20    :  #include "syzygy/core/unittest_util.h"
  21    :  #include "syzygy/pe/decomposer.h"
  22    :  #include "syzygy/pe/unittest_util.h"
  23    :  
  24    :  namespace pe {
  25    :  
  26    :  using block_graph::BlockGraph;
  27    :  using core::RelativeAddress;
  28    :  
  29    :  namespace {
  30    :  
  31    :  class ImageLayoutTest: public testing::PELibUnitTest {
  32    :    // Insert your customizations here.
  33    :  };
  34    :  
  35    :  // Returns true if the given sections are equivalent, ignoring the data_size.
  36    :  bool SectionsAreEqual(const ImageLayout::SectionInfo& a,
  37  E :                        const ImageLayout::SectionInfo& b) {
  38    :    return a.name == b.name && a.addr == b.addr && a.size == b.size &&
  39  E :        a.characteristics == b.characteristics;
  40  E :  }
  41    :  
  42    :  }  // namespace
  43    :  
  44  E :  TEST_F(ImageLayoutTest, BuildCanonicalImageLayout) {
  45  E :    base::FilePath image_path(testing::GetExeRelativePath(testing::kTestDllName));
  46  E :    PEFile image_file;
  47    :  
  48  E :    ASSERT_TRUE(image_file.Init(image_path));
  49    :  
  50  E :    Decomposer decomposer(image_file);
  51  E :    BlockGraph block_graph;
  52  E :    ImageLayout image_layout(&block_graph);
  53  E :    ASSERT_TRUE(decomposer.Decompose(&image_layout));
  54    :  
  55  E :    ImageLayout canonical_image_layout(&block_graph);
  56  E :    EXPECT_TRUE(BuildCanonicalImageLayout(&canonical_image_layout));
  57    :  
  58    :    // We expect the sections to be in the same order and have the same size and
  59    :    // addresses. We do not check the data size, as it is effectively impossible
  60    :    // for us to perform the exact same logic as the original toolchain here, and
  61    :    // we often see different results. In fact, we are more aggressive at trimming
  62    :    // NULLS from the end of a section, especially when it comes to .relocs.
  63    :    EXPECT_EQ(image_layout.sections.size(),
  64  E :              canonical_image_layout.sections.size());
  65  E :    for (size_t i = 0; i < image_layout.sections.size(); ++i) {
  66    :      EXPECT_TRUE(SectionsAreEqual(image_layout.sections[i],
  67  E :                                   canonical_image_layout.sections[i]));
  68  E :    }
  69    :  
  70    :    BlockGraph::AddressSpace::RangeMapConstIter block_it1 =
  71  E :        image_layout.blocks.begin();
  72    :    BlockGraph::AddressSpace::RangeMapConstIter block_it2 =
  73  E :        canonical_image_layout.blocks.begin();
  74  E :    while (true) {
  75  E :      if (block_it1 == image_layout.blocks.end())
  76  E :        break;
  77  E :      if (block_it2 == canonical_image_layout.blocks.end())
  78  i :        break;
  79    :  
  80    :      // We expect the same block to be mapped to the same position in each
  81    :      // image.
  82  E :      EXPECT_EQ(block_it1->first, block_it2->first);
  83  E :      EXPECT_EQ(block_it1->second, block_it2->second);
  84    :  
  85  E :      ++block_it1;
  86  E :      ++block_it2;
  87  E :    }
  88  E :  }
  89    :  
  90  E :  TEST_F(ImageLayoutTest, CopyImageLayoutWithoutPadding) {
  91  E :    base::FilePath image_path(testing::GetExeRelativePath(testing::kTestDllName));
  92  E :    PEFile image_file;
  93    :  
  94  E :    ASSERT_TRUE(image_file.Init(image_path));
  95    :  
  96  E :    Decomposer decomposer(image_file);
  97  E :    BlockGraph block_graph;
  98  E :    ImageLayout orig_image_layout(&block_graph);
  99  E :    ASSERT_TRUE(decomposer.Decompose(&orig_image_layout));
 100    :  
 101  E :    ImageLayout image_layout(&block_graph);
 102  E :    EXPECT_TRUE(CopyImageLayoutWithoutPadding(orig_image_layout, &image_layout));
 103    :  
 104  E :    EXPECT_LT(image_layout.blocks.size(), orig_image_layout.blocks.size());
 105  E :    EXPECT_EQ(image_layout.blocks.size(), block_graph.blocks().size());
 106    :  
 107    :    BlockGraph::BlockMap::const_iterator block_it =
 108  E :        block_graph.blocks().begin();
 109  E :    for (; block_it != block_graph.blocks().end(); ++block_it) {
 110  E :      const BlockGraph::Block* block = &(block_it->second);
 111  E :      EXPECT_EQ(0u, block->attributes() & BlockGraph::PADDING_BLOCK);
 112  E :    }
 113  E :  }
 114    :  
 115    :  }  // namespace pe