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