llvm/unittests/ObjCopy/ObjCopyTest.cpp

   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 }