xref: /llvm-project/llvm/unittests/ObjCopy/ObjCopyTest.cpp (revision dedf006fa9045404b948297bd496524fd04054ff)
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.
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>>
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.
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.
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 
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 
208 TEST(CopySimpleInMemoryFile, COFF) {
209   SCOPED_TRACE("CopySimpleInMemoryFileCOFF");
210 
211   copySimpleInMemoryFileImpl(SimpleFileCOFFYAML,
212                              [](const Binary &File) { return File.isCOFF(); });
213 }
214 
215 TEST(CopySimpleInMemoryFile, ELF) {
216   SCOPED_TRACE("CopySimpleInMemoryFileELF");
217 
218   copySimpleInMemoryFileImpl(SimpleFileELFYAML,
219                              [](const Binary &File) { return File.isELF(); });
220 }
221 
222 TEST(CopySimpleInMemoryFile, MachO) {
223   SCOPED_TRACE("CopySimpleInMemoryFileMachO");
224 
225   copySimpleInMemoryFileImpl(SimpleFileMachOYAML,
226                              [](const Binary &File) { return File.isMachO(); });
227 }
228 
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 
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 
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 
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 
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 
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 
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 
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 
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 
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 
368 TEST(RemoveSectionByPattern, COFF) {
369   SCOPED_TRACE("removeSectionByPatternCOFF");
370 
371   removeSectionByPatternImpl(
372       SimpleFileCOFFYAML, [](const Binary &File) { return File.isCOFF(); },
373       "\\.text*", ".text");
374 }
375 
376 TEST(RemoveSectionByPattern, ELF) {
377   SCOPED_TRACE("removeSectionByPatternELF");
378 
379   removeSectionByPatternImpl(
380       SimpleFileELFYAML, [](const Binary &File) { return File.isELF(); },
381       "\\.text*", ".text");
382 }
383 
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 
392 TEST(RemoveSectionByPattern, Wasm) {
393   SCOPED_TRACE("removeSectionByPatternWasm");
394 
395   removeSectionByPatternImpl(
396       SimpleFileWasmYAML, [](const Binary &File) { return File.isWasm(); },
397       "text*", "text");
398 }
399