1 : // Copyright 2013 Google Inc. All Rights Reserved.
2 : //
3 : // Licensed under the Apache License, Version 2.0 (the "License");
4 : // you may not use this file except in compliance with the License.
5 : // You may obtain a copy of the License at
6 : //
7 : // http://www.apache.org/licenses/LICENSE-2.0
8 : //
9 : // Unless required by applicable law or agreed to in writing, software
10 : // distributed under the License is distributed on an "AS IS" BASIS,
11 : // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 : // See the License for the specific language governing permissions and
13 : // limitations under the License.
14 :
15 : #include "syzygy/pe/pe_coff_file.h"
16 :
17 : #include "base/native_library.h"
18 : #include "base/path_service.h"
19 : #include "base/files/file_path.h"
20 : #include "base/strings/string_util.h"
21 : #include "gmock/gmock.h"
22 : #include "gtest/gtest.h"
23 : #include "syzygy/core/unittest_util.h"
24 : #include "syzygy/pe/unittest_util.h"
25 :
26 : namespace pe {
27 :
28 : namespace {
29 :
30 : using core::AbsoluteAddress;
31 : using core::FileOffsetAddress;
32 : using core::RelativeAddress;
33 :
34 : template <typename AddressSpaceTraits>
35 : class TestPECoffFile : public PECoffFile<AddressSpaceTraits> {
36 : public:
37 : // Partial initialization, for common parts.
38 E : bool Init(const base::FilePath& path, bool has_pe_headers) {
39 E : if (!PECoffFile::Init(path))
40 i : return false;
41 :
42 E : FileOffsetAddress file_header_start(0);
43 E : if (has_pe_headers) {
44 : // Read NT header position at 0x3C according to the spec.
45 E : uint32 nt_header_pos = 0;
46 E : if (!ReadAt(0x3C, &nt_header_pos, sizeof(nt_header_pos)))
47 i : return false;
48 :
49 : file_header_start.set_value(nt_header_pos +
50 E : offsetof(IMAGE_NT_HEADERS, FileHeader));
51 : }
52 :
53 E : bool success = ReadCommonHeaders(file_header_start);
54 E : if (success)
55 E : success = ReadSections();
56 :
57 E : return success;
58 E : }
59 : };
60 :
61 : template <typename AddressType, size_t shift_base, size_t header_base>
62 : struct ShiftedAddressSpaceTraits {
63 : typedef typename AddressType AddressType;
64 : typedef unsigned int SizeType;
65 :
66 : static const size_t kShiftBase = shift_base;
67 :
68 E : static const AddressType invalid_address() {
69 E : return AddressType::kInvalidAddress;
70 E : }
71 :
72 E : static const AddressType header_address() {
73 E : return AddressType(header_base);
74 E : }
75 :
76 E : static AddressType GetSectionAddress(const IMAGE_SECTION_HEADER& header) {
77 : // VirtualAddress is not necessarily 0 for object files, but it
78 : // should, and MSVC should produce 0; we use that info as
79 : // a heuristic to decide on the address.
80 E : if (header.VirtualAddress == 0) {
81 : if ((header.Characteristics & IMAGE_SCN_CNT_UNINITIALIZED_DATA) != 0 &&
82 E : header.PointerToRawData == 0) {
83 : // Unmapped section.
84 E : return invalid_address();
85 i : } else {
86 E : return AddressType(kShiftBase + header.PointerToRawData);
87 : }
88 i : } else {
89 E : return AddressType(kShiftBase + header.VirtualAddress);
90 : }
91 E : }
92 :
93 E : static SizeType GetSectionSize(const IMAGE_SECTION_HEADER& header) {
94 E : return 1;
95 E : }
96 : };
97 :
98 : typedef ShiftedAddressSpaceTraits<AbsoluteAddress, 0x10000, 0x4>
99 : ShiftedVirtualAddressTraits;
100 : typedef ShiftedAddressSpaceTraits<FileOffsetAddress, 0x100, 0x100>
101 : ShiftedFileOffsetAddressTraits;
102 :
103 : class PECoffFileTest : public testing::PELibUnitTest {
104 : typedef testing::PELibUnitTest Super;
105 :
106 : public:
107 E : PECoffFileTest() {
108 E : }
109 :
110 E : virtual void SetUp() override {
111 E : Super::SetUp();
112 :
113 E : test_dll_path_ = testing::GetExeRelativePath(testing::kTestDllName);
114 : test_dll_coff_obj_path_ =
115 E : testing::GetExeTestDataRelativePath(testing::kTestDllCoffObjName);
116 : test_dll_ltcg_obj_path_ =
117 E : testing::GetExeTestDataRelativePath(testing::kTestDllLtcgObjName);
118 E : }
119 :
120 : protected:
121 E : bool InitImages() {
122 : return pe_image_file_.Init(test_dll_path_, true) &&
123 E : coff_image_file_.Init(test_dll_coff_obj_path_, false);
124 E : }
125 :
126 : base::FilePath test_dll_path_;
127 : base::FilePath test_dll_coff_obj_path_;
128 : base::FilePath test_dll_ltcg_obj_path_;
129 :
130 : TestPECoffFile<ShiftedVirtualAddressTraits> pe_image_file_;
131 : TestPECoffFile<ShiftedFileOffsetAddressTraits> coff_image_file_;
132 :
133 : private:
134 : DISALLOW_COPY_AND_ASSIGN(PECoffFileTest);
135 : };
136 :
137 : bool IsCoffSectionMapped(
138 : const TestPECoffFile<ShiftedFileOffsetAddressTraits>& image_file,
139 E : size_t i) {
140 E : const IMAGE_SECTION_HEADER* hdr = image_file.section_header(i);
141 : return hdr != NULL &&
142 : (ShiftedFileOffsetAddressTraits::GetSectionAddress(*hdr) !=
143 E : ShiftedFileOffsetAddressTraits::invalid_address());
144 E : }
145 :
146 : } // namespace
147 :
148 E : TEST_F(PECoffFileTest, Init) {
149 E : EXPECT_TRUE(pe_image_file_.file_header() == NULL);
150 E : EXPECT_TRUE(pe_image_file_.section_headers() == NULL);
151 E : EXPECT_TRUE(coff_image_file_.file_header() == NULL);
152 E : EXPECT_TRUE(coff_image_file_.section_headers() == NULL);
153 :
154 E : ASSERT_TRUE(pe_image_file_.Init(test_dll_path_, true));
155 E : ASSERT_TRUE(coff_image_file_.Init(test_dll_coff_obj_path_, false));
156 :
157 E : EXPECT_TRUE(pe_image_file_.file_header() != NULL);
158 E : EXPECT_TRUE(pe_image_file_.section_headers() != NULL);
159 E : EXPECT_TRUE(coff_image_file_.file_header() != NULL);
160 E : EXPECT_TRUE(coff_image_file_.section_headers() != NULL);
161 :
162 E : EXPECT_EQ(pe_image_file_.file_header()->Machine, IMAGE_FILE_MACHINE_I386);
163 E : EXPECT_EQ(coff_image_file_.file_header()->Machine, IMAGE_FILE_MACHINE_I386);
164 :
165 E : EXPECT_TRUE(pe_image_file_.file_header()->SizeOfOptionalHeader != 0);
166 E : EXPECT_TRUE(coff_image_file_.file_header()->PointerToSymbolTable != 0);
167 E : }
168 :
169 E : TEST_F(PECoffFileTest, FailOnAnonymousObject) {
170 E : ASSERT_FALSE(coff_image_file_.Init(test_dll_ltcg_obj_path_, false));
171 E : }
172 :
173 : // Compare header data obtained from different methods.
174 E : TEST_F(PECoffFileTest, ReadFileHeader) {
175 E : ASSERT_TRUE(InitImages());
176 E : ASSERT_TRUE(coff_image_file_.file_header() != NULL);
177 :
178 E : IMAGE_FILE_HEADER header = {};
179 : ASSERT_TRUE(coff_image_file_.ReadImage(coff_image_file_.header_address(),
180 E : &header, sizeof(header)));
181 : EXPECT_TRUE(memcmp(static_cast<void*>(&header),
182 : static_cast<const void*>(coff_image_file_.file_header()),
183 E : sizeof(header)) == 0);
184 : uint8* ptr = coff_image_file_.GetImageData(coff_image_file_.header_address(),
185 E : sizeof(header));
186 E : ASSERT_TRUE(ptr != NULL);
187 : EXPECT_TRUE(memcmp(static_cast<void*>(&header),
188 : static_cast<void*>(ptr),
189 E : sizeof(header)) == 0);
190 E : }
191 :
192 E : TEST_F(PECoffFileTest, Contains) {
193 E : ASSERT_TRUE(InitImages());
194 E : EXPECT_TRUE(pe_image_file_.Contains(pe_image_file_.header_address(), 1));
195 E : EXPECT_FALSE(pe_image_file_.Contains(pe_image_file_.header_address() - 1, 1));
196 :
197 : // Should be a gap in the address space before the big shift of
198 : // ShiftedVirtualAddressTraits::kShiftBase.
199 : EXPECT_FALSE(pe_image_file_.Contains(
200 E : AbsoluteAddress(ShiftedVirtualAddressTraits::kShiftBase - 1), 1));
201 :
202 E : EXPECT_TRUE(coff_image_file_.Contains(coff_image_file_.header_address(), 1));
203 : EXPECT_FALSE(coff_image_file_.Contains(
204 E : coff_image_file_.header_address() - 1, 1));
205 E : }
206 :
207 E : TEST_F(PECoffFileTest, GetSectionIndex) {
208 E : ASSERT_TRUE(InitImages());
209 :
210 E : size_t num_sections = pe_image_file_.file_header()->NumberOfSections;
211 E : for (size_t i = 0; i < num_sections; ++i) {
212 : AbsoluteAddress section_start(
213 : ShiftedVirtualAddressTraits::GetSectionAddress(
214 E : *pe_image_file_.section_header(i)));
215 E : EXPECT_EQ(i, pe_image_file_.GetSectionIndex(section_start, 1));
216 E : }
217 :
218 E : AbsoluteAddress off_by_one(ShiftedVirtualAddressTraits::kShiftBase - 1);
219 E : EXPECT_EQ(kInvalidSection, pe_image_file_.GetSectionIndex(off_by_one, 1));
220 :
221 E : num_sections = coff_image_file_.file_header()->NumberOfSections;
222 E : for (size_t i = 0; i < num_sections; ++i) {
223 E : if (IsCoffSectionMapped(coff_image_file_, i)) {
224 : FileOffsetAddress section_start(
225 : ShiftedFileOffsetAddressTraits::GetSectionAddress(
226 E : *coff_image_file_.section_header(i)));
227 E : EXPECT_EQ(i, coff_image_file_.GetSectionIndex(section_start, 1));
228 : }
229 E : }
230 E : }
231 :
232 E : TEST_F(PECoffFileTest, GetSectionHeader) {
233 E : ASSERT_TRUE(InitImages());
234 :
235 E : size_t num_sections = pe_image_file_.file_header()->NumberOfSections;
236 E : for (size_t i = 0; i < num_sections; ++i) {
237 : AbsoluteAddress section_start(
238 : ShiftedVirtualAddressTraits::GetSectionAddress(
239 E : *pe_image_file_.section_header(i)));
240 : EXPECT_EQ(pe_image_file_.section_header(i),
241 E : pe_image_file_.GetSectionHeader(section_start, 1));
242 E : }
243 :
244 E : num_sections = coff_image_file_.file_header()->NumberOfSections;
245 E : for (size_t i = 0; i < num_sections; ++i) {
246 E : if (IsCoffSectionMapped(coff_image_file_, i)) {
247 : FileOffsetAddress section_start(
248 : ShiftedFileOffsetAddressTraits::GetSectionAddress(
249 E : *coff_image_file_.section_header(i)));
250 : EXPECT_EQ(coff_image_file_.section_header(i),
251 E : coff_image_file_.GetSectionHeader(section_start, 1));
252 : }
253 E : }
254 E : }
255 :
256 E : TEST_F(PECoffFileTest, GetImageData) {
257 E : ASSERT_TRUE(InitImages());
258 E : ASSERT_TRUE(pe_image_file_.file_header() != NULL);
259 E : ASSERT_TRUE(coff_image_file_.file_header() != NULL);
260 :
261 E : size_t num_sections = pe_image_file_.file_header()->NumberOfSections;
262 E : for (size_t i = 0; i < num_sections; ++i) {
263 : AbsoluteAddress section_start(
264 : ShiftedVirtualAddressTraits::GetSectionAddress(
265 E : *pe_image_file_.section_header(i)));
266 E : EXPECT_TRUE(pe_image_file_.GetImageData(section_start, 1) != NULL);
267 E : }
268 :
269 E : num_sections = coff_image_file_.file_header()->NumberOfSections;
270 E : for (size_t i = 0; i < num_sections; ++i) {
271 E : if (IsCoffSectionMapped(coff_image_file_, i)) {
272 : FileOffsetAddress section_start(
273 : ShiftedFileOffsetAddressTraits::GetSectionAddress(
274 E : *coff_image_file_.section_header(i)));
275 E : EXPECT_TRUE(coff_image_file_.GetImageData(section_start, 1) != NULL);
276 : }
277 E : }
278 :
279 : // Get arbitrary image data spanning across sections. Expect this to fail
280 : // with GetImageData, but to succeed with GetImageDataFromFileOffset.
281 E : AbsoluteAddress header_abs(ShiftedVirtualAddressTraits::header_address());
282 : FileOffsetAddress header_off(
283 E : ShiftedFileOffsetAddressTraits::header_address());
284 : size_t len = (ShiftedFileOffsetAddressTraits::GetSectionAddress(
285 : *coff_image_file_.section_header(0)) -
286 : header_off) +
287 E : 1;
288 E : EXPECT_FALSE(pe_image_file_.GetImageData(header_abs, len));
289 E : EXPECT_TRUE(pe_image_file_.GetImageDataByFileOffset(header_off, len));
290 E : }
291 :
292 : } // namespace pe
|