1 // Copyright 2018 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
5 #include "chrome/chrome_elf/pe_image_safe/pe_image_safe.h"
6
7 #include "base/files/file_util.h"
8 #include "base/path_service.h"
9 #include "testing/gtest/include/gtest/gtest.h"
10
11 namespace pe_image_safe {
12 namespace {
13
14 constexpr wchar_t kPeFile[] = L"chrome_elf.dll";
15 constexpr DWORD kPageSize = 4096;
16
17 struct CompareData {
18 bool bits_64;
19 DWORD size_of_image;
20 DWORD time_date_stamp;
21 PIMAGE_DOS_HEADER dos_header;
22 PIMAGE_NT_HEADERS nt_headers;
23 PIMAGE_FILE_HEADER file_header;
24 PIMAGE_OPTIONAL_HEADER optional_header;
25 };
26
27 // Raw collection of some PE header data, without using pe_image_safe.
28 // This function assumes full headers from a legitimate PE image.
GetComparisonData(char * buffer,CompareData * data)29 bool GetComparisonData(char* buffer, CompareData* data) {
30 data->dos_header = reinterpret_cast<PIMAGE_DOS_HEADER>(buffer);
31 if (data->dos_header->e_magic != IMAGE_DOS_SIGNATURE)
32 return false;
33
34 data->nt_headers = reinterpret_cast<PIMAGE_NT_HEADERS>(
35 reinterpret_cast<char*>(data->dos_header) + data->dos_header->e_lfanew);
36 if (data->nt_headers->Signature != IMAGE_NT_SIGNATURE)
37 return false;
38
39 data->file_header = &data->nt_headers->FileHeader;
40 data->time_date_stamp = data->nt_headers->FileHeader.TimeDateStamp;
41 data->optional_header = &data->nt_headers->OptionalHeader;
42
43 if (data->optional_header->Magic == IMAGE_NT_OPTIONAL_HDR64_MAGIC) {
44 PIMAGE_OPTIONAL_HEADER64 optional_header =
45 reinterpret_cast<PIMAGE_OPTIONAL_HEADER64>(
46 &data->nt_headers->OptionalHeader);
47 data->bits_64 = true;
48 data->size_of_image = optional_header->SizeOfImage;
49 } else {
50 PIMAGE_OPTIONAL_HEADER32 optional_header =
51 reinterpret_cast<PIMAGE_OPTIONAL_HEADER32>(
52 &data->nt_headers->OptionalHeader);
53 data->bits_64 = false;
54 data->size_of_image = optional_header->SizeOfImage;
55 }
56
57 return true;
58 }
59
60 //------------------------------------------------------------------------------
61
TEST(PEImageSafe,SanityTest)62 TEST(PEImageSafe, SanityTest) {
63 // Open and read in a PE file.
64 base::FilePath pe_path;
65 EXPECT_TRUE(base::PathService::Get(base::DIR_EXE, &pe_path));
66 pe_path = pe_path.Append(kPeFile);
67
68 base::File file(pe_path, base::File::FLAG_OPEN | base::File::FLAG_READ);
69 ASSERT_TRUE(file.IsValid());
70
71 std::vector<char> buffer;
72 buffer.resize(kPageSize);
73 ASSERT_EQ(file.Read(0, &buffer[0], kPageSize), static_cast<int>(kPageSize));
74 file.Close();
75
76 // Grab some key data out of the pe headers first, NOT using pe_image_safe.
77 CompareData data_for_comparison = {};
78 ASSERT_TRUE(GetComparisonData(buffer.data(), &data_for_comparison));
79
80 // Apply scaffolding.
81 PEImageSafe pe_image(buffer.data(), static_cast<DWORD>(buffer.size()));
82 PIMAGE_DOS_HEADER dos_header = pe_image.GetDosHeader();
83 EXPECT_TRUE(dos_header);
84 EXPECT_EQ(dos_header, data_for_comparison.dos_header);
85
86 PIMAGE_FILE_HEADER file_header = pe_image.GetFileHeader();
87 EXPECT_TRUE(file_header);
88 EXPECT_EQ(file_header, data_for_comparison.file_header);
89 EXPECT_EQ(file_header->TimeDateStamp, data_for_comparison.time_date_stamp);
90
91 BYTE* optional_header = pe_image.GetOptionalHeader();
92 EXPECT_TRUE(optional_header);
93 EXPECT_EQ(optional_header,
94 reinterpret_cast<BYTE*>(data_for_comparison.optional_header));
95
96 pe_image_safe::ImageBitness bitness = pe_image.GetImageBitness();
97 EXPECT_NE(bitness, pe_image_safe::ImageBitness::kUnknown);
98 EXPECT_TRUE(((bitness == pe_image_safe::ImageBitness::k64) &&
99 data_for_comparison.bits_64) ||
100 ((bitness == pe_image_safe::ImageBitness::k32) &&
101 !data_for_comparison.bits_64));
102
103 if (bitness == pe_image_safe::ImageBitness::k64) {
104 PIMAGE_OPTIONAL_HEADER64 optional_header64 =
105 reinterpret_cast<PIMAGE_OPTIONAL_HEADER64>(optional_header);
106 EXPECT_EQ(optional_header64->SizeOfImage,
107 data_for_comparison.size_of_image);
108 } else {
109 PIMAGE_OPTIONAL_HEADER32 optional_header32 =
110 reinterpret_cast<PIMAGE_OPTIONAL_HEADER32>(optional_header);
111 EXPECT_EQ(optional_header32->SizeOfImage,
112 data_for_comparison.size_of_image);
113 }
114 }
115
116 } // namespace
117 } // namespace pe_image_safe
118