Index: llvm/include/llvm/Object/MutableELFObject.h =================================================================== --- llvm/include/llvm/Object/MutableELFObject.h +++ llvm/include/llvm/Object/MutableELFObject.h @@ -63,6 +63,17 @@ operator const Elf_Sym &() const { return Header; } }; +template class MutableELFSegment { + using Elf_Phdr = Elf_Phdr_Impl; + +public: + Elf_Phdr Header; + + MutableELFSegment(Elf_Phdr Header) : Header(Header) {} + + 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. @@ -197,6 +208,14 @@ : 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) { @@ -209,11 +228,13 @@ using Elf_Shdr = typename ELFT::Shdr; using Elf_Ehdr = typename ELFT::Ehdr; + 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; @@ -232,6 +253,8 @@ return ELFObjectFile::getSymbol(Sym); } + const Elf_Phdr *getSegment(DataRefImpl Sec) const { return &Segments[Sec.p]; } + /// 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 @@ -253,6 +276,8 @@ : DynSymbols; } + MutableELFSegment &makeMutableSegment(uint64_t Seg); + static DataRefImpl toDataRef(uintptr_t Ptr); static DataRefImpl toDataRef(uint32_t A, uint32_t B); @@ -260,12 +285,43 @@ 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 Object == Other.Object && Index == 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 Object->Segments.getEffectiveIndex(Index); + } + + void moveNext() { Index = Object->Segments.getNextIndex(Index); } + + const Elf_Phdr &getHeader() const { return Object->Segments[Index]; } + }; + + using segment_iterator = content_iterator; + section_iterator section_begin() const override { return section_iterator( SectionRef(toDataRef(Sections.getFirstIndex()), this)); @@ -301,6 +357,18 @@ this)); } + segment_iterator segment_begin() { + return segment_iterator(SegmentRef(this)); + } + + segment_iterator segment_end() { + return segment_iterator(SegmentRef(this, Segments.getEndIndex())); + } + + 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()); @@ -320,6 +388,11 @@ *Name); } + /// Returns a mutable reference to a segment. + Expected &> getMutableSegment(SegmentRef Segment) { + return makeMutableSegment(Segment.getIndex()); + } + /// Returns a mutable reference to the section pointed to by Sec. A possible /// usage to change all sections with alignment of 0 to 1 could be: /// @code{.cpp} @@ -347,6 +420,11 @@ Sections.add(std::move(Sec)); } + /// Add segment \param Seg. + void addSegment(MutableELFSegment &&Seg) { + Segments.add(std::move(Seg)); + } + /// Removes a symbol. void removeSymbol(SymbolRef Sym) { assert(Sections.getOriginal(Sym.getRawDataRefImpl().d.a).sh_type == @@ -444,6 +522,16 @@ return Ref; } +template +MutableELFSegment & +MutableELFObject::makeMutableSegment(uint64_t Seg) { + if (MutableELFSegment *Phdr = Segments.getIfNew(Seg)) + return *Phdr; + + const Elf_Phdr &Phdr = Segments[Seg]; + return Segments.makeMutable(Seg, Phdr); +} + } // namespace object } // namespace llvm Index: llvm/unittests/Object/MutableELFObjectTest.cpp =================================================================== --- llvm/unittests/Object/MutableELFObjectTest.cpp +++ llvm/unittests/Object/MutableELFObjectTest.cpp @@ -736,3 +736,75 @@ ASSERT_THAT_EXPECTED(StartAddr, Succeeded()); EXPECT_EQ(*StartAddr, 5u); } + +// 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 + - Type: 0x12345 + FileSize: 0 + MemSize: 100)"); + + ASSERT_THAT_EXPECTED(ErrOrObj, Succeeded()); + auto *ELFObjFile = dyn_cast>(ErrOrObj->get()); + ASSERT_TRUE(ELFObjFile); + MutableELFObject MutableObject(std::move(*ELFObjFile)); + + EXPECT_EQ( + std::distance(MutableObject.segment_begin(), MutableObject.segment_end()), + 2); + + 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); + + auto SecondSegment = ++MutableObject.segment_begin(); + EXPECT_EQ(SecondSegment->getHeader().p_type, 0x12345u); + EXPECT_EQ(SecondSegment->getHeader().p_filesz, 0u); + EXPECT_EQ(SecondSegment->getHeader().p_memsz, 100u); +} + +// Test mutating program headers. +TEST(MutableELFObject, MutateProgramHeaders) { + 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_vaddr, 0x1000u); + Expected &> MutSegmentOrErr = + MutableObject.getMutableSegment(*FirstSegment); + ASSERT_THAT_EXPECTED(MutSegmentOrErr, Succeeded()); + MutableELFSegment &MutSegment = *MutSegmentOrErr; + MutSegment.Header.p_vaddr = 0x1100; + EXPECT_EQ(FirstSegment->getHeader().p_vaddr, 0x1100u); +}