chrome/common/safe_browsing/pe_image_reader_win_unittest.cc

   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 <windows.h>
   6 #include <wintrust.h>
   7
   8 #include "base/files/file_path.h"
   9 #include "base/files/memory_mapped_file.h"
  10 #include "base/path_service.h"
  11 #include "chrome/common/chrome_paths.h"
  12 #include "chrome/common/safe_browsing/pe_image_reader_win.h"
  13 #include "testing/gmock/include/gmock/gmock.h"
  14 #include "testing/gtest/include/gtest/gtest.h"
  15
  16 using ::testing::_;
  17 using ::testing::Gt;
  18 using ::testing::NotNull;
  19 using ::testing::Return;
  20 using ::testing::StrictMock;
  21
  22 struct TestData {
  23   const char* filename;
  24   safe_browsing::PeImageReader::WordSize word_size;
  25   WORD machine_identifier;
  26   WORD optional_header_size;
  27   size_t number_of_sections;
  28   size_t number_of_debug_entries;
  29 };
  30
  31 // A test fixture parameterized on test data containing the name of a PE image
  32 // to parse and the expected values to be read from it. The file is read from
  33 // the src/chrome/test/data/safe_browsing directory.
  34 class PeImageReaderTest : public testing::TestWithParam<const TestData*> {
  35  protected:
  36   PeImageReaderTest() : expected_data_(GetParam()) {}
  37
  38   void SetUp() override {
  39     ASSERT_TRUE(PathService::Get(chrome::DIR_TEST_DATA, &data_file_path_));
  40     data_file_path_ = data_file_path_.AppendASCII("safe_browsing");
  41     data_file_path_ = data_file_path_.AppendASCII(expected_data_->filename);
  42
  43     ASSERT_TRUE(data_file_.Initialize(data_file_path_));
  44
  45     ASSERT_TRUE(image_reader_.Initialize(data_file_.data(),
  46                                          data_file_.length()));
  47   }
  48
  49   const TestData* expected_data_;
  50   base::FilePath data_file_path_;
  51   base::MemoryMappedFile data_file_;
  52   safe_browsing::PeImageReader image_reader_;
  53 };
  54
  55 TEST_P(PeImageReaderTest, GetWordSize) {
  56   EXPECT_EQ(expected_data_->word_size, image_reader_.GetWordSize());
  57 }
  58
  59 TEST_P(PeImageReaderTest, GetDosHeader) {
  60   const IMAGE_DOS_HEADER* dos_header = image_reader_.GetDosHeader();
  61   ASSERT_NE(reinterpret_cast<const IMAGE_DOS_HEADER*>(NULL), dos_header);
  62   EXPECT_EQ(IMAGE_DOS_SIGNATURE, dos_header->e_magic);
  63 }
  64
  65 TEST_P(PeImageReaderTest, GetCoffFileHeader) {
  66   const IMAGE_FILE_HEADER* file_header = image_reader_.GetCoffFileHeader();
  67   ASSERT_NE(reinterpret_cast<const IMAGE_FILE_HEADER*>(NULL), file_header);
  68   EXPECT_EQ(expected_data_->machine_identifier, file_header->Machine);
  69   EXPECT_EQ(expected_data_->optional_header_size,
  70             file_header->SizeOfOptionalHeader);
  71 }
  72
  73 TEST_P(PeImageReaderTest, GetOptionalHeaderData) {
  74   size_t optional_header_size = 0;
  75   const uint8_t* optional_header_data =
  76       image_reader_.GetOptionalHeaderData(&optional_header_size);
  77   ASSERT_NE(reinterpret_cast<const uint8_t*>(NULL), optional_header_data);
  78   EXPECT_EQ(expected_data_->optional_header_size, optional_header_size);
  79 }
  80
  81 TEST_P(PeImageReaderTest, GetNumberOfSections) {
  82   EXPECT_EQ(expected_data_->number_of_sections,
  83             image_reader_.GetNumberOfSections());
  84 }
  85
  86 TEST_P(PeImageReaderTest, GetSectionHeaderAt) {
  87   size_t number_of_sections = image_reader_.GetNumberOfSections();
  88   for (size_t i = 0; i < number_of_sections; ++i) {
  89     const IMAGE_SECTION_HEADER* section_header =
  90         image_reader_.GetSectionHeaderAt(i);
  91     ASSERT_NE(reinterpret_cast<const IMAGE_SECTION_HEADER*>(NULL),
  92               section_header);
  93   }
  94 }
  95
  96 TEST_P(PeImageReaderTest, InitializeFailTruncatedFile) {
  97   // Compute the size of all headers through the section headers.
  98   const IMAGE_SECTION_HEADER* last_section_header =
  99       image_reader_.GetSectionHeaderAt(image_reader_.GetNumberOfSections() - 1);
 100   const uint8_t* headers_end =
 101       reinterpret_cast<const uint8_t*>(last_section_header) +
 102       sizeof(*last_section_header);
 103   size_t header_size = headers_end - data_file_.data();
 104   safe_browsing::PeImageReader short_reader;
 105
 106   // Initialize should succeed when all headers are present.
 107   EXPECT_TRUE(short_reader.Initialize(data_file_.data(), header_size));
 108
 109   // But fail if anything is missing.
 110   for (size_t i = 0; i < header_size; ++i) {
 111     EXPECT_FALSE(short_reader.Initialize(data_file_.data(), i));
 112   }
 113 }
 114
 115 TEST_P(PeImageReaderTest, GetExportSection) {
 116   size_t section_size = 0;
 117   const uint8_t* export_section = image_reader_.GetExportSection(&section_size);
 118   ASSERT_NE(reinterpret_cast<const uint8_t*>(NULL), export_section);
 119   EXPECT_NE(0U, section_size);
 120 }
 121
 122 TEST_P(PeImageReaderTest, GetNumberOfDebugEntries) {
 123   EXPECT_EQ(expected_data_->number_of_debug_entries,
 124             image_reader_.GetNumberOfDebugEntries());
 125 }
 126
 127 TEST_P(PeImageReaderTest, GetDebugEntry) {
 128   size_t number_of_debug_entries = image_reader_.GetNumberOfDebugEntries();
 129   for (size_t i = 0; i < number_of_debug_entries; ++i) {
 130     const uint8_t* raw_data = NULL;
 131     size_t raw_data_size = 0;
 132     const IMAGE_DEBUG_DIRECTORY* entry =
 133         image_reader_.GetDebugEntry(i, &raw_data, &raw_data_size);
 134     EXPECT_NE(reinterpret_cast<const IMAGE_DEBUG_DIRECTORY*>(NULL), entry);
 135     EXPECT_NE(reinterpret_cast<const uint8_t*>(NULL), raw_data);
 136     EXPECT_NE(0U, raw_data_size);
 137   }
 138 }
 139
 140 namespace {
 141
 142 const TestData kTestData[] = {
 143   {
 144     "module_with_exports_x86.dll",
 145     safe_browsing::PeImageReader::WORD_SIZE_32,
 146     IMAGE_FILE_MACHINE_I386,
 147     sizeof(IMAGE_OPTIONAL_HEADER32),
 148     4,
 149     1,
 150   }, {
 151     "module_with_exports_x64.dll",
 152     safe_browsing::PeImageReader::WORD_SIZE_64,
 153     IMAGE_FILE_MACHINE_AMD64,
 154     sizeof(IMAGE_OPTIONAL_HEADER64),
 155     5,
 156     1,
 157   },
 158 };
 159
 160 }  // namespace
 161
 162 INSTANTIATE_TEST_CASE_P(WordSize32,
 163                         PeImageReaderTest,
 164                         testing::Values(&kTestData[0]));
 165 INSTANTIATE_TEST_CASE_P(WordSize64,
 166                         PeImageReaderTest,
 167                         testing::Values(&kTestData[1]));
 168
 169 // An object exposing a PeImageReader::EnumCertificatesCallback that invokes a
 170 // virtual OnCertificate() method. This method is suitable for mocking in tests.
 171 class CertificateReceiver {
 172  public:
 173   void* AsContext() { return this; }
 174   static bool OnCertificateCallback(uint16_t revision,
 175                                     uint16_t certificate_type,
 176                                     const uint8_t* certificate_data,
 177                                     size_t certificate_data_size,
 178                                     void* context) {
 179     return reinterpret_cast<CertificateReceiver*>(context)->OnCertificate(
 180         revision, certificate_type, certificate_data, certificate_data_size);
 181   }
 182
 183  protected:
 184   CertificateReceiver() {}
 185   virtual ~CertificateReceiver() {}
 186   virtual bool OnCertificate(uint16_t revision,
 187                              uint16_t certificate_type,
 188                              const uint8_t* certificate_data,
 189                              size_t certificate_data_size) = 0;
 190 };
 191
 192 class MockCertificateReceiver : public CertificateReceiver {
 193  public:
 194   MockCertificateReceiver() {}
 195   MOCK_METHOD4(OnCertificate, bool(uint16_t, uint16_t, const uint8_t*, size_t));
 196
 197  private:
 198   DISALLOW_COPY_AND_ASSIGN(MockCertificateReceiver);
 199 };
 200
 201 struct CertificateTestData {
 202   const char* filename;
 203   int num_signers;
 204 };
 205
 206 // A test fixture parameterized on test data containing the name of a PE image
 207 // to parse and the expected values to be read from it. The file is read from
 208 // the src/chrome/test/data/safe_browsing/download_protection directory.
 209 class PeImageReaderCertificateTest
 210     : public testing::TestWithParam<const CertificateTestData*> {
 211  protected:
 212   PeImageReaderCertificateTest() : expected_data_(GetParam()) {}
 213
 214   void SetUp() override {
 215     ASSERT_TRUE(PathService::Get(chrome::DIR_TEST_DATA, &data_file_path_));
 216     data_file_path_ = data_file_path_.AppendASCII("safe_browsing");
 217     data_file_path_ = data_file_path_.AppendASCII("download_protection");
 218     data_file_path_ = data_file_path_.AppendASCII(expected_data_->filename);
 219     ASSERT_TRUE(data_file_.Initialize(data_file_path_));
 220     ASSERT_TRUE(image_reader_.Initialize(data_file_.data(),
 221                                          data_file_.length()));
 222   }
 223
 224   const CertificateTestData* expected_data_;
 225   base::FilePath data_file_path_;
 226   base::MemoryMappedFile data_file_;
 227   safe_browsing::PeImageReader image_reader_;
 228 };
 229
 230 TEST_P(PeImageReaderCertificateTest, EnumCertificates) {
 231   StrictMock<MockCertificateReceiver> receiver;
 232   if (expected_data_->num_signers) {
 233     EXPECT_CALL(receiver, OnCertificate(WIN_CERT_REVISION_2_0,
 234                                         WIN_CERT_TYPE_PKCS_SIGNED_DATA,
 235                                         NotNull(),
 236                                         Gt(0U)))
 237         .Times(expected_data_->num_signers)
 238         .WillRepeatedly(Return(true));
 239   }
 240   EXPECT_TRUE(image_reader_.EnumCertificates(
 241       &CertificateReceiver::OnCertificateCallback, receiver.AsContext()));
 242 }
 243
 244 TEST_P(PeImageReaderCertificateTest, AbortEnum) {
 245   StrictMock<MockCertificateReceiver> receiver;
 246   if (expected_data_->num_signers) {
 247     // Return false for the first cert, thereby stopping the enumeration.
 248     EXPECT_CALL(receiver, OnCertificate(_, _, _, _)).WillOnce(Return(false));
 249     EXPECT_FALSE(image_reader_.EnumCertificates(
 250         &CertificateReceiver::OnCertificateCallback, receiver.AsContext()));
 251   } else {
 252     // An unsigned file always reports true with no invocations of the callback.
 253     EXPECT_TRUE(image_reader_.EnumCertificates(
 254         &CertificateReceiver::OnCertificateCallback, receiver.AsContext()));
 255   }
 256 }
 257
 258 namespace {
 259
 260 const CertificateTestData kCertificateTestData[] = {
 261   {
 262     "signed.exe",
 263     1,
 264   }, {
 265     "unsigned.exe",
 266     0,
 267   }, {
 268     "wow_helper.exe",
 269     1,
 270   }, {
 271     "signed_twice.exe",
 272     2,
 273   },
 274 };
 275
 276 }  // namespace
 277
 278 INSTANTIATE_TEST_CASE_P(SignedExe,
 279                         PeImageReaderCertificateTest,
 280                         testing::Values(&kCertificateTestData[0]));
 281 INSTANTIATE_TEST_CASE_P(UnsignedExe,
 282                         PeImageReaderCertificateTest,
 283                         testing::Values(&kCertificateTestData[1]));
 284 INSTANTIATE_TEST_CASE_P(WowHelperExe,
 285                         PeImageReaderCertificateTest,
 286                         testing::Values(&kCertificateTestData[2]));
 287 INSTANTIATE_TEST_CASE_P(SignedTwiceExe,
 288                         PeImageReaderCertificateTest,
 289                         testing::Values(&kCertificateTestData[3]));