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

[abrachet]

Summary

Added support for segments

[abrachet]

Rebase

[jhenderson]

How about some testing for changing program header properties, e.g. size or address?

llvm/include/llvm/Object/MutableELFObject.h

210

Any particular reason you've moved this line?

221

This should be in one of the earlier changes probably, if it can be.

245–246

Did you mean to remove this?

305

What is the point of this and the less-than operator?

306

Just do:

return Object == Other.Object && Index == Other.Index;

308

Blank line before this. Also, is it clang-formatted?

319

What's the point of this function?

321–327

I'm not convinced this patch should attempt to add sections to segments. The reason for this is because they're not an intrinsic part of segments, but rather an artefact of the layout of an ELF object. Relatedly, this doesn't handle segments within other segments, so feels incomplete. I'd defer it to a patch that handles laying out an ELF.

357–358

Did you mean to remove this? Looks like a bad rebase to me.

362–363

Did you mean to remove this? Looks like a bad rebase to me.

394–397

Did you mean to remove this? Looks like a bad rebase to me.

415

Comment?

Also Seg -> Segment

llvm/unittests/Object/MutableELFObjectTest.cpp

438

Seems unrelated?

493

Seems unrelated?

503

Seems unrelated?

751

I'd have at least 2 different segments here and test that both can be read in.

[abrachet]

Addressed comments

[abrachet]

Ahh sorry about all those unrelated changes! Not sure what happened there probably a bad rebase like you said.

llvm/include/llvm/Object/MutableELFObject.h

305

They are needed by content_iterator same with moveNext

321–327

Ok, I'll move it out to a part 7.

[MaskRay]

llvm/include/llvm/Object/MutableELFObject.h

72

The ctor is not really necessary.

211

This should be added to Part I D64281.

302

You can use uint64_t Index = 0; (did you mean size_t Index = 0;?) and delete the ctor.

424

Delete &&. See my comment in Part 4 D66062.

[jhenderson]

llvm/include/llvm/Object/MutableELFObject.h

527

Can this be inlined into getMutableSegment?

528

Is there a particular reason this implementation requires getIfNew, whereas the symbol and section versions don't?

llvm/unittests/Object/MutableELFObjectTest.cpp

808

Changing the address of a segment is very unlikely to ever actually happen, I'd think. I'd change the p_offset field instead.

[abrachet]

Addressed comments

[abrachet]

llvm/include/llvm/Object/MutableELFObject.h

72

It eventually gets called by an emplace_back

211

I had this in an earlier patch but then moved it out because it wasn't needed until now. I don't have a preference either way which patch it goes in.

302

Hmm am I missing something on how to construct objects? C++ doesn't have the nice C syntax like SegmentRef S = {.Object = nullptr, .Index = 0}; Is there a way to leave out these boiler plate constructors that I don't know of?

528

Yes, although I'm not sure if you'll like it :) Creating a MutableELFSegment is more expensive than the others because it takes a list of all sections that reside in it. So using getIfNew allows to not have to calculate that if it is already been made mutable then it just returns a mutable reference to a MutableELFSegment. I say you might not like that because what I described doesn't happen in this patch but only will in Part 7.

[abrachet]

Use the right diff (!)

[jhenderson]

Posting this here since I've not seen a patch for Part 7 yet. I may have said this elsewhere already, but I don't think it's a good idea to have sections as members of segments. I think you'll find the layout algorithm easier to work with when using a parent relationship, rather than a child one. Note that this is how llvm-objcopy is already written. That should allow you to handle multiple levels of nesting/overlapping etc without too much difficulty. That would also, I believe, make the Segment and Section implementations much closer to each other.

llvm/include/llvm/Object/MutableELFObject.h

302

Do you ever need the default value of Index? It seems to me to be unnecessary.

[abrachet]

Added parent segment to both section and segment types and methods to find parent segments

[abrachet]

That should allow you to handle multiple levels of nesting/overlapping etc without too much difficulty.

Hmm if a segment had a list of its child segments, then couldn't layout be done recursively on each child segment? Would this not be more straight forward?

I think what I currently have should move more in the direction you are suggesting. I have a concern with my current implementation, its if I should use the original program headers or the updated ones. It's possible to change a segments vaddr, although you have mentioned unlikely. This would be problematic to use sectionInSegment after changing the program headers because then it would be obviously wrong. So right now I use the original program headers. Also something to look out for, I use parent segments as indexes not pointers and those are ssize_t not size_t because the alternative of using one past the last segment like an end iterator was too ugly than using -1 to denote no parent. I could also use size_t's max I suppose.

[jhenderson]

In D66070#1640579, @abrachet wrote:

That should allow you to handle multiple levels of nesting/overlapping etc without too much difficulty.

Hmm if a segment had a list of its child segments, then couldn't layout be done recursively on each child segment? Would this not be more straight forward?

What happens if a segment is inside another segment? You only want to lay it out if it hasn't been laid out already, so you'd have to find a way to track that. On the other hand, if you only lay out things that have no parent (and then lay out the things with parents based on their parent's layout), you don't have this issue. Keeping track of children is not necessarily bad, as it allows you to lay the children out at the same time as the parent, but it's not as important (you can always do two passes for example).

I think what I currently have should move more in the direction you are suggesting. I have a concern with my current implementation, its if I should use the original program headers or the updated ones. It's possible to change a segments vaddr, although you have mentioned unlikely. This would be problematic to use sectionInSegment after changing the program headers because then it would be obviously wrong. So right now I use the original program headers.

Using the original program headers is probably the right thing to do, as that ensures the original relative layout is preserved.

Also something to look out for, I use parent segments as indexes not pointers and those are ssize_t not size_t because the alternative of using one past the last segment like an end iterator was too ugly than using -1 to denote no parent. I could also use size_t's max I suppose.

See my inline suggestion about using Optional. size_t's max would be my next suggestion.

llvm/include/llvm/Object/MutableELFObject.h

268–269

How about using Optional<size_t> instead of -1? I think that would more clearly convey the meaning and would avoid the weirdness of using ssize_t everywhere.

273

Could you not just move SegmentRef's definition earlier?

276

I'd keep this and findSectionsParentSegment() next to each other, and I'd seriously consider just calling them both findParentSegment().

551–552

Something you may wish to consider is effectively flattening the nesting, so that only top-level segments are parents, i.e. if a segment has a parent, then the parent cannot have a parent itself. This may make the layout algorithm simpler, but could probably be left until you're at that point.

Same comment applies to sections.

558

Is this identical to the llvm-objcopy version? If not, could you highlight how it's different, please.

582

Delete the blank line here.

You need to be careful. It looks to me like a segment could end up as its own parent here, if I'm not misreading the code.

Also, you should consider taking alignment requirements into account, in the event that two segments have the same size and offset. See https://bugs.llvm.org/show_bug.cgi?id=42963.

Also, consider the following bug in your sectionWithinSegment code: https://bugs.llvm.org/show_bug.cgi?id=41126.

[abrachet]

Use Optional<size_t> to represents parent segments. Addressed other comments.

[abrachet]

llvm/include/llvm/Object/MutableELFObject.h

268–269

Thats much better :)

551–552

Something like if (segmentWithinSegment(Child, Phdr) && findParentSegment(Phdr) != None)?

Same comment applies to sections.

I think I am misunderstanding something, do sections have anything to do whether a segment can have a parent? There is no overlapping sections I thought so I don't think this relates to that. Sorry about the confusion!

582

You need to be careful. It looks to me like a segment could end up as its own parent here, if I'm not misreading the code.

Indeed, thanks fixed that.

Also, you should consider taking alignment requirements into account, in the event that two segments have the same size and offset. See https://bugs.llvm.org/show_bug.cgi?id=42963.

I have added this it happens on line 941 of the unit test, is that what you mean?

Also, consider the following bug in your sectionWithinSegment code: https://bugs.llvm.org/show_bug.cgi?id=41126.

I think I have this:

|-Segment-|
    |-Section-|

tested and it properly says that the section is in the segment. There was no change needed for that.
To get yaml2obj to output that I just made the sections size larger than the segment and put the section in the segment. I had trouble with the other two examples in that bug report because Offset: was not recognized for sections? So I don't know how to test these two cases:

    |-Segment-|
|-Section-|

  |-Segment-|
|---Section---|

but I know for sure they will not be reported as being in the segment.

[jhenderson]

Could you write some comments in the test code that show what the cases are that you are testing, similar to how you've written it in your inline comments. This will make it easier to understand what is going on.

llvm/include/llvm/Object/MutableELFObject.h

551–552

I'm thinking that if the segment layout looks something like the following:

| Section1 | Section2 | Section3 | Section4 |
| Segment1            | Segment2 |
| Segment3                                  |

Everything should have Segment3 as their parent (except Segment3 itself of course).

If you don't flatten the tree, then sections 1, 2, and 3 will have parents of Segment 1 or 2 which will make it harder to perform layout, as you effectively have to chase the parents pointer to find the top-level element to determine where to move them.

I think I am misunderstanding something, do sections have anything to do whether a segment can have a parent? There is no overlapping sections I thought so I don't think this relates to that. Sorry about the confusion!

Yes, I think you misunderstood. I was just referring to the fact that the parent of a section should be the top-most parent in the tree.

557

except it uses

558

and not

582

Indeed, thanks fixed that.

I think you've missed an issue - what happens if two segments have the same offset, size and alignment? In that case, neither will be treated as a parent of the other, and the algorithm falls apart.

I had trouble with the other two examples in that bug report because Offset: was not recognized for sections? So I don't know how to test these two cases:

I think there's a ShOffset field available to set the claimed offset of a section (the actual data will still be written where it would be). The alternative to that is to effectively force the section's offset by setting the AddressAlign and Address fields, so that yaml2obj has to place it at a certain position. You'll see quite a few examples of this within tests. Once you have forced a section to a specific position, you can then just set the ProgramHeader fields accordingly.

583

I'd rename Parent to PotentialParent and Child to Segment, because the Parent isn't necessarily a Parent yet (and therefore the Child isn't necessarily a Child of that segment).

[abrachet]

Addressed review comments

[abrachet]

llvm/include/llvm/Object/MutableELFObject.h

582

I think you've missed an issue - what happens if two segments have the same offset, size and alignment? In that case, neither will be treated as a parent of the other, and the algorithm falls apart.

Ok, now it is finally fixed. It isn't very pretty because I need to pass the indexes of each segment but I just choose whichever segment is first to be the parent.

[jhenderson]

I'm out of time for today on this review. I'll come back to it tomorrow.

llvm/include/llvm/Object/MutableELFObject.h

529–530

A comment briefly explaining this would be helpful.

535–538

Some code comments explaining what is going on here would be useful.

547–548

An explanation with some comments as to why you can have this early out would make this clearer.

587

ChildOffset -> SegmentIndex
ParentOffset -> ParentIndex

Write a comment at the start of this function outlining the behaviour, please.

llvm/unittests/Object/MutableELFObjectTest.cpp

900–901

doesnt -> doesn't

But really, I'm a little confused what this comment is trying to say. I think you want to say something about offsets, because it's confusing what you're talking about in this comment. It appears you're saying 4 doesn't come before 5, but it's earlier in the program header list, which seems to contradict what you're saying.

llvm/include/llvm/Object/MutableELFObject.h

template <typename ELFT> struct MutableELFSymbol {
   std::string Name;
 
   MutableELFSymbol(const Elf_Sym &Header, StringRef Name)
       : Header(Header), Name(Name) {}
 
   operator const Elf_Sym &() const { return Header; }
 };
 
+template <typename ELFT> class MutableELFSegment {
+  using Elf_Phdr = Elf_Phdr_Impl<ELFT>;
+
+public:
+  Elf_Phdr Header;
+
+  MutableELFSegment(Elf_Phdr Header) : Header(Header) {}
+
+  operator const Elf_Phdr &() const { return Header; }
+};
+
 template <typename ELFT> class MutableELFObject : public ELFObjectFile<ELFT> {
   /// This class is used for a 'copy on write' effect with tables in an ELF
   /// object file.
   ///
   /// The table holds entries of two different but related types. OrigType
   /// represents the object type in the file originaly while NewType is the
   /// type to create on modification of the entry. It is required that
   /// NewType have an `operator const OrigType &()`.
public:
     const NewType *getConstIfNew(uint64_t Index) const {
       assert(Index < Mappings.size() && "Out of bounds");
       assert(Mappings[Index].Type != MappingType::Removed);
       return Mappings[Index].Type == MappingType::New
                  ? &NewValues[Mappings[Index]]
                  : 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 Mapping.Type != MappingType::Removed;
       });
     }
 
     size_t originalSize() const { return OriginalValues.size(); }
   };
 
   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_Sym, MutableELFSymbol<ELFT>>;
 
   Elf_Ehdr FileHeader;
+  MutableTable<Elf_Phdr, MutableELFSegment<ELFT>> Segments;
   MutableTable<Elf_Shdr, MutableELFSection<ELFT>> Sections;
   MutableSymbolTable Symbols;
   MutableSymbolTable DynSymbols;
 
   /// 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.
   const Elf_Shdr *getSection(DataRefImpl Sec) const override {
     return &Sections[Sec.p];
   }
 
   const Elf_Sym *getSymbol(DataRefImpl Sym) const override {
     const MutableSymbolTable &Table = getSymbolTable(Sym);
     if (const MutableELFSymbol<ELFT> *SymOrNull = Table.getConstIfNew(Sym.d.b))
       return &SymOrNull->Header;
     return ELFObjectFile<ELFT>::getSymbol(Sym);
   }
 
+  const Elf_Phdr *getSegment(DataRefImpl Sec) const { return &Segments[Sec.p]; }
+
   /// moveSectionNext must be overriden because the section_iterators in
   /// MutableELFObject work differently than in ELFObjectFile. In this class,
   /// sections are iterated with their index, not address in the file, which
   /// allows use with MutableTable.
   void moveSectionNext(DataRefImpl &Sec) const override;
   Expected<StringRef> getSectionName(DataRefImpl Sec) const override;
   uint64_t getSectionIndex(DataRefImpl Sec) const override;
   Expected<ArrayRef<uint8_t>>
   getSectionContents(DataRefImpl Sec) const override;
 
   void moveSymbolNext(DataRefImpl &Sym) const override;
   Expected<StringRef> getSymbolName(DataRefImpl Sym) const override;
   ArrayRef<Elf_Sym> findSymbolTable(uint64_t ShType) const;
   uint32_t findSectionOfType(uint64_t ShType) const;
 
   const MutableSymbolTable &getSymbolTable(DataRefImpl Sym) const {
     return Sections.getOriginal(Sym.d.a).sh_type == ELF::SHT_SYMTAB
                ? Symbols
                : DynSymbols;
   }
 
+  MutableELFSegment<ELFT> &makeMutableSegment(uint64_t Seg);
+
   static DataRefImpl toDataRef(uintptr_t Ptr);
   static DataRefImpl toDataRef(uint32_t A, uint32_t B);
 
 public:
   explicit MutableELFObject(ELFObjectFile<ELFT> &&B)
       : ELFObjectFile<ELFT>(std::move(B)),
         FileHeader(*reinterpret_cast<const Elf_Ehdr *>(this->base())),
+        Segments(ArrayRef<Elf_Phdr>(reinterpret_cast<const Elf_Phdr *>(
+                                        this->base() + getHeader().e_phoff),
+                                    getHeader().e_phnum)),
         Sections(ArrayRef<Elf_Shdr>(reinterpret_cast<const Elf_Shdr *>(
                                         this->base() + getHeader().e_shoff),
                                     getHeader().e_shnum)),
         Symbols(findSymbolTable(ELF::SHT_SYMTAB)),
         DynSymbols(findSymbolTable(ELF::SHT_DYNSYM)) {}
 
+  class SegmentRef {
+    MutableELFObject<ELFT> *Object;
+    uint64_t Index;
+
+  public:
+    SegmentRef(MutableELFObject<ELFT> *Object, uint64_t Index = 0)
+        : 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;
+    }
+
+    uint64_t getIndex() const {
+      return Object->Segments.getEffectiveIndex(Index);
+    }
+
+    void moveNext() { Index = Object->Segments.getNextIndex(Index); }
+
+    const Elf_Phdr &getHeader() const { return Object->Segments[Index]; }
+  };
+
+  using segment_iterator = content_iterator<SegmentRef>;
+
   section_iterator section_begin() const override {
     return section_iterator(
         SectionRef(toDataRef(Sections.getFirstIndex()), this));
   }
 
   section_iterator section_end() const override {
     return section_iterator(
         SectionRef(toDataRef(Sections.getEndIndex()), this));
   }
 
   basic_symbol_iterator symbol_begin() const override {
public:
   }
 
   elf_symbol_iterator dynamic_symbol_end() const override {
     return basic_symbol_iterator(SymbolRef(
         toDataRef(findSectionOfType(ELF::SHT_DYNSYM), DynSymbols.getEndIndex()),
         this));
   }
 
+  segment_iterator segment_begin() {
+    return segment_iterator(SegmentRef(this));
+  }
+
+  segment_iterator segment_end() {
+    return segment_iterator(SegmentRef(this, Segments.getEndIndex()));
+  }
+
+  iterator_range<segment_iterator> segments() {
+    return iterator_range<segment_iterator>(segment_begin(), segment_end());
+  }
+
   /// Returns a mutable reference to the symbol at the specified index.
   Expected<MutableELFSymbol<ELFT> &> getMutableSymbol(SymbolRef Sym) {
     Expected<StringRef> Name = getSymbolName(Sym.getRawDataRefImpl());
     if (!Name)
       return Name.takeError();
     return Symbols.makeMutable(Sym.getRawDataRefImpl().d.b,
                                Symbols[Sym.getRawDataRefImpl().d.b], *Name);
   }
 
   /// Returns a mutable reference to the dynamic symbol at the specified index.
   Expected<MutableELFSymbol<ELFT> &> getMutableDynamicSymbol(SymbolRef Sym) {
     Expected<StringRef> Name = getSymbolName(Sym.getRawDataRefImpl());
     if (!Name)
       return Name.takeError();
     return DynSymbols.makeMutable(Sym.getRawDataRefImpl().d.b,
                                   DynSymbols[Sym.getRawDataRefImpl().d.b],
                                   *Name);
   }
 
+  /// Returns a mutable reference to a segment.
+  Expected<MutableELFSegment<ELFT> &> getMutableSegment(SegmentRef Segment) {
+    return makeMutableSegment(Segment.getIndex());
+  }
+
   /// Returns a mutable reference to the section pointed to by Sec. A possible
   /// usage to change all sections with alignment of 0 to 1 could be:
   /// @code{.cpp}
   /// for (const SectionRef &Sec : MutObj.sections())
   ///   if (!Sec.getAlignment()) {
   ///     auto MutSecOrErr = MutObj.getMutableSection(Sec);
   ///     if (!MutSecOrErr)
   ///       handleError(MutSecOrErr.takeError());
   ///     MutSecOrErr->Header.sh_addrallgin = 1;
   ///   }
   /// @endcode
   Expected<MutableELFSection<ELFT> &> getMutableSection(SectionRef Sec) {
     const Elf_Shdr_Impl<ELFT> &Header = Sections[Sec.getRawDataRefImpl().p];
     Expected<StringRef> Name = getSectionName(Sec.getRawDataRefImpl());
     if (!Name)
       return Name.takeError();
     return Sections.makeMutable(Sec.getRawDataRefImpl().p, Header, *Name, this);
   }
 
   /// Add symbol \param Sym.
   void addSymbol(MutableELFSymbol<ELFT> &&Sym) { Symbols.add(std::move(Sym)); }
 
   /// Add section \param Sec.
   void addSection(MutableELFSection<ELFT> &&Sec) {
     Sections.add(std::move(Sec));
   }
 
+  /// Add segment \param Seg.
+  void addSegment(MutableELFSegment<ELFT> &&Seg) {
+    Segments.add(std::move(Seg));
+  }
+
   /// Removes a symbol.
   void removeSymbol(SymbolRef Sym) {
     assert(Sections.getOriginal(Sym.getRawDataRefImpl().d.a).sh_type ==
                ELF::SHT_SYMTAB &&
            "Not pointing to symbol table");
     Symbols.remove(Sym.getRawDataRefImpl().d.b);
   }
 
   /// Removes a section.
   void removeSection(SectionRef Sec) {
     Sections.remove(Sec.getRawDataRefImpl().p);
   }
 
   const Elf_Ehdr &getHeader() const override { return FileHeader; }
   Elf_Ehdr &getMutableHeader() { return FileHeader; }
 };
 
 template <typename ELFT>
 void MutableELFObject<ELFT>::moveSectionNext(DataRefImpl &Sec) const {
   Sec.p = Sections.getNextIndex(Sec.p);
 }
 

 Expected<ArrayRef<uint8_t>>
 MutableELFObject<ELFT>::getSectionContents(DataRefImpl Sec) const {
   if (const MutableELFSection<ELFT> *SecOrNull = Sections.getConstIfNew(Sec.p))
     return ArrayRef<uint8_t>(SecOrNull->Data.data(), SecOrNull->Header.sh_size);
   return ELFObjectFile<ELFT>::getSectionContents(Sec);
 }
 
 template <typename ELFT>
 void MutableELFObject<ELFT>::moveSymbolNext(DataRefImpl &Sym) const {
   const MutableSymbolTable &SymbolTable = getSymbolTable(Sym);
   Sym.d.b = SymbolTable.getNextIndex(Sym.d.b);
 }
 
 template <typename ELFT>
 Expected<StringRef>
 MutableELFObject<ELFT>::getSymbolName(DataRefImpl Sym) const {
   const MutableSymbolTable &SymbolTable = getSymbolTable(Sym);
   if (const MutableELFSymbol<ELFT> *SymOrNull =
           SymbolTable.getConstIfNew(Sym.d.b))
     return SymOrNull->Name;

 template <typename ELFT>
 DataRefImpl MutableELFObject<ELFT>::toDataRef(uint32_t A, uint32_t B) {
   DataRefImpl Ref;
   Ref.d.a = A;
   Ref.d.b = B;
   return Ref;
 }
 
+template <typename ELFT>
+MutableELFSegment<ELFT> &
+MutableELFObject<ELFT>::makeMutableSegment(uint64_t Seg) {
+  if (MutableELFSegment<ELFT> *Phdr = Segments.getIfNew(Seg))
+    return *Phdr;
+
+  const Elf_Phdr &Phdr = Segments[Seg];
+  return Segments.makeMutable(Seg, Phdr);
+}
+
 } // namespace object
 } // namespace llvm
 
 #endif // LLVM_OBJECT_MUTABLEELFOBJECT_H

llvm/unittests/Object/MutableELFObjectTest.cpp

DynamicSymbols:
   EXPECT_EQ(std::distance(MutableObject.dynamic_symbol_begin(),
                           MutableObject.dynamic_symbol_end()),
             3);
 
   auto FirstSym = ++MutableObject.symbol_begin();
   Expected<StringRef> NameOrErr = symbol_iterator(FirstSym)->getName();
   ASSERT_THAT_EXPECTED(NameOrErr, Succeeded());
   EXPECT_EQ(*NameOrErr, "first");
   MutableObject.removeSymbol(*FirstSym);
 
   auto NewFirstSym = ++MutableObject.symbol_begin();
   auto NewNameOrErr = symbol_iterator(NewFirstSym)->getName();
   ASSERT_THAT_EXPECTED(NewNameOrErr, Succeeded());
   EXPECT_EQ(*NewNameOrErr, "second");
 
   EXPECT_EQ(
       std::distance(MutableObject.symbol_begin(), MutableObject.symbol_end()),
EXPECT_EQ(
       std::distance(MutableObject.section_begin(), MutableObject.section_end()),
       6);
 
   auto Iter = ++MutableObject.section_begin();
   StringRef Name;
   ASSERT_FALSE(Iter->getName(Name));
   EXPECT_EQ(Name, ".remove_me");
 
   MutableObject.removeSection(*Iter);
 
   EXPECT_EQ(
       std::distance(MutableObject.section_begin(), MutableObject.section_end()),
       5);
 
   Iter = ++MutableObject.section_begin();
   ASSERT_FALSE(Iter->getName(Name));
   EXPECT_EQ(Name, ".remove_me_too");
 
   MutableObject.removeSection(*Iter);
 
   EXPECT_EQ(
       std::distance(MutableObject.section_begin(), MutableObject.section_end()),
       4);
 
   Iter = MutableObject.section_begin();
   EXPECT_EQ(Iter->getIndex(), 0u);
 
FileHeader:
   Header.e_flags = 100;
   EXPECT_EQ(MutableObject.getPlatformFlags(), 100u);
 
   Header.e_entry = 5;
   Expected<uint64_t> StartAddr = MutableObject.getStartAddress();
   ASSERT_THAT_EXPECTED(StartAddr, Succeeded());
   EXPECT_EQ(*StartAddr, 5u);
 }
+
+// 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
+  - Type:       0x12345
+    FileSize:   0
+    MemSize:    100)");
+
+  ASSERT_THAT_EXPECTED(ErrOrObj, Succeeded());
+  auto *ELFObjFile = dyn_cast<ELFObjectFile<ELF64LE>>(ErrOrObj->get());
+  ASSERT_TRUE(ELFObjFile);
+  MutableELFObject<ELF64LE> 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<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(std::move(*ELFObjFile));
+
+  auto FirstSegment = MutableObject.segment_begin();
+  EXPECT_EQ(FirstSegment->getHeader().p_vaddr, 0x1000u);
+  Expected<MutableELFSegment<ELF64LE> &> MutSegmentOrErr =
+      MutableObject.getMutableSegment(*FirstSegment);
+  ASSERT_THAT_EXPECTED(MutSegmentOrErr, Succeeded());
+  MutableELFSegment<ELF64LE> &MutSegment = *MutSegmentOrErr;
+  MutSegment.Header.p_vaddr = 0x1100;
+  EXPECT_EQ(FirstSegment->getHeader().p_vaddr, 0x1100u);
+}