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