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 }; uint64_t Index; MappedType Type; @@ -106,6 +111,7 @@ /// operator OrigType to make the proper conversion. const OrigType &operator[](uint64_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,56 @@ /// index. const OrigType &getOriginal(uint64_t Index) const { assert(Index < OriginalValues.size() && "Out of bounds"); + assert(Mappings[Index].Type != MappingType::Removed); 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. + uint64_t getNextIndex(uint64_t CurrentIndex) const { + assert(CurrentIndex < Mappings.size()); + return *llvm::find_if(seq(CurrentIndex + 1, (uint64_t)Mappings.size()), + [this](uint64_t Index) { + return Mappings[Index].Type != + MappingType::Removed; + }); + } + + /// Returns the first index which has not been removed. + uint64_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. + uint64_t getEffectiveIndex(uint64_t Index) const { + uint64_t Ret = 0; + for (uint64_t I = 0; I < Index; ++I) + if (Mappings[I].Type != MappingType::Removed) + ++Ret; + 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. + 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(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(uint64_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]]; NewValues.emplace_back(Arguments...); @@ -136,13 +183,18 @@ /// has had makeMutable called on it. Otherwise this method returns nullptr. const NewType *getConstIfNew(uint64_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 size() const { + return llvm::count_if(Mappings, [](MappingType &Mapping) { + return Mapping.Type != MappingType::Removed; + }); + } size_t originalSize() const { return OriginalValues.size(); } }; @@ -185,6 +237,7 @@ 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 +261,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)); } @@ -273,16 +332,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 @@ -301,6 +375,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,27 @@ using namespace object; using namespace yaml; +template +bool MatchesRange( + Iter Begin, Iter End, std::initializer_list List, + function_ref Extractor = [](Iter It) { return *It; }) { + if (MatchLen) + if ((size_t)std::distance(Begin, End) != List.size()) + return false; + + for (T Value : List) { + if (Value != Extractor(Begin)) + return false; + + ++Begin; + if (Begin == End) + break; + } + + return true; +} + // Test that when no modifications have been made SectionRef's methods work // the same in both ELFObjectFile and MutableELFObject. TEST(MutableELFObject, NoChange) { @@ -394,3 +415,145 @@ 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)); + + // 4 additional sections, index 0, .symtab, .strtab and .shstrtab are added + // by yaml2obj. + EXPECT_EQ( + std::distance(MutableObject.section_begin(), MutableObject.section_end()), + 6); + + bool Match = MatchesRange( + MutableObject.section_begin(), MutableObject.section_end(), + {"", ".remove_me", ".remove_me_too", ".symtab", ".strtab", ".shstrtab"}, + [](section_iterator Sec) { + StringRef Name; + Sec->getName(Name); + return Name; + }); + EXPECT_TRUE(Match); + + auto Iter = ++MutableObject.section_begin(); + StringRef Name; + ASSERT_FALSE(Iter->getName(Name)); + EXPECT_EQ(Name, ".remove_me"); + + MutableObject.removeSection(*Iter); + + EXPECT_EQ( + std::distance(MutableObject.section_begin(), MutableObject.section_end()), + 5); + + Iter = ++MutableObject.section_begin(); + ASSERT_FALSE(Iter->getName(Name)); + EXPECT_EQ(Name, ".remove_me_too"); + + MutableObject.removeSection(*Iter); + + EXPECT_EQ( + std::distance(MutableObject.section_begin(), MutableObject.section_end()), + 4); + + Iter = MutableObject.section_begin(); + EXPECT_EQ(Iter->getIndex(), 0u); + + ASSERT_FALSE(Iter->getName(Name)); + EXPECT_EQ(Name, ""); + ++Iter; + EXPECT_EQ(Iter->getIndex(), 1u); + ASSERT_FALSE(Iter->getName(Name)); + EXPECT_EQ(Name, ".symtab"); + ++Iter; + EXPECT_EQ(Iter->getIndex(), 2u); + ASSERT_FALSE(Iter->getName(Name)); + EXPECT_EQ(Name, ".strtab"); + ++Iter; + EXPECT_EQ(Iter->getIndex(), 3u); + ASSERT_FALSE(Iter->getName(Name)); + EXPECT_EQ(Name, ".shstrtab"); + ++Iter; + EXPECT_EQ(MutableObject.section_end(), Iter); +}