249 changes: 249 additions & 0 deletions llvm/unittests/ObjCopy/ObjCopyTest.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -116,3 +116,252 @@ TEST(CopySimpleInMemoryFile, Wasm) {
)",
[](const Binary &File) { return File.isWasm(); });
}

enum Action : uint8_t { AddSection, UpdateSection };

void addOrUpdateSectionToFileImpl(
const char *YamlCreationString,
std::function<bool(const Binary &File)> IsValidFormat,
StringRef NewSectionName, StringRef NewSectionData, Action SectionAction) {
auto ErrHandler = [&](const Twine &Msg) { FAIL() << "Error: " << Msg; };

// Create Object file from YAML description.
SmallVector<char> Storage;
std::unique_ptr<ObjectFile> Obj =
yaml2ObjectFile(Storage, YamlCreationString, ErrHandler);
ASSERT_TRUE(Obj);
ASSERT_TRUE(IsValidFormat(*Obj));

std::unique_ptr<MemoryBuffer> NewSectionBuffer =
MemoryBuffer::getMemBuffer(NewSectionData, NewSectionName, false);
std::string Name;
if (Obj->isMachO())
Name = "__TEXT," + NewSectionName.str();
else
Name = NewSectionName.str();

ConfigManager Config;
Config.Common.OutputFilename = "a.out";
if (SectionAction == AddSection)
Config.Common.AddSection.push_back({Name, std::move(NewSectionBuffer)});
else
Config.Common.UpdateSection.push_back({Name, std::move(NewSectionBuffer)});

// Call executeObjcopyOnBinary()
SmallVector<char> DataVector;
raw_svector_ostream OutStream(DataVector);
Error Err = objcopy::executeObjcopyOnBinary(Config, *Obj.get(), OutStream);
ASSERT_FALSE(std::move(Err));

MemoryBufferRef Buffer(StringRef(DataVector.data(), DataVector.size()),
Config.Common.OutputFilename);

// Check copied file.
Expected<std::unique_ptr<Binary>> Result = createBinary(Buffer);
ASSERT_THAT_EXPECTED(Result, Succeeded());
ASSERT_TRUE(IsValidFormat(**Result));
ASSERT_TRUE((*Result)->isObject());

// Check that copied file has the new section.
bool HasNewSection = false;
for (const object::SectionRef &Sect :
static_cast<ObjectFile *>((*Result).get())->sections()) {
Expected<StringRef> SectNameOrErr = Sect.getName();
ASSERT_THAT_EXPECTED(SectNameOrErr, Succeeded());

if (*SectNameOrErr == NewSectionName) {
HasNewSection = true;
Expected<StringRef> SectionData = Sect.getContents();
ASSERT_THAT_EXPECTED(SectionData, Succeeded());
EXPECT_TRUE(Sect.getSize() == NewSectionData.size());
EXPECT_TRUE(memcmp(SectionData->data(), NewSectionData.data(),
NewSectionData.size()) == 0);
break;
}
}
EXPECT_TRUE(HasNewSection);
}

TEST(AddSection, COFF) {
SCOPED_TRACE("addSectionToFileCOFF");

addOrUpdateSectionToFileImpl(
R"(
--- !COFF
header:
Machine: IMAGE_FILE_MACHINE_AMD64
Characteristics: [ ]
sections:
- Name: .text
Characteristics: [ ]
Alignment: 4
SectionData: E800000000C3C3C3
symbols:
...
)",
[](const Binary &File) { return File.isCOFF(); }, ".foo", "1234",
AddSection);
}

TEST(AddSection, ELF) {
SCOPED_TRACE("addSectionToFileELF");

addOrUpdateSectionToFileImpl(
R"(
--- !ELF
FileHeader:
Class: ELFCLASS64
Data: ELFDATA2LSB
Type: ET_REL)",
[](const Binary &File) { return File.isELF(); }, ".foo", "1234",
AddSection);
}

TEST(AddSection, MachO) {
SCOPED_TRACE("addSectionToFileMachO");

addOrUpdateSectionToFileImpl(
R"(
--- !mach-o
FileHeader:
magic: 0xFEEDFACF
cputype: 0x01000007
cpusubtype: 0x80000003
filetype: 0x00000001
ncmds: 1
sizeofcmds: 152
flags: 0x00002000
reserved: 0x00000000
LoadCommands:
- cmd: LC_SEGMENT_64
cmdsize: 152
segname: __TEXT
vmaddr: 0
vmsize: 4
fileoff: 184
filesize: 4
maxprot: 7
initprot: 7
nsects: 1
flags: 0
Sections:
- sectname: __text
segname: __TEXT
addr: 0x0000000000000000
content: 'AABBCCDD'
size: 4
offset: 184
align: 0
reloff: 0x00000000
nreloc: 0
flags: 0x80000400
reserved1: 0x00000000
reserved2: 0x00000000
reserved3: 0x00000000
...
)",
[](const Binary &File) { return File.isMachO(); }, "__foo", "1234",
AddSection);
}

TEST(AddSection, Wasm) {
SCOPED_TRACE("addSectionToFileWasm");

addOrUpdateSectionToFileImpl(
R"(
--- !WASM
FileHeader:
Version: 0x00000001
...
)",
[](const Binary &File) { return File.isWasm(); }, ".foo", "1234",
AddSection);
}

TEST(UpdateSection, COFF) {
SCOPED_TRACE("updateSectionToFileCOFF");

addOrUpdateSectionToFileImpl(
R"(
--- !COFF
header:
Machine: IMAGE_FILE_MACHINE_AMD64
Characteristics: [ ]
sections:
- Name: .foo
Characteristics: [ ]
Alignment: 4
SectionData: E800000000C3C3C3
symbols:
...
)",
[](const Binary &File) { return File.isCOFF(); }, ".foo", "1234",
UpdateSection);
}

TEST(UpdateSection, ELF) {
SCOPED_TRACE("updateSectionToFileELF");

addOrUpdateSectionToFileImpl(
R"(
--- !ELF
FileHeader:
Class: ELFCLASS64
Data: ELFDATA2LSB
Type: ET_REL
Sections:
- Name: .foo
Type: SHT_PROGBITS
Flags: [ SHF_ALLOC ]
Content: "12345678"
)",
[](const Binary &File) { return File.isELF(); }, ".foo", "1234",
UpdateSection);
}

TEST(UpdateSection, MachO) {
SCOPED_TRACE("updateSectionToFileMachO");

addOrUpdateSectionToFileImpl(
R"(
--- !mach-o
FileHeader:
magic: 0xFEEDFACF
cputype: 0x01000007
cpusubtype: 0x80000003
filetype: 0x00000001
ncmds: 1
sizeofcmds: 152
flags: 0x00002000
reserved: 0x00000000
LoadCommands:
- cmd: LC_SEGMENT_64
cmdsize: 152
segname: __TEXT
vmaddr: 0
vmsize: 4
fileoff: 184
filesize: 4
maxprot: 7
initprot: 7
nsects: 1
flags: 0
Sections:
- sectname: __foo
segname: __TEXT
addr: 0x0000000000000000
content: 'AABBCCDD'
size: 4
offset: 184
align: 0
reloff: 0x00000000
nreloc: 0
flags: 0x80000400
reserved1: 0x00000000
reserved2: 0x00000000
reserved3: 0x00000000
...
)",
[](const Binary &File) { return File.isMachO(); }, "__foo", "1234",
UpdateSection);
}