Index: llvm/include/llvm/Object/MutableELFObject.h =================================================================== --- llvm/include/llvm/Object/MutableELFObject.h +++ llvm/include/llvm/Object/MutableELFObject.h @@ -35,6 +35,10 @@ Header.sh_size); } + MutableELFSection(const Elf_Shdr &Header, StringRef Name, + OwningArrayRef &&Data) + : Header(Header), Name(Name), Data(std::move(Data)) {} + void setData(ArrayRef Ref) { Data = OwningArrayRef(Ref); Header.sh_size = Data.size(); @@ -158,12 +162,18 @@ /// Removes the entry at \param Index. void remove(size_t Index) { - assert(Index < OriginalValues.size() && "Out of bounds"); + assert(Index < Mappings.size() && "Out of bounds"); assert(Mappings[Index].Type != MappingType::Removed && "Entry already removed"); Mappings[Index].Type = MappingType::Removed; } + /// Add's \param Value to the end of the table. + void add(NewType &&Value) { + Mappings.emplace_back(NewValues.size(), MappingType::New); + NewValues.emplace_back(std::move(Value)); + } + /// If the entry at index Index has already been made mutable, this returns /// a reference to that. Otherwise, this replaces the current entry at the /// specified index with a NewType constructued with Arguments. @@ -326,6 +336,10 @@ return Sections.makeMutable(Sec.getRawDataRefImpl().p, Header, *Name, this); } + void addSymbol(MutableELFSymbol Sym) { Symbols.add(std::move(Sym)); } + + void addSection(MutableELFSection Sec) { Sections.add(std::move(Sec)); } + /// Removes a symbol. void removeSymbol(SymbolRef Sym) { assert(Sections.getOriginal(Sym.getRawDataRefImpl().d.a).sh_type == Index: llvm/unittests/Object/MutableELFObjectTest.cpp =================================================================== --- llvm/unittests/Object/MutableELFObjectTest.cpp +++ llvm/unittests/Object/MutableELFObjectTest.cpp @@ -29,15 +29,19 @@ return Ret; } -StringRef getSectionName(const SectionRef &Sec) { - Expected NameOrErr = Sec.getName(); +template StringRef getName(const T &Named) { + Expected NameOrErr = Named.getName(); if (!NameOrErr) { consumeError(NameOrErr.takeError()); - return "Couldn't get section name"; + return "Couldn't get name"; } return *NameOrErr; } +StringRef getSectionName(const SectionRef &Sec) { return getName(Sec); } + +StringRef getSymbolName(const SymbolRef &Sym) { return getName(Sym); } + // Test that when no modifications have been made SectionRef's methods work // the same in both ELFObjectFile and MutableELFObject. TEST(MutableELFObject, NoChange) { @@ -545,3 +549,180 @@ ++Iter; EXPECT_EQ(MutableObject.section_end(), Iter); } + +TEST(MutableELFObject, AddSectionAndSymbols) { + SmallString<0> Storage; + Expected> ErrOrObj = yaml2ObjectFile(Storage, R"( +--- !ELF +FileHeader: + Class: ELFCLASS64 + Data: ELFDATA2LSB + Type: ET_REL + Machine: EM_X86_64)"); + + ASSERT_THAT_EXPECTED(ErrOrObj, Succeeded()); + auto *ELFObjFile = dyn_cast>(ErrOrObj->get()); + ASSERT_TRUE(ELFObjFile); + MutableELFObject MutableObject(std::move(*ELFObjFile)); + + std::vector SectionNames = + collect(MutableObject.sections(), getSectionName); + EXPECT_THAT(SectionNames, + ::testing::ElementsAre("", ".symtab", ".strtab", ".shstrtab")); + + Elf_Shdr_Impl Header; + Header.sh_size = 64; + + MutableObject.addSection({Header, ".newsec", OwningArrayRef(64)}); + + SectionNames = collect(MutableObject.sections(), getSectionName); + EXPECT_THAT(SectionNames, ::testing::ElementsAre("", ".symtab", ".strtab", + ".shstrtab", ".newsec")); + + auto Iter = MutableObject.section_begin(); + std::advance(Iter, 4); + + EXPECT_EQ(Iter->getSize(), 64u); + EXPECT_EQ(getSectionName(*Iter), ".newsec"); + + // yaml2obj creates the index 0 symbol even when no other symbols exist. + EXPECT_EQ( + std::distance(MutableObject.symbol_begin(), MutableObject.symbol_end()), + 1); + std::vector SymbolNames = + collect(MutableObject.symbols(), getSymbolName); + EXPECT_THAT(SymbolNames, ::testing::ElementsAre("")); + + Elf_Sym_Impl Sym; + Sym.st_shndx = ELF::SHN_ABS; + Sym.st_value = 1234; + Sym.st_size = 2; + MutableObject.addSymbol({Sym, "newsym"}); + + SymbolNames = collect(MutableObject.symbols(), getSymbolName); + EXPECT_THAT(SymbolNames, ::testing::ElementsAre("", "newsym")); + + ELFSymbolRef NewSym = *symbol_iterator(++MutableObject.symbol_begin()); + EXPECT_EQ(NewSym.getValue(), 1234u); + EXPECT_EQ(NewSym.getSize(), 2u); + Expected NameOrErr = NewSym.getName(); + ASSERT_THAT_EXPECTED(NameOrErr, Succeeded()); + EXPECT_EQ(*NameOrErr, "newsym"); +} + +// Add many symbols. +TEST(MutableELFObject, AddMany) { + SmallString<0> Storage; + Expected> ErrOrObj = yaml2ObjectFile(Storage, R"( +--- !ELF +FileHeader: + Class: ELFCLASS64 + Data: ELFDATA2LSB + Type: ET_REL + Machine: EM_X86_64)"); + + ASSERT_THAT_EXPECTED(ErrOrObj, Succeeded()); + auto *ELFObjFile = dyn_cast>(ErrOrObj->get()); + ASSERT_TRUE(ELFObjFile); + MutableELFObject MutableObject(std::move(*ELFObjFile)); + + EXPECT_EQ( + std::distance(MutableObject.symbol_begin(), MutableObject.symbol_end()), + 1); + + Elf_Sym_Impl Sym; + for (size_t I = 0; I < 100; ++I) + MutableObject.addSymbol({Sym, ""}); + + EXPECT_EQ( + std::distance(MutableObject.symbol_begin(), MutableObject.symbol_end()), + 101); +} + +// Add and then remove a symbol. +TEST(MutableELFObject, AddRemoveSymbol) { + SmallString<0> Storage; + Expected> ErrOrObj = yaml2ObjectFile(Storage, R"( +--- !ELF +FileHeader: + Class: ELFCLASS64 + Data: ELFDATA2LSB + Type: ET_REL + Machine: EM_X86_64)"); + + ASSERT_THAT_EXPECTED(ErrOrObj, Succeeded()); + auto *ELFObjFile = dyn_cast>(ErrOrObj->get()); + ASSERT_TRUE(ELFObjFile); + MutableELFObject MutableObject(std::move(*ELFObjFile)); + + EXPECT_EQ( + std::distance(MutableObject.symbol_begin(), MutableObject.symbol_end()), + 1); + + Elf_Sym_Impl Sym; + MutableObject.addSymbol({Sym, ""}); + + EXPECT_EQ( + std::distance(MutableObject.symbol_begin(), MutableObject.symbol_end()), + 2); + + basic_symbol_iterator NewSym = ++MutableObject.symbol_begin(); + MutableObject.removeSymbol(*NewSym); + + EXPECT_EQ( + std::distance(MutableObject.symbol_begin(), MutableObject.symbol_end()), + 1); +} + +// Add and then remove sections. +TEST(MutableELFObject, AddRemoveSection) { + SmallString<0> Storage; + Expected> ErrOrObj = yaml2ObjectFile(Storage, R"( +--- !ELF +FileHeader: + Class: ELFCLASS64 + Data: ELFDATA2LSB + Type: ET_REL + Machine: EM_X86_64)"); + + ASSERT_THAT_EXPECTED(ErrOrObj, Succeeded()); + auto *ELFObjFile = dyn_cast>(ErrOrObj->get()); + ASSERT_TRUE(ELFObjFile); + MutableELFObject MutableObject(std::move(*ELFObjFile)); + + auto Distance = [&MutableObject] { + return std::distance(MutableObject.section_begin(), + MutableObject.section_end()); + }; + + auto IterAtIndex = [&MutableObject, Distance](unsigned Index) { + assert(Index < Distance()); + auto Iter = MutableObject.section_begin(); + std::advance(Iter, Index); + return Iter; + }; + + EXPECT_EQ(Distance(), 4); + + Elf_Shdr_Impl Header; + Header.sh_size = 2; + MutableObject.addSection({Header, "", OwningArrayRef(0)}); + EXPECT_EQ(Distance(), 5); + EXPECT_EQ(IterAtIndex(4)->getSize(), 2u); + + MutableObject.removeSection(*(IterAtIndex(1))); + EXPECT_EQ(Distance(), 4); + EXPECT_EQ(IterAtIndex(3)->getSize(), 2u); + + Header.sh_size = 4; + MutableObject.addSection({Header, "", OwningArrayRef(0)}); + EXPECT_EQ(Distance(), 5); + EXPECT_EQ(IterAtIndex(4)->getSize(), 4u); + + MutableObject.addSection({Header, "", OwningArrayRef(0)}); + MutableObject.addSection({Header, "", OwningArrayRef(0)}); + Header.sh_size = 100; + MutableObject.addSection({Header, "", OwningArrayRef(0)}); + EXPECT_EQ(Distance(), 8); + EXPECT_EQ(IterAtIndex(7)->getSize(), 100u); +}