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