1    :  // Copyright 2013 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/coff_file_writer.h"
  16    :  
  17    :  #include <cstring>
  18    :  
  19    :  #include "base/path_service.h"
  20    :  #include "base/files/file_util.h"
  21    :  #include "gmock/gmock.h"
  22    :  #include "gtest/gtest.h"
  23    :  #include "syzygy/block_graph/typed_block.h"
  24    :  #include "syzygy/core/unittest_util.h"
  25    :  #include "syzygy/pe/coff_decomposer.h"
  26    :  #include "syzygy/pe/coff_file.h"
  27    :  #include "syzygy/pe/pe_utils.h"
  28    :  #include "syzygy/pe/unittest_util.h"
  29    :  
  30    :  namespace pe {
  31    :  
  32    :  using block_graph::BlockGraph;
  33    :  using block_graph::ConstTypedBlock;
  34    :  using core::RelativeAddress;
  35    :  
  36    :  typedef testing::ApplicationTestBase CoffFileWriterTest;
  37    :  
  38  E :  TEST_F(CoffFileWriterTest, RedecomposeAfterWrite) {
  39    :    // Compute paths.
  40  E :    base::FilePath temp_dir;
  41  E :    ASSERT_NO_FATAL_FAILURE(CreateTemporaryDir(&temp_dir));
  42    :  
  43    :    base::FilePath image_path_0(
  44  E :        testing::GetExeTestDataRelativePath(testing::kTestDllCoffObjName));
  45  E :    base::FilePath image_path_1(temp_dir.Append(testing::kTestDllName));
  46    :  
  47    :    // Decompose the original image.
  48  E :    CoffFile image_file_0;
  49  E :    ASSERT_TRUE(image_file_0.Init(image_path_0));
  50    :  
  51  E :    CoffDecomposer decomposer_0(image_file_0);
  52  E :    block_graph::BlockGraph block_graph_0;
  53  E :    pe::ImageLayout image_layout_0(&block_graph_0);
  54  E :    ASSERT_TRUE(decomposer_0.Decompose(&image_layout_0));
  55    :  
  56    :    // Write temporary image file.
  57  E :    CoffFileWriter writer_0(&image_layout_0);
  58  E :    ASSERT_TRUE(writer_0.WriteImage(image_path_1));
  59    :  
  60    :    // Redecompose.
  61  E :    CoffFile image_file_1;
  62  E :    ASSERT_TRUE(image_file_1.Init(image_path_1));
  63    :  
  64  E :    CoffDecomposer decomposer_1(image_file_1);
  65  E :    block_graph::BlockGraph block_graph_1;
  66  E :    pe::ImageLayout image_layout_1(&block_graph_1);
  67  E :    ASSERT_TRUE(decomposer_1.Decompose(&image_layout_1));
  68    :  
  69    :    // Compare the results of the two decompositions.
  70  E :    EXPECT_EQ(image_layout_0.sections, image_layout_1.sections);
  71  E :    EXPECT_EQ(image_layout_0.blocks.size(), image_layout_1.blocks.size());
  72    :  
  73  E :    ConstTypedBlock<IMAGE_FILE_HEADER> file_header_0;
  74  E :    ConstTypedBlock<IMAGE_FILE_HEADER> file_header_1;
  75    :    BlockGraph::Block* headers_block_0 =
  76  E :        image_layout_0.blocks.GetBlockByAddress(RelativeAddress(0));
  77    :    BlockGraph::Block* headers_block_1 =
  78  E :        image_layout_1.blocks.GetBlockByAddress(RelativeAddress(0));
  79  E :    ASSERT_TRUE(headers_block_0 != NULL);
  80  E :    ASSERT_TRUE(headers_block_1 != NULL);
  81  E :    ASSERT_TRUE(file_header_0.Init(0, headers_block_0));
  82  E :    ASSERT_TRUE(file_header_1.Init(0, headers_block_1));
  83    :    EXPECT_TRUE(std::memcmp(&file_header_0[0], &file_header_1[0],
  84  E :                            sizeof(file_header_0[0])) == 0);
  85  E :  }
  86    :  
  87    :  }  // namespace pe