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 12struct 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. 24class 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 45TEST_P(PeImageReaderTest, GetWordSize) { 46 EXPECT_EQ(expected_data_->word_size, image_reader_.GetWordSize()); 47} 48 49TEST_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 55TEST_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 63TEST_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 71TEST_P(PeImageReaderTest, GetNumberOfSections) { 72 EXPECT_EQ(expected_data_->number_of_sections, 73 image_reader_.GetNumberOfSections()); 74} 75 76TEST_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 86TEST_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 105TEST_P(PeImageReaderTest, GetExportSection) { 106 size_t section_size = 0; 107 const uint8_t* export_section = image_reader_.GetExportSection(§ion_size); 108 ASSERT_NE(reinterpret_cast<const uint8_t*>(NULL), export_section); 109 EXPECT_NE(0U, section_size); 110} 111 112TEST_P(PeImageReaderTest, GetNumberOfDebugEntries) { 113 EXPECT_EQ(expected_data_->number_of_debug_entries, 114 image_reader_.GetNumberOfDebugEntries()); 115} 116 117TEST_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 130namespace { 131 132const 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 152INSTANTIATE_TEST_CASE_P(WordSize32, 153 PeImageReaderTest, 154 testing::Values(&kTestData[0])); 155INSTANTIATE_TEST_CASE_P(WordSize64, 156 PeImageReaderTest, 157 testing::Values(&kTestData[1])); 158