Index: llvm/include/llvm/Object/MutableELFObject.h =================================================================== --- llvm/include/llvm/Object/MutableELFObject.h +++ llvm/include/llvm/Object/MutableELFObject.h @@ -61,6 +61,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. @@ -198,16 +209,26 @@ : nullptr; } + NewType *getIfNew(size_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; + } + size_t originalSize() const { return OriginalValues.size(); } }; 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; @@ -226,6 +247,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 @@ -247,6 +270,8 @@ : DynSymbols; } + MutableELFSegment &makeMutableSegment(uint64_t Seg); + static DataRefImpl toDataRef(uintptr_t Ptr); static DataRefImpl toDataRef(uint32_t A, uint32_t B); @@ -254,12 +279,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; + size_t Index = 0; + + public: + SegmentRef(MutableELFObject *Object, size_t Index) + : 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; + } + + size_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)); @@ -295,6 +351,16 @@ this)); } + segment_iterator segment_begin() { return segment_iterator({this, 0}); } + + segment_iterator segment_end() { + return segment_iterator({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()); @@ -314,6 +380,15 @@ *Name); } + /// Returns a mutable reference to a segment. + Expected &> getMutableSegment(SegmentRef Segment) { + if (MutableELFSegment *Phdr = Segments.getIfNew(Segment.getIndex())) + return *Phdr; + + const Elf_Phdr &Phdr = Segments[Segment.getIndex()]; + return Segments.makeMutable(Segment.getIndex(), Phdr); + } + /// 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} @@ -338,6 +413,8 @@ void addSection(MutableELFSection Sec) { Sections.add(std::move(Sec)); } + 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 == @@ -437,6 +514,10 @@ return Ref; } +template +MutableELFSegment & +MutableELFObject::makeMutableSegment(uint64_t Seg) {} + } // namespace object } // namespace llvm Index: llvm/unittests/Object/MutableELFObjectTest.cpp =================================================================== --- llvm/unittests/Object/MutableELFObjectTest.cpp +++ llvm/unittests/Object/MutableELFObjectTest.cpp @@ -759,3 +759,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 ] + Offset: 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_offset, 0x1000u); + Expected &> MutSegmentOrErr = + MutableObject.getMutableSegment(*FirstSegment); + ASSERT_THAT_EXPECTED(MutSegmentOrErr, Succeeded()); + MutableELFSegment &MutSegment = *MutSegmentOrErr; + MutSegment.Header.p_offset = 0x1100; + EXPECT_EQ(FirstSegment->getHeader().p_offset, 0x1100u); +}