chore: shorten change

RaisinTen · RaisinTen · commit 2b7cbdd9fdca · 2023-03-02T21:29:05.000+05:30
Signed-off-by: Darshan Sen &lt;raisinten@gmail.com&gt;
diff --git a/postject-api.h b/postject-api.h
@@ -44,55 +44,11 @@ static inline bool postject_has_resource() {
 }
 
 #if defined(__linux__)
-struct postject__dl_iterate_phdr_data {
-  void* data;
-  size_t size;
-  const char* elf_section_name;
-  size_t elf_section_name_size;
-};
-
 static int postject__dl_iterate_phdr_callback(struct dl_phdr_info* info,
                                               size_t size,
-                                              void* opaque) {
-  struct postject__dl_iterate_phdr_data* data =
-      (struct postject__dl_iterate_phdr_data*)opaque;
-
-  // iterate program headers
-  for (unsigned i = 0; i < info->dlpi_phnum; i++) {
-    // skip everything but notes
-    if (info->dlpi_phdr[i].p_type != PT_NOTE)
-      continue;
-
-    // note segment starts at base address + segment virtual address
-    uintptr_t pos = (info->dlpi_addr + info->dlpi_phdr[i].p_vaddr);
-    uintptr_t end = (pos + info->dlpi_phdr[i].p_memsz);
-
-    // iterate through segment until we reach the end
-    while (pos < end) {
-      if (pos + sizeof(ElfW(Nhdr)) > end) {
-        break;  // invalid
-      }
-
-      ElfW(Nhdr)* note = (ElfW(Nhdr)*)(uintptr_t)pos;
-      if (note->n_namesz != 0 && note->n_descsz != 0 &&
-          strncmp((char*)(pos + sizeof(ElfW(Nhdr))), data->elf_section_name,
-                  data->elf_section_name_size) == 0) {
-        data->size = note->n_descsz;
-        // advance past note header and aligned name
-        // to get to description data
-        data->data = (void*)((uintptr_t)note + sizeof(ElfW(Nhdr)) +
-                             roundup(note->n_namesz, 4));
-        // found the note, so terminate the search by returning 1
-        return 1;
-      }
-
-      pos += (sizeof(ElfW(Nhdr)) + roundup(note->n_namesz, 4) +
-              roundup(note->n_descsz, 4));
-    }
-  }
-
-  // wasn't able to find the note in the main executable program headers, so
-  // terminate the search
+                                              void* data) {
+  // Snag the dl_phdr_info struct for the main program, then stop iterating
+  *((struct dl_phdr_info*)data) = *info;
   return 1;
 }
 #endif
@@ -167,14 +123,49 @@ static const void* postject_find_resource(
     name = options->elf_section_name;
   }
 
-  struct postject__dl_iterate_phdr_data data;
-  data.data = NULL;
-  data.size = 0;
-  data.elf_section_name = name;
-  data.elf_section_name_size = sizeof(name);
-  dl_iterate_phdr(postject__dl_iterate_phdr_callback, &data);
-  *size = data.size;
-  return data.data;
+  struct dl_phdr_info main_program_info;
+  dl_iterate_phdr(postject__dl_iterate_phdr_callback, &main_program_info);
+
+  uintptr_t p = (uintptr_t)main_program_info.dlpi_phdr;
+  size_t n = main_program_info.dlpi_phnum;
+  uintptr_t base_addr = main_program_info.dlpi_addr;
+
+  // iterate program header
+  for (; n > 0; n--, p += sizeof(ElfW(Phdr))) {
+    ElfW(Phdr)* phdr = (ElfW(Phdr)*)p;
+
+    // skip everything but notes
+    if (phdr->p_type != PT_NOTE) {
+      continue;
+    }
+
+    // note segment starts at base address + segment virtual address
+    uintptr_t pos = (base_addr + phdr->p_vaddr);
+    uintptr_t end = (pos + phdr->p_memsz);
+
+    // iterate through segment until we reach the end
+    while (pos < end) {
+      if (pos + sizeof(ElfW(Nhdr)) > end) {
+        break;  // invalid
+      }
+
+      ElfW(Nhdr)* note = (ElfW(Nhdr)*)(uintptr_t)pos;
+      if (note->n_namesz != 0 && note->n_descsz != 0 &&
+          strncmp((char*)(pos + sizeof(ElfW(Nhdr))), (char*)name,
+                  sizeof(name)) == 0) {
+        *size = note->n_descsz;
+        // advance past note header and aligned name
+        // to get to description data
+        return (void*)((uintptr_t)note + sizeof(ElfW(Nhdr)) +
+                       roundup(note->n_namesz, 4));
+      }
+
+      pos += (sizeof(ElfW(Nhdr)) + roundup(note->n_namesz, 4) +
+              roundup(note->n_descsz, 4));
+    }
+  }
+  return NULL;
+
 #elif defined(_WIN32)
   void* ptr = NULL;
   char* resource_name = NULL;
