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