Index: llvm/include/llvm/Object/MutableELFObject.h =================================================================== --- llvm/include/llvm/Object/MutableELFObject.h +++ llvm/include/llvm/Object/MutableELFObject.h @@ -17,6 +17,7 @@ namespace object { template class MutableELFObject; +template class MutableELFSegment; /// This class is used by MutableELFObject to provide a mutable version of an /// ELFObjectFile Section. @@ -25,10 +26,12 @@ Elf_Shdr Header; std::string Name; + Optional ParentIndex; OwningArrayRef Data; - MutableELFSection(const Elf_Shdr &Header, StringRef Name, void *Base) - : Header(Header), Name(Name), + MutableELFSection(const Elf_Shdr &Header, StringRef Name, + Optional ParentIndex, void *Base) + : Header(Header), Name(Name), ParentIndex(ParentIndex), Data(OwningArrayRef(Header.sh_size)) { std::memcpy(Data.data(), reinterpret_cast(Base) + Header.sh_offset, @@ -61,6 +64,19 @@ operator const Elf_Sym &() const { return Header; } }; +template class MutableELFSegment { + using Elf_Phdr = Elf_Phdr_Impl; + +public: + Elf_Phdr Header; + Optional ParentIndex; + + MutableELFSegment(Elf_Phdr Header, Optional ParentIndex) + : Header(Header), ParentIndex(ParentIndex) {} + + 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. @@ -203,15 +219,43 @@ 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; +public: + class SegmentRef { + const MutableELFObject *Object; + size_t Index; + + public: + SegmentRef(const 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 Index; } + + void moveNext() { ++Index; } + + const Elf_Phdr &getHeader() const { return Object->Segments[Index]; } + }; + +private: /// Many getSection* methods in ELFObjectFile use getSection to get the /// the header associated with that section. This override returns a valid /// section header whether the section has been modified or not. @@ -249,17 +293,25 @@ static DataRefImpl toDataRef(uintptr_t Ptr); static DataRefImpl toDataRef(uint32_t A, uint32_t B); + static bool segmentWithinSegment(const Elf_Phdr &Child, size_t ChildIndex, + const Elf_Phdr &Parent, size_t ParentIndex); + 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)) {} + using segment_iterator = content_iterator; + section_iterator section_begin() const override { return section_iterator( SectionRef(toDataRef(Sections.getFirstIndex()), this)); @@ -295,6 +347,18 @@ this)); } + segment_iterator segment_begin() { + return segment_iterator(SegmentRef(this, 0)); + } + + 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()); @@ -331,7 +395,15 @@ Expected Name = getSectionName(Sec.getRawDataRefImpl()); if (!Name) return Name.takeError(); - return Sections.makeMutable(Sec.getRawDataRefImpl().p, Header, *Name, this); + return Sections.makeMutable(Sec.getRawDataRefImpl().p, Header, *Name, + findParentSegment(Sec), this); + } + + /// Returns a mutable reference to a segment. + Expected &> getMutableSegment(SegmentRef Segment) { + const Elf_Phdr &Phdr = Segments[Segment.getIndex()]; + return Segments.makeMutable(Segment.getIndex(), Phdr, + findParentSegment(Segment)); } void addSymbol(MutableELFSymbol Sym) { Symbols.add(std::move(Sym)); } @@ -347,6 +419,25 @@ Symbols.remove(Sym.getRawDataRefImpl().d.b); } + bool sectionInSegment(SectionRef Sec, SegmentRef Segment) const { + Optional Seg = findParentSegment(Sec); + if (!Seg) + return false; + return *Seg == Segment.getIndex(); + } + + bool segmentInSegment(SegmentRef First, SegmentRef Second) const { + Optional Segment = findParentSegment(First); + if (!Segment) + return false; + return *Segment == Second.getIndex(); + } + + /// Returns None if the section has no parent segment, otherwise returns the + /// index of the parent. + Optional findParentSegment(SectionRef Sec) const; + Optional findParentSegment(SegmentRef Segment) const; + /// Removes a section. void removeSection(SectionRef Sec) { assert(Sec.getRawDataRefImpl().p && "Cannot remove index 0 section"); @@ -437,6 +528,81 @@ return Ref; } +template +Optional +MutableELFObject::findParentSegment(SectionRef Sec) const { + uint64_t Index = Sec.getRawDataRefImpl().p; + if (const MutableELFSection *Section = Sections.getConstIfNew(Index)) + return Section->ParentIndex; + + const Elf_Shdr &Shdr = Sections[Index]; + for (size_t I = 0, E = Segments.originalSize(); I < E; ++I) { + const Elf_Phdr &Phdr = Segments.getOriginal(I); + if (sectionWithinSegment(Shdr, Phdr)) { + Optional ParentsParent = findParentSegment(SegmentRef(this, I)); + return ParentsParent ? *ParentsParent : I; + } + } + return None; +} + +template +Optional +MutableELFObject::findParentSegment(SegmentRef Segment) const { + size_t Index = Segment.getIndex(); + if (const MutableELFSegment *Segment = Segments.getConstIfNew(Index)) + return Segment->ParentIndex; + + const Elf_Phdr &Child = Segments[Index]; + for (size_t I = 0, E = Segments.originalSize(); I < E; ++I) { + const Elf_Phdr &PotentialParent = Segments.getOriginal(I); + if (segmentWithinSegment(Child, Index, PotentialParent, I)) { + Optional ParentsParent = findParentSegment(SegmentRef(this, I)); + return ParentsParent ? *ParentsParent : I; + } + } + return None; +} + +// This is identical to the same function in llvm-objcopy except it uses ELFT +// types and not the llvm-objcopy ones. +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; +} + +template +bool MutableELFObject::segmentWithinSegment( + const typename ELFT::Phdr &Segment, size_t ChildOffset, + const typename ELFT::Phdr &PotentialParent, size_t ParentOffset) { + if (Segment.p_offset == PotentialParent.p_offset && + Segment.p_filesz == PotentialParent.p_filesz) { + if (PotentialParent.p_align == Segment.p_align) + return ParentOffset < ChildOffset; + return PotentialParent.p_align > Segment.p_align; + } + + return PotentialParent.p_offset <= Segment.p_offset && + PotentialParent.p_offset + PotentialParent.p_filesz > Segment.p_offset; +} + } // namespace object } // namespace llvm Index: llvm/unittests/Object/MutableELFObjectTest.cpp =================================================================== --- llvm/unittests/Object/MutableELFObjectTest.cpp +++ llvm/unittests/Object/MutableELFObjectTest.cpp @@ -759,3 +759,287 @@ 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); +} + +// Test sectionInSegment method. +TEST(MutableELFObject, SectionInSegment) { + 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 + Flags: [ SHF_ALLOC ] + Type: SHT_NOBITS + - Name: goes_past_third + Type: 0 + Size: 0x1000 +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 + - Type: 0 + FileSize: 0x100 + Sections: + - Section: goes_past_third)"); + + 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()), + 3); + segment_iterator FirstSegment = MutableObject.segment_begin(); + segment_iterator SecondSegment = ++MutableObject.segment_begin(); + segment_iterator ThirdSegment = ++(++MutableObject.segment_begin()); + section_iterator TextSection = ++MutableObject.section_begin(); + section_iterator DataSection = ++(++MutableObject.section_begin()); + section_iterator PastThirdSection = ++(++(++MutableObject.section_begin())); + EXPECT_TRUE(MutableObject.sectionInSegment(*TextSection, *FirstSegment)); + EXPECT_FALSE(MutableObject.sectionInSegment(*DataSection, *FirstSegment)); + EXPECT_TRUE(MutableObject.sectionInSegment(*DataSection, *SecondSegment)); + EXPECT_TRUE(MutableObject.sectionInSegment(*PastThirdSection, *ThirdSegment)); +} + +// Test when a segment is in a segment. +TEST(MutableELFObject, SegmentInSegment) { + SmallString<0> Storage; + Expected> ErrOrObj = yaml2ObjectFile(Storage, R"( +--- !ELF +FileHeader: + Class: ELFCLASS64 + Data: ELFDATA2LSB + Type: ET_EXEC + Machine: EM_X86_64 +ProgramHeaders: + - Type: 0 + Offset: 0x1000 + FileSize: 0x1000 + - Type: 1 + Offset: 0x1010 + - Type: 2 + Offset: 0x3000 + FileSize: 0x100 + Align: 0x100 + - Type: 3 + Offset: 0x3000 + FileSize: 0x100 + Align: 1 + - Type: 4 # Reorder 4 and 5 so 4 doesnt come before and become + Offset: 0x4001 # the parent just because it came first in the loop. + FileSize: 0x100 + - Type: 5 # These segments can be visualized roughly like this: + Offset: 0x4000 # |- 6 -| + FileSize: 0x1000 # |---- 4 ----| + - Type: 6 # |---------- 5 --------| + Offset: 0x4002 + FileSize: 0x10)"); + + 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()), + 7); + + auto GetSegmentAtIndex = [&MutableObject](size_t Index) { + segment_iterator It = MutableObject.segment_begin(); + std::advance(It, Index); + return It; + }; + segment_iterator FirstSegment = GetSegmentAtIndex(0); + segment_iterator SecondSegment = GetSegmentAtIndex(1); + segment_iterator ThirdSegment = GetSegmentAtIndex(2); + segment_iterator FourthSegment = GetSegmentAtIndex(3); + EXPECT_TRUE(MutableObject.segmentInSegment(*SecondSegment, *FirstSegment)); + EXPECT_FALSE(MutableObject.segmentInSegment(*FirstSegment, *FirstSegment)) + << "Segment should not be its own parent"; + EXPECT_TRUE(MutableObject.segmentInSegment(*FourthSegment, *ThirdSegment)) + << "Segments with same offset and size should choose larger alignment to " + "be the parent"; + EXPECT_FALSE(MutableObject.segmentInSegment(*ThirdSegment, *FourthSegment)); + + Optional ParentOf6 = + MutableObject.findParentSegment(*GetSegmentAtIndex(6)); + ASSERT_NE(ParentOf6, None); + EXPECT_EQ(*ParentOf6, 5u) + << "Segment 6 fits in 4, but the most parental segment is 5"; + Optional ParentOf4 = + MutableObject.findParentSegment(*GetSegmentAtIndex(4)); + ASSERT_NE(ParentOf4, None); + EXPECT_EQ(*ParentOf4, 5u) << "Segment 4's parent should be 5"; + Optional ParentOf5 = + MutableObject.findParentSegment(*GetSegmentAtIndex(5)); + EXPECT_EQ(ParentOf5, None) << "Segment 5 should have no parent"; +} + +TEST(MutableELFObject, SameSegmentOneParent) { + SmallString<0> Storage; + Expected> ErrOrObj = yaml2ObjectFile(Storage, R"( +--- !ELF +FileHeader: + Class: ELFCLASS64 + Data: ELFDATA2LSB + Type: ET_EXEC + Machine: EM_X86_64 +ProgramHeaders: + - Type: 0 + Offset: 0x100 + FileSize: 0x100 + Align: 0x100 + - Type: 0 + Offset: 0x100 + FileSize: 0x100 + Align: 0x100)"); + + 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 First = MutableObject.segment_begin(); + segment_iterator Second = ++MutableObject.segment_begin(); + + Optional FirstParent = MutableObject.findParentSegment(*First); + EXPECT_EQ(FirstParent, None) + << "First segment shouldn't have a parent because it is first"; + Optional SecondParent = MutableObject.findParentSegment(*Second); + ASSERT_NE(SecondParent, None); + EXPECT_EQ(*SecondParent, 0u); +} + +TEST(MutableELFObject, TrueSectionParent) { + SmallString<0> Storage; + Expected> ErrOrObj = yaml2ObjectFile(Storage, R"( +--- !ELF +FileHeader: + Class: ELFCLASS64 + Data: ELFDATA2LSB + Type: ET_EXEC + Machine: EM_X86_64 +Sections: + - Name: 1 + Type: 0 + ShOffset: 1010 + Size: 50 + - Name: 2 + Type: 0 +ProgramHeaders: + - Type: 0 # Put fake parent first. + Offset: 1000 + FileSize: 1000 + - Type: 1 + Offset: 0 + FileSize: 10000)"); + + ASSERT_THAT_EXPECTED(ErrOrObj, Succeeded()); + auto *ELFObjFile = dyn_cast>(ErrOrObj->get()); + ASSERT_TRUE(ELFObjFile); + MutableELFObject MutableObject(std::move(*ELFObjFile)); + + section_iterator First = ++MutableObject.section_begin(); + Optional ParentSegment = MutableObject.findParentSegment(*First); + ASSERT_NE(ParentSegment, None); + EXPECT_EQ(*ParentSegment, 1u); +}