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. |
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, Yaml: YamlCreationString, ErrHandler); |
111 | if (!Obj) |
112 | return createError(Err: "could not create ObjectFile from yaml description" ); |
113 | |
114 | if (!IsValidFormat(*Obj)) |
115 | return createError(Err: "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>> |
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, Out&: 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(Source: Buffer); |
135 | |
136 | // Check copied file. |
137 | if (!Result) |
138 | return Result; |
139 | |
140 | if (!IsValidFormat(**Result)) |
141 | return createError(Err: "wrong file format" ); |
142 | |
143 | if (!(*Result)->isObject()) |
144 | return createError(Err: "binary is not object file" ); |
145 | |
146 | return Result; |
147 | } |
148 | |
149 | // \returns true if specified \p File has a section named \p 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. |
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 | |
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, In&: *Obj.get(), DataVector, IsValidFormat); |
204 | ASSERT_THAT_EXPECTED(Result, Succeeded()); |
205 | } |
206 | |
207 | TEST(CopySimpleInMemoryFile, COFF) { |
208 | SCOPED_TRACE("CopySimpleInMemoryFileCOFF" ); |
209 | |
210 | copySimpleInMemoryFileImpl(YamlCreationString: SimpleFileCOFFYAML, |
211 | IsValidFormat: [](const Binary &File) { return File.isCOFF(); }); |
212 | } |
213 | |
214 | TEST(CopySimpleInMemoryFile, ELF) { |
215 | SCOPED_TRACE("CopySimpleInMemoryFileELF" ); |
216 | |
217 | copySimpleInMemoryFileImpl(YamlCreationString: SimpleFileELFYAML, |
218 | IsValidFormat: [](const Binary &File) { return File.isELF(); }); |
219 | } |
220 | |
221 | TEST(CopySimpleInMemoryFile, MachO) { |
222 | SCOPED_TRACE("CopySimpleInMemoryFileMachO" ); |
223 | |
224 | copySimpleInMemoryFileImpl(YamlCreationString: SimpleFileMachOYAML, |
225 | IsValidFormat: [](const Binary &File) { return File.isMachO(); }); |
226 | } |
227 | |
228 | TEST(CopySimpleInMemoryFile, Wasm) { |
229 | SCOPED_TRACE("CopySimpleInMemoryFileWasm" ); |
230 | |
231 | copySimpleInMemoryFileImpl(YamlCreationString: SimpleFileWasmYAML, |
232 | IsValidFormat: [](const Binary &File) { return File.isWasm(); }); |
233 | } |
234 | |
235 | enum Action : uint8_t { AddSection, UpdateSection }; |
236 | |
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(InputData: NewSectionData, BufferName: NewSectionName, RequiresNullTerminator: 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(Elt: {Name, std::move(NewSectionBuffer)}); |
262 | else |
263 | Config.Common.UpdateSection.push_back(Elt: {Name, std::move(NewSectionBuffer)}); |
264 | |
265 | // Call executeObjcopyOnBinary() |
266 | SmallVector<char> DataVector; |
267 | Expected<std::unique_ptr<Binary>> Result = |
268 | callObjCopy(Config, In&: *Obj.get(), DataVector, IsValidFormat); |
269 | ASSERT_THAT_EXPECTED(Result, Succeeded()); |
270 | |
271 | // Check that copied file has the new section. |
272 | checkSectionData(File&: *static_cast<ObjectFile *>((*Result).get()), SectionName: NewSectionName, |
273 | SectionData: NewSectionData); |
274 | } |
275 | |
276 | TEST(AddSection, COFF) { |
277 | SCOPED_TRACE("addSectionToFileCOFF" ); |
278 | |
279 | addOrUpdateSectionToFileImpl( |
280 | YamlCreationString: SimpleFileCOFFYAML, IsValidFormat: [](const Binary &File) { return File.isCOFF(); }, |
281 | NewSectionName: ".foo" , NewSectionData: "1234" , SectionAction: AddSection); |
282 | } |
283 | |
284 | TEST(AddSection, ELF) { |
285 | SCOPED_TRACE("addSectionToFileELF" ); |
286 | |
287 | addOrUpdateSectionToFileImpl( |
288 | YamlCreationString: SimpleFileELFYAML, IsValidFormat: [](const Binary &File) { return File.isELF(); }, |
289 | NewSectionName: ".foo" , NewSectionData: "1234" , SectionAction: AddSection); |
290 | } |
291 | |
292 | TEST(AddSection, MachO) { |
293 | SCOPED_TRACE("addSectionToFileMachO" ); |
294 | |
295 | addOrUpdateSectionToFileImpl( |
296 | YamlCreationString: SimpleFileMachOYAML, IsValidFormat: [](const Binary &File) { return File.isMachO(); }, |
297 | NewSectionName: "__foo" , NewSectionData: "1234" , SectionAction: AddSection); |
298 | } |
299 | |
300 | TEST(AddSection, Wasm) { |
301 | SCOPED_TRACE("addSectionToFileWasm" ); |
302 | |
303 | addOrUpdateSectionToFileImpl( |
304 | YamlCreationString: SimpleFileWasmYAML, IsValidFormat: [](const Binary &File) { return File.isWasm(); }, |
305 | NewSectionName: ".foo" , NewSectionData: "1234" , SectionAction: AddSection); |
306 | } |
307 | |
308 | TEST(UpdateSection, COFF) { |
309 | SCOPED_TRACE("updateSectionToFileCOFF" ); |
310 | |
311 | addOrUpdateSectionToFileImpl( |
312 | YamlCreationString: SimpleFileCOFFYAML, IsValidFormat: [](const Binary &File) { return File.isCOFF(); }, |
313 | NewSectionName: ".text" , NewSectionData: "1234" , SectionAction: UpdateSection); |
314 | } |
315 | |
316 | TEST(UpdateSection, ELF) { |
317 | SCOPED_TRACE("updateSectionToFileELF" ); |
318 | |
319 | addOrUpdateSectionToFileImpl( |
320 | YamlCreationString: SimpleFileELFYAML, IsValidFormat: [](const Binary &File) { return File.isELF(); }, |
321 | NewSectionName: ".text" , NewSectionData: "1234" , SectionAction: UpdateSection); |
322 | } |
323 | |
324 | TEST(UpdateSection, MachO) { |
325 | SCOPED_TRACE("updateSectionToFileMachO" ); |
326 | |
327 | addOrUpdateSectionToFileImpl( |
328 | YamlCreationString: SimpleFileMachOYAML, IsValidFormat: [](const Binary &File) { return File.isMachO(); }, |
329 | NewSectionName: "__text" , NewSectionData: "1234" , SectionAction: UpdateSection); |
330 | } |
331 | |
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 | Pattern: SectionWildcard, MS: objcopy::MatchStyle::Wildcard, |
350 | ErrorCallback: [](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, In&: *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 | |
367 | TEST(RemoveSectionByPattern, COFF) { |
368 | SCOPED_TRACE("removeSectionByPatternCOFF" ); |
369 | |
370 | removeSectionByPatternImpl( |
371 | YamlCreationString: SimpleFileCOFFYAML, IsValidFormat: [](const Binary &File) { return File.isCOFF(); }, |
372 | SectionWildcard: "\\.text*" , SectionName: ".text" ); |
373 | } |
374 | |
375 | TEST(RemoveSectionByPattern, ELF) { |
376 | SCOPED_TRACE("removeSectionByPatternELF" ); |
377 | |
378 | removeSectionByPatternImpl( |
379 | YamlCreationString: SimpleFileELFYAML, IsValidFormat: [](const Binary &File) { return File.isELF(); }, |
380 | SectionWildcard: "\\.text*" , SectionName: ".text" ); |
381 | } |
382 | |
383 | TEST(RemoveSectionByPattern, MachO) { |
384 | SCOPED_TRACE("removeSectionByPatternMachO" ); |
385 | |
386 | removeSectionByPatternImpl( |
387 | YamlCreationString: SimpleFileMachOYAML, IsValidFormat: [](const Binary &File) { return File.isMachO(); }, |
388 | SectionWildcard: "__TEXT,__text*" , SectionName: "__text" ); |
389 | } |
390 | |
391 | TEST(RemoveSectionByPattern, Wasm) { |
392 | SCOPED_TRACE("removeSectionByPatternWasm" ); |
393 | |
394 | removeSectionByPatternImpl( |
395 | YamlCreationString: SimpleFileWasmYAML, IsValidFormat: [](const Binary &File) { return File.isWasm(); }, |
396 | SectionWildcard: "text*" , SectionName: "text" ); |
397 | } |
398 | |