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