tmp fix for p_align issue

src/patcherex2/components/allocation_managers/allocation_manager.py

@@ -154,7 +154,7 @@ def _find_in_mapped_blocks(
     def _create_new_mapped_block(
         self, size: int, flag=MemoryFlag.RWX, align=0x1
     ) -> bool:
-        # map 0x1000 bytes # TODO: currently we won't use available file/mem blocks, instead we create new one at the end of the file
+        # TODO: currently we won't use available file/mem blocks, instead we create new one at the end of the file
         file_addr = None
         mem_addr = None
         for block in self.blocks[FileBlock]:
@@ -163,8 +163,19 @@ def _create_new_mapped_block(
                 block.addr += 0x2000
         for block in self.blocks[MemoryBlock]:
             if block.size == -1:
-                # mem_addr % 0x1000 should equal to file_addr % 0x1000 TODO
-                mem_addr = block.addr + (file_addr % 0x1000)
+                # NOTE: mem_addr % p_align should equal to file_addr % p_align
+                # Check `man elf` and search for `p_align` for more information
+                # FIXME: shouldn't do any assumption on component type, reimpl in a better way
+                # FIXME: even worse, importing ELF will cause circular import
+                # TODO: consider merge allocation_manager and binfmt_tool into one component
+                if self.p.binfmt_tool.__class__.__name__ == "ELF":
+                    max_seg_align = max(
+                        [segment["p_align"] for segment in self.p.binfmt_tool._segments]
+                        + [0]
+                    )
+                    mem_addr = block.addr + (file_addr % max_seg_align)
+                else:
+                    mem_addr = block.addr + (file_addr % 0x1000)
                 block.addr = mem_addr + 0x2000
         if file_addr and mem_addr:
             self.add_block(

src/patcherex2/components/binfmt_tools/elf.py

@@ -185,7 +185,9 @@ def _add_end_of_file_block(self) -> None:
         self.p.allocation_manager.add_block(block)
 
         addr = load_segments[-1]["p_vaddr"] + load_segments[-1]["p_memsz"]
-        addr = (addr + 0xFFF) & ~0xFFF  # round up to 0x1000
+        # TODO: should we use the max_align of the segments?
+        max_align = max([segment["p_align"] for segment in self._segments] + [0])
+        addr = (addr + (max_align - 1)) & ~(max_align - 1)  # round up to max_align
         block = MemoryBlock(addr, -1)
         self.p.allocation_manager.add_block(block)
 
@@ -342,6 +344,7 @@ def finalize(self) -> None:
             self._segments.append(phdr_load_segment)
 
             # magic
+            # TODO: should we use the max_align of the segments?
             load_segments_rounded = []
             first_load_segment = None
             for segment in self._segments:
@@ -350,24 +353,24 @@ def finalize(self) -> None:
                         first_load_segment = segment
                     load_segments_rounded.append(
                         (
-                            # start of the segment, round down to multiple of 0x1000
+                            # start of the segment, round down to multiple of max_align
                             (segment["p_vaddr"] - first_load_segment["p_vaddr"])
                             - (
                                 (segment["p_vaddr"] - first_load_segment["p_vaddr"])
-                                % 0x1000
+                                % max_align
                             ),
-                            # end of the segment, round up to multiple of 0x1000
+                            # end of the segment, round up to multiple of max_align
                             int(
                                 (
                                     segment["p_vaddr"]
                                     + segment["p_memsz"]
                                     - first_load_segment["p_vaddr"]
-                                    + 0x1000
+                                    + max_align
                                     - 1
                                 )
-                                / 0x1000
+                                / max_align
                             )
-                            * 0x1000,
+                            * max_align,
                         )
                     )
             load_segments_rounded = sorted(load_segments_rounded, key=lambda x: x[0])
@@ -402,9 +405,10 @@ def finalize(self) -> None:
             for prev_seg, next_seg in zip(
                 load_segments_rounded[:-1], load_segments_rounded[1:]
             ):
+                # TODO: should we use the max_align of the segments?
                 potential_base = (
-                    max(prev_seg[1], self.p.binfmt_tool.file_size) + 0xFFF
-                ) & ~0xFFF  # round up to 0x1000
+                    max(prev_seg[1], self.p.binfmt_tool.file_size) + (max_align - 1)
+                ) & ~(max_align - 1)  # round up to max_align
                 if next_seg[0] - potential_base > self._elf.header["e_phentsize"] * len(
                     self._segments
                 ):  # if there is space between segments, put phdr here
@@ -417,10 +421,11 @@ def finalize(self) -> None:
             # this is to workaround a weird issue in the dynamic linker of glibc
             # we want to make sure p_vaddr (phdr_start) == p_offset (len(ncontent))
             if phdr_start <= self.p.binfmt_tool.file_size:
+                # TODO: should we use the max_align of the segments?
                 # p_vaddr <= p_offset: pad the file (p_offset) to page size, and let p_vaddr = p_offset
                 self.p.binfmt_tool.file_size = (
-                    self.p.binfmt_tool.file_size + 0xFFF
-                ) & ~0xFFF  # round up to 0x1000
+                    self.p.binfmt_tool.file_size + (max_align - 1)
+                ) & ~(max_align - 1)  # round up to max_align
                 phdr_start = self.p.binfmt_tool.file_size
 
             # update phdr segment and its corresponding load segment