pe_image_reader_unittest.cc 9.9 KB

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