Index: llvm/include/llvm/Object/MutableELFObject.h =================================================================== --- llvm/include/llvm/Object/MutableELFObject.h +++ llvm/include/llvm/Object/MutableELFObject.h @@ -35,6 +35,12 @@ Header.sh_size); } + MutableELFSection(const Elf_Shdr &Header, StringRef Name, + OwningArrayRef &&Data) + : Header(Header), Name(Name), Data(std::move(Data)) {} + + MutableELFSection(MutableELFSection &&Other) = default; + void setData(ArrayRef Ref) { Data = OwningArrayRef(Ref); Header.sh_size = Data.size(); @@ -126,13 +132,9 @@ return OriginalValues[Index]; } - /// Gets the next index that hasn't been removed after \param CurrentIndex. - /// If there is not an index after \param CurrentIndex which has not been - /// removed, this returns the index past the last entry. + /// Gets next index that hasn't been removed after \param CurrentIndex. uint64_t getNextIndex(uint64_t CurrentIndex) const { - assert(CurrentIndex < Mappings.size()); - iterator_range> Range = - seq(CurrentIndex + 1, (uint64_t)Mappings.size()); + auto Range = seq(CurrentIndex + 1, (uint64_t)Mappings.size()); return *std::find_if(Range.begin(), Range.end(), [this](uint64_t Index) { return Mappings[Index].Type != MappingType::Removed; }); @@ -153,18 +155,24 @@ return Ret; } - /// Returns the index of the last element. This is different to size() - /// because size() returns the number of entries that haven't been removed. + /// different to size() because size() returns the number of entries that + /// haven't been removed. size_t getEndIndex() const { return Mappings.size(); } /// Removes the entry at \param Index. void remove(uint64_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) { + NewValues.emplace_back(std::move(Value)); + Mappings.emplace_back(NewValues.size() - 1, MappingType::New); + } + /// 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. @@ -333,18 +341,24 @@ return Sections.makeMutable(Sec.getRawDataRefImpl().p, Header, *Name, this); } + /// Add symbol \param Sym. + void addSymbol(MutableELFSymbol &&Sym) { Symbols.add(std::move(Sym)); } + + /// Add section \param Sec. + 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 == ELF::SHT_SYMTAB && "Not pointing to symbol table"); - assert(Sym.getRawDataRefImpl().d.b && "Cannot remove index 0 symbol"); Symbols.remove(Sym.getRawDataRefImpl().d.b); } /// Removes a section. void removeSection(SectionRef Sec) { - assert(Sec.getRawDataRefImpl().p && "Cannot remove index 0 section"); Sections.remove(Sec.getRawDataRefImpl().p); } }; Index: llvm/unittests/Object/MutableELFObjectTest.cpp =================================================================== --- llvm/unittests/Object/MutableELFObjectTest.cpp +++ llvm/unittests/Object/MutableELFObjectTest.cpp @@ -349,7 +349,7 @@ ASSERT_TRUE(ELFObjFile); MutableELFObject MutableObject(std::move(*ELFObjFile)); - symbol_iterator TestSym(++MutableObject.symbol_begin()); + auto TestSym = symbol_iterator(++MutableObject.symbol_begin()); Expected &> MutSymOrErr = MutableObject.getMutableSymbol(*TestSym); ASSERT_THAT_EXPECTED(MutSymOrErr, Succeeded()); @@ -387,7 +387,7 @@ ASSERT_TRUE(ELFObjFile); MutableELFObject MutableObject(std::move(*ELFObjFile)); - symbol_iterator DynSym(++MutableObject.dynamic_symbol_begin()); + auto DynSym = symbol_iterator(++MutableObject.dynamic_symbol_begin()); EXPECT_EQ(DynSym->getValue(), 0x1234U); EXPECT_TRUE(DynSym->getFlags() & SymbolRef::SF_Global); auto NameOrErr = DynSym->getName(); @@ -480,7 +480,7 @@ MutableELFObject MutableObject(std::move(*ELFObjFile)); // 4 additional sections, index 0, .symtab, .strtab and .shstrtab are added - // by yaml2obj. + // yaml2obj. EXPECT_EQ( std::distance(MutableObject.section_begin(), MutableObject.section_end()), 6); @@ -526,3 +526,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)); + + // yaml2obj creates the .symtab, .strtab and .shstrtab sections by default. + EXPECT_EQ( + std::distance(MutableObject.section_begin(), MutableObject.section_end()), + 4); + + Elf_Shdr_Impl Header; + Header.sh_size = 64; + + MutableObject.addSection({Header, ".newsec", OwningArrayRef(64)}); + + EXPECT_EQ( + std::distance(MutableObject.section_begin(), MutableObject.section_end()), + 5); + + auto Iter = MutableObject.section_begin(); + std::advance(Iter, 4); + + EXPECT_EQ(Iter->getSize(), 64u); + StringRef Name; + EXPECT_FALSE(Iter->getName(Name)); + EXPECT_EQ(Name, ".newsec"); + + // yaml2obj creates the index 0 symbol even when no other symbols exist. + EXPECT_EQ( + std::distance(MutableObject.symbol_begin(), MutableObject.symbol_end()), + 1); + + Elf_Sym_Impl Sym; + Sym.st_shndx = ELF::SHN_ABS; + Sym.st_value = 1234; + Sym.st_size = 2; + MutableObject.addSymbol({Sym, "newsym"}); + + EXPECT_EQ( + std::distance(MutableObject.symbol_begin(), MutableObject.symbol_end()), + 2); + + 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); +}