Index: llvm/include/llvm/Object/MutableELFObject.h =================================================================== --- llvm/include/llvm/Object/MutableELFObject.h +++ llvm/include/llvm/Object/MutableELFObject.h @@ -11,6 +11,7 @@ #include "llvm/ADT/ArrayRef.h" #include "llvm/Object/ELFObjectFile.h" +#include namespace llvm { namespace object { @@ -78,6 +79,28 @@ return Ref; } +template +static bool 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; +} + template struct MutableELFSection { Elf_Shdr_Impl Header; std::string Name; @@ -95,10 +118,78 @@ } }; +template class MutableELFSegment { + std::set Sections; + const MutableELFObject *ObjFile; + typename ELFT::Off OriginalOffset; + +public: + Elf_Phdr_Impl Header; + ArrayRef Contents; + + MutableELFSegment(uint64_t ToCopy, const MutableELFObject *ObjFile) + : ObjFile(ObjFile) { + const Elf_Ehdr_Impl &EHeader = + *reinterpret_cast *>(ObjFile->base()); + assert(ToCopy < EHeader.e_phnum && "Copying segment that doesn't exist"); + Header = reinterpret_cast *>( + ObjFile->base() + EHeader.e_phoff)[ToCopy]; + OriginalOffset = Header.p_offset; + Contents = ArrayRef( + reinterpret_cast(ObjFile->base() + Header.p_offset), + Header.p_filesz); + + ArrayRef> SecHeaders( + reinterpret_cast *>(ObjFile->base() + + EHeader.e_shoff), + EHeader.e_shnum); + uint64_t Index = 0; + for (const auto &Sec : SecHeaders) { + if (sectionWithinSegment(Sec, Header)) + Sections.insert(Index); + ++Index; + } + } + + void addSection(section_iterator 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); + } + + bool sectionInSegment(section_iterator Sec) const { + return Sections.find(Sec->getRawDataRefImpl().p) != Sections.end(); + } + + class iterator : public std::set::iterator { + using Base = std::set::iterator; + const MutableELFObject *ObjFile; + + public: + iterator(Base B, const MutableELFObject *ObjFile) + : Base(B), ObjFile(ObjFile) {} + + SectionRef operator*() const { + return SectionRef(toDataRef(Base::operator*()), ObjFile); + } + }; + + iterator begin() { return iterator(Sections.begin(), ObjFile); } + + iterator end() { return iterator(Sections.end(), ObjFile); } +}; + template class MutableELFObject : public ELFObjectFile { friend struct MutableELFSection; + friend class MutableELFSegment; MutableRange> Sections; + std::vector> Segments; protected: using MappingType = @@ -129,10 +220,15 @@ bool isBerkeleyData(DataRefImpl Sec) const override; public: - MutableELFObject(ELFObjectFile &B) + explicit MutableELFObject(ELFObjectFile &B) : ELFObjectFile(std::move(B)), Sections(B.section_begin(), B.section_end(), - [&](SectionRef Ref) { return Ref.getRawDataRefImpl().p; }) {} + [&](SectionRef Ref) { return Ref.getRawDataRefImpl().p; }) { + const Elf_Ehdr_Impl *Header = + reinterpret_cast *>(this->base()); + for (int I = 0; I < Header->e_phnum; ++I) + Segments.emplace_back(I, this); + } section_iterator section_begin() const override { return section_iterator(SectionRef(toDataRef(0), this)); @@ -142,6 +238,19 @@ return section_iterator(SectionRef(toDataRef(Sections.size()), this)); } + section_iterator findSectionByName(StringRef Name) const; + + using segment_iterator = + typename std::vector>::iterator; + + segment_iterator segment_begin() { return Segments.begin(); } + + segment_iterator segment_end() { return Segments.end(); } + + iterator_range segments() { + return iterator_range(segment_begin(), segment_end()); + } + MutableELFSection *getMutableSection(section_iterator Sec) { uintptr_t Index = Sec->getRawDataRefImpl().p; return Sections.makeMutable(Index, this); @@ -230,6 +339,20 @@ return ELFObjectFile::isBerkeleyData(getSectionRef(Sec)); } +template +section_iterator +MutableELFObject::findSectionByName(StringRef Name) const { + const section_iterator End = section_end(); + for (section_iterator Iter = section_begin(); Iter != End; ++Iter) { + StringRef SecName; + if (Iter->getName(SecName)) + return section_end(); + if (SecName == Name) + return Iter; + } + return section_end(); +} + } // namespace object } // namespace llvm Index: llvm/unittests/Object/MutableELFObjectTest.cpp =================================================================== --- llvm/unittests/Object/MutableELFObjectTest.cpp +++ llvm/unittests/Object/MutableELFObjectTest.cpp @@ -18,6 +18,14 @@ using namespace object; using namespace yaml; +template static ptrdiff_t iterRangeLength(const T &Range) { + return std::distance(Range.begin(), Range.end()); +} + +template static ptrdiff_t iterRangeLength(T &Range) { + return std::distance(Range.begin(), Range.end()); +} + // Change a sections name and test that SectionRef::getName() returns the new // name. TEST(MutableELFObject, ChangeSectionName) { @@ -96,7 +104,7 @@ MutableELFObject MutableObject(*ELFObjFile); ptrdiff_t NumSections = - std::distance(MutableObject.section_begin(), MutableObject.section_end()); + std::distance(MutableObject.section_begin(), MutableObject.section_end()); auto FirstSec = ++MutableObject.section_begin(); Expected Contents = FirstSec->getContents(); @@ -159,9 +167,9 @@ const ObjectFile &ObjFile = *ErrOrObj->get(); ptrdiff_t ObjFileSecs = - std::distance(ELFObjFile->section_begin(), ELFObjFile->section_end()); + std::distance(ELFObjFile->section_begin(), ELFObjFile->section_end()); ptrdiff_t MutObjSecs = - std::distance(MutableObject.section_begin(), MutableObject.section_end()); + std::distance(MutableObject.section_begin(), MutableObject.section_end()); EXPECT_EQ(ObjFileSecs, MutObjSecs); auto TestSections = [](SectionRef ObjFile, SectionRef MutObj) { @@ -180,8 +188,138 @@ EXPECT_EQ_ITER(isStripped); EXPECT_EQ_ITER(isText); EXPECT_EQ_ITER(isVirtual); +#undef EXPECT_EQ_ITER }; for (const auto &Tuple : zip(MutableObject.sections(), ObjFile.sections())) TestSections(std::get<0>(Tuple), std::get<1>(Tuple)); } + +// Test findSectionByName method +TEST(MutableELFObject, FindSectionByName) { + 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 ] + AddressAlign: 0x1000 + Size: 24)"); + + ASSERT_THAT_EXPECTED(ErrOrObj, Succeeded()); + auto *ELFObjFile = dyn_cast>(ErrOrObj->get()); + ASSERT_TRUE(ELFObjFile); + MutableELFObject MutableObject(*ELFObjFile); + + section_iterator Section1 = ++MutableObject.section_begin(); + section_iterator FoundSection = MutableObject.findSectionByName(".text"); + EXPECT_EQ(Section1, FoundSection); + + section_iterator DoesntExist = + MutableObject.findSectionByName("doesnt exist"); + EXPECT_EQ(DoesntExist, MutableObject.section_end()); +} + +// 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(*ELFObjFile); + + using segment_iterator = MutableELFObject::segment_iterator; + + EXPECT_EQ(iterRangeLength(MutableObject.segments()), 1) + << "More segments than expected"; + segment_iterator FirstSegment = MutableObject.segment_begin(); + EXPECT_TRUE(FirstSegment->Header.p_flags & (ELF::PF_R | ELF::PF_X)); + EXPECT_FALSE(FirstSegment->Header.p_flags & ELF::PF_W); + EXPECT_EQ(FirstSegment->Header.p_vaddr, 0x1000U); + EXPECT_EQ(FirstSegment->Header.p_paddr, 0x2000U); + EXPECT_EQ(FirstSegment->Header.p_type, ELF::PT_LOAD); +} + +// 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(*ELFObjFile); + + using segment_iterator = MutableELFObject::segment_iterator; + + EXPECT_EQ(iterRangeLength(MutableObject.segments()), 2); + segment_iterator FirstSegment = MutableObject.segment_begin(); + segment_iterator SecondSegment = FirstSegment + 1; + + section_iterator TextSection = MutableObject.findSectionByName(".text"); + section_iterator RoDataSection = MutableObject.findSectionByName(".data"); + section_iterator BssSection = MutableObject.findSectionByName(".bss"); + EXPECT_TRUE(FirstSegment->sectionInSegment(TextSection)); + EXPECT_FALSE(FirstSegment->sectionInSegment(RoDataSection)); + EXPECT_FALSE(FirstSegment->sectionInSegment(BssSection)); + + EXPECT_FALSE(SecondSegment->sectionInSegment(TextSection)); + EXPECT_TRUE(SecondSegment->sectionInSegment(RoDataSection)); + EXPECT_FALSE(SecondSegment->sectionInSegment(BssSection)); + + EXPECT_EQ(iterRangeLength(*FirstSegment), 1); + EXPECT_EQ(iterRangeLength(*SecondSegment), 1); + + EXPECT_FALSE(SecondSegment->sectionInSegment(BssSection)); + SecondSegment->addSection(BssSection); + EXPECT_TRUE(SecondSegment->sectionInSegment(BssSection)); +}