chrome/common/safe_browsing/pe_image_reader_win_unittest.cc - chromium/src - Git at Google

 // Copyright 2014 The Chromium Authors. All rights reserved.
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.

 #include <windows.h>
 #include <stddef.h>
 #include <stdint.h>
 #include <wintrust.h>

 #include "base/files/file_path.h"
 #include "base/files/memory_mapped_file.h"
 #include "base/macros.h"
 #include "base/path_service.h"
 #include "chrome/common/chrome_paths.h"
 #include "chrome/common/safe_browsing/pe_image_reader_win.h"
 #include "testing/gmock/include/gmock/gmock.h"
 #include "testing/gtest/include/gtest/gtest.h"

 using ::testing::_;
 using ::testing::Gt;
 using ::testing::NotNull;
 using ::testing::Return;
 using ::testing::StrictMock;

 struct TestData {
   const char* filename;
   safe_browsing::PeImageReader::WordSize word_size;
   WORD machine_identifier;
   WORD optional_header_size;
   size_t number_of_sections;
   size_t number_of_debug_entries;
 };

 // A test fixture parameterized on test data containing the name of a PE image
 // to parse and the expected values to be read from it. The file is read from
 // the src/chrome/test/data/safe_browsing directory.
 class PeImageReaderTest : public testing::TestWithParam<const TestData*> {
  protected:
   PeImageReaderTest() : expected_data_(GetParam()) {}

   void SetUp() override {
     ASSERT_TRUE(
         base::PathService::Get(chrome::DIR_TEST_DATA, &data_file_path_));
     data_file_path_ = data_file_path_.AppendASCII("safe_browsing");
     data_file_path_ = data_file_path_.AppendASCII(expected_data_->filename);

     ASSERT_TRUE(data_file_.Initialize(data_file_path_));

     ASSERT_TRUE(image_reader_.Initialize(data_file_.data(),
                                          data_file_.length()));
   }

   const TestData* expected_data_;
   base::FilePath data_file_path_;
   base::MemoryMappedFile data_file_;
   safe_browsing::PeImageReader image_reader_;
 };

 TEST_P(PeImageReaderTest, GetWordSize) {
   EXPECT_EQ(expected_data_->word_size, image_reader_.GetWordSize());
 }

 TEST_P(PeImageReaderTest, GetDosHeader) {
   const IMAGE_DOS_HEADER* dos_header = image_reader_.GetDosHeader();
   ASSERT_NE(reinterpret_cast<const IMAGE_DOS_HEADER*>(NULL), dos_header);
   EXPECT_EQ(IMAGE_DOS_SIGNATURE, dos_header->e_magic);
 }

 TEST_P(PeImageReaderTest, GetCoffFileHeader) {
   const IMAGE_FILE_HEADER* file_header = image_reader_.GetCoffFileHeader();
   ASSERT_NE(reinterpret_cast<const IMAGE_FILE_HEADER*>(NULL), file_header);
   EXPECT_EQ(expected_data_->machine_identifier, file_header->Machine);
   EXPECT_EQ(expected_data_->optional_header_size,
             file_header->SizeOfOptionalHeader);
 }

 TEST_P(PeImageReaderTest, GetOptionalHeaderData) {
   size_t optional_header_size = 0;
   const uint8_t* optional_header_data =
       image_reader_.GetOptionalHeaderData(&optional_header_size);
   ASSERT_NE(reinterpret_cast<const uint8_t*>(NULL), optional_header_data);
   EXPECT_EQ(expected_data_->optional_header_size, optional_header_size);
 }

 TEST_P(PeImageReaderTest, GetNumberOfSections) {
   EXPECT_EQ(expected_data_->number_of_sections,
             image_reader_.GetNumberOfSections());
 }

 TEST_P(PeImageReaderTest, GetSectionHeaderAt) {
   size_t number_of_sections = image_reader_.GetNumberOfSections();
   for (size_t i = 0; i < number_of_sections; ++i) {
     const IMAGE_SECTION_HEADER* section_header =
         image_reader_.GetSectionHeaderAt(i);
     ASSERT_NE(reinterpret_cast<const IMAGE_SECTION_HEADER*>(NULL),
               section_header);
   }
 }

 TEST_P(PeImageReaderTest, InitializeFailTruncatedFile) {
   // Compute the size of all headers through the section headers.
   const IMAGE_SECTION_HEADER* last_section_header =
       image_reader_.GetSectionHeaderAt(image_reader_.GetNumberOfSections() - 1);
   const uint8_t* headers_end =
       reinterpret_cast<const uint8_t*>(last_section_header) +
       sizeof(*last_section_header);
   size_t header_size = headers_end - data_file_.data();
   safe_browsing::PeImageReader short_reader;

   // Initialize should succeed when all headers are present.
   EXPECT_TRUE(short_reader.Initialize(data_file_.data(), header_size));

   // But fail if anything is missing.
   for (size_t i = 0; i < header_size; ++i) {
     EXPECT_FALSE(short_reader.Initialize(data_file_.data(), i));
   }
 }

 TEST_P(PeImageReaderTest, GetExportSection) {
   size_t section_size = 0;
   const uint8_t* export_section = image_reader_.GetExportSection(&section_size);
   ASSERT_NE(reinterpret_cast<const uint8_t*>(NULL), export_section);
   EXPECT_NE(0U, section_size);
 }

 TEST_P(PeImageReaderTest, GetNumberOfDebugEntries) {
   EXPECT_EQ(expected_data_->number_of_debug_entries,
             image_reader_.GetNumberOfDebugEntries());
 }

 TEST_P(PeImageReaderTest, GetDebugEntry) {
   size_t number_of_debug_entries = image_reader_.GetNumberOfDebugEntries();
   for (size_t i = 0; i < number_of_debug_entries; ++i) {
     const uint8_t* raw_data = NULL;
     size_t raw_data_size = 0;
     const IMAGE_DEBUG_DIRECTORY* entry =
         image_reader_.GetDebugEntry(i, &raw_data, &raw_data_size);
     EXPECT_NE(reinterpret_cast<const IMAGE_DEBUG_DIRECTORY*>(NULL), entry);
     EXPECT_NE(reinterpret_cast<const uint8_t*>(NULL), raw_data);
     EXPECT_NE(0U, raw_data_size);
   }
 }

 namespace {

 const TestData kTestData[] = {
   {
     "module_with_exports_x86.dll",
     safe_browsing::PeImageReader::WORD_SIZE_32,
     IMAGE_FILE_MACHINE_I386,
     sizeof(IMAGE_OPTIONAL_HEADER32),
     4,
     1,
   }, {
     "module_with_exports_x64.dll",
     safe_browsing::PeImageReader::WORD_SIZE_64,
     IMAGE_FILE_MACHINE_AMD64,
     sizeof(IMAGE_OPTIONAL_HEADER64),
     5,
     1,
   },
 };

 }  // namespace

 INSTANTIATE_TEST_SUITE_P(WordSize32,
                          PeImageReaderTest,
                          testing::Values(&kTestData[0]));
 INSTANTIATE_TEST_SUITE_P(WordSize64,
                          PeImageReaderTest,
                          testing::Values(&kTestData[1]));

 // An object exposing a PeImageReader::EnumCertificatesCallback that invokes a
 // virtual OnCertificate() method. This method is suitable for mocking in tests.
 class CertificateReceiver {
  public:
   void* AsContext() { return this; }
   static bool OnCertificateCallback(uint16_t revision,
                                     uint16_t certificate_type,
                                     const uint8_t* certificate_data,
                                     size_t certificate_data_size,
                                     void* context) {
     return reinterpret_cast<CertificateReceiver*>(context)->OnCertificate(
         revision, certificate_type, certificate_data, certificate_data_size);
   }

  protected:
   CertificateReceiver() {}
   virtual ~CertificateReceiver() {}
   virtual bool OnCertificate(uint16_t revision,
                              uint16_t certificate_type,
                              const uint8_t* certificate_data,
                              size_t certificate_data_size) = 0;
 };

 class MockCertificateReceiver : public CertificateReceiver {
  public:
   MockCertificateReceiver() {}
   MOCK_METHOD4(OnCertificate, bool(uint16_t, uint16_t, const uint8_t*, size_t));

  private:
   DISALLOW_COPY_AND_ASSIGN(MockCertificateReceiver);
 };

 struct CertificateTestData {
   const char* filename;
   int num_signers;
 };

 // A test fixture parameterized on test data containing the name of a PE image
 // to parse and the expected values to be read from it. The file is read from
 // the src/chrome/test/data/safe_browsing/download_protection directory.
 class PeImageReaderCertificateTest
     : public testing::TestWithParam<const CertificateTestData*> {
  protected:
   PeImageReaderCertificateTest() : expected_data_(GetParam()) {}

   void SetUp() override {
     ASSERT_TRUE(
         base::PathService::Get(chrome::DIR_TEST_DATA, &data_file_path_));
     data_file_path_ = data_file_path_.AppendASCII("safe_browsing");
     data_file_path_ = data_file_path_.AppendASCII("download_protection");
     data_file_path_ = data_file_path_.AppendASCII(expected_data_->filename);
     ASSERT_TRUE(data_file_.Initialize(data_file_path_));
     ASSERT_TRUE(image_reader_.Initialize(data_file_.data(),
                                          data_file_.length()));
   }

   const CertificateTestData* expected_data_;
   base::FilePath data_file_path_;
   base::MemoryMappedFile data_file_;
   safe_browsing::PeImageReader image_reader_;
 };

 TEST_P(PeImageReaderCertificateTest, EnumCertificates) {
   StrictMock<MockCertificateReceiver> receiver;
   if (expected_data_->num_signers) {
     EXPECT_CALL(receiver, OnCertificate(WIN_CERT_REVISION_2_0,
                                         WIN_CERT_TYPE_PKCS_SIGNED_DATA,
                                         NotNull(),
                                         Gt(0U)))
         .Times(expected_data_->num_signers)
         .WillRepeatedly(Return(true));
   }
   EXPECT_TRUE(image_reader_.EnumCertificates(
       &CertificateReceiver::OnCertificateCallback, receiver.AsContext()));
 }

 TEST_P(PeImageReaderCertificateTest, AbortEnum) {
   StrictMock<MockCertificateReceiver> receiver;
   if (expected_data_->num_signers) {
     // Return false for the first cert, thereby stopping the enumeration.
     EXPECT_CALL(receiver, OnCertificate(_, _, _, _)).WillOnce(Return(false));
     EXPECT_FALSE(image_reader_.EnumCertificates(
         &CertificateReceiver::OnCertificateCallback, receiver.AsContext()));
   } else {
     // An unsigned file always reports true with no invocations of the callback.
     EXPECT_TRUE(image_reader_.EnumCertificates(
         &CertificateReceiver::OnCertificateCallback, receiver.AsContext()));
   }
 }

 namespace {

 const CertificateTestData kCertificateTestData[] = {
   {
     "signed.exe",
     1,
   }, {
     "unsigned.exe",
     0,
   }, {
     "disable_outdated_build_detector.exe",
     1,
   }, {
     "signed_twice.exe",
     2,
   },
 };

 }  // namespace

 INSTANTIATE_TEST_SUITE_P(SignedExe,
                          PeImageReaderCertificateTest,
                          testing::Values(&kCertificateTestData[0]));
 INSTANTIATE_TEST_SUITE_P(UnsignedExe,
                          PeImageReaderCertificateTest,
                          testing::Values(&kCertificateTestData[1]));
 INSTANTIATE_TEST_SUITE_P(DisableOutdatedBuildDetectorExe,
                          PeImageReaderCertificateTest,
                          testing::Values(&kCertificateTestData[2]));
 INSTANTIATE_TEST_SUITE_P(SignedTwiceExe,
                          PeImageReaderCertificateTest,
                          testing::Values(&kCertificateTestData[3]));
	// Copyright 2014 The Chromium Authors. All rights reserved.
	// Use of this source code is governed by a BSD-style license that can be
	// found in the LICENSE file.

	#include <windows.h>
	#include <stddef.h>
	#include <stdint.h>
	#include <wintrust.h>

	#include "base/files/file_path.h"
	#include "base/files/memory_mapped_file.h"
	#include "base/macros.h"
	#include "base/path_service.h"
	#include "chrome/common/chrome_paths.h"
	#include "chrome/common/safe_browsing/pe_image_reader_win.h"
	#include "testing/gmock/include/gmock/gmock.h"
	#include "testing/gtest/include/gtest/gtest.h"

	using ::testing::_;
	using ::testing::Gt;
	using ::testing::NotNull;
	using ::testing::Return;
	using ::testing::StrictMock;

	struct TestData {
	const char* filename;
	safe_browsing::PeImageReader::WordSize word_size;
	WORD machine_identifier;
	WORD optional_header_size;
	size_t number_of_sections;
	size_t number_of_debug_entries;
	};

	// A test fixture parameterized on test data containing the name of a PE image
	// to parse and the expected values to be read from it. The file is read from
	// the src/chrome/test/data/safe_browsing directory.
	class PeImageReaderTest : public testing::TestWithParam<const TestData*> {
	protected:
	PeImageReaderTest() : expected_data_(GetParam()) {}

	void SetUp() override {
	ASSERT_TRUE(
	base::PathService::Get(chrome::DIR_TEST_DATA, &data_file_path_));
	data_file_path_ = data_file_path_.AppendASCII("safe_browsing");
	data_file_path_ = data_file_path_.AppendASCII(expected_data_->filename);

	ASSERT_TRUE(data_file_.Initialize(data_file_path_));

	ASSERT_TRUE(image_reader_.Initialize(data_file_.data(),
	data_file_.length()));
	}

	const TestData* expected_data_;
	base::FilePath data_file_path_;
	base::MemoryMappedFile data_file_;
	safe_browsing::PeImageReader image_reader_;
	};

	TEST_P(PeImageReaderTest, GetWordSize) {
	EXPECT_EQ(expected_data_->word_size, image_reader_.GetWordSize());
	}

	TEST_P(PeImageReaderTest, GetDosHeader) {
	const IMAGE_DOS_HEADER* dos_header = image_reader_.GetDosHeader();
	ASSERT_NE(reinterpret_cast<const IMAGE_DOS_HEADER*>(NULL), dos_header);
	EXPECT_EQ(IMAGE_DOS_SIGNATURE, dos_header->e_magic);
	}

	TEST_P(PeImageReaderTest, GetCoffFileHeader) {
	const IMAGE_FILE_HEADER* file_header = image_reader_.GetCoffFileHeader();
	ASSERT_NE(reinterpret_cast<const IMAGE_FILE_HEADER*>(NULL), file_header);
	EXPECT_EQ(expected_data_->machine_identifier, file_header->Machine);
	EXPECT_EQ(expected_data_->optional_header_size,
	file_header->SizeOfOptionalHeader);
	}

	TEST_P(PeImageReaderTest, GetOptionalHeaderData) {
	size_t optional_header_size = 0;
	const uint8_t* optional_header_data =
	image_reader_.GetOptionalHeaderData(&optional_header_size);
	ASSERT_NE(reinterpret_cast<const uint8_t*>(NULL), optional_header_data);
	EXPECT_EQ(expected_data_->optional_header_size, optional_header_size);
	}

	TEST_P(PeImageReaderTest, GetNumberOfSections) {
	EXPECT_EQ(expected_data_->number_of_sections,
	image_reader_.GetNumberOfSections());
	}

	TEST_P(PeImageReaderTest, GetSectionHeaderAt) {
	size_t number_of_sections = image_reader_.GetNumberOfSections();
	for (size_t i = 0; i < number_of_sections; ++i) {
	const IMAGE_SECTION_HEADER* section_header =
	image_reader_.GetSectionHeaderAt(i);
	ASSERT_NE(reinterpret_cast<const IMAGE_SECTION_HEADER*>(NULL),
	section_header);
	}
	}

	TEST_P(PeImageReaderTest, InitializeFailTruncatedFile) {
	// Compute the size of all headers through the section headers.
	const IMAGE_SECTION_HEADER* last_section_header =
	image_reader_.GetSectionHeaderAt(image_reader_.GetNumberOfSections() - 1);
	const uint8_t* headers_end =
	reinterpret_cast<const uint8_t*>(last_section_header) +
	sizeof(*last_section_header);
	size_t header_size = headers_end - data_file_.data();
	safe_browsing::PeImageReader short_reader;

	// Initialize should succeed when all headers are present.
	EXPECT_TRUE(short_reader.Initialize(data_file_.data(), header_size));

	// But fail if anything is missing.
	for (size_t i = 0; i < header_size; ++i) {
	EXPECT_FALSE(short_reader.Initialize(data_file_.data(), i));
	}
	}

	TEST_P(PeImageReaderTest, GetExportSection) {
	size_t section_size = 0;
	const uint8_t* export_section = image_reader_.GetExportSection(&section_size);
	ASSERT_NE(reinterpret_cast<const uint8_t*>(NULL), export_section);
	EXPECT_NE(0U, section_size);
	}

	TEST_P(PeImageReaderTest, GetNumberOfDebugEntries) {
	EXPECT_EQ(expected_data_->number_of_debug_entries,
	image_reader_.GetNumberOfDebugEntries());
	}

	TEST_P(PeImageReaderTest, GetDebugEntry) {
	size_t number_of_debug_entries = image_reader_.GetNumberOfDebugEntries();
	for (size_t i = 0; i < number_of_debug_entries; ++i) {
	const uint8_t* raw_data = NULL;
	size_t raw_data_size = 0;
	const IMAGE_DEBUG_DIRECTORY* entry =
	image_reader_.GetDebugEntry(i, &raw_data, &raw_data_size);
	EXPECT_NE(reinterpret_cast<const IMAGE_DEBUG_DIRECTORY*>(NULL), entry);
	EXPECT_NE(reinterpret_cast<const uint8_t*>(NULL), raw_data);
	EXPECT_NE(0U, raw_data_size);
	}
	}

	namespace {

	const TestData kTestData[] = {
	{
	"module_with_exports_x86.dll",
	safe_browsing::PeImageReader::WORD_SIZE_32,
	IMAGE_FILE_MACHINE_I386,
	sizeof(IMAGE_OPTIONAL_HEADER32),
	4,
	1,
	}, {
	"module_with_exports_x64.dll",
	safe_browsing::PeImageReader::WORD_SIZE_64,
	IMAGE_FILE_MACHINE_AMD64,
	sizeof(IMAGE_OPTIONAL_HEADER64),
	5,
	1,
	},
	};

	} // namespace

	INSTANTIATE_TEST_SUITE_P(WordSize32,
	PeImageReaderTest,
	testing::Values(&kTestData[0]));
	INSTANTIATE_TEST_SUITE_P(WordSize64,
	PeImageReaderTest,
	testing::Values(&kTestData[1]));

	// An object exposing a PeImageReader::EnumCertificatesCallback that invokes a
	// virtual OnCertificate() method. This method is suitable for mocking in tests.
	class CertificateReceiver {
	public:
	void* AsContext() { return this; }
	static bool OnCertificateCallback(uint16_t revision,
	uint16_t certificate_type,
	const uint8_t* certificate_data,
	size_t certificate_data_size,
	void* context) {
	return reinterpret_cast<CertificateReceiver*>(context)->OnCertificate(
	revision, certificate_type, certificate_data, certificate_data_size);
	}

	protected:
	CertificateReceiver() {}
	virtual ~CertificateReceiver() {}
	virtual bool OnCertificate(uint16_t revision,
	uint16_t certificate_type,
	const uint8_t* certificate_data,
	size_t certificate_data_size) = 0;
	};

	class MockCertificateReceiver : public CertificateReceiver {
	public:
	MockCertificateReceiver() {}
	MOCK_METHOD4(OnCertificate, bool(uint16_t, uint16_t, const uint8_t*, size_t));

	private:
	DISALLOW_COPY_AND_ASSIGN(MockCertificateReceiver);
	};

	struct CertificateTestData {
	const char* filename;
	int num_signers;
	};

	// A test fixture parameterized on test data containing the name of a PE image
	// to parse and the expected values to be read from it. The file is read from
	// the src/chrome/test/data/safe_browsing/download_protection directory.
	class PeImageReaderCertificateTest
	: public testing::TestWithParam<const CertificateTestData*> {
	protected:
	PeImageReaderCertificateTest() : expected_data_(GetParam()) {}

	void SetUp() override {
	ASSERT_TRUE(
	base::PathService::Get(chrome::DIR_TEST_DATA, &data_file_path_));
	data_file_path_ = data_file_path_.AppendASCII("safe_browsing");
	data_file_path_ = data_file_path_.AppendASCII("download_protection");
	data_file_path_ = data_file_path_.AppendASCII(expected_data_->filename);
	ASSERT_TRUE(data_file_.Initialize(data_file_path_));
	ASSERT_TRUE(image_reader_.Initialize(data_file_.data(),
	data_file_.length()));
	}

	const CertificateTestData* expected_data_;
	base::FilePath data_file_path_;
	base::MemoryMappedFile data_file_;
	safe_browsing::PeImageReader image_reader_;
	};

	TEST_P(PeImageReaderCertificateTest, EnumCertificates) {
	StrictMock<MockCertificateReceiver> receiver;
	if (expected_data_->num_signers) {
	EXPECT_CALL(receiver, OnCertificate(WIN_CERT_REVISION_2_0,
	WIN_CERT_TYPE_PKCS_SIGNED_DATA,
	NotNull(),
	Gt(0U)))
	.Times(expected_data_->num_signers)
	.WillRepeatedly(Return(true));
	}
	EXPECT_TRUE(image_reader_.EnumCertificates(
	&CertificateReceiver::OnCertificateCallback, receiver.AsContext()));
	}

	TEST_P(PeImageReaderCertificateTest, AbortEnum) {
	StrictMock<MockCertificateReceiver> receiver;
	if (expected_data_->num_signers) {
	// Return false for the first cert, thereby stopping the enumeration.
	EXPECT_CALL(receiver, OnCertificate(_, _, _, _)).WillOnce(Return(false));
	EXPECT_FALSE(image_reader_.EnumCertificates(
	&CertificateReceiver::OnCertificateCallback, receiver.AsContext()));
	} else {
	// An unsigned file always reports true with no invocations of the callback.
	EXPECT_TRUE(image_reader_.EnumCertificates(
	&CertificateReceiver::OnCertificateCallback, receiver.AsContext()));
	}
	}

	namespace {

	const CertificateTestData kCertificateTestData[] = {
	{
	"signed.exe",
	1,
	}, {
	"unsigned.exe",
	0,
	}, {
	"disable_outdated_build_detector.exe",
	1,
	}, {
	"signed_twice.exe",
	2,
	},
	};

	} // namespace

	INSTANTIATE_TEST_SUITE_P(SignedExe,
	PeImageReaderCertificateTest,
	testing::Values(&kCertificateTestData[0]));
	INSTANTIATE_TEST_SUITE_P(UnsignedExe,
	PeImageReaderCertificateTest,
	testing::Values(&kCertificateTestData[1]));
	INSTANTIATE_TEST_SUITE_P(DisableOutdatedBuildDetectorExe,
	PeImageReaderCertificateTest,
	testing::Values(&kCertificateTestData[2]));
	INSTANTIATE_TEST_SUITE_P(SignedTwiceExe,
	PeImageReaderCertificateTest,
	testing::Values(&kCertificateTestData[3]));