[Sanitizers][Darwin] Correct iterating of MachO load commands

[wrotki]

The condition to stop iterating so far was to look for load command cmd field == 0. The iteration would continue past the commands area, and would finally find lc->cmd ==0, if lucky. Or crash with bus error, if out of luck.

Correcting this by limiting the number of iterations to the count specified in mach_header(_64) ncmds field.

rdar://143903403

[llvmbot]

@llvm/pr-subscribers-compiler-rt-sanitizer

Author: Mariusz Borsa (wrotki)

Changes

The condition to stop iterating so far was to look for load command cmd field == 0. The iteration would continue past the commands area, and would finally find lc->cmd ==0, if lucky. Or crash with bus error, if out of luck.

Correcting this by limiting the number of iterations to the count specified in mach_header(_64) ncmds field.

rdar://143903403

---

Full diff: https://github.com/llvm/llvm-project/pull/130161.diff

2 Files Affected:

• (modified) compiler-rt/lib/sanitizer_common/sanitizer_procmaps_mac.cpp (+29-8)
• (modified) compiler-rt/lib/sanitizer_common/tests/sanitizer_procmaps_mac_test.cpp (+20-5)

diff --git a/compiler-rt/lib/sanitizer_common/sanitizer_procmaps_mac.cpp b/compiler-rt/lib/sanitizer_common/sanitizer_procmaps_mac.cpp
index 5ff8d1832556f..3b67382db2add 100644
--- a/compiler-rt/lib/sanitizer_common/sanitizer_procmaps_mac.cpp
+++ b/compiler-rt/lib/sanitizer_common/sanitizer_procmaps_mac.cpp
@@ -334,9 +334,22 @@ static const load_command *NextCommand(const load_command *lc) {
   return (const load_command *)((const char *)lc + lc->cmdsize);
 }
 
-static void FindUUID(const load_command *first_lc, u8 *uuid_output) {
-  for (const load_command *lc = first_lc; lc->cmd != 0; lc = NextCommand(lc)) {
-    if (lc->cmd != LC_UUID) continue;
+#ifdef MH_MAGIC_64
+static const size_t header_size = sizeof(mach_header_64);
+#else
+static const size_t header_size = sizeof(mach_header);
+#endif
+
+static void FindUUID(const load_command *first_lc, const mach_header *hdr,
+                     u8 *uuid_output) {
+  unsigned short curcmd = 0;
+  for (const load_command *lc = first_lc; curcmd < hdr->ncmds;
+       curcmd++, lc = NextCommand(lc)) {
+
+    CHECK_LT((const char*)lc, (const char*)hdr + header_size + hdr->sizeofcmds);
+
+    if (lc->cmd != LC_UUID)
+      continue;
 
     const uuid_command *uuid_lc = (const uuid_command *)lc;
     const uint8_t *uuid = &uuid_lc->uuid[0];
@@ -345,9 +358,16 @@ static void FindUUID(const load_command *first_lc, u8 *uuid_output) {
   }
 }
 
-static bool IsModuleInstrumented(const load_command *first_lc) {
-  for (const load_command *lc = first_lc; lc->cmd != 0; lc = NextCommand(lc)) {
-    if (lc->cmd != LC_LOAD_DYLIB) continue;
+static bool IsModuleInstrumented(const load_command *first_lc,
+                                 const mach_header *hdr) {
+  unsigned short curcmd = 0;
+  for (const load_command *lc = first_lc; curcmd < hdr->ncmds;
+       curcmd++, lc = NextCommand(lc)) {
+
+    CHECK_LT((const char*)lc, (const char*)hdr + header_size + hdr->sizeofcmds);
+
+    if (lc->cmd != LC_LOAD_DYLIB)
+      continue;
 
     const dylib_command *dylib_lc = (const dylib_command *)lc;
     uint32_t dylib_name_offset = dylib_lc->dylib.name.offset;
@@ -394,9 +414,10 @@ bool MemoryMappingLayout::Next(MemoryMappedSegment *segment) {
         }
       }
       FindUUID((const load_command *)data_.current_load_cmd_addr,
-               data_.current_uuid);
+               hdr, data_.current_uuid);
       data_.current_instrumented = IsModuleInstrumented(
-          (const load_command *)data_.current_load_cmd_addr);
+          (const load_command *)data_.current_load_cmd_addr,
+          hdr);
     }
 
     while (data_.current_load_cmd_count > 0) {
diff --git a/compiler-rt/lib/sanitizer_common/tests/sanitizer_procmaps_mac_test.cpp b/compiler-rt/lib/sanitizer_common/tests/sanitizer_procmaps_mac_test.cpp
index f622bba246309..b8f887a2fdd99 100644
--- a/compiler-rt/lib/sanitizer_common/tests/sanitizer_procmaps_mac_test.cpp
+++ b/compiler-rt/lib/sanitizer_common/tests/sanitizer_procmaps_mac_test.cpp
@@ -37,10 +37,12 @@ class MemoryMappingLayoutMock final : public MemoryMappingLayout {
     .uuid = {}
   };
 
-  static constexpr char dylib_name[] = "libclang_rt.\0\0\0"; // 8 bytes aligned, padded with zeros per loader.h
+  static constexpr char libclang_rt_dylib_name[]    = "libclang_rt.\0\0\0"; // 8 bytes aligned, padded with zeros per loader.h
+  static constexpr char uninstrumented_dylib_name[] = "uninst___rt.\0\0\0"; // 8 bytes aligned, padded with zeros per loader.h
+
   static constexpr dylib_command mock_dylib_command = {
     .cmd = LC_LOAD_DYLIB,
-    .cmdsize = sizeof(dylib_command) + sizeof(dylib_name),
+    .cmdsize = sizeof(dylib_command) + sizeof(libclang_rt_dylib_name),
     .dylib = {
       .name = {
         .offset = sizeof(dylib_command)
@@ -59,7 +61,7 @@ class MemoryMappingLayoutMock final : public MemoryMappingLayout {
   std::vector<unsigned char> mock_header;
 
 public:
-  MemoryMappingLayoutMock(): MemoryMappingLayout(false) {
+  MemoryMappingLayoutMock(bool instrumented): MemoryMappingLayout(false) {
     EXPECT_EQ(mock_uuid_command.cmdsize % 8, 0u);
     EXPECT_EQ(mock_dylib_command.cmdsize % 8, 0u);
 
@@ -89,6 +91,7 @@ class MemoryMappingLayoutMock final : public MemoryMappingLayout {
     copy((unsigned char *)&mock_dylib_command,
       ((unsigned char *)&mock_dylib_command) + sizeof(dylib_command), // as mock_dylib_command.cmdsize contains the following string
       back_inserter(mock_header));
+    const char (&dylib_name)[16] = instrumented ? libclang_rt_dylib_name : uninstrumented_dylib_name;
     copy((unsigned char *)dylib_name,
       ((unsigned char *)dylib_name) + sizeof(dylib_name),
       back_inserter(mock_header));
@@ -120,8 +123,20 @@ class MemoryMappingLayoutMock final : public MemoryMappingLayout {
   }
 };
 
-TEST(MemoryMappingLayout, Next) {
-  __sanitizer::MemoryMappingLayoutMock memory_mapping;
+TEST(MemoryMappingLayout, NextInstrumented) {
+  __sanitizer::MemoryMappingLayoutMock memory_mapping(true);
+  __sanitizer::MemoryMappedSegment segment;
+  size_t size = memory_mapping.SizeOfLoadCommands();
+  while (memory_mapping.Next(&segment)) {
+    size_t offset = memory_mapping.CurrentLoadCommandOffset();
+    EXPECT_LE(offset, size);
+  }
+  size_t final_offset = memory_mapping.CurrentLoadCommandOffset();
+  EXPECT_EQ(final_offset, size); // All commands processed, no more, no less
+}
+
+TEST(MemoryMappingLayout, NextUnInstrumented) {
+  __sanitizer::MemoryMappingLayoutMock memory_mapping(true);
   __sanitizer::MemoryMappedSegment segment;
   size_t size = memory_mapping.SizeOfLoadCommands();
   while (memory_mapping.Next(&segment)) {

[yln]

Unable to give real feedback, but your explanation and the change look sensible to me.
2 small nits.

LGTM, thanks Mariusz!