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,40 @@ 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(section_iterator Sec) const { + return Sections.find(Sec->getRawDataRefImpl().p) != Sections.end(); + } + bool hasSection(SectionRef Sec) const { + return Sections.find(Sec.getRawDataRefImpl().p) != Sections.end(); + } + + void addSection(section_iterator Sec) { + Sections.insert(Sec->getRawDataRefImpl().p); + } + void addSection(SectionRef Sec) { + Sections.insert(Sec.getRawDataRefImpl().p); + } + + void removeSection(section_iterator 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 +236,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 +254,15 @@ size_t originalSize() const { return OriginalValues.size(); } }; - using Elf_Shdr = Elf_Shdr_Impl; using Elf_Ehdr = Elf_Ehdr_Impl; + using Elf_Shdr = Elf_Shdr_Impl; + using Elf_Phdr = Elf_Phdr_Impl; using Elf_Sym = Elf_Sym_Impl; using MutableSymbolTable = MutableTable>; Elf_Ehdr FileHeader; + MutableTable> Segments; MutableTable> Sections; MutableSymbolTable Symbols; MutableSymbolTable DynSymbols; @@ -252,19 +298,65 @@ : 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); + } + bool hasSection(section_iterator Sec) const { return hasSection(*Sec); } + + void addSection(SectionRef Sec) { + return Object->segmentAddSection(Index, Sec.getRawDataRefImpl().p); + } + void addSection(section_iterator Sec) { return addSection(*Sec); } + }; + + using segment_iterator = content_iterator; + section_iterator section_begin() const override { return section_iterator( SectionRef(toDataRef(Sections.getFirstIndex()), this)); @@ -301,6 +393,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()); @@ -348,6 +452,14 @@ return getMutableSection(*Sec); } + Expected &> getMutableSegment(SegmentRef Seg) { + return makeMutableSegment(Seg.getIndex()); + } + + Expected &> getMutableSegment(segment_iterator Seg) { + return getMutableSegment(*Seg); + } + /// Add symbol \param Sym. MutableELFSymbol &addSymbol(MutableELFSymbol &&Sym) { return Symbols.add(std::move(Sym)); @@ -478,6 +590,55 @@ return Sections.makeMutable(Sec.getRawDataRefImpl().p, Header, *Name, this); } +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 Elf_Shdr_Impl &Sec, const Elf_Phdr_Impl &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) { @@ -584,3 +588,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)); +}