1    :  // Copyright 2014 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/transforms/pe_remove_empty_sections_transform.h"
  16    :  
  17    :  #include "gtest/gtest.h"
  18    :  #include "syzygy/pe/decomposer.h"
  19    :  #include "syzygy/pe/pe_utils.h"
  20    :  #include "syzygy/pe/unittest_util.h"
  21    :  
  22    :  namespace pe {
  23    :  namespace transforms {
  24    :  
  25    :  namespace {
  26    :  
  27    :  using block_graph::BlockGraph;
  28    :  typedef block_graph::BlockGraph::Section Section;
  29    :  
  30    :  // _asm ret
  31    :  const uint8 kCodeRet[] = { 0xC3 };
  32    :  
  33    :  // Dummy data.
  34    :  const uint8 kData[] = { 0x01, 0x02, 0x03, 0x04 };
  35    :  
  36    :  class PERemoveEmptySectionsTransformTest : public testing::Test {
  37    :   public:
  38    :    PERemoveEmptySectionsTransformTest()
  39    :        : image_layout_(&block_graph_),
  40  E :          block_header_(NULL) {
  41  E :    }
  42    :  
  43  E :    virtual void SetUp() override {
  44    :      // Set the block graph type to PE_IMAGE.
  45  E :      block_graph_.set_image_format(BlockGraph::PE_IMAGE);
  46    :  
  47    :      // Create the text section.
  48  E :      BlockGraph::Section* section_text = block_graph_.AddSection(".text", 0);
  49  E :      ASSERT_TRUE(section_text != NULL);
  50  E :      pe::ImageLayout::SectionInfo section_info_text = {};
  51  E :      section_info_text.name = section_text->name();
  52  E :      section_info_text.addr = core::RelativeAddress(0x1000);
  53  E :      section_info_text.size = 0x1000;
  54  E :      section_info_text.data_size = 0x1000;
  55  E :      image_layout_.sections.push_back(section_info_text);
  56    :  
  57    :      // Create the unused section.
  58  E :      BlockGraph::Section* section_unused = block_graph_.AddSection(".unused", 0);
  59  E :      ASSERT_TRUE(section_unused != NULL);
  60  E :      pe::ImageLayout::SectionInfo section_info_unused = {};
  61  E :      section_info_unused.name = section_unused->name();
  62  E :      section_info_unused.addr = core::RelativeAddress(0x2000);
  63  E :      section_info_unused.size = 0x1000;
  64  E :      section_info_unused.data_size = 0x1000;
  65  E :      image_layout_.sections.push_back(section_info_unused);
  66    :  
  67    :      // Create the DOS header.
  68    :      block_header_ = block_graph_.AddBlock(BlockGraph::DATA_BLOCK,
  69    :                                            sizeof(kData),
  70  E :                                            "header");
  71  E :      ASSERT_TRUE(block_header_ != NULL);
  72    :  
  73    :      // Create main function block.
  74    :      block_main_ =
  75  E :          block_graph_.AddBlock(BlockGraph::CODE_BLOCK, sizeof(kCodeRet), "main");
  76  E :      DCHECK_NE(reinterpret_cast<BlockGraph::Block*>(NULL), block_main_);
  77  E :      block_main_->SetData(kCodeRet, sizeof(kCodeRet));
  78  E :      block_main_->SetLabel(0, "main", BlockGraph::CODE_LABEL);
  79    :  
  80    :  
  81    :      // Put blocks into text section.
  82  E :      block_header_->set_section(section_text->id());
  83  E :      block_main_->set_section(section_text->id());
  84  E :    }
  85    :  
  86    :    PEFile pe_file_;
  87    :    testing::DummyTransformPolicy policy_;
  88    :    BlockGraph block_graph_;
  89    :    ImageLayout image_layout_;
  90    :    BlockGraph::Block* block_header_;
  91    :    BlockGraph::Block* block_main_;
  92    :  };
  93    :  
  94    :  }  // namespace
  95    :  
  96    :  
  97  E :  TEST_F(PERemoveEmptySectionsTransformTest, RemoveSection) {
  98    :    // Validate that both section exists.
  99  E :    Section* section_text = block_graph_.FindSection(".text");
 100  E :    Section* section_ununsed = block_graph_.FindSection(".unused");
 101  E :    ASSERT_TRUE(section_text != NULL);
 102  E :    ASSERT_TRUE(section_ununsed != NULL);
 103    :  
 104    :    // Apply the transform.
 105  E :    PERemoveEmptySectionsTransform transform;
 106    :    EXPECT_TRUE(block_graph::ApplyBlockGraphTransform(
 107  E :        &transform, &policy_, &block_graph_, block_header_));
 108    :  
 109    :    // Validate that .text still exists.
 110  E :    EXPECT_EQ(section_text, block_graph_.FindSection(".text"));
 111    :  
 112    :    // Validate that the unused section no longer exists.
 113  E :    section_ununsed = block_graph_.FindSection(".unused");
 114  E :    EXPECT_TRUE(section_ununsed == NULL);
 115  E :  }
 116    :  
 117    :  }  // namespace transforms
 118    :  }  // namespace pe