Draft: Speed up SegmentInjectorModifier

ofrak_core/CHANGELOG.md

@@ -19,6 +19,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/)
 - Add YAFFS2 filesystem support ([#739](https://github.com/redballoonsecurity/ofrak/pull/739))
 
 ### Changed
+- Optimize `SegmentInjectorModifier` for bulk segment injection ([#685](https://github.com/redballoonsecurity/ofrak/pull/685))
 - Resources tagged as `FlashResource` now have their `FlashAttributes` inferred automatically by the new `FlashGeometryHeuristicAnalyzer` and become unpackable via `FlashResourceUnpacker` without user-supplied geometry. Workflows that previously attached a hand-crafted `FlashAttributes` should verify the inferred geometry matches. If no standard geometry fits the image, the analyzer logs a warning and returns no attributes so other analyzers (e.g. `BinwalkAnalyzer`) can still run ([#737](https://github.com/redballoonsecurity/ofrak/pull/737))
 - Remove test dependencies that are already in the global `requirements-dev.txt` ([#695](https://github.com/redballoonsecurity/ofrak/pull/695))
 

ofrak_core/src/ofrak/core/patch_maker/modifiers.py

@@ -13,7 +13,6 @@
 from ofrak.component.modifier import Modifier
 from ofrak.core.architecture import ProgramAttributes
 from ofrak.core.complex_block import ComplexBlock
-from ofrak.core.injector import BinaryInjectorModifier, BinaryInjectorModifierConfig
 from ofrak.core.memory_region import MemoryRegion
 from ofrak.core.patch_maker.linkable_binary import LinkableBinary
 from ofrak.core.program import Program
@@ -22,6 +21,7 @@
 from ofrak.service.resource_service_i import ResourceFilter, ResourceSort, ResourceSortDirection
 from ofrak_type.memory_permissions import MemoryPermissions
 from ofrak_type.error import NotFoundError
+from ofrak_type.range import Range
 
 LOGGER = logging.getLogger(__file__)
 
@@ -240,16 +240,16 @@ async def modify(self, resource: Resource, config: SegmentInjectorModifierConfig
             )
         )
 
-        injection_tasks: List[Tuple[Resource, BinaryInjectorModifierConfig]] = []
+        patches: List[Tuple[int, bytes]] = []  # (root_offset, data)
 
         for segment, segment_data in config.segments_and_data:
             if segment.length == 0 or not segment.is_allocated:
                 continue
-            if segment.length > 0:
-                LOGGER.debug(
-                    f"    Segment {segment.segment_name} - {segment.length} "
-                    f"bytes @ {hex(segment.vm_address)}",
-                )
+
+            LOGGER.debug(
+                f"    Segment {segment.segment_name} - {segment.length} "
+                f"bytes @ {hex(segment.vm_address)}",
+            )
 
             if segment.segment_name.startswith(".rela"):
                 continue
@@ -273,8 +273,9 @@ async def modify(self, resource: Resource, config: SegmentInjectorModifierConfig
 
             region_mapped_to_data = region.resource.get_data_id() is not None
             if region_mapped_to_data:
-                patches = [(segment.vm_address, segment_data)]
-                injection_tasks.append((region.resource, BinaryInjectorModifierConfig(patches)))
+                range_in_root = await region.resource.get_data_range_within_root()
+                offset = range_in_root.start + segment.vm_address - region.virtual_address
+                patches.append((offset, segment_data))
             else:
                 if segment.is_bss:
                     # uninitialized section like .bss mapped to arbitrary memory range without corresponding
@@ -285,23 +286,35 @@ async def modify(self, resource: Resource, config: SegmentInjectorModifierConfig
                     f"{hex(segment.vm_address)} is not mapped to data"
                 )
 
-        for injected_resource, injection_config in injection_tasks:
-            result = await injected_resource.run(BinaryInjectorModifier, injection_config)
-            # The above can patch data of any of injected_resources' descendants or ancestors
-            # We don't want to delete injected_resources or its ancestors, so subtract them from the
-            # set of patched resources
-            patched_descendants = result.resources_modified.difference(
-                {
-                    r.get_id()
-                    for r in await injected_resource.get_ancestors(
-                        ResourceFilter(include_self=True)
-                    )
-                }
-            )
-            to_delete = [
-                r for r in await resource.get_descendants() if r.get_id() in patched_descendants
-            ]
-            await asyncio.gather(*(r.delete() for r in to_delete))
+        if not patches:
+            return
+
+        # Predict stale descendants from patch ranges before modifying
+        patch_ranges = [Range.from_size(offset, len(data)) for offset, data in patches]
+        all_descendants = list(await resource.get_descendants())
+
+        async def _get_data_range(r: Resource) -> Optional[Range]:
+            if r.get_data_id() is None:
+                return None
+            try:
+                return await r.get_data_range_within_root()
+            except Exception:
+                return None
+
+        desc_ranges = await asyncio.gather(*(_get_data_range(r) for r in all_descendants))
+        stale = [
+            r
+            for r, rng in zip(all_descendants, desc_ranges)
+            if rng is not None and any(rng.overlaps(pr) for pr in patch_ranges)
+        ]
+
+        # Apply all patches in a single bulk operation
+        resource_data = bytearray(await resource.get_data())
+        for offset, data in patches:
+            resource_data[offset : offset + len(data)] = data
+        resource.queue_patch(Range.from_size(0, await resource.get_data_length()), bytes(resource_data))
+
+        await asyncio.gather(*(r.delete() for r in stale))
 
 
 @dataclass

ofrak_core/version.py

@@ -1 +1 @@
-VERSION = "3.4.0rc12"
+VERSION = "3.4.0rc13"