Index: llvm/include/llvm/Object/MutableELFObject.h =================================================================== --- llvm/include/llvm/Object/MutableELFObject.h +++ llvm/include/llvm/Object/MutableELFObject.h @@ -68,7 +68,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 +81,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 +110,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 +121,55 @@ /// 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 non remove index after \param CurrentIndex. + uint64_t getNextIndex(uint64_t CurrentIndex) const { + // <= so that Index can reach an end iterator state. + while (Mappings[++CurrentIndex].Type == MappingType::Removed && + CurrentIndex <= Mappings.size()) + ; + return CurrentIndex; + } + + /// Returns the first index which has not been removed. + uint64_t getFirstIndex() const { + for (uint64_t I = 0; I != Mappings.size(); ++I) + if (Mappings[I].Type != MappingType::Removed) + return I; + return Mappings.size(); + } + + /// Counts the number of non removed entries between 0 and \param Index. + uint64_t getRelativeIndex(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 than size + /// because size returns the number of valid (non removed) entries. + size_t getLastIndex() 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 +181,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 +235,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,32 +259,39 @@ 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.getLastIndex()), 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.getLastIndex()), + 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()), - this)); + return basic_symbol_iterator( + SymbolRef(toDataRef(findSectionOfType(ELF::SHT_DYNSYM), + DynSymbols.getLastIndex()), + this)); } /// Returns a mutable reference to the symbol at the specified index. @@ -273,16 +331,35 @@ 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"); + Symbols.remove(Sym.getRawDataRefImpl().d.b); + } + + /// Removes a symbol. + void removeSymbol(symbol_iterator Sym) { removeSymbol(*Sym); } + + /// Removes a section. + void removeSection(SectionRef Sec) { + Sections.remove(Sec.getRawDataRefImpl().p); + } + + /// Removes a section. + void removeSection(section_iterator Sec) { removeSection(*Sec); } }; 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.getRelativeIndex(Sec.p); } template @@ -301,6 +378,11 @@ return ELFObjectFile::getSectionContents(Sec); } +template +void MutableELFObject::moveSymbolNext(DataRefImpl &Sym) const { + Sym.d.b = Symbols.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 @@ -398,3 +398,111 @@ 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)"); + + ASSERT_THAT_EXPECTED(ErrOrObj, Succeeded()); + auto *ELFObjFile = dyn_cast>(ErrOrObj->get()); + ASSERT_TRUE(ELFObjFile); + MutableELFObject MutableObject(std::move(*ELFObjFile)); + + auto Distance = + std::distance(MutableObject.symbol_begin(), MutableObject.symbol_end()); + EXPECT_EQ(Distance, 3); + + auto FirstSym = ++MutableObject.symbol_begin(); + auto 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"); + + Distance = + std::distance(MutableObject.symbol_begin(), MutableObject.symbol_end()); + EXPECT_EQ(Distance, 2); +} + +// 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_EQ( + std::distance(MutableObject.section_begin(), MutableObject.section_end()), + 6); + + 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); + + // At this point Iter is on a removed index, it is invalid and any method + // should abort (from an assert). If for some reason tests are being run with + // assertions disabled, these ASSERTs should not be run. +#ifndef NDEBUG + ASSERT_DEATH(Iter->getName(Name), ".*"); + ASSERT_DEATH(Iter->getSize(), ".*"); + ASSERT_DEATH(Iter->isData(), ".*"); +#endif + + 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(), 1u); + + ASSERT_FALSE(Iter->getName(Name)); + EXPECT_NE(Name, ".remove_me"); + EXPECT_NE(Name, ".remove_me_too"); +}