[Object] Create MutableELFObject Class for Doing Mutations on ELFObjectFiles [Part 2]

[abrachet]

Summary

Added support for symbols to MutableELFObject.

[abrachet]

Added unit tests and increased functionality.

[abrachet]

Made segment support similar to sections from ObjectFile for more continuity.

[abrachet]

rebase

[abrachet]

rebase

[grimar]

Few comments/suggestions are below.

llvm/include/llvm/Object/MutableELFObject.h

213

Mappings.erase(Mappings.begin() + Index)?

Perhaps also worth adding an assert like was done above.

218

Mappings.push_back({NewValues.size() - 1, MappingType::New});

243

You should probably at least add an assert checking ShType is one of the symbol table types.

250

return {};
(Its a bit confusing to see None which is normally used for returning the empty Optional<>)

253

It could be findSectionIndex, since you nowere check that ShType is a symbol table.

270

What about implementing const version using non-const? I.e. call getWhichTable from getWhichTable const.

297

I think when you have a symbol iterator, you usually know if it is a regular or dynamic
symbol. I am not sure it worth to allow mixing them. I.e. I am thinking about making 2 different functions,
like removeSymbol and removeDynamicSymbol.

I am not feeling strong here, though it seems to be less error prone.

[abrachet]

Addressed review comments

[abrachet]

Thanks for the comments @grimar!

llvm/include/llvm/Object/MutableELFObject.h

218

I'll do you one better, emplace_back :)

270

How do I do this without const_casting this? Is that how you do it? This is very useful because this duplicating methods like this has been bugging me recently. Thanks!

297

I have added these, I don't feel strongly either. FWIW, all of the getSymbol* methods in ELFObjectFile work on either symbol type so I duplicated that here. I'll keep this as not done so others who have an opinion stronger than ours can chime in.

[grimar]

llvm/include/llvm/Object/MutableELFObject.h

270

I think there is no way to do this without const_cast, but this is a normal/common way.

[jhenderson]

I don't think you want to be modifying the number of dynamic symbols at all. That would require you to make other changes to loadable code, which would impact addresses, and effectively require you to completely relink the program.

You may however still need to modify the contents of the existing dynamic symbols (specifically you may need to be able to update their section indexes).

llvm/include/llvm/Object/MutableELFObject.h

255

How about findSectionOfType?

270

It doesn't look like you use the non-const version anywhere, so it can be deleted, right?

[abrachet]

Added remove to MutableTable
Added getMutableSymbol to MutableELFObject
Added unit tests for older changes

[abrachet]

Remove removeDynSymbol from MutableELFObject

[abrachet]

In D65367#1621113, @jhenderson wrote:

I don't think you want to be modifying the number of dynamic symbols at all. That would require you to make other changes to loadable code, which would impact addresses, and effectively require you to completely relink the program.

You may however still need to modify the contents of the existing dynamic symbols (specifically you may need to be able to update their section indexes).

Ok I can just not add an addDynSymbol and removeDynSymbol methods then.

[jhenderson]

Not had time to look at everything in this patch, but it feels like you don't want removal functionality to be added in this patch at the same time as basic symbol functionality. That should be a different patch, and can then include both symbol and section removal. What do you think?

By the way, the overall approach with removal seems sound.

llvm/include/llvm/Object/MutableELFObject.h

210–211

This comment doesn't document what the function does. It talks about class implementation details, which doesn't really make sense.

212

Assert that the Index hasn't already been removed.

213

"and only the equivalent to push_back, add." sounds like an unfinished sentence to me.

213–214

I'm inclined to think that getNextIndex shouldn't provide you the first index under any circumstance. You have getFirstNonRemoved instead.

215

No need for doxygen-style comments in the middle of a function.

Your opening comment should however explain what happens when the last index is passed in/reached by the loop.

216

This comment should refer to the parameter e.g. "Adds \param New to the back of the list."

216

I'd call this lastIndex, as this doesn't return an iterator.

element, this -> element. This

217

Don't use "New" as a variable name, since it won't work with the upcoming change in variable naming rules.

222

Comment?

230

How about getFirstIndex, to match getNextIndex?

Also this function needs a comment.

231

I'd put the I == end() clause in here, rather than in the body of the loop and just return I instead of the llvm_unreachable statement.

[abrachet]

Moved remove and add out of this patch

[jhenderson]

llvm/include/llvm/Object/MutableELFObject.h

212

be, entries -> be: entries

215

Does this belong to a later patch?

243

I'd be inclined to just call this getSymbolTable.

288–291

This change seems to be unrelated?

Avoid refactors of code which don't have anything to do with the current change, and put them in their own change.

319–321

This overload seems a bit superfluous to me.

324

Call this getMutableDynamicSymbol

333–335

This overload seems a bit superfluous to me.

496

Why is this looping over only the original sections?

llvm/unittests/Object/MutableELFObjectTest.cpp

171

Looks to me like you need a checking first of all counting the number of symbols. Also, why do you need to construct a new symbol_iterator, rather than using the output of ++MutableObject.symbol_begin() directly?

178

FirstSym and SecondSym?

Also, same comment as above.

186

"no change public methods" what are those? :)

I think you need to rewrite this comment.

216

I think the same comments made to patch 1 apply here too, regarding the std::move-ing of ELFObjFile making this reference invalid.

225

I think less confusing names might be FromBase and FromMutable or something like that. My first reaction was "why is a file a SymbolRef?"

253

This loop needs a comment explaining its purpose and the consumeError call.

263

Is this comment correct?

284

Same comments as earlier.

290

This seems a bit redundant...

Same goes for the other changes below. You're basically testing that C++ works...

I think you want to be fetching the symbol again from one or both of the object interfaces (base or sub-class) and showing that the value has been updated.

302

Don't abbreviate names unnecessarily: use "BasicDynamicSymbol"

322

Same comments as above.

328

My gut reaction was "where are the equivalent dynamic symbol tests for the earlier regular symbol cases?" Perhaps worth a comment to that effect somewhere, perhaps making it clear that the earlier cases test common code.

[rupprecht]

llvm/include/llvm/Object/MutableELFObject.h

158

using Elf_Sym = typename ELFT::Sym;

llvm/unittests/Object/MutableELFObjectTest.cpp

294–298

What happens to the string table in this case?

The names for symbols in the symbol table is backed by the string table, so if you change the name from "test" to "new_name", the string table needs to be rewritten for that, but it doesn't looks like there's a corresponding MutableTable StringTable to back that change.

[abrachet]

Addressed review comments

[abrachet]

llvm/include/llvm/Object/MutableELFObject.h

215

No it is needed. I explain this in an answer to your question here

496

Because ELFObjectFile::getSymbol() (and also directly many other symbol related methods) needs the DataRef to have an index to the original section, and this is how to differentiate between dynamic and regular symbols. I will think more about this, but this was the least disruptive way to inherit from ELFObjectFile I believe. It's worth noting that this is a private method.

llvm/unittests/Object/MutableELFObjectTest.cpp

171

SymbolicFile::symbol_begin() returns a basic_symbol_iterator pointing to BasicSymbolRefs whose only meaningful method for testing is getFlags(), but symbol_iterator points to SymbolRefs which has getValue and getName.

290

Oops! They are more meaningful now :)

294–298

Nothing happens for now. The idea, and this is the same as objcopy::elf::Object, is that it gets updated only when doing final layout.

[jhenderson]

llvm/unittests/Object/MutableELFObjectTest.cpp

171

Okay. Please write this as:

symbol_iterator FirstSym(++MutableObject.symbol_begin());

189–191

I'm not sure you need this here?

194

Perhaps "in SymbolRef's public methods beteween..." to avoid using "between" twice.

257

ElfObjFile might be invalid. You will need to make another copy of the input, like in patch 1.

[abrachet]

Addressed review comments

[jhenderson]

llvm/unittests/Object/MutableELFObjectTest.cpp

286

You've made a change to address my comment for this style above. Please make the same change throughout. It wasn't a local issue for my earlier point :)

[abrachet]

Addressed comments

[abrachet]

llvm/unittests/Object/MutableELFObjectTest.cpp

286

Sorry about that, silly mistake.

[rupprecht]

I don't see anything major left, but please wait for James to take another look

llvm/include/llvm/Object/MutableELFObject.h

503–516

I don't think these methods have anything to do with MutableELFObject, so they should moved Object/SymbolicFile.h so others can use it. (A FIXME comment is fine for now).

llvm/unittests/Object/MutableELFObjectTest.cpp

294–298

I suppose it depends on how the writer/finalizer for MutableElfObject looks like -- a naive implementation that just iterates over sections and writes out the contents will lose the renamed values, because those will be privately owned by MutableELFSymbol, and not reflected in the (non-mutated) string table.

[jhenderson]

LGTM.

[MaskRay]

llvm/include/llvm/Object/MutableELFObject.h

497

uint32_t I = 0, E = Sections.originalSize(); I != E; ++I)

https://llvm.org/docs/CodingStandards.html#don-t-evaluate-end-every-time-through-a-loop

Not a strong opinion though.

[abrachet]

Addressed review comments

[abrachet]

llvm/include/llvm/Object/MutableELFObject.h

497

Good catch! Thanks

llvm/include/llvm/Object/MutableELFObject.h

 //===-- MutableELFObject.h --------------------------------------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_OBJECT_MUTABLEELFOBJECT_H
 #define LLVM_OBJECT_MUTABLEELFOBJECT_H
 
 #include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/Sequence.h"
 #include "llvm/Object/ELFObjectFile.h"
+#include <set>
 
 namespace llvm {
 namespace object {
 
-// T is a wrapper type around Iterable::value_type and must have a constructor
-// taking an Iterable::value_type.
-template <typename T> class MutableRange {
-public:
-  struct MappingType {
-    uintptr_t Ptr;
-    bool New;
 
-    MappingType(bool New, uintptr_t Ptr) : Ptr(Ptr), New(New) {}
-  };
 
+// T is a wrapper type around Iterable::value_type and must have a constructor
+// taking an Iterable::value_type.
+template <typename T> class MutableRange {
+public:
   using iterator = typename std::vector<MappingType>::iterator;
   using value_type = MappingType;
 
 private:
   std::vector<MappingType> Mappings;
   std::vector<T> NewValues;
 
 public:
   template <typename Iterable>
   MutableRange(
       Iterable Range,
       function_ref<uintptr_t(typename Iterable::iterator::value_type)> Extract)
       : MutableRange(Range.begin(), Range.end(), Extract) {}
 
   template <typename Iter>
   MutableRange(Iter Begin, Iter End,
                function_ref<uintptr_t(typename Iter::value_type)> Extract) {
     std::transform(Begin, End, std::back_inserter(Mappings),
                    [&Extract](typename Iter::value_type Value) {
                      return MappingType(false, Extract(Value));
                    });
   }
 
+  MutableRange(function_ref<void(std::vector<MappingType> &)> MappingCreator) {
+    MappingCreator(Mappings);
+  }
+
   MappingType operator[](uint64_t Index) const { return Mappings[Index]; }
 
   size_t size() const { return Mappings.size(); }
 
   iterator begin() { return Mappings.begin(); }
   iterator end() { return Mappings.end(); }
 
   const T *getNew(uint64_t Index) const { return &NewValues[Index]; }

 template <typename ELFT> class MutableELFObject;
 
 static inline DataRefImpl toDataRef(uintptr_t Ptr) {
   DataRefImpl Ref;
   Ref.p = Ptr;
   return Ref;
 }
 
+template <typename ELFT>
+static bool sectionWithinSegment(const Elf_Shdr_Impl<ELFT> &Sec,
+                                 const Elf_Phdr_Impl<ELFT> &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 <typename ELFT> class SegmentRef {
+  uint64_t Index;
+  MutableELFObject<ELFT> *OwningObject;
+
+public:
+  SegmentRef(uintptr_t Index, MutableELFObject<ELFT> *OwningObject)
+      : Index(Index), OwningObject(OwningObject) {}
+
+  bool operator==(const SegmentRef &Other) const;
+  bool operator!=(const SegmentRef &Other) const;
+
+  void moveNext();
+
+  uint64_t getIndex() const { return Index; }
+  const MutableELFObject<ELFT> *getObject() const { return OwningObject; }
+
+  uint64_t getType() const;
+  uint64_t getOffset() const;
+  uint64_t getVirtualAddress() const;
+  uint64_t getPhysicalAddress() const;
+  uint64_t getFileSize() const;
+  uint64_t getMemSize() const;
+  uint64_t getFlags() const;
+  uint64_t getAlign() const;
+
+  const Elf_Phdr_Impl<ELFT> &getHeader() const;
+
+  bool hasSection(section_iterator Sec) const;
+  bool hasSection(uint64_t Sec) const;
+};
+
 template <typename ELFT> struct MutableELFSection {
   Elf_Shdr_Impl<ELFT> Header;
   std::string Name;
   OwningArrayRef<uint8_t> Data;
 
   MutableELFSection(uintptr_t ToCopy, const MutableELFObject<ELFT> *ObjFile)
       : Header(*ObjFile->getSection(toDataRef(ToCopy))),
         Data(OwningArrayRef<uint8_t>(Header.sh_size)) {
     ::memcpy(Data.data(), ObjFile->base() + Header.sh_offset, Header.sh_size);
   }
 
   void setData(ArrayRef<uint8_t> Ref) {
     Data = OwningArrayRef<uint8_t>(Ref);
     Header.sh_size = Data.size();
   }
 };
 
+template <typename ELFT> class MutableELFSegment {
+  std::set<uint64_t> Sections;
+  const MutableELFObject<ELFT> *ObjFile;
+
+public:
+  Elf_Phdr_Impl<ELFT> Header;
+
+  MutableELFSegment(uintptr_t ToCopy, const MutableELFObject<ELFT> *ObjFile)
+      : ObjFile(ObjFile),
+        Header(*reinterpret_cast<const Elf_Phdr_Impl<ELFT> *>(ToCopy)) {
+    for (const auto &Sec : ObjFile->sections())
+      if (sectionWithinSegment(ObjFile->getSectionHeader(Sec), Header))
+        Sections.insert(Sec.getIndex());
+  }
+
+  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 hasSection(section_iterator Sec) const {
+    return Sections.find(Sec->getRawDataRefImpl().p) != Sections.end();
+  }
+
+  bool hasSection(uint64_t Sec) const {
+    return Sections.find(Sec) != Sections.end();
+  }
+
+  class iterator : public std::set<uint64_t>::iterator {
+    using Base = std::set<uint64_t>::iterator;
+    const MutableELFObject<ELFT> *ObjFile;
+
+  public:
+    iterator(Base B, const MutableELFObject<ELFT> *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 <typename ELFT> class MutableELFObject : public ELFObjectFile<ELFT> {
   friend struct MutableELFSection<ELFT>;
+  friend class MutableELFSegment<ELFT>;
+  friend class SegmentRef<ELFT>;
 
   MutableRange<MutableELFSection<ELFT>> Sections;
+  MutableRange<MutableELFSegment<ELFT>> Segments;
 
 protected:
-  using MappingType =
-      typename MutableRange<MutableELFSection<ELFT>>::MappingType;
+  const Elf_Ehdr_Impl<ELFT> &header() {
+    return *reinterpret_cast<const Elf_Ehdr_Impl<ELFT> *>(this->base());
+  }
 
   // Returns DataRef with pointer to the correct section header.
   DataRefImpl getSectionRef(DataRefImpl Sec) const {
     MappingType Mapping = Sections[Sec.p];
     return Mapping.New ? toDataRef(reinterpret_cast<uintptr_t>(
                              &Sections.getNew(Mapping.Ptr)->Header))
                        : toDataRef(Mapping.Ptr);
   }
 
+  const Elf_Shdr_Impl<ELFT> &getSectionHeader(uint64_t Index) const;
+  const Elf_Shdr_Impl<ELFT> &getSectionHeader(section_iterator Sec) const;
+
   void moveSectionNext(DataRefImpl &Sec) const override;
   Expected<StringRef> getSectionName(DataRefImpl Sec) const override;
   uint64_t getSectionAddress(DataRefImpl Sec) const override;
   uint64_t getSectionIndex(DataRefImpl Sec) const override;
   uint64_t getSectionSize(DataRefImpl Sec) const override;
   Expected<ArrayRef<uint8_t>>
   getSectionContents(DataRefImpl Sec) const override;
   uint64_t getSectionAlignment(DataRefImpl Sec) const override;
   bool isSectionCompressed(DataRefImpl Sec) const override;
   bool isSectionText(DataRefImpl Sec) const override;
   bool isSectionData(DataRefImpl Sec) const override;
   bool isSectionBSS(DataRefImpl Sec) const override;
   bool isSectionVirtual(DataRefImpl Sec) const override;
   bool isBerkeleyText(DataRefImpl Sec) const override;
   bool isBerkeleyData(DataRefImpl Sec) const override;
 
+  uint64_t getSegmentType(uint64_t Index) const;
+  uint64_t getSegmentOffset(uint64_t Index) const;
+  uint64_t getSegmentVirtualAddress(uint64_t Index) const;
+  uint64_t getSegmentPhysicalAddress(uint64_t Index) const;
+  uint64_t getSegmentFileSize(uint64_t Index) const;
+  uint64_t getSegmentMemSize(uint64_t Index) const;
+  uint64_t getSegmentFlags(uint64_t Index) const;
+  uint64_t getSegmentAlign(uint64_t Index) const;
+  const Elf_Phdr_Impl<ELFT> &getSegmentHeader(uint64_t Index) const;
+  Elf_Phdr_Impl<ELFT> &getSegmentHeader(uint64_t Index);
+  bool segmentHasSection(uint64_t Segment, uint64_t Section) const;
+
 public:
-  MutableELFObject(ELFObjectFile<ELFT> &B)
+  explicit MutableELFObject(ELFObjectFile<ELFT> &B)
       : ELFObjectFile<ELFT>(std::move(B)),
         Sections(B.section_begin(), B.section_end(),
-                 [&](SectionRef Ref) { return Ref.getRawDataRefImpl().p; }) {}
+                 [&](SectionRef Ref) { return Ref.getRawDataRefImpl().p; }),
+        Segments([this](std::vector<MappingType> &Vec) {
+          const Elf_Ehdr_Impl<ELFT> &Header =
+              *reinterpret_cast<const Elf_Ehdr_Impl<ELFT> *>(this->base());
+          const Elf_Phdr_Impl<ELFT> *Phdrs =
+              reinterpret_cast<const Elf_Phdr_Impl<ELFT> *>(this->base() +
+                                                            Header.e_phoff);
+          auto Sequence = seq(0ULL, (uint64_t)Header.e_phnum);
+          std::transform(Sequence.begin(), Sequence.end(),
+                         std::back_inserter(Vec), [Phdrs](int Index) {
+                           return MappingType(
+                               /*New=*/false,
+                               reinterpret_cast<uintptr_t>(Phdrs + Index));
+                         });
+        }) {}
 
   section_iterator section_begin() const override {
     return section_iterator(SectionRef(toDataRef(0), this));
   }
 
   section_iterator section_end() const override {
     return section_iterator(SectionRef(toDataRef(Sections.size()), this));
   }
 
+  section_iterator findSectionByName(StringRef Name) const;
+
   MutableELFSection<ELFT> *getMutableSection(section_iterator Sec) {
     uintptr_t Index = Sec->getRawDataRefImpl().p;
     return Sections.makeMutable(Index, this);
   }
+
+  using segment_iterator = content_iterator<SegmentRef<ELFT>>;
+
+  segment_iterator segment_begin() {
+    return segment_iterator(SegmentRef<ELFT>(0, this));
+  }
+
+  segment_iterator segment_end() {
+    return segment_iterator(SegmentRef<ELFT>(header().e_phnum, this));
+  }
+
+  iterator_range<segment_iterator> segments() {
+    return iterator_range<segment_iterator>(segment_begin(), segment_end());
+  }
+
+  MutableELFSegment<ELFT> *getMutableSegment(segment_iterator Seg) {
+    return Segments.makeMutable(Seg->getIndex(), this);
+  }
 };
 
+template <typename ELFT> void SegmentRef<ELFT>::moveNext() { ++Index; }
+
+template <typename ELFT>
+bool SegmentRef<ELFT>::operator==(const SegmentRef<ELFT> &Other) const {
+  return Index == Other.Index && OwningObject == Other.OwningObject;
+}
+
+template <typename ELFT>
+bool SegmentRef<ELFT>::operator!=(const SegmentRef<ELFT> &Other) const {
+  return Index != Other.Index || OwningObject != Other.OwningObject;
+}
+
+template <typename ELFT> uint64_t SegmentRef<ELFT>::getType() const {
+  return OwningObject->getSegmentType(Index);
+}
+
+template <typename ELFT> uint64_t SegmentRef<ELFT>::getOffset() const {
+  return OwningObject->getSegmentOffset(Index);
+}
+
+template <typename ELFT> uint64_t SegmentRef<ELFT>::getVirtualAddress() const {
+  return OwningObject->getSegmentVirtualAddress(Index);
+}
+
+template <typename ELFT> uint64_t SegmentRef<ELFT>::getPhysicalAddress() const {
+  return OwningObject->getSegmentPhysicalAddress(Index);
+}
+
+template <typename ELFT> uint64_t SegmentRef<ELFT>::getFileSize() const {
+  return OwningObject->getSegmentFileSize(Index);
+}
+
+template <typename ELFT> uint64_t SegmentRef<ELFT>::getMemSize() const {
+  return OwningObject->getSegmentMemSize(Index);
+}
+
+template <typename ELFT> uint64_t SegmentRef<ELFT>::getFlags() const {
+  return OwningObject->getSegmentFlags(Index);
+}
+
+template <typename ELFT> uint64_t SegmentRef<ELFT>::getAlign() const {
+  return OwningObject->getSegmentAlign(Index);
+}
+
+template <typename ELFT>
+const Elf_Phdr_Impl<ELFT> &SegmentRef<ELFT>::getHeader() const {
+  return OwningObject->getSegmentHeader(Index);
+}
+
+template <typename ELFT>
+bool SegmentRef<ELFT>::hasSection(section_iterator Sec) const {
+  return OwningObject->segmentHasSection(Index, Sec->getRawDataRefImpl().p);
+}
+
+template <typename ELFT> bool SegmentRef<ELFT>::hasSection(uint64_t Sec) const {
+  return OwningObject->segmentHasSection(Index, Sec);
+}
+
+template <typename ELFT>
+const Elf_Shdr_Impl<ELFT> &
+MutableELFObject<ELFT>::getSectionHeader(uint64_t Index) const {
+  MappingType Mapping = Sections[Index];
+  if (Mapping.New)
+    return Sections.getNew(Mapping.Ptr)->Header;
+  return *this->getSection(toDataRef(Mapping.Ptr));
+}
+
+template <typename ELFT>
+const Elf_Shdr_Impl<ELFT> &
+MutableELFObject<ELFT>::getSectionHeader(section_iterator Sec) const {
+  return getSectionHeader(Sec->getRawDataRefImpl().p);
+}
+
 template <typename ELFT>
 void MutableELFObject<ELFT>::moveSectionNext(DataRefImpl &Sec) const {
   ++Sec.p;
 }
 
 template <typename ELFT>
 Expected<StringRef>
 MutableELFObject<ELFT>::getSectionName(DataRefImpl Sec) const {
bool MutableELFObject<ELFT>::isBerkeleyText(DataRefImpl Sec) const {
   return ELFObjectFile<ELFT>::isBerkeleyText(getSectionRef(Sec));
 }
 
 template <typename ELFT>
 bool MutableELFObject<ELFT>::isBerkeleyData(DataRefImpl Sec) const {
   return ELFObjectFile<ELFT>::isBerkeleyData(getSectionRef(Sec));
 }
 
+template <typename ELFT>
+uint64_t MutableELFObject<ELFT>::getSegmentType(uint64_t Index) const {
+  return getSegmentHeader(Index).p_type;
+}
+
+template <typename ELFT>
+uint64_t MutableELFObject<ELFT>::getSegmentFlags(uint64_t Index) const {
+  return getSegmentHeader(Index).p_flags;
+}
+
+template <typename ELFT>
+uint64_t MutableELFObject<ELFT>::getSegmentOffset(uint64_t Index) const {
+  return getSegmentHeader(Index).p_offset;
+}
+
+template <typename ELFT>
+uint64_t
+MutableELFObject<ELFT>::getSegmentVirtualAddress(uint64_t Index) const {
+  return getSegmentHeader(Index).p_vaddr;
+}
+
+template <typename ELFT>
+uint64_t
+MutableELFObject<ELFT>::getSegmentPhysicalAddress(uint64_t Index) const {
+  return getSegmentHeader(Index).p_paddr;
+}
+
+template <typename ELFT>
+uint64_t MutableELFObject<ELFT>::getSegmentFileSize(uint64_t Index) const {
+  return getSegmentHeader(Index).p_filesz;
+}
+
+template <typename ELFT>
+uint64_t MutableELFObject<ELFT>::getSegmentMemSize(uint64_t Index) const {
+  return getSegmentHeader(Index).p_memsz;
+}
+
+template <typename ELFT>
+uint64_t MutableELFObject<ELFT>::getSegmentAlign(uint64_t Index) const {
+  return getSegmentHeader(Index).p_align;
+}
+
+template <typename ELFT>
+bool MutableELFObject<ELFT>::segmentHasSection(uint64_t Segment,
+                                               uint64_t Section) const {
+  MappingType Mapping = Segments[Segment];
+  if (Mapping.New)
+    return Segments.getNew(Mapping.Ptr)->hasSection(Section);
+  return sectionWithinSegment(getSectionHeader(Section),
+                              getSegmentHeader(Segment));
+}
+
+template <typename ELFT>
+const Elf_Phdr_Impl<ELFT> &
+MutableELFObject<ELFT>::getSegmentHeader(uint64_t Index) const {
+  MappingType Mapping = Segments[Index];
+  if (Mapping.New)
+    return Segments.getNew(Index)->Header;
+  return *reinterpret_cast<const Elf_Phdr_Impl<ELFT> *>(Mapping.Ptr);
+}
+
+template <typename ELFT>
+section_iterator
+MutableELFObject<ELFT>::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
 
 #endif // LLVM_OBJECT_MUTABLEELFOBJECT_H

llvm/unittests/Object/MutableELFObjectTest.cpp

 #include "llvm/Support/Error.h"
 #include "llvm/Testing/Support/Error.h"
 #include "gtest/gtest.h"
 
 using namespace llvm;
 using namespace object;
 using namespace yaml;
 
+template <typename T> static ptrdiff_t iterRangeLength(const T &Range) {
+  return std::distance(Range.begin(), Range.end());
+}
+
+template <typename T> 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) {
   SmallString<0> Storage;
   Expected<std::unique_ptr<ObjectFile>> ErrOrObj = yaml2ObjectFile(Storage, R"(
 --- !ELF
 FileHeader:
   Class:     ELFCLASS64
- Name:    .text
     Content: "DEADBEEF")");
 
   ASSERT_THAT_EXPECTED(ErrOrObj, Succeeded());
   auto *ELFObjFile = dyn_cast<ELFObjectFile<ELF64LE>>(ErrOrObj->get());
   ASSERT_TRUE(ELFObjFile);
   MutableELFObject<ELF64LE> MutableObject(*ELFObjFile);
 
   ptrdiff_t NumSections =
       std::distance(MutableObject.section_begin(), MutableObject.section_end());
 
   auto FirstSec = ++MutableObject.section_begin();
   Expected<StringRef> Contents = FirstSec->getContents();
   ASSERT_THAT_EXPECTED(Contents, Succeeded());
 
   EXPECT_EQ(*Contents, "\xDE\xAD\xBE\xEF");
   EXPECT_EQ(FirstSec->getSize(), Contents->size());
 
Sections:
 
   ASSERT_THAT_EXPECTED(ErrOrObj, Succeeded());
   auto *ELFObjFile = dyn_cast<ELFObjectFile<ELF64LE>>(ErrOrObj->get());
   ASSERT_TRUE(ELFObjFile);
   MutableELFObject<ELF64LE> MutableObject(*ELFObjFile);
   const ObjectFile &ObjFile = *ErrOrObj->get();
 
   ptrdiff_t ObjFileSecs =
       std::distance(ELFObjFile->section_begin(), ELFObjFile->section_end());
   ptrdiff_t MutObjSecs =
       std::distance(MutableObject.section_begin(), MutableObject.section_end());
   EXPECT_EQ(ObjFileSecs, MutObjSecs);
 
   auto TestSections = [](SectionRef ObjFile, SectionRef MutObj) {
 #define EXPECT_EQ_ITER(getMethod)                                              \
   EXPECT_EQ(ObjFile.getMethod(), MutObj.getMethod())
     EXPECT_EQ_ITER(getAddress);
     EXPECT_EQ_ITER(getAlignment);
     EXPECT_EQ_ITER(getIndex);
     EXPECT_EQ_ITER(getSize);
     EXPECT_EQ_ITER(isBerkeleyData);
     EXPECT_EQ_ITER(isBerkeleyText);
     EXPECT_EQ_ITER(isBitcode);
     EXPECT_EQ_ITER(isBSS);
     EXPECT_EQ_ITER(isCompressed);
     EXPECT_EQ_ITER(isData);
     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<std::unique_ptr<ObjectFile>> 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<ELFObjectFile<ELF64LE>>(ErrOrObj->get());
+  ASSERT_TRUE(ELFObjFile);
+  MutableELFObject<ELF64LE> 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<std::unique_ptr<ObjectFile>> 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<ELFObjectFile<ELF64LE>>(ErrOrObj->get());
+  ASSERT_TRUE(ELFObjFile);
+  MutableELFObject<ELF64LE> MutableObject(*ELFObjFile);
+
+  auto FirstSegment = MutableObject.segment_begin();
+  EXPECT_EQ(FirstSegment->getType(), ELF::PT_LOAD);
+  EXPECT_TRUE(FirstSegment->getFlags() & (ELF::PF_X | ELF::PF_R));
+  EXPECT_FALSE(FirstSegment->getFlags() & ELF::PF_W);
+  EXPECT_EQ(FirstSegment->getVirtualAddress(), 0x1000U);
+  EXPECT_EQ(FirstSegment->getPhysicalAddress(), 0x2000U);
+
+  EXPECT_EQ(iterRangeLength(MutableObject.segments()), 1);
+}
+
+// Test
+TEST(MutableELFObject, SectionWithinSegment) {
+  SmallString<0> Storage;
+  Expected<std::unique_ptr<ObjectFile>> 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<ELFObjectFile<ELF64LE>>(ErrOrObj->get());
+  ASSERT_TRUE(ELFObjFile);
+  MutableELFObject<ELF64LE> MutableObject(*ELFObjFile);
+
+  using segment_iterator = MutableELFObject<ELF64LE>::segment_iterator;
+
+  EXPECT_EQ(iterRangeLength(MutableObject.segments()), 2);
+  segment_iterator FirstSegment = MutableObject.segment_begin();
+  segment_iterator SecondSegment = ++MutableObject.segment_begin();
+
+  section_iterator End = MutableObject.section_end();
+  section_iterator TextSection = MutableObject.findSectionByName(".text");
+  EXPECT_NE(TextSection, End);
+  section_iterator DataSection = MutableObject.findSectionByName(".data");
+  EXPECT_NE(DataSection, End);
+  section_iterator BssSection = MutableObject.findSectionByName(".bss");
+  EXPECT_NE(BssSection, End);
+  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<std::unique_ptr<ObjectFile>> 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<ELFObjectFile<ELF64LE>>(ErrOrObj->get());
+  ASSERT_TRUE(ELFObjFile);
+  MutableELFObject<ELF64LE> MutableObject(*ELFObjFile);
+
+  using segment_iterator = MutableELFObject<ELF64LE>::segment_iterator;
+
+  segment_iterator SecondSegment = ++MutableObject.segment_begin();
+  section_iterator DataSection = MutableObject.findSectionByName(".data");
+  section_iterator BssSection = MutableObject.findSectionByName(".bss");
+  auto *MutableSegment = MutableObject.getMutableSegment(SecondSegment);
+  EXPECT_TRUE(MutableSegment->hasSection(DataSection));
+  MutableSegment->addSection(BssSection);
+  EXPECT_TRUE(MutableSegment->hasSection(BssSection));
+  EXPECT_TRUE(SecondSegment->hasSection(BssSection));
+}