Index: llvm/include/llvm/Object/MutableELFObject.h =================================================================== --- llvm/include/llvm/Object/MutableELFObject.h +++ llvm/include/llvm/Object/MutableELFObject.h @@ -10,7 +10,9 @@ #define LLVM_OBJECT_MUTABLEELFOBJECT_H #include "llvm/ADT/ArrayRef.h" +#include "llvm/ADT/Sequence.h" #include "llvm/Object/ELFObjectFile.h" +#include namespace llvm { namespace object { @@ -62,6 +64,34 @@ operator const Elf_Sym &() const { return Header; } }; +template class MutableELFSegment { + using Elf_Phdr = Elf_Phdr_Impl; + std::set Sections; + +public: + Elf_Phdr Header; + + MutableELFSegment(Elf_Phdr Header, std::vector Sections) + : Sections(Sections.begin(), Sections.end()), Header(Header) {} + + bool hasSection(SectionRef Sec) const { + return Sections.find(Sec.getRawDataRefImpl().p) != Sections.end(); + } + + void addSection(SectionRef Sec) { + Sections.insert(Sec.getRawDataRefImpl().p); + } + + void removeSection(SectionRef Sec) { + uintptr_t Index = Sec.getRawDataRefImpl().p; + assert(Sections.find(Index) != Sections.end() && + "Removing section not previously in segment"); + Sections.erase(Index); + } + + operator const Elf_Phdr &() const { return Header; } +}; + template class MutableELFObject : public ELFObjectFile { /// This class is used for a 'copy on write' effect with tables in an ELF /// object file. @@ -200,9 +230,17 @@ : nullptr; } + NewType *getIfNew(uint64_t Index) { + 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 llvm::count_if(Mappings, [](MappingType &Mapping) { + return llvm::count_if(Mappings, [](const MappingType &Mapping) { return Mapping.Type != MappingType::Removed; }); } @@ -210,13 +248,15 @@ size_t originalSize() const { return OriginalValues.size(); } }; - using Elf_Shdr = typename ELFT::Shdr; using Elf_Ehdr = typename ELFT::Ehdr; + using Elf_Shdr = typename ELFT::Shdr; + using Elf_Phdr = typename ELFT::Phdr; using Elf_Sym = typename ELFT::Sym; using MutableSymbolTable = MutableTable>; Elf_Ehdr FileHeader; + MutableTable> Segments; MutableTable> Sections; MutableSymbolTable Symbols; MutableSymbolTable DynSymbols; @@ -248,7 +288,6 @@ getSectionContents(DataRefImpl Sec) const override; void moveSymbolNext(DataRefImpl &Sym) const override; - const Elf_Sym *getSymbol(DataRefImpl Sym) const override; Expected getSymbolName(DataRefImpl Sym) const override; ArrayRef findSymbolTable(uint64_t ShType) const; uint32_t findSectionOfType(uint64_t ShType) const; @@ -259,19 +298,63 @@ : DynSymbols; } + const Elf_Phdr *getSegment(DataRefImpl Sec) const { return &Segments[Sec.p]; } + void segmentAddSection(uint64_t Seg, uint64_t Sec); + bool segmentHasSection(uint64_t Seg, uint64_t Sec) const; + MutableELFSegment &makeMutableSegment(uint64_t Seg); + static DataRefImpl toDataRef(uintptr_t Ptr); static DataRefImpl toDataRef(uint32_t A, uint32_t B); + static bool sectionWithinSegment(const Elf_Shdr &Sec, const Elf_Phdr &Seg); public: explicit MutableELFObject(ELFObjectFile &&B) : ELFObjectFile(std::move(B)), FileHeader(*reinterpret_cast(this->base())), + Segments(ArrayRef(reinterpret_cast( + this->base() + getHeader().e_phoff), + getHeader().e_phnum)), Sections(ArrayRef(reinterpret_cast( this->base() + getHeader().e_shoff), getHeader().e_shnum)), Symbols(findSymbolTable(ELF::SHT_SYMTAB)), DynSymbols(findSymbolTable(ELF::SHT_DYNSYM)) {} + class SegmentRef { + MutableELFObject *Object; + uint64_t Index; + + public: + SegmentRef(MutableELFObject *Object, uint64_t Index = 0) + : Object(Object), Index(Index) {} + + bool operator==(const SegmentRef &Other) const { + return std::tie(Object, Index) == std::tie(Other.Object, Other.Index); + } + bool operator<(const SegmentRef &Other) const { + assert(Object == Other.Object && "Not from the same object"); + return Index < Other.Index; + } + + uint64_t getIndex() const { return Index; } + + void moveNext() { ++Index; } + + const Elf_Phdr &getHeader() const { return Object->Segments[Index]; } + + MutableELFSegment &makeMutable() {} + + bool hasSection(SectionRef Sec) const { + return Object->segmentHasSection(Index, Sec.getRawDataRefImpl().p); + } + + void addSection(SectionRef Sec) { + return Object->segmentAddSection(Index, Sec.getRawDataRefImpl().p); + } + }; + + using segment_iterator = content_iterator; + section_iterator section_begin() const override { return section_iterator( SectionRef(toDataRef(Sections.getFirstIndex()), this)); @@ -308,6 +391,18 @@ this)); } + segment_iterator segment_begin() { + return segment_iterator(SegmentRef(this)); + } + + segment_iterator segment_end() { + return segment_iterator(SegmentRef(this, Segments.getLastIndex())); + } + + iterator_range segments() { + return iterator_range(segment_begin(), segment_end()); + } + /// Returns a mutable reference to the symbol at the specified index. Expected &> getMutableSymbol(SymbolRef Sym) { Expected Name = getSymbolName(Sym.getRawDataRefImpl()); @@ -346,6 +441,10 @@ return Sections.makeMutable(Sec.getRawDataRefImpl().p, Header, *Name, this); } + Expected &> getMutableSegment(SegmentRef Seg) { + return makeMutableSegment(Seg.getIndex()); + } + /// Add symbol \param Sym. MutableELFSymbol &addSymbol(MutableELFSymbol &&Sym) { return Symbols.add(std::move(Sym)); @@ -364,17 +463,11 @@ 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); } - Elf_Ehdr &getHeader() { return FileHeader; } }; @@ -409,15 +502,6 @@ Sym.d.b = Symbols.getNextIndex(Sym.d.b); } -template -const Elf_Sym_Impl * -MutableELFObject::getSymbol(DataRefImpl Sym) const { - const MutableSymbolTable &Table = getWhichTable(Sym); - if (const MutableELFSymbol *SymOrNull = Table.getConstIfNew(Sym.d.b)) - return &SymOrNull->Header; - return ELFObjectFile::getSymbol(Sym); -} - template Expected MutableELFObject::getSymbolName(DataRefImpl Sym) const { @@ -466,6 +550,55 @@ return Ref; } +template +void MutableELFObject::segmentAddSection(uint64_t Seg, uint64_t Sec) { + makeMutableSegment(Seg).addSection(Sec); +} + +template +bool MutableELFObject::segmentHasSection(uint64_t Seg, + uint64_t Sec) const { + if (const MutableELFSegment *Phdr = Segments.getConstIfNew(Seg)) + return Phdr->hasSection(SectionRef(toDataRef(Sec), this)); + return sectionWithinSegment(Sections[Sec], Segments[Seg]); +} + +template +MutableELFSegment & +MutableELFObject::makeMutableSegment(uint64_t Seg) { + if (MutableELFSegment *Phdr = Segments.getIfNew(Seg)) + return *Phdr; + + std::vector SectionsInSeg; + const Elf_Phdr &Phdr = Segments[Seg]; + for (uint64_t Index = 0; Index < Sections.size(); ++Index) + if (sectionWithinSegment(*getSection(toDataRef(Index)), Phdr)) + SectionsInSeg.push_back(Index); + return Segments.makeMutable(Seg, Phdr, SectionsInSeg); +} + +template +bool MutableELFObject::sectionWithinSegment( + const typename ELFT::Shdr &Sec, const typename ELFT::Phdr &Seg) { + uint64_t SecSize = Sec.sh_size ? Sec.sh_size : 1; + + if (Sec.sh_type == ELF::SHT_NOBITS) { + if (!(Sec.sh_flags & ELF::SHF_ALLOC)) + return false; + + bool SectionIsTLS = Sec.sh_flags & ELF::SHF_TLS; + bool SegmentIsTLS = Seg.p_type == ELF::PT_TLS; + if (SectionIsTLS != SegmentIsTLS) + return false; + + return Seg.p_vaddr <= Sec.sh_addr && + Seg.p_vaddr + Seg.p_memsz >= Sec.sh_addr + SecSize; + } + + return Seg.p_offset <= Sec.sh_offset && + Seg.p_offset + Seg.p_filesz >= Sec.sh_offset + SecSize; +} + } // namespace object } // namespace llvm Index: llvm/unittests/Object/MutableELFObjectTest.cpp =================================================================== --- llvm/unittests/Object/MutableELFObjectTest.cpp +++ llvm/unittests/Object/MutableELFObjectTest.cpp @@ -19,6 +19,10 @@ using namespace object; using namespace yaml; +template ptrdiff_t distance(T &Range) { + return std::distance(Range.begin(), Range.end()); +} + // Test that when no modifications have been made SectionRef's methods work // the same in both ELFObjectFile and MutableELFObject. TEST(MutableELFObject, NoChange) { @@ -432,7 +436,7 @@ auto NameOrErr = symbol_iterator(FirstSym)->getName(); ASSERT_THAT_EXPECTED(NameOrErr, Succeeded()); EXPECT_EQ(*NameOrErr, "first"); - MutableObject.removeSymbol(FirstSym); + MutableObject.removeSymbol(*FirstSym); auto NewFirstSym = ++MutableObject.symbol_begin(); auto NewNameOrErr = symbol_iterator(NewFirstSym)->getName(); @@ -474,7 +478,7 @@ ASSERT_FALSE(Iter->getName(Name)); EXPECT_EQ(Name, ".remove_me"); - MutableObject.removeSection(Iter); + MutableObject.removeSection(*Iter); EXPECT_EQ( std::distance(MutableObject.section_begin(), MutableObject.section_end()), @@ -493,7 +497,7 @@ ASSERT_FALSE(Iter->getName(Name)); EXPECT_EQ(Name, ".remove_me_too"); - MutableObject.removeSection(Iter); + MutableObject.removeSection(*Iter); EXPECT_EQ( std::distance(MutableObject.section_begin(), MutableObject.section_end()), @@ -586,3 +590,149 @@ Header.e_type = ELF::ET_EXEC; EXPECT_EQ(MutableObject.getHeader().e_type, ELF::ET_EXEC); } + +// Test MutableELFObject constructor properly reads program headers. +TEST(MutableELFObject, ProgramHeaders) { + SmallString<0> Storage; + Expected> ErrOrObj = yaml2ObjectFile(Storage, R"( +--- !ELF +FileHeader: + Class: ELFCLASS64 + Data: ELFDATA2LSB + Type: ET_EXEC + Machine: EM_X86_64 +ProgramHeaders: + - Type: PT_LOAD + Flags: [ PF_X, PF_R ] + VAddr: 0x1000 + PAddr: 0x2000)"); + + ASSERT_THAT_EXPECTED(ErrOrObj, Succeeded()); + auto *ELFObjFile = dyn_cast>(ErrOrObj->get()); + ASSERT_TRUE(ELFObjFile); + MutableELFObject MutableObject(std::move(*ELFObjFile)); + + auto FirstSegment = MutableObject.segment_begin(); + EXPECT_EQ(FirstSegment->getHeader().p_type, ELF::PT_LOAD); + EXPECT_TRUE(FirstSegment->getHeader().p_flags & (ELF::PF_X | ELF::PF_R)); + EXPECT_FALSE(FirstSegment->getHeader().p_flags & ELF::PF_W); + EXPECT_EQ(FirstSegment->getHeader().p_vaddr, 0x1000U); + EXPECT_EQ(FirstSegment->getHeader().p_paddr, 0x2000U); + + EXPECT_EQ( + std::distance(MutableObject.segment_begin(), MutableObject.segment_end()), + 1); +} + +TEST(MutableELFObject, SectionWithinSegment) { + SmallString<0> Storage; + Expected> ErrOrObj = yaml2ObjectFile(Storage, R"( +--- !ELF +FileHeader: + Class: ELFCLASS64 + Data: ELFDATA2LSB + Type: ET_EXEC + Machine: EM_X86_64 +Sections: + - Name: .text + Type: SHT_PROGBITS + Flags: [ SHF_ALLOC, SHF_EXECINSTR ] + Size: 24 + - Name: .data + Flags: [ SHF_ALLOC ] + Type: SHT_PROGBITS + Size: 4064 + Content: "DEADBEEF" + - Name: .bss + Type: SHT_NOBITS + Flags: [ SHF_ALLOC ] +ProgramHeaders: + - Type: PT_LOAD + Flags: [ PF_X, PF_R ] + VAddr: 0x1000 + Sections: + - Section: .text + - Type: PT_LOAD + Flags: [ PF_X, PF_W, PF_R ] + VAddr: 0x2000 + Sections: + - Section: .data)"); + + ASSERT_THAT_EXPECTED(ErrOrObj, Succeeded()); + auto *ELFObjFile = dyn_cast>(ErrOrObj->get()); + ASSERT_TRUE(ELFObjFile); + MutableELFObject MutableObject(std::move(*ELFObjFile)); + + using segment_iterator = MutableELFObject::segment_iterator; + + EXPECT_EQ( + std::distance(MutableObject.segment_begin(), MutableObject.segment_end()), + 2); + segment_iterator FirstSegment = MutableObject.segment_begin(); + segment_iterator SecondSegment = ++MutableObject.segment_begin(); + section_iterator TextSection = ++MutableObject.section_begin(); + section_iterator DataSection = ++(++MutableObject.section_begin()); + section_iterator BSSSection = ++(++(++MutableObject.section_begin())); + EXPECT_TRUE(FirstSegment->hasSection(*TextSection)); + EXPECT_FALSE(FirstSegment->hasSection(*DataSection)); + EXPECT_FALSE(FirstSegment->hasSection(*BSSSection)); + + EXPECT_FALSE(SecondSegment->hasSection(*TextSection)); + EXPECT_TRUE(SecondSegment->hasSection(*DataSection)); + EXPECT_FALSE(SecondSegment->hasSection(*BSSSection)); +} + +// Test adding and removing sections from segments +TEST(MutableELFObject, AddRemoveSections) { + SmallString<0> Storage; + Expected> ErrOrObj = yaml2ObjectFile(Storage, R"( +--- !ELF +FileHeader: + Class: ELFCLASS64 + Data: ELFDATA2LSB + Type: ET_EXEC + Machine: EM_X86_64 +Sections: + - Name: .text + Type: SHT_PROGBITS + Flags: [ SHF_ALLOC, SHF_EXECINSTR ] + Size: 24 + - Name: .data + Flags: [ SHF_ALLOC ] + Type: SHT_PROGBITS + Size: 4064 + Content: "DEADBEEF" + - Name: .bss + Type: SHT_NOBITS + Flags: [ SHF_ALLOC ] +ProgramHeaders: + - Type: PT_LOAD + Flags: [ PF_X, PF_R ] + VAddr: 0x1000 + Sections: + - Section: .text + - Type: PT_LOAD + Flags: [ PF_X, PF_W, PF_R ] + VAddr: 0x2000 + Sections: + - Section: .data)"); + + ASSERT_THAT_EXPECTED(ErrOrObj, Succeeded()); + auto *ELFObjFile = dyn_cast>(ErrOrObj->get()); + ASSERT_TRUE(ELFObjFile); + MutableELFObject MutableObject(std::move(*ELFObjFile)); + + using segment_iterator = MutableELFObject::segment_iterator; + + segment_iterator SecondSegment = ++MutableObject.segment_begin(); + section_iterator TextSection = ++MutableObject.section_begin(); + section_iterator DataSection = ++(++MutableObject.section_begin()); + + auto SegOrErr = MutableObject.getMutableSegment(*SecondSegment); + ASSERT_THAT_EXPECTED(SegOrErr, Succeeded()); + MutableELFSegment &MutSecond = *SegOrErr; + EXPECT_TRUE(MutSecond.hasSection(*DataSection)); + EXPECT_FALSE(MutSecond.hasSection(*TextSection)); + MutSecond.addSection(*TextSection); + EXPECT_TRUE(MutSecond.hasSection(*TextSection)); +}