MachO::binary: add extend_section (#1137)

* MachO::Binary::add_section: respect alignment of empty sections

* MachO/layout_check: relax offset continuity requirement for __text section

Sections in MachO are prepended, but __text section is not moved.
Due to this there might be redundant gap between __text section and a
section that comes before __text.

Note: we don't want to alter distance between __DATA and __text, because
there might be position relative references in metadata.

* tests/macho: use lief.MachO.check_layout

* MachO::binary: add extend_section

This commit adds member function MachO::Binary::extend_section that is
able to extend sections from the first segment in MachO.

* macho/test_builder: add test_extend_section

api/python/lief/MachO/__init__.pyi

@@ -476,6 +476,8 @@ class Binary(lief.Binary):
     @overload
     def add_section(self, section: Section) -> Section: ...
 
+    def extend_section(self, section: Section, size: int) -> bool: ...
+
     def add_library(self, library_name: str) -> LoadCommand: ...
 
     def get(self, type: LoadCommand.TYPE) -> LoadCommand: ...

api/python/src/MachO/objects/pyBinary.cpp

@@ -630,6 +630,11 @@ void create<Binary>(nb::module_& m) {
         "section"_a,
         nb::rv_policy::reference_internal)
 
+    .def("extend_section",
+        nb::overload_cast<Section&, size_t>(&Binary::extend_section),
+        "Extend the **content** of the given " RST_CLASS_REF(lief.MachO.Section) " by ``size``"_doc,
+        "section"_a, "size"_a)
+
     .def("add_library",
         nb::overload_cast<const std::string&>(&Binary::add_library),
         "Add a new library dependency"_doc,

include/LIEF/MachO/Binary.hpp

@@ -410,6 +410,12 @@ class LIEF_API Binary : public LIEF::Binary  {
   /// Extend the **content** of the given SegmentCommand
   bool extend_segment(const SegmentCommand& segment, size_t size);
 
+  /// Extend the **content** of the given Section.
+  /// @note This method may create a gap between the current section and the next one, if `size`
+  ///       is not multiple of the maximum alignment of sections before the current one.
+  /// @note This method works only with sections that belong to the first segment.
+  bool extend_section(Section& section, size_t size);
+
   /// Remove the `PIE` flag
   bool disable_pie();
 

src/MachO/Binary.cpp

@@ -883,7 +883,7 @@ ok_error_t Binary::shift(size_t value) {
   // Segment that wraps this load command table
   SegmentCommand* load_cmd_segment = segment_from_offset(loadcommands_end);
   if (load_cmd_segment == nullptr) {
-    LIEF_WARN("Can't find segment associated with last load command");
+    LIEF_ERR("Can't find segment associated with load command space");
     return make_error_code(lief_errors::file_format_error);
   }
   LIEF_DEBUG("LC Table wrapped by {} / End offset: 0x{:x} (size: {:x})",
@@ -1229,6 +1229,88 @@ bool Binary::extend_segment(const SegmentCommand& segment, size_t size) {
   return true;
 }
 
+bool Binary::extend_section(Section& section, size_t size) {
+  // All sections must keep their requested alignment.
+  //
+  // As per current implementation of `shift` method, space is allocated between
+  // the last load command and the first section by shifting everything to the "right".
+  // After that we shift `section` and all other sections that come before it to the "left",
+  // so that we create a gap of at least `size` wide after the current `section`.
+  // Finally, we assign new size to the `section`.
+  //
+  // Let's say we are extending section S.
+  // There might be sections P that come prior S, and there might be sections A that come after S.
+  // We try to keep relative relationships between sections in groups P and A,
+  // such that relative offsets from one section to another one are unchanged,
+  // however preserving the same relationship between sections from different groups is impossible.
+  // We achieve this by shifting P and S to the left by size rounded up to the maximum common alignment factor.
+
+  const uint64_t loadcommands_start = is64_ ? sizeof(details::mach_header_64) :
+                                              sizeof(details::mach_header);
+  const uint64_t loadcommands_end = loadcommands_start + header().sizeof_cmds();
+  SegmentCommand* load_cmd_segment = segment_from_offset(loadcommands_end);
+  if (load_cmd_segment == nullptr) {
+    LIEF_ERR("Can't find segment associated with load command space");
+    return false;
+  }
+
+  if (section.segment() != load_cmd_segment) {
+    LIEF_ERR("Can't extend section that belongs to segment '{}' which is not the first one", section.segment_name());
+    return false;
+  }
+
+  // Select sections that we need to shift.
+  sections_cache_t sections_to_shift;
+  for (Section& s : sections()) {
+    if (s.offset() > section.offset() || s.offset() == 0) {
+      continue;
+    }
+    sections_to_shift.push_back(&s);
+  }
+  assert(!sections_to_shift.empty());
+
+  // Stable-sort by end_offset, preserving original order.
+  std::stable_sort(sections_to_shift.begin(), sections_to_shift.end(),
+            [](const Section* a, const Section* b) {
+              return a->offset() + a->size() > b->offset() + b->size();
+            });
+
+  // If we are extending an empty section, then there may be many sections at this offset,
+  // we do not want to shift sections which were added after the current one.
+  // Such that we preserve order of sections in which they were added to the binary.
+  if (section.size() == 0) {
+    auto it = std::find(sections_to_shift.begin(), sections_to_shift.end(), &section);
+    assert(it != sections_to_shift.end());
+    sections_to_shift.erase(std::next(it), sections_to_shift.end());
+  }
+
+  // Find maximum alignment
+  auto it_maxa = std::max_element(sections_to_shift.begin(), sections_to_shift.end(),
+            [](const Section* a, const Section* b) {
+              return a->alignment() < b->alignment();
+            });
+  const size_t max_alignment = 1 << (*it_maxa)->alignment();
+
+  // Resize command space, if needed.
+  const size_t shift_value = align(size, max_alignment);
+  if (auto result = ensure_command_space(shift_value); is_err(result)) {
+    LIEF_ERR("Failed to ensure command space {}: {}", shift_value, to_string(get_error(result)));
+    return false;
+  }
+  available_command_space_ -= shift_value;
+
+  // Shift selected sections to allocate requested space for `section`.
+  for (Section* s : sections_to_shift) {
+    s->offset(s->offset() - shift_value);
+    s->address(s->address() - shift_value);
+  }
+
+  // Extend the given `section`.
+  section.size(section.size() + size);
+
+  return true;
+}
+
 void Binary::remove_section(const std::string& name, bool clear) {
   Section* sec_to_delete = get_section(name);
   if (sec_to_delete == nullptr) {
@@ -1338,7 +1420,7 @@ Section* Binary::add_section(const SegmentCommand& segment, const Section& secti
     // Section offset is not defined: we need to allocate space enough to fit its content.
     const size_t hdr_size = is64_ ? sizeof(details::section_64) :
                                     sizeof(details::section_32);
-    const size_t alignment = content.empty() ? 0 : 1 << section.alignment();
+    const size_t alignment = 1 << section.alignment();
     const size_t needed_size = hdr_size + content.size() + alignment;
 
     // Request size with a gap of alignment, so we would have enough room

src/MachO/layout_check.cpp

@@ -845,6 +845,17 @@ bool LayoutChecker::check_section_contiguity() {
     // Skip this check for object file
     return true;
   }
+
+  // LIEF allocates space for new/extended sections between the last load command
+  // and the first section in the first segment.
+  // We are not willing to change the distance between the end of the `__text` section
+  // and start of `__DATA` segment, keeping `__text` section in a "fixed" position.
+  // Due to above there might happen a "reverse" alignment gap between `__text` section
+  // and a section that was allocated in front of it.
+  auto is_gap_reversed = [](const Section* LHS, const Section* RHS) {
+    return align_down(RHS->offset() - LHS->size(), LHS->alignment());
+  };
+
   for (const SegmentCommand& segment : binary.segments()) {
     std::vector<const Section*> sections_vec;
     auto sections = segment.sections();
@@ -874,28 +885,28 @@ bool LayoutChecker::check_section_contiguity() {
               });
 
     uint64_t next_expected_offset = sections_vec[0]->offset();
-    const uint32_t alignment = 1 << sections_vec[0]->alignment();
-    if (sections_vec[0]->offset() % (1 << sections_vec[0]->alignment()) != 0) {
-      return error("section '{}' offset (0x{:04x}) is misaligned (align=0x{:04x})",
-                   sections_vec[0]->name(), sections_vec[0]->offset(),
-                   alignment);
-    }
-
-    for (const Section* section : sections_vec) {
+    for (auto it = sections_vec.begin(); it != sections_vec.end(); ++it) {
+      const Section* section = *it;
       const uint32_t alignment = 1 << section->alignment();
-      if ((section->offset() % (1 << section->alignment())) != 0) {
-        return error("section '{}' offset (0x{:04x}) is misaligned (align=0x{:04x})",
+      if ((section->offset() % alignment) != 0) {
+        return error("section '{}' offset ({:#06x}) is misaligned (align={:#06x})",
                      section->name(), section->virtual_address(),
                      alignment);
       }
 
-      next_expected_offset = align(next_expected_offset, 1 << section->alignment());
+      next_expected_offset = align(next_expected_offset, alignment);
 
       if (section->offset() != next_expected_offset) {
-        return error("section '{}' is not at the expected offset: 0x{:04x} "
-                     "(expected: 0x{:04x}, alignment: 0x{:04x})",
-                     section->name(), section->offset(), next_expected_offset,
-                     section->alignment());
+        auto message = fmt::format("section '{}' is not at the expected offset: {:#06x} "
+                                   "(expected={:#06x}, align={:#06x})",
+                                   section->name(), section->offset(),
+                                   next_expected_offset, alignment);
+        if (it != sections_vec.begin() && is_gap_reversed(*std::prev(it), *it) && section->name() == "__text") {
+          LIEF_WARN("Permitting section gap which could be caused by LIEF add_section/extend_section: {}", message);
+          next_expected_offset = section->offset();
+        } else {
+          return error(message);
+        }
       }
       next_expected_offset += section->size();
     }

tests/macho/test_builder.py

@@ -159,25 +159,6 @@ def test_extend_cmd(tmp_path):
 
     assert new[lief.MachO.LoadCommand.TYPE.UUID].size == original_size + 0x4000
 
-def verify_sections_continuity(binary):
-    # Check that sections are continuous and non-overlapping in each segment
-    bad_section_types = [
-            lief.MachO.Section.TYPE.ZEROFILL,
-            lief.MachO.Section.TYPE.THREAD_LOCAL_ZEROFILL,
-            lief.MachO.Section.TYPE.GB_ZEROFILL,
-            ]
-    for segment in binary.segments:
-        if len(segment.sections) == 0:
-            continue
-        sections = sorted([s for s in segment.sections if s.type not in bad_section_types], key=lambda s: s.offset)
-        next_expected_offset = sections[0].offset
-        assert sections[0].offset % (1 << sections[0].alignment) == 0
-        for section in sections:
-            next_expected_offset = align_to(next_expected_offset, 1 << section.alignment)
-            assert section.offset % (1 << section.alignment) == 0
-            assert section.offset == next_expected_offset
-            next_expected_offset += section.size
-
 def test_add_section_id(tmp_path):
     bin_path = pathlib.Path(get_sample("MachO/MachO64_x86-64_binary_id.bin"))
     original = lief.MachO.parse(bin_path.as_posix()).at(0)
@@ -188,7 +169,8 @@ def test_add_section_id(tmp_path):
         section = lief.MachO.Section(f"__lief_{i}", [0x90] * 0x100)
         original.add_section(section)
 
-    verify_sections_continuity(original)
+    checked, err = lief.MachO.check_layout(original)
+    assert checked, err
 
     assert original.virtual_size % original.page_size == 0
 
@@ -205,6 +187,46 @@ def test_add_section_id(tmp_path):
         print(stdout)
         assert re.search(r'uid=', stdout) is not None
 
+def test_extend_section(tmp_path):
+    bin_path = pathlib.Path(get_sample("MachO/MachO64_x86-64_binary_id.bin"))
+    original = lief.MachO.parse(bin_path.as_posix()).at(0)
+    output = f"{tmp_path}/test_extend_section.bin"
+
+    text_segment = original.get_segment("__TEXT")
+    def make_section(name, alignment):
+        section = lief.MachO.Section(f"__lief_{alignment}")
+        section.alignment = alignment
+        return section
+
+    sections = [
+        make_section("__lief_8", 8),
+        make_section("__lief_7", 7),
+        make_section("__lief_6", 6),
+        make_section("__lief_5", 5),
+        make_section("__lief_4", 4),
+        make_section("__lief_3", 3),
+        make_section("__lief_2", 2),
+        make_section("__lief_1", 1),
+        make_section("__lief_0", 0),
+        make_section("__lief_00", 0),
+    ]
+    sections = [original.add_section(text_segment, s) for s in sections]
+
+    checked, err = lief.MachO.check_layout(original)
+    assert checked, err
+
+    for i, section in enumerate(sections):
+        assert original.extend_section(section, 1 << section.alignment)
+
+    checked, err = lief.MachO.check_layout(original)
+    assert checked, err
+
+    original.write(output)
+    new = lief.MachO.parse(output).at(0)
+
+    checked, err = lief.MachO.check_layout(new)
+    assert checked, err
+
 @pytest.mark.skipif(is_github_ci(), reason="sshd does not work on Github Action")
 def test_add_section_ssh(tmp_path):
     bin_path = pathlib.Path(get_sample("MachO/MachO64_x86-64_binary_sshd.bin"))