Additional fixes for #1130

lief-project · Dec 9, 2024 · 52a5d5d · 52a5d5d
1 parent f1b3586
commit 52a5d5d
Show file tree

Hide file tree

Showing 4 changed files with 44 additions and 3 deletions.
diff --git a/include/LIEF/MachO/Binary.hpp b/include/LIEF/MachO/Binary.hpp
@@ -944,6 +944,9 @@ class LIEF_API Binary : public LIEF::Binary  {
   void sort_segments();
   void refresh_seg_offset();
 
+  /// Check if the given segment can go in the offset_seg_ cache
+  static LIEF_LOCAL bool can_cache_segment(const SegmentCommand& segment);
+
   private:
   /// Default constructor
   LIEF_LOCAL Binary();

diff --git a/src/MachO/Binary.cpp b/src/MachO/Binary.cpp
@@ -600,7 +600,9 @@ void Binary::sort_segments() {
   for (auto it = start; it != end; ++it) {
     SegmentCommand& seg = *(*it)->as<SegmentCommand>();
     seg.index_ = segments_.size();
-    offset_seg_[seg.file_offset()] = &seg;
+    if (can_cache_segment(seg)) {
+      offset_seg_[seg.file_offset()] = &seg;
+    }
     segments_.push_back(&seg);
   }
 }
@@ -2336,7 +2338,7 @@ ExportInfo* Binary::add_exported_function(uint64_t address, const std::string& n
 void Binary::refresh_seg_offset() {
   offset_seg_.clear();
   for (SegmentCommand* segment : segments_) {
-    if (segment->file_offset_ == 0) {
+    if (!can_cache_segment(*segment)) {
       continue;
     }
     offset_seg_[segment->file_offset()] = segment;
@@ -2385,6 +2387,21 @@ Binary::stub_iterator Binary::symbol_stubs() const {
   return make_range(std::move(begin), std::move(end));
 }
 
+bool Binary::can_cache_segment(const SegmentCommand& segment) {
+  if (segment.file_offset() > 0 && segment.file_size() > 0) {
+    return true;
+  }
+
+  if (segment.name() == "__TEXT") {
+    // In some cases (c.f. <samples>/MachO/issue_1130.macho)
+    // the __TEXT segment can have a file_size set to 0 while it is logically
+    // revelant to cache it
+    return true;
+  }
+
+  return false;
+}
+
 Binary::~Binary() = default;
 
 std::ostream& Binary::print(std::ostream& os) const {

diff --git a/src/MachO/BinaryParser.tcc b/src/MachO/BinaryParser.tcc
@@ -279,7 +279,7 @@ ok_error_t BinaryParser::parse_load_commands() {
           segment->index_ = binary_->segments_.size();
           binary_->segments_.push_back(segment);
 
-          if (segment->file_size() > 0) {
+          if (Binary::can_cache_segment(*segment)) {
             binary_->offset_seg_[segment->file_offset()] = segment;
           }
 

diff --git a/tests/macho/test_library_injection.py b/tests/macho/test_library_injection.py
@@ -5,6 +5,7 @@
 import pathlib
 import re
 import pytest
+from pathlib import Path
 from subprocess import Popen
 from utils import is_osx, get_sample, is_apple_m1
 
@@ -128,3 +129,23 @@ def test_all_x86_64(tmp_path):
     print(stdout)
 
     assert re.search(r'CTOR CALLED', stdout) is not None
+
+
+@pytest.mark.parametrize("sample", [
+    "MachO/MachO64_x86-64_binary_sshd.bin",
+    "MachO/issue_1130.macho",
+])
+def test_segment_caching(tmp_path: Path, sample):
+    bin_path = Path(get_sample(sample))
+    original = lief.MachO.parse(bin_path.as_posix()).at(0)
+    output = tmp_path / bin_path.name
+    library_path = "/private/var/folders/vb/jj4r3nc1657b19v26p3kpclc0000gp/T/pytest-of-github-runner/pytest-18/test_ssh0/libexample.dylib"
+
+    original.add_library(library_path)
+
+    original.remove_signature()
+    original.write(output.as_posix())
+    new = lief.MachO.parse(output).at(0)
+
+    checked, err = lief.MachO.check_layout(new)
+    assert checked, err