Home | History | Annotate | Download | only in safe_browsing 
      1 // Copyright 2014 The Chromium Authors. All rights reserved.
      2 // Use of this source code is governed by a BSD-style license that can be
      3 // found in the LICENSE file.
      4 
      5 #include "base/files/file_path.h"
      6 #include "base/files/memory_mapped_file.h"
      7 #include "base/path_service.h"
      8 #include "chrome/browser/safe_browsing/pe_image_reader_win.h"
      9 #include "chrome/common/chrome_paths.h"
     10 #include "testing/gtest/include/gtest/gtest.h"
     11 
     12 struct TestData {
     13   const char* filename;
     14   safe_browsing::PeImageReader::WordSize word_size;
     15   WORD machine_identifier;
     16   WORD optional_header_size;
     17   size_t number_of_sections;
     18   size_t number_of_debug_entries;
     19 };
     20 
     21 // A test fixture parameterized on test data containing the name of a PE image
     22 // to parse and the expected values to be read from it. The file is read from
     23 // the src/chrome/test/data/safe_browsing directory.
     24 class PeImageReaderTest : public testing::TestWithParam<const TestData*> {
     25  protected:
     26   PeImageReaderTest() : expected_data_(GetParam()) {}
     27 
     28   virtual void SetUp() OVERRIDE {
     29     ASSERT_TRUE(PathService::Get(chrome::DIR_TEST_DATA, &data_file_path_));
     30     data_file_path_ = data_file_path_.AppendASCII("safe_browsing");
     31     data_file_path_ = data_file_path_.AppendASCII(expected_data_->filename);
     32 
     33     ASSERT_TRUE(data_file_.Initialize(data_file_path_));
     34 
     35     ASSERT_TRUE(image_reader_.Initialize(data_file_.data(),
     36                                          data_file_.length()));
     37   }
     38 
     39   const TestData* expected_data_;
     40   base::FilePath data_file_path_;
     41   base::MemoryMappedFile data_file_;
     42   safe_browsing::PeImageReader image_reader_;
     43 };
     44 
     45 TEST_P(PeImageReaderTest, GetWordSize) {
     46   EXPECT_EQ(expected_data_->word_size, image_reader_.GetWordSize());
     47 }
     48 
     49 TEST_P(PeImageReaderTest, GetDosHeader) {
     50   const IMAGE_DOS_HEADER* dos_header = image_reader_.GetDosHeader();
     51   ASSERT_NE(reinterpret_cast<const IMAGE_DOS_HEADER*>(NULL), dos_header);
     52   EXPECT_EQ(IMAGE_DOS_SIGNATURE, dos_header->e_magic);
     53 }
     54 
     55 TEST_P(PeImageReaderTest, GetCoffFileHeader) {
     56   const IMAGE_FILE_HEADER* file_header = image_reader_.GetCoffFileHeader();
     57   ASSERT_NE(reinterpret_cast<const IMAGE_FILE_HEADER*>(NULL), file_header);
     58   EXPECT_EQ(expected_data_->machine_identifier, file_header->Machine);
     59   EXPECT_EQ(expected_data_->optional_header_size,
     60             file_header->SizeOfOptionalHeader);
     61 }
     62 
     63 TEST_P(PeImageReaderTest, GetOptionalHeaderData) {
     64   size_t optional_header_size = 0;
     65   const uint8_t* optional_header_data =
     66       image_reader_.GetOptionalHeaderData(&optional_header_size);
     67   ASSERT_NE(reinterpret_cast<const uint8_t*>(NULL), optional_header_data);
     68   EXPECT_EQ(expected_data_->optional_header_size, optional_header_size);
     69 }
     70 
     71 TEST_P(PeImageReaderTest, GetNumberOfSections) {
     72   EXPECT_EQ(expected_data_->number_of_sections,
     73             image_reader_.GetNumberOfSections());
     74 }
     75 
     76 TEST_P(PeImageReaderTest, GetSectionHeaderAt) {
     77   size_t number_of_sections = image_reader_.GetNumberOfSections();
     78   for (size_t i = 0; i < number_of_sections; ++i) {
     79     const IMAGE_SECTION_HEADER* section_header =
     80         image_reader_.GetSectionHeaderAt(i);
     81     ASSERT_NE(reinterpret_cast<const IMAGE_SECTION_HEADER*>(NULL),
     82               section_header);
     83   }
     84 }
     85 
     86 TEST_P(PeImageReaderTest, InitializeFailTruncatedFile) {
     87   // Compute the size of all headers through the section headers.
     88   const IMAGE_SECTION_HEADER* last_section_header =
     89       image_reader_.GetSectionHeaderAt(image_reader_.GetNumberOfSections() - 1);
     90   const uint8_t* headers_end =
     91       reinterpret_cast<const uint8_t*>(last_section_header) +
     92       sizeof(*last_section_header);
     93   size_t header_size = headers_end - data_file_.data();
     94   safe_browsing::PeImageReader short_reader;
     95 
     96   // Initialize should succeed when all headers are present.
     97   EXPECT_TRUE(short_reader.Initialize(data_file_.data(), header_size));
     98 
     99   // But fail if anything is missing.
    100   for (size_t i = 0; i < header_size; ++i) {
    101     EXPECT_FALSE(short_reader.Initialize(data_file_.data(), i));
    102   }
    103 }
    104 
    105 TEST_P(PeImageReaderTest, GetExportSection) {
    106   size_t section_size = 0;
    107   const uint8_t* export_section = image_reader_.GetExportSection(&section_size);
    108   ASSERT_NE(reinterpret_cast<const uint8_t*>(NULL), export_section);
    109   EXPECT_NE(0U, section_size);
    110 }
    111 
    112 TEST_P(PeImageReaderTest, GetNumberOfDebugEntries) {
    113   EXPECT_EQ(expected_data_->number_of_debug_entries,
    114             image_reader_.GetNumberOfDebugEntries());
    115 }
    116 
    117 TEST_P(PeImageReaderTest, GetDebugEntry) {
    118   size_t number_of_debug_entries = image_reader_.GetNumberOfDebugEntries();
    119   for (size_t i = 0; i < number_of_debug_entries; ++i) {
    120     const uint8_t* raw_data = NULL;
    121     size_t raw_data_size = 0;
    122     const IMAGE_DEBUG_DIRECTORY* entry =
    123         image_reader_.GetDebugEntry(i, &raw_data, &raw_data_size);
    124     EXPECT_NE(reinterpret_cast<const IMAGE_DEBUG_DIRECTORY*>(NULL), entry);
    125     EXPECT_NE(reinterpret_cast<const uint8_t*>(NULL), raw_data);
    126     EXPECT_NE(0U, raw_data_size);
    127   }
    128 }
    129 
    130 namespace {
    131 
    132 const TestData kTestData[] = {
    133   {
    134     "module_with_exports_x86.dll",
    135     safe_browsing::PeImageReader::WORD_SIZE_32,
    136     IMAGE_FILE_MACHINE_I386,
    137     sizeof(IMAGE_OPTIONAL_HEADER32),
    138     4,
    139     1,
    140   }, {
    141     "module_with_exports_x64.dll",
    142     safe_browsing::PeImageReader::WORD_SIZE_64,
    143     IMAGE_FILE_MACHINE_AMD64,
    144     sizeof(IMAGE_OPTIONAL_HEADER64),
    145     5,
    146     1,
    147   },
    148 };
    149 
    150 }  // namespace
    151 
    152 INSTANTIATE_TEST_CASE_P(WordSize32,
    153                         PeImageReaderTest,
    154                         testing::Values(&kTestData[0]));
    155 INSTANTIATE_TEST_CASE_P(WordSize64,
    156                         PeImageReaderTest,
    157                         testing::Values(&kTestData[1]));
    158