Index: llvm/tools/llvm-objcopy/MutableObject/MutableObject.h =================================================================== --- /dev/null +++ llvm/tools/llvm-objcopy/MutableObject/MutableObject.h @@ -0,0 +1,427 @@ +#include "../Buffer.h" +#include "llvm/ADT/ArrayRef.h" +#include "llvm/ADT/Optional.h" +#include "llvm/ADT/SmallVector.h" +#include "llvm/Object/ELF.h" +#include "llvm/Object/ELFObjectFile.h" +#include "llvm/Object/ELFTypes.h" +#include "llvm/Object/ObjectFile.h" +#include "llvm/Support/Errc.h" +#include "llvm/Support/Error.h" +#include "llvm/Support/Casting.h" +#include +#include +#include +#include + +using namespace llvm; +using namespace object; + +template class MutableELFObject; + +class SectionBase { +protected: + ArrayRef Data; + StringRef Name; + +public: + SectionBase() = default; + SectionBase(const SectionBase &) = default; + SectionBase(SectionBase &&) = default; + SectionBase(ArrayRef Data, StringRef Name) : Data(Data), Name(Name) {} + virtual ~SectionBase() {} + + StringRef name() const { return Name; } + ArrayRef data() const { return Data; } + + virtual SectionBase *clone(Optional NewName = None) const { + auto Ptr = new SectionBase(*this); + if (NewName) + Ptr->Name = *NewName; + return Ptr; + } + + virtual SectionBase *clone(ArrayRef NewData) const { + auto Ptr = new SectionBase(*this); + Ptr->Data = NewData; + return Ptr; + } +}; + +class MutableObject { +// Access modifiers are pretty permissive throughout for now because +// I can't be bothered to make getters. +protected: + using SecPtr = std::unique_ptr; + using key_t = uint64_t; + + std::unordered_map UpdatedSections; + + std::vector Sections; + + friend class section_range; + + struct SectionRef { + uint64_t Index : 63; + bool New : 1; + + SectionRef(uint64_t Index) : Index(Index), New(true) {} + SectionRef(uint64_t Index, bool Updated) : Index(Index), New(Updated) {} + + operator unsigned long() { + return Index; + } + + }; + + + class section_range { + friend class MutableObject; + MutableObject &MutObject; + + public: + + section_range(MutableObject &MutObject) : MutObject(MutObject) {} + + class section_iter { + friend class section_range; + MutableObject &MutObject; + // Start at max and call operator++ in case section 0 gets modified + // modifing sh_size if more than SHN_LORESERVE sections for example + uint64_t Index = UINT64_MAX; + bool InUpdatedSection = false; + + SectionBase *Current = nullptr; + + explicit section_iter(MutableObject &MutObject) : MutObject(MutObject) { + operator++(); + } + + static section_iter createEnd(MutableObject &MutObject) { + section_iter Ret(MutObject); + Ret.Index = UINT64_MAX; + Ret.Current = nullptr; + return Ret; + } + + public: + + SectionRef getRef() { + return SectionRef(Index, InUpdatedSection); + } + + bool operator==(const section_iter &Other) { + return &Other.MutObject == &MutObject && + Other.Index == Index && Other.Current == Current; + } + + bool operator!=(const section_iter &Other) { + return !operator==(Other); + } + + SectionBase &operator*() { + assert(Current && "invalid section"); + return *Current; + } + + SectionBase *operator->() { + return Current; + } + + section_iter operator++() { + auto End = MutObject.UpdatedSections.end(); + for (;;) { + ++Index; + auto Found = MutObject.UpdatedSections.find(Index); + if (Found != End) { + // nullptr in UpdatedSections means it was removed, skip it. + if (!Found->second.get()) + continue; + + InUpdatedSection = true; + Current = Found->second.get(); + return *this; + } + + InUpdatedSection = false; + + // Index is past the original number of seections and no longer + // found in UpdatedSections so we are done. Create end() state. + if (Index >= MutObject.Sections.size()) { + Index = UINT64_MAX; + Current = nullptr; + return *this; + } + + Current = MutObject.Sections[Index].get(); + return *this; + } + llvm_unreachable("impossible to exit loop without returning"); + } + }; + + section_iter begin() { + return section_iter(MutObject); + } + + section_iter end() { + return section_iter::createEnd(MutObject); + } + }; + +public: + MutableObject() = default; + MutableObject(const MutableObject &) = default; + MutableObject(MutableObject &&) = default; + virtual ~MutableObject() {} + virtual Error init() = 0; + + virtual Error writeToDisk(objcopy::Buffer &Buff) = 0; + + virtual const uint8_t *base() const = 0; + + Error renameSection(SectionRef Index, StringRef NewName) { + if (Index >= Sections.size()) + return createError("section doesn't exist"); + UpdatedSections[Index] = + std::unique_ptr(Sections[Index]->clone(NewName)); + return Error::success(); + } + + Error removeSection(SectionRef Index) { + if (Index >= Sections.size()) + return createError("section doesn't exist"); + + //if (Index.New) return createError("can't remove new section, for now"); + + UpdatedSections[Index] = nullptr; + + return Error::success(); + } + + Error replaceSection(SectionRef Index, SecPtr NewSection) { + if (Index >= Sections.size()) + return createError("section doesn't exist"); + UpdatedSections[Index] = std::move(NewSection); + return Error::success(); + } + + Optional getSection(SectionRef Index) const { + // Searching if the key exists so is_contained doesn't work. + auto Found = UpdatedSections.find(Index); + if (Found != UpdatedSections.end()) { + if (Found->second.get()) + return Found->second.get(); + return None; + } + if (Index > Sections.size()) return None; + return Sections[Index].get(); + } + + Optional findSectionByName(StringRef Name) /* const */ { + auto Iter = sections().begin(); + const auto End = sections().end(); + for (; Iter != End; ++Iter) + if (Iter->name() == Name) + return Iter.getRef(); + + return None; + } + + section_range sections() { + return section_range(*this); + } +}; + +template +class ELFSection : public SectionBase { + Elf_Shdr_Base &Shdr; + +public: + friend class MutableELFObject; + + ELFSection(const Elf_Shdr_Base &Shdr, MutableELFObject &Object) + : SectionBase(ArrayRef(Object.base() + Shdr.sh_offset, Shdr.sh_size), + Object.nameFromShstrtab(Shdr.sh_name)), Shdr(const_cast&>(Shdr)) {} + + ELFSection(const ELFSection &) = default; + ELFSection(ELFSection &&) = default; + + ELFSection *clone(Optional NewName = None) const override { + auto *Ptr = new ELFSection(*this); + if (NewName) + Ptr->Name = *NewName; + return Ptr; + } + + ELFSection *clone(ArrayRef NewData) const override { + auto *Ptr = new ELFSection(*this); + Ptr->Shdr.sh_size = static_cast(NewData.size()); + Ptr->Data = NewData; + return Ptr; + } + + static bool classof(const SectionBase *Sec) { + return true; + } +}; + +template +class MutableELFObject : public MutableObject { + ELFObjectFile &Object; + ArrayRef Shstrtab; + +public: + MutableELFObject(ELFObjectFile &Object) : Object(Object) {} + + Error init() override; + + Error createNewShstrtab(); + Expected layoutSections(typename ELFT::Ehdr &Header); + Error writeToDisk(objcopy::Buffer &Buff) override; + + const uint8_t *base() const override { + return Object.base(); + } + + static StringRef nameFromStrtab(uint64_t Index, ArrayRef Strtab) { + if (!Strtab.data()) + return StringRef(); + + assert(Index < Strtab.size()); + return reinterpret_cast(Strtab.data()) + Index; + } + + StringRef nameFromShstrtab(uint64_t Index) const { + return nameFromStrtab(Index, Shstrtab); + } +}; + +template +Error MutableELFObject::init() { + auto &Header = *Object.getELFFile()->getHeader(); + auto SectionsOrErr = Object.getELFFile()->sections(); + if (!SectionsOrErr) + return SectionsOrErr.takeError(); + auto SectionArr = SectionsOrErr.get(); + Shstrtab = ArrayRef(base() + SectionArr[Header.e_shstrndx].sh_offset, + SectionArr[Header.e_shstrndx].sh_size); + + Sections.reserve(Header.e_shnum); + for (const auto &Shdr : SectionArr) { + // Later construct different classes derived from ELF section but + // for now just create an ELFSection. + switch (Shdr.sh_type) { + default: + Sections.emplace_back(make_unique>(Shdr, *this)); + } + } + + return Error::success(); +} + +template +Error MutableELFObject::createNewShstrtab() { + std::vector NewData {0}; + + auto PushBackStr = [&NewData] (std::string Str) -> size_t { + size_t Index = NewData.size(); + for (auto c : Str) + NewData.push_back(c); + NewData.push_back(0); + return Index; + }; + + // Here the sections sh_name field is changed which goes against leaving + // sections immutable, only updating, but this should be ok because their name + // is still the same, just no longer pointing to the same strtab. + for (auto &Sec : sections()) { + size_t Shstrndx = PushBackStr(Sec.name()); + auto *Sect = dyn_cast>(&Sec); + assert(Sect); + Sect->Shdr.sh_name = Shstrndx; + } + + assert(!NewData.back()); + + Optional Shstrtab = findSectionByName(".shstrtab"); + if (!Shstrtab) + return createError("Couldn't find .shstrtab"); + auto OldShstrtab = getSection(Shstrtab.getValue()).getValue(); + + OwningArrayRef OwnedArray = ArrayRef(NewData); + std::unique_ptr NewShstrtab(OldShstrtab->clone(OwnedArray)); + if (Error E = replaceSection(Shstrtab.getValue(), std::move(NewShstrtab))) + return E; + return Error::success(); +} + +template +Expected MutableELFObject::layoutSections(typename ELFT::Ehdr &Header) { + size_t NumSections = 0; + + size_t CurrentOffset = sizeof(Header); + + auto GetNextAligned = [&CurrentOffset](unsigned Alignment) -> uint64_t { + if (Alignment == 0) + Alignment = 1; + return (CurrentOffset + Alignment - 1) & ~(Alignment - 1); + }; + + for (auto &Sec : sections()) { + if (!NumSections++) + continue; + + auto *ElfSection = dyn_cast>(&Sec); + assert(ElfSection); + + CurrentOffset = GetNextAligned(ElfSection->Shdr.sh_addralign); + ElfSection->Shdr.sh_offset = CurrentOffset; + CurrentOffset += ElfSection->Shdr.sh_size; + } + + Header.e_shnum = NumSections; + Header.e_shoff = CurrentOffset; + + return CurrentOffset + sizeof(typename ELFT::Shdr) * NumSections; +} + +template +Error MutableELFObject::writeToDisk(objcopy::Buffer &Buff) { + auto ConstHeader = Object.getELFFile()->getHeader(); + typename ELFT::Ehdr Header = *ConstHeader; + if (Error E = createNewShstrtab()) + return E; + + Expected FileSizeOrErr = layoutSections(Header); + if (!FileSizeOrErr) + return FileSizeOrErr.takeError(); + + if (Error E = Buff.allocate(*FileSizeOrErr)) + return E; + + uint64_t Shstrtabndx = 0; + for (const auto &Sec : sections()) { + if (Sec.name() == ".shstrtab") + break; + Shstrtabndx++; + } + + Header.e_shstrndx = Shstrtabndx; + + uint8_t *Start = Buff.getBufferStart(); + memcpy(Start, &Header, sizeof(Header)); + auto *SectionArr = + reinterpret_cast(Start + Header.e_shoff); + uint64_t SectionCount = 0; + for (auto &Sec : sections()) { + auto *ElfSection = dyn_cast>(&Sec); + assert(ElfSection); + memcpy(SectionArr + SectionCount, &ElfSection->Shdr, sizeof(ElfSection->Shdr)); + SectionCount++; + memcpy(Start + ElfSection->Shdr.sh_offset, Sec.data().data(), Sec.data().size()); + + } + + return Error::success(); +} Index: llvm/tools/llvm-objcopy/llvm-objcopy.cpp =================================================================== --- llvm/tools/llvm-objcopy/llvm-objcopy.cpp +++ llvm/tools/llvm-objcopy/llvm-objcopy.cpp @@ -46,6 +46,8 @@ #include #include +#include "MutableObject/MutableObject.h" + namespace llvm { namespace objcopy { @@ -54,6 +56,7 @@ LLVM_ATTRIBUTE_NORETURN void error(Twine Message) { WithColor::error(errs(), ToolName) << Message << "\n"; + errs().flush(); exit(1); } @@ -264,9 +267,54 @@ return Error::success(); } +static Error testing(StringRef Filename) { + Expected> BinaryOrErr = + object::createBinary(Filename); + if (!BinaryOrErr) + return createFileError(Filename, BinaryOrErr.takeError()); + + auto *ELFBinary = dyn_cast>(BinaryOrErr->getBinary()); + if (!ELFBinary) + return createError("not an ELF file"); + + auto MutObject = MutableELFObject(*ELFBinary); + if (Error E = MutObject.init()) + return E; + + for (const auto &Sec : MutObject.sections()) + outs() << Sec.name() << "\n"; + + auto Comment = MutObject.findSectionByName(".comment"); + if (!Comment) + return createError("couldn't find '.comment'"); + + if (Error E = MutObject.removeSection(*Comment)) + return E; + + for (const auto &Sec : MutObject.sections()) + outs() << Sec.name() << "\n"; + + FileBuffer Buff("out.o"); + if (Error E = MutObject.writeToDisk(Buff)) + return E; + + if (Error E = Buff.commit()) + return E; + + return Error::success(); +} + int main(int argc, char **argv) { InitLLVM X(argc, argv); ToolName = argv[0]; + if (!strcmp(argv[1], "TEST")) { + if (Error E = testing(argv[2])) { + logAllUnhandledErrors(std::move(E), WithColor::error(errs(), ToolName)); + return 1; + } + return 0; + } + bool IsStrip = sys::path::stem(ToolName).contains("strip"); Expected DriverConfig = IsStrip ? parseStripOptions(makeArrayRef(argv + 1, argc), reportWarning)