Index: llvm/include/llvm/Object/MutableELFObject.h =================================================================== --- llvm/include/llvm/Object/MutableELFObject.h +++ llvm/include/llvm/Object/MutableELFObject.h @@ -10,6 +10,7 @@ #define LLVM_OBJECT_MUTABLEELFOBJECT_H #include "llvm/ADT/ArrayRef.h" +#include "llvm/ADT/Sequence.h" #include "llvm/Object/ELFObjectFile.h" namespace llvm { @@ -68,7 +69,11 @@ /// The table keeps a list of mappings, these mappings can have one of two /// states either Original or New. In the case of Original the index /// associated with the mapping is into the original table in the file. For - /// a New mapping, the index is into the NewValues vector. + /// a New mapping, the index is into the NewValues vector. A Removed mapping + /// means that an entry has been removed and should be skipped over when + /// iterating over the MutableTable. Notably removed values stay in the table + /// to preserve the index of all other other entries. Similarly insertion is + /// not supported. /// /// The design of MutableTable allows fewer copies to be made than there /// would otherwise need to be: entries with no modifications never get @@ -77,7 +82,7 @@ /// allowed by the object file standard. template class MutableTable { struct MappingType { - enum MappedType { Original, New }; + enum MappedType { Original, New, Removed }; size_t Index; MappedType Type; @@ -106,6 +111,7 @@ /// operator OrigType to make the proper conversion. const OrigType &operator[](size_t Index) const { assert(Index < Mappings.size() && "Out of bounds"); + assert(Mappings[Index].Type != MappingType::Removed); if (Mappings[Index].Type == MappingType::New) return static_cast(NewValues[Mappings[Index]]); return OriginalValues[Mappings[Index]]; @@ -116,15 +122,54 @@ /// index. const OrigType &getOriginal(size_t Index) const { assert(Index < OriginalValues.size() && "Out of bounds"); + assert(Mappings[Index].Type != MappingType::Removed); return OriginalValues[Index]; } + /// Gets the index of the next entry that hasn't been removed after \param + /// CurrentIndex. If there is no entry after \param CurrentIndex which has + /// not been removed, this returns the index one past the last entry. + size_t getNextIndex(size_t CurrentIndex) const { + assert(CurrentIndex < Mappings.size()); + return *llvm::find_if( + seq(CurrentIndex + 1, Mappings.size()), [this](size_t Index) { + return Mappings[Index].Type != MappingType::Removed; + }); + } + + /// Returns the first index which has not been removed. + size_t getFirstIndex() const { + return Mappings[0].Type != MappingType::Removed ? 0 : getNextIndex(0); + } + + /// Counts the number of entries that haven't been removed between 0 and + /// \param Index. + size_t getEffectiveIndex(size_t Index) const { + size_t Ret = 0; + for (size_t I = 0; I < Index; ++I) + if (Mappings[I].Type != MappingType::Removed) + ++Ret; + return Ret; + } + + /// Returns the index of the last element. + size_t getEndIndex() const { return Mappings.size(); } + + /// Removes the entry at \param Index. + void remove(size_t Index) { + assert(Index < OriginalValues.size() && "Out of bounds"); + assert(Mappings[Index].Type != MappingType::Removed && + "Entry already removed"); + Mappings[Index].Type = MappingType::Removed; + } + /// 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. template NewType &makeMutable(size_t Index, Args &&... Arguments) { assert(Index < Mappings.size() && "Out of bounds"); + assert(Mappings[Index].Type != MappingType::Removed); if (Mappings[Index].Type == MappingType::New) return NewValues[Mappings[Index]]; Mappings[Index] = MappingType(NewValues.size(), MappingType::New); @@ -136,14 +181,12 @@ /// has had makeMutable called on it. Otherwise this method returns nullptr. const NewType *getConstIfNew(size_t Index) const { assert(Index < Mappings.size() && "Out of bounds"); + assert(Mappings[Index].Type != MappingType::Removed); return Mappings[Index].Type == MappingType::New ? &NewValues[Mappings[Index]] : nullptr; } - /// Return the number of elements in the table. - size_t size() const { return Mappings.size(); } - size_t originalSize() const { return OriginalValues.size(); } }; @@ -168,10 +211,6 @@ return &Sections[Sec.p]; } - /// moveSectionNext must be overridden because MutableELFObject::section_begin - /// works differently than in ELFObjectFile. In this class, sections are - /// iterated with their index, not address in the file, which allows use with - /// MutableTable. const Elf_Sym *getSymbol(DataRefImpl Sym) const override { const MutableSymbolTable &Table = getSymbolTable(Sym); if (const MutableELFSymbol *SymOrNull = Table.getConstIfNew(Sym.d.b)) @@ -179,12 +218,17 @@ return ELFObjectFile::getSymbol(Sym); } + /// moveSectionNext must be overriden because the section_iterators in + /// MutableELFObject work differently than in ELFObjectFile. In this class, + /// sections are iterated with their index, not address in the file, which + /// allows use with MutableTable. void moveSectionNext(DataRefImpl &Sec) const override; Expected getSectionName(DataRefImpl Sec) const override; uint64_t getSectionIndex(DataRefImpl Sec) const override; Expected> getSectionContents(DataRefImpl Sec) const override; + void moveSymbolNext(DataRefImpl &Sym) const override; Expected getSymbolName(DataRefImpl Sym) const override; ArrayRef findSymbolTable(uint64_t ShType) const; uint32_t findSectionOfType(uint64_t ShType) const; @@ -208,31 +252,37 @@ DynSymbols(findSymbolTable(ELF::SHT_DYNSYM)) {} section_iterator section_begin() const override { - return section_iterator(SectionRef(toDataRef(0), this)); + return section_iterator( + SectionRef(toDataRef(Sections.getFirstIndex()), this)); } section_iterator section_end() const override { - return section_iterator(SectionRef(toDataRef(Sections.size()), this)); + return section_iterator( + SectionRef(toDataRef(Sections.getEndIndex()), this)); } basic_symbol_iterator symbol_begin() const override { - return basic_symbol_iterator( - SymbolRef(toDataRef(findSectionOfType(ELF::SHT_SYMTAB), 0), this)); + return basic_symbol_iterator(SymbolRef( + toDataRef(findSectionOfType(ELF::SHT_SYMTAB), Symbols.getFirstIndex()), + this)); } basic_symbol_iterator symbol_end() const override { return basic_symbol_iterator(SymbolRef( - toDataRef(findSectionOfType(ELF::SHT_SYMTAB), Symbols.size()), this)); + toDataRef(findSectionOfType(ELF::SHT_SYMTAB), Symbols.getEndIndex()), + this)); } elf_symbol_iterator dynamic_symbol_begin() const override { return basic_symbol_iterator( - SymbolRef(toDataRef(findSectionOfType(ELF::SHT_DYNSYM), 0), this)); + SymbolRef(toDataRef(findSectionOfType(ELF::SHT_DYNSYM), + DynSymbols.getFirstIndex()), + this)); } elf_symbol_iterator dynamic_symbol_end() const override { return basic_symbol_iterator(SymbolRef( - toDataRef(findSectionOfType(ELF::SHT_DYNSYM), DynSymbols.size()), + toDataRef(findSectionOfType(ELF::SHT_DYNSYM), DynSymbols.getEndIndex()), this)); } @@ -274,16 +324,31 @@ return Name.takeError(); return Sections.makeMutable(Sec.getRawDataRefImpl().p, Header, *Name, this); } + + /// 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); + } }; template void MutableELFObject::moveSectionNext(DataRefImpl &Sec) const { - ++Sec.p; + Sec.p = Sections.getNextIndex(Sec.p); } template uint64_t MutableELFObject::getSectionIndex(DataRefImpl Sec) const { - return Sec.p; + return Sections.getEffectiveIndex(Sec.p); } template @@ -302,6 +367,12 @@ return ELFObjectFile::getSectionContents(Sec); } +template +void MutableELFObject::moveSymbolNext(DataRefImpl &Sym) const { + const MutableSymbolTable &SymbolTable = getSymbolTable(Sym); + Sym.d.b = SymbolTable.getNextIndex(Sym.d.b); +} + template Expected MutableELFObject::getSymbolName(DataRefImpl Sym) const { Index: llvm/unittests/Object/MutableELFObjectTest.cpp =================================================================== --- llvm/unittests/Object/MutableELFObjectTest.cpp +++ llvm/unittests/Object/MutableELFObjectTest.cpp @@ -19,6 +19,25 @@ using namespace object; using namespace yaml; +template +std::vector collect( + Range TheRange, + function_ref Extractor = + [](auto It) { return *It; }) { + std::vector Ret; + llvm::transform(TheRange, std::back_inserter(Ret), Extractor); + return Ret; +} + +StringRef getSectionName(const SectionRef &Sec) { + Expected NameOrErr = Sec.getName(); + if (!NameOrErr) { + consumeError(NameOrErr.takeError()); + return "Couldn't get section name"; + } + return *NameOrErr; +} + // Test that when no modifications have been made SectionRef's methods work // the same in both ELFObjectFile and MutableELFObject. TEST(MutableELFObject, NoChange) { @@ -114,38 +133,18 @@ ASSERT_TRUE(ELFObjFile); MutableELFObject MutableObject(std::move(*ELFObjFile)); - auto compareSectionName = [](section_iterator Iter, const char *Name) { - auto NameOrErr = Iter->getName(); - ASSERT_THAT_EXPECTED(NameOrErr, Succeeded()); - EXPECT_EQ(*NameOrErr, Name); - }; - - ptrdiff_t NumSections = - std::distance(MutableObject.section_begin(), MutableObject.section_end()); - ASSERT_EQ(NumSections, 7); - - auto Iter = MutableObject.section_begin(); - compareSectionName(Iter, nullptr); - - compareSectionName(++Iter, ".sec0"); - compareSectionName(++Iter, ".sec1"); - compareSectionName(++Iter, ".sec2"); + EXPECT_THAT(collect(MutableObject.sections(), getSectionName), + ::testing::ElementsAre("", ".sec0", ".sec1", ".sec2", ".symtab", + ".strtab", ".shstrtab")); - Iter = MutableObject.section_begin(); + section_iterator Iter = MutableObject.section_begin(); std::advance(Iter, 2); auto MutSectionOrErr = MutableObject.getMutableSection(*Iter); - ASSERT_EQ(std::distance(MutableObject.section_begin(), Iter), 2); - ASSERT_EQ( - std::distance(MutableObject.section_begin(), MutableObject.section_end()), - 7); ASSERT_THAT_EXPECTED(MutSectionOrErr, Succeeded()); MutSectionOrErr->Name = ".new_name"; - - Iter = MutableObject.section_begin(); - compareSectionName(Iter, nullptr); - compareSectionName(++Iter, ".sec0"); - compareSectionName(++Iter, ".new_name"); - compareSectionName(++Iter, ".sec2"); + EXPECT_THAT(collect(MutableObject.sections(), getSectionName), + ::testing::ElementsAre("", ".sec0", ".new_name", ".sec2", + ".symtab", ".strtab", ".shstrtab")); } // Test MutableELFSection::setData(). @@ -394,3 +393,119 @@ ASSERT_THAT_EXPECTED(NameOrErr, Succeeded()); EXPECT_EQ(*NameOrErr, "test"); } + +// Test removeSymbol method. +TEST(MutableELFObject, RemoveSymbols) { + SmallString<0> Storage; + Expected> ErrOrObj = yaml2ObjectFile(Storage, R"( +--- !ELF +FileHeader: + Class: ELFCLASS64 + Data: ELFDATA2LSB + Type: ET_REL + Machine: EM_X86_64 +Symbols: + - Name: first + Index: SHN_ABS + Value: 0x1234 + Binding: STB_LOCAL + - Name: second + Index: SHN_ABS + Value: 0 + Binding: STB_GLOBAL +DynamicSymbols: + - Name: a + - Name: b)"); + + 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()), + 3); + + EXPECT_EQ(std::distance(MutableObject.dynamic_symbol_begin(), + MutableObject.dynamic_symbol_end()), + 3); + + auto FirstSym = ++MutableObject.symbol_begin(); + Expected NameOrErr = symbol_iterator(FirstSym)->getName(); + ASSERT_THAT_EXPECTED(NameOrErr, Succeeded()); + EXPECT_EQ(*NameOrErr, "first"); + MutableObject.removeSymbol(*FirstSym); + + auto NewFirstSym = ++MutableObject.symbol_begin(); + auto NewNameOrErr = symbol_iterator(NewFirstSym)->getName(); + ASSERT_THAT_EXPECTED(NewNameOrErr, Succeeded()); + EXPECT_EQ(*NewNameOrErr, "second"); + + EXPECT_EQ( + std::distance(MutableObject.symbol_begin(), MutableObject.symbol_end()), + 2); + EXPECT_EQ(std::distance(MutableObject.dynamic_symbol_begin(), + MutableObject.dynamic_symbol_end()), + 3); + + MutableObject.removeSymbol(*NewFirstSym); + EXPECT_EQ( + std::distance(MutableObject.symbol_begin(), MutableObject.symbol_end()), + 1); + EXPECT_EQ(std::distance(MutableObject.dynamic_symbol_begin(), + MutableObject.dynamic_symbol_end()), + 3); +} + +// Test removeSection method. +TEST(MutableELFObject, RemoveSections) { + SmallString<0> Storage; + Expected> ErrOrObj = yaml2ObjectFile(Storage, R"( +--- !ELF +FileHeader: + Class: ELFCLASS64 + Data: ELFDATA2LSB + Type: ET_REL + Machine: EM_X86_64 +Sections: + - Name: .remove_me + Type: SHT_PROGBITS + - Name: .remove_me_too + Type: SHT_PROGBITS)"); + + ASSERT_THAT_EXPECTED(ErrOrObj, Succeeded()); + auto *ELFObjFile = dyn_cast>(ErrOrObj->get()); + ASSERT_TRUE(ELFObjFile); + MutableELFObject MutableObject(std::move(*ELFObjFile)); + + EXPECT_THAT(collect(MutableObject.sections(), getSectionName), + ::testing::ElementsAre("", ".remove_me", ".remove_me_too", + ".symtab", ".strtab", ".shstrtab")); + + auto Iter = ++MutableObject.section_begin(); + Expected NameOrErr = Iter->getName(); + ASSERT_THAT_EXPECTED(NameOrErr, Succeeded()); + EXPECT_EQ(*NameOrErr, ".remove_me"); + MutableObject.removeSection(*Iter); + + EXPECT_THAT(collect(MutableObject.sections(), getSectionName), + ::testing::ElementsAre("", ".remove_me_too", ".symtab", ".strtab", + ".shstrtab")); + + Iter = ++MutableObject.section_begin(); + MutableObject.removeSection(*Iter); + + EXPECT_THAT(collect(MutableObject.sections(), getSectionName), + ::testing::ElementsAre("", ".symtab", ".strtab", ".shstrtab")); + + Iter = MutableObject.section_begin(); + EXPECT_EQ(Iter->getIndex(), 0u); + ++Iter; + EXPECT_EQ(Iter->getIndex(), 1u); + ++Iter; + EXPECT_EQ(Iter->getIndex(), 2u); + ++Iter; + EXPECT_EQ(Iter->getIndex(), 3u); + ++Iter; + EXPECT_EQ(MutableObject.section_end(), Iter); +}