1 //===- ObjCopyTest.cpp ----------------------------------------------------===//
2 //
3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4 // See https://llvm.org/LICENSE.txt for license information.
5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6 //
7 //===----------------------------------------------------------------------===//
8
9 #include "llvm/ObjCopy/ObjCopy.h"
10 #include "llvm/ObjCopy/ConfigManager.h"
11 #include "llvm/Object/ObjectFile.h"
12 #include "llvm/ObjectYAML/yaml2obj.h"
13 #include "llvm/Support/Error.h"
14 #include "llvm/Support/FileUtilities.h"
15 #include "llvm/Testing/Support/Error.h"
16 #include "gtest/gtest.h"
17
18 using namespace llvm;
19 using namespace object;
20 using namespace objcopy;
21 using namespace yaml;
22
23 const char *SimpleFileCOFFYAML = R"(
24 --- !COFF
25 header:
26 Machine: IMAGE_FILE_MACHINE_AMD64
27 Characteristics: [ ]
28 sections:
29 - Name: .text
30 Characteristics: [ ]
31 Alignment: 4
32 SectionData: E800000000C3C3C3
33 symbols:
34 ...
35 )";
36
37 const char *SimpleFileELFYAML = R"(
38 --- !ELF
39 FileHeader:
40 Class: ELFCLASS64
41 Data: ELFDATA2LSB
42 Type: ET_REL
43 Sections:
44 - Name: .text
45 Type: SHT_PROGBITS
46 Flags: [ SHF_ALLOC ]
47 Content: "12345678"
48 )";
49
50 const char *SimpleFileMachOYAML = R"(
51 --- !mach-o
52 FileHeader:
53 magic: 0xFEEDFACF
54 cputype: 0x01000007
55 cpusubtype: 0x80000003
56 filetype: 0x00000001
57 ncmds: 1
58 sizeofcmds: 152
59 flags: 0x00002000
60 reserved: 0x00000000
61 LoadCommands:
62 - cmd: LC_SEGMENT_64
63 cmdsize: 152
64 segname: __TEXT
65 vmaddr: 0
66 vmsize: 4
67 fileoff: 184
68 filesize: 4
69 maxprot: 7
70 initprot: 7
71 nsects: 1
72 flags: 0
73 Sections:
74 - sectname: __text
75 segname: __TEXT
76 addr: 0x0000000000000000
77 content: 'AABBCCDD'
78 size: 4
79 offset: 184
80 align: 0
81 reloff: 0x00000000
82 nreloc: 0
83 flags: 0x80000400
84 reserved1: 0x00000000
85 reserved2: 0x00000000
86 reserved3: 0x00000000
87 ...
88 )";
89
90 const char *SimpleFileWasmYAML = R"(
91 --- !WASM
92 FileHeader:
93 Version: 0x00000001
94 Sections:
95 - Type: CUSTOM
96 Name: text
97 Payload: ABC123
98 ...
99 )";
100
101 // Create ObjectFile from \p YamlCreationString and do validation using \p
102 // IsValidFormat checker. \p Storage is a storage for data. \returns created
103 // ObjectFile.
createObjectFileFromYamlDescription(const char * YamlCreationString,SmallVector<char> & Storage,function_ref<bool (const Binary & File)> IsValidFormat)104 Expected<std::unique_ptr<ObjectFile>> createObjectFileFromYamlDescription(
105 const char *YamlCreationString, SmallVector<char> &Storage,
106 function_ref<bool(const Binary &File)> IsValidFormat) {
107 auto ErrHandler = [&](const Twine &Msg) { FAIL() << "Error: " << Msg; };
108
109 std::unique_ptr<ObjectFile> Obj =
110 yaml2ObjectFile(Storage, YamlCreationString, ErrHandler);
111 if (!Obj)
112 return createError("could not create ObjectFile from yaml description");
113
114 if (!IsValidFormat(*Obj))
115 return createError("wrong file format");
116
117 return std::move(Obj);
118 }
119
120 // Call objcopy::executeObjcopyOnBinary for \p Config and \p In. \p DataVector
121 // is a holder for data. \returns Binary for copied data.
122 Expected<std::unique_ptr<Binary>>
callObjCopy(ConfigManager & Config,object::Binary & In,SmallVector<char> & DataVector,function_ref<bool (const Binary & File)> IsValidFormat)123 callObjCopy(ConfigManager &Config, object::Binary &In,
124 SmallVector<char> &DataVector,
125 function_ref<bool(const Binary &File)> IsValidFormat) {
126 raw_svector_ostream OutStream(DataVector);
127
128 if (Error Err = objcopy::executeObjcopyOnBinary(Config, In, OutStream))
129 return std::move(Err);
130
131 MemoryBufferRef Buffer(StringRef(DataVector.data(), DataVector.size()),
132 Config.Common.OutputFilename);
133
134 Expected<std::unique_ptr<Binary>> Result = createBinary(Buffer);
135
136 // Check copied file.
137 if (!Result)
138 return Result;
139
140 if (!IsValidFormat(**Result))
141 return createError("wrong file format");
142
143 if (!(*Result)->isObject())
144 return createError("binary is not object file");
145
146 return Result;
147 }
148
149 // \returns true if specified \p File has a section named \p SectionName.
hasSection(ObjectFile & File,StringRef SectionName)150 bool hasSection(ObjectFile &File, StringRef SectionName) {
151 for (const object::SectionRef &Sec : File.sections()) {
152 Expected<StringRef> CurSecNameOrErr = Sec.getName();
153 if (!CurSecNameOrErr)
154 continue;
155
156 if (*CurSecNameOrErr == SectionName)
157 return true;
158 }
159
160 return false;
161 }
162
163 // Check that specified \p File has a section \p SectionName and its data
164 // matches \p SectionData.
checkSectionData(ObjectFile & File,StringRef SectionName,StringRef SectionData)165 void checkSectionData(ObjectFile &File, StringRef SectionName,
166 StringRef SectionData) {
167 for (const object::SectionRef &Sec : File.sections()) {
168 Expected<StringRef> CurSecNameOrErr = Sec.getName();
169 ASSERT_THAT_EXPECTED(CurSecNameOrErr, Succeeded());
170
171 if (*CurSecNameOrErr == SectionName) {
172 Expected<StringRef> CurSectionData = Sec.getContents();
173 ASSERT_THAT_EXPECTED(CurSectionData, Succeeded());
174 EXPECT_TRUE(Sec.getSize() == SectionData.size());
175 EXPECT_TRUE(memcmp(CurSectionData->data(), SectionData.data(),
176 SectionData.size()) == 0);
177 return;
178 }
179 }
180
181 // Section SectionName must be presented.
182 EXPECT_TRUE(false);
183 }
184
copySimpleInMemoryFileImpl(const char * YamlCreationString,function_ref<bool (const Binary & File)> IsValidFormat)185 void copySimpleInMemoryFileImpl(
186 const char *YamlCreationString,
187 function_ref<bool(const Binary &File)> IsValidFormat) {
188 SCOPED_TRACE("copySimpleInMemoryFileImpl");
189
190 // Create Object file from YAML description.
191 SmallVector<char> Storage;
192 Expected<std::unique_ptr<ObjectFile>> Obj =
193 createObjectFileFromYamlDescription(YamlCreationString, Storage,
194 IsValidFormat);
195 ASSERT_THAT_EXPECTED(Obj, Succeeded());
196
197 ConfigManager Config;
198 Config.Common.OutputFilename = "a.out";
199
200 // Call executeObjcopyOnBinary()
201 SmallVector<char> DataVector;
202 Expected<std::unique_ptr<Binary>> Result =
203 callObjCopy(Config, *Obj.get(), DataVector, IsValidFormat);
204 ASSERT_THAT_EXPECTED(Result, Succeeded());
205 }
206
TEST(CopySimpleInMemoryFile,COFF)207 TEST(CopySimpleInMemoryFile, COFF) {
208 SCOPED_TRACE("CopySimpleInMemoryFileCOFF");
209
210 copySimpleInMemoryFileImpl(SimpleFileCOFFYAML,
211 [](const Binary &File) { return File.isCOFF(); });
212 }
213
TEST(CopySimpleInMemoryFile,ELF)214 TEST(CopySimpleInMemoryFile, ELF) {
215 SCOPED_TRACE("CopySimpleInMemoryFileELF");
216
217 copySimpleInMemoryFileImpl(SimpleFileELFYAML,
218 [](const Binary &File) { return File.isELF(); });
219 }
220
TEST(CopySimpleInMemoryFile,MachO)221 TEST(CopySimpleInMemoryFile, MachO) {
222 SCOPED_TRACE("CopySimpleInMemoryFileMachO");
223
224 copySimpleInMemoryFileImpl(SimpleFileMachOYAML,
225 [](const Binary &File) { return File.isMachO(); });
226 }
227
TEST(CopySimpleInMemoryFile,Wasm)228 TEST(CopySimpleInMemoryFile, Wasm) {
229 SCOPED_TRACE("CopySimpleInMemoryFileWasm");
230
231 copySimpleInMemoryFileImpl(SimpleFileWasmYAML,
232 [](const Binary &File) { return File.isWasm(); });
233 }
234
235 enum Action : uint8_t { AddSection, UpdateSection };
236
addOrUpdateSectionToFileImpl(const char * YamlCreationString,function_ref<bool (const Binary & File)> IsValidFormat,StringRef NewSectionName,StringRef NewSectionData,Action SectionAction)237 void addOrUpdateSectionToFileImpl(
238 const char *YamlCreationString,
239 function_ref<bool(const Binary &File)> IsValidFormat,
240 StringRef NewSectionName, StringRef NewSectionData, Action SectionAction) {
241 SCOPED_TRACE("addOrUpdateSectionToFileImpl");
242
243 // Create Object file from YAML description.
244 SmallVector<char> Storage;
245 Expected<std::unique_ptr<ObjectFile>> Obj =
246 createObjectFileFromYamlDescription(YamlCreationString, Storage,
247 IsValidFormat);
248 ASSERT_THAT_EXPECTED(Obj, Succeeded());
249
250 std::unique_ptr<MemoryBuffer> NewSectionBuffer =
251 MemoryBuffer::getMemBuffer(NewSectionData, NewSectionName, false);
252 std::string Name;
253 if ((*Obj)->isMachO())
254 Name = "__TEXT," + NewSectionName.str();
255 else
256 Name = NewSectionName.str();
257
258 ConfigManager Config;
259 Config.Common.OutputFilename = "a.out";
260 if (SectionAction == AddSection)
261 Config.Common.AddSection.push_back({Name, std::move(NewSectionBuffer)});
262 else
263 Config.Common.UpdateSection.push_back({Name, std::move(NewSectionBuffer)});
264
265 // Call executeObjcopyOnBinary()
266 SmallVector<char> DataVector;
267 Expected<std::unique_ptr<Binary>> Result =
268 callObjCopy(Config, *Obj.get(), DataVector, IsValidFormat);
269 ASSERT_THAT_EXPECTED(Result, Succeeded());
270
271 // Check that copied file has the new section.
272 checkSectionData(*static_cast<ObjectFile *>((*Result).get()), NewSectionName,
273 NewSectionData);
274 }
275
TEST(AddSection,COFF)276 TEST(AddSection, COFF) {
277 SCOPED_TRACE("addSectionToFileCOFF");
278
279 addOrUpdateSectionToFileImpl(
280 SimpleFileCOFFYAML, [](const Binary &File) { return File.isCOFF(); },
281 ".foo", "1234", AddSection);
282 }
283
TEST(AddSection,ELF)284 TEST(AddSection, ELF) {
285 SCOPED_TRACE("addSectionToFileELF");
286
287 addOrUpdateSectionToFileImpl(
288 SimpleFileELFYAML, [](const Binary &File) { return File.isELF(); },
289 ".foo", "1234", AddSection);
290 }
291
TEST(AddSection,MachO)292 TEST(AddSection, MachO) {
293 SCOPED_TRACE("addSectionToFileMachO");
294
295 addOrUpdateSectionToFileImpl(
296 SimpleFileMachOYAML, [](const Binary &File) { return File.isMachO(); },
297 "__foo", "1234", AddSection);
298 }
299
TEST(AddSection,Wasm)300 TEST(AddSection, Wasm) {
301 SCOPED_TRACE("addSectionToFileWasm");
302
303 addOrUpdateSectionToFileImpl(
304 SimpleFileWasmYAML, [](const Binary &File) { return File.isWasm(); },
305 ".foo", "1234", AddSection);
306 }
307
TEST(UpdateSection,COFF)308 TEST(UpdateSection, COFF) {
309 SCOPED_TRACE("updateSectionToFileCOFF");
310
311 addOrUpdateSectionToFileImpl(
312 SimpleFileCOFFYAML, [](const Binary &File) { return File.isCOFF(); },
313 ".text", "1234", UpdateSection);
314 }
315
TEST(UpdateSection,ELF)316 TEST(UpdateSection, ELF) {
317 SCOPED_TRACE("updateSectionToFileELF");
318
319 addOrUpdateSectionToFileImpl(
320 SimpleFileELFYAML, [](const Binary &File) { return File.isELF(); },
321 ".text", "1234", UpdateSection);
322 }
323
TEST(UpdateSection,MachO)324 TEST(UpdateSection, MachO) {
325 SCOPED_TRACE("updateSectionToFileMachO");
326
327 addOrUpdateSectionToFileImpl(
328 SimpleFileMachOYAML, [](const Binary &File) { return File.isMachO(); },
329 "__text", "1234", UpdateSection);
330 }
331
removeSectionByPatternImpl(const char * YamlCreationString,function_ref<bool (const Binary & File)> IsValidFormat,StringRef SectionWildcard,StringRef SectionName)332 void removeSectionByPatternImpl(
333 const char *YamlCreationString,
334 function_ref<bool(const Binary &File)> IsValidFormat,
335 StringRef SectionWildcard, StringRef SectionName) {
336 SCOPED_TRACE("removeSectionByPatternImpl");
337
338 // Create Object file from YAML description.
339 SmallVector<char> Storage;
340 Expected<std::unique_ptr<ObjectFile>> Obj =
341 createObjectFileFromYamlDescription(YamlCreationString, Storage,
342 IsValidFormat);
343 ASSERT_THAT_EXPECTED(Obj, Succeeded());
344
345 // Check that section is present.
346 EXPECT_TRUE(hasSection(**Obj, SectionName));
347
348 Expected<NameOrPattern> Pattern = objcopy::NameOrPattern::create(
349 SectionWildcard, objcopy::MatchStyle::Wildcard,
350 [](Error Err) { return Err; });
351
352 ConfigManager Config;
353 Config.Common.OutputFilename = "a.out";
354 EXPECT_THAT_ERROR(Config.Common.ToRemove.addMatcher(std::move(Pattern)),
355 Succeeded());
356
357 SmallVector<char> DataVector;
358 Expected<std::unique_ptr<Binary>> Result =
359 callObjCopy(Config, *Obj.get(), DataVector, IsValidFormat);
360 ASSERT_THAT_EXPECTED(Result, Succeeded());
361
362 // Check that section was removed.
363 EXPECT_FALSE(
364 hasSection(*static_cast<ObjectFile *>((*Result).get()), SectionName));
365 }
366
TEST(RemoveSectionByPattern,COFF)367 TEST(RemoveSectionByPattern, COFF) {
368 SCOPED_TRACE("removeSectionByPatternCOFF");
369
370 removeSectionByPatternImpl(
371 SimpleFileCOFFYAML, [](const Binary &File) { return File.isCOFF(); },
372 "\\.text*", ".text");
373 }
374
TEST(RemoveSectionByPattern,ELF)375 TEST(RemoveSectionByPattern, ELF) {
376 SCOPED_TRACE("removeSectionByPatternELF");
377
378 removeSectionByPatternImpl(
379 SimpleFileELFYAML, [](const Binary &File) { return File.isELF(); },
380 "\\.text*", ".text");
381 }
382
TEST(RemoveSectionByPattern,MachO)383 TEST(RemoveSectionByPattern, MachO) {
384 SCOPED_TRACE("removeSectionByPatternMachO");
385
386 removeSectionByPatternImpl(
387 SimpleFileMachOYAML, [](const Binary &File) { return File.isMachO(); },
388 "__TEXT,__text*", "__text");
389 }
390
TEST(RemoveSectionByPattern,Wasm)391 TEST(RemoveSectionByPattern, Wasm) {
392 SCOPED_TRACE("removeSectionByPatternWasm");
393
394 removeSectionByPatternImpl(
395 SimpleFileWasmYAML, [](const Binary &File) { return File.isWasm(); },
396 "text*", "text");
397 }
398