FastSegmentInjectorModifier

disassemblers/ofrak_ghidra/tests/test_ghidra_program_analyzer.py

@@ -211,5 +211,5 @@ async def _make_dummy_program(resource: Resource, arch_info):
 
     await resource.run(
         SegmentInjectorModifier,
-        SegmentInjectorModifierConfig.from_fem(fem),
+        SegmentInjectorModifierConfig.from_fems([fem]),
     )

ofrak_core/CHANGELOG.md

@@ -6,6 +6,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/)
 ## [Unreleased](https://github.com/redballoonsecurity/ofrak/tree/master)
 
 ### Added
+- Add FastSegmentInjectorModifier for optimized bulk segment injection ([#685](https://github.com/redballoonsecurity/ofrak/pull/685))
 - Add `FlashGeometryHeuristicAnalyzer` that infers `FlashAttributes` (page/OOB geometry) for raw NAND dumps tagged as `FlashResource`, using YAFFS2, Linux MTD large-page OOB, small-page OOB density, and exact Hamming ECC verification as signals ([#737](https://github.com/redballoonsecurity/ofrak/pull/737))
 - Add `FlashFieldType.SPARE` and `FlashSpareAreaResource` so `FlashOobResourceUnpacker` can preserve OOB bytes without ECC/Checksum computation, and extend `FlashLogicalDataResourcePacker` to consume the sibling `FlashSpareAreaResource` so the original OOB layout is reconstructed verbatim during repacking ([#738](https://github.com/redballoonsecurity/ofrak/pull/738))
 - Add cramfs unpacker and packer, using `fsck.cramfs` and `mkfs.cramfs` from `util-linux` 2.38.1 ([#719](https://github.com/redballoonsecurity/ofrak/pull/719))

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

@@ -3,7 +3,7 @@
 import os
 import tempfile312 as tempfile
 from dataclasses import dataclass
-from typing import Dict, List, Optional, Tuple, Type, Union, cast
+from typing import Dict, Iterable, List, Optional, Tuple, Type, Union, cast
 
 from ofrak_patch_maker.model import PatchRegionConfig, FEM
 from ofrak_patch_maker.patch_maker import PatchMaker
@@ -22,6 +22,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__)
 
@@ -211,6 +212,18 @@ def from_fem(fem: FEM) -> "SegmentInjectorModifierConfig":
             extracted_segments.append((segment, segment_data))
         return SegmentInjectorModifierConfig(tuple(extracted_segments))
 
+    @staticmethod
+    def from_fems(fems: Iterable[FEM]) -> "SegmentInjectorModifierConfig":
+        """
+        Automatically build a config from a list of FEMs by extracting each segment's bytes and metadata.
+        """
+        extracted_segments: List[Tuple[Segment, bytes]] = []
+        for fem in fems:
+            extracted_segments.extend(
+                list(SegmentInjectorModifierConfig.from_fem(fem).segments_and_data)
+            )
+        return SegmentInjectorModifierConfig(tuple(extracted_segments))
+
 
 class SegmentInjectorModifier(Modifier[SegmentInjectorModifierConfig]):
     """
@@ -304,6 +317,88 @@ async def modify(self, resource: Resource, config: SegmentInjectorModifierConfig
             await asyncio.gather(*(r.delete() for r in to_delete))
 
 
+class FastSegmentInjectorModifier(Modifier[SegmentInjectorModifierConfig]):
+    """
+    Optimized segment injection modifier that applies multiple segment patches in a single bulk
+    operation. Unlike SegmentInjectorModifier which runs BinaryInjectorModifier separately for each
+    segment, this modifier collects all segment patches and applies them to the root resource data
+    in one pass, improving performance for multi-segment injections.
+
+    **WARNING: This modifier does NOT delete stale descendant resources.** After patching, the
+    resource tree may contain outdated descendants that no longer match the modified data. This is a
+    critical difference from SegmentInjectorModifier, which automatically cleans up invalidated
+    descendants. The resource tree may be in an inconsistent state after using this modifier.
+
+    This modifier skips .bss, .rela, and .got segments, and applies all valid segment patches
+    directly to the program resource without invoking individual injection modifiers. Use when
+    applying complex patches with many segments.
+    """
+
+    targets = (Program,)
+
+    async def modify(self, resource: Resource, config: SegmentInjectorModifierConfig) -> None:
+        segments_to_skip = [
+            ".bss",
+            ".rela",
+            ".got",
+        ]
+        sorted_regions = list(
+            await resource.get_descendants_as_view(
+                MemoryRegion,
+                r_filter=ResourceFilter(include_self=True, tags=(MemoryRegion,)),
+                r_sort=ResourceSort(
+                    attribute=MemoryRegion.Size,
+                    direction=ResourceSortDirection.DESCENDANT,
+                ),
+            )
+        )
+
+        patches_to_resource: List[Tuple[int, bytes]] = list()
+
+        for segment, segment_data in config.segments_and_data:
+            # Skip bss, rela, got, etc
+            if segment.length == 0 or segment.vm_address == 0:
+                continue
+            if any([segment.segment_name.startswith(prefix) for prefix in segments_to_skip]):
+                continue
+
+            LOGGER.info(
+                f'    Injecting segment "{segment.segment_name}": {segment.length} '
+                f"bytes @ {hex(segment.vm_address)}",
+            )
+
+            # Find and validate target region
+            target_region = MemoryRegion.get_mem_region_with_vaddr_from_sorted(
+                segment.vm_address, sorted_regions
+            )
+            if target_region is None:
+                raise ValueError(
+                    f"Cannot inject patch because the memory region at vaddr "
+                    f"{hex(segment.vm_address)} is None"
+                )
+            if target_region.resource.get_data_id() is None:
+                raise ValueError(
+                    f"Cannot inject patch because the memory region at vaddr "
+                    f"{hex(segment.vm_address)} is dataless"
+                )
+
+            # Calculate the offset in the root resource
+            range_in_root = await target_region.resource.get_data_range_within_root()
+            offset = range_in_root.start + segment.vm_address - target_region.virtual_address
+            patches_to_resource.append((offset, segment_data))
+
+        # Apply segment patches in-memory
+        resource_data = bytearray(await resource.get_data())
+        for offset, segment_data in patches_to_resource:
+            resource_data[offset : offset + len(segment_data)] = segment_data
+
+        # Patch OFRAK resource
+        resource.queue_patch(
+            Range.from_size(0, await resource.get_data_length()),
+            bytes(resource_data),
+        )
+
+
 @dataclass
 class FunctionReplacementModifierConfig(ComponentConfig):
     """

ofrak_core/tests/components/test_patch_maker_component.py

@@ -39,8 +39,10 @@
     FunctionReplacementModifier,
     SegmentInjectorModifierConfig,
     SegmentInjectorModifier,
+    FastSegmentInjectorModifier,
     SourceBundle,
 )
+from ofrak_patch_maker.model import FEM, LinkedExecutable
 from ofrak_patch_maker.toolchain.model import (
     CompilerOptimizationLevel,
     BinFileType,
@@ -330,3 +332,77 @@ async def test_segment_injector_deletes_patched_descendants(ofrak_context: OFRAK
 
     # check that resources have been deleted
     assert results.resources_deleted == expected_deleted_ids
+
+
+async def test_fast_segment_injector_applies_patches(ofrak_context: OFRAKContext):
+    """
+    Tests that FastSegmentInjectorModifier applies segment patches correctly (REQ6.1).
+
+    Unlike SegmentInjectorModifier, this does NOT delete patched descendants.
+    """
+    root_resource = await ofrak_context.create_root_resource_from_file(ARM32_PROGRAM_PATH)
+    await root_resource.unpack_recursively()
+
+    main_start = 0x8068
+    main_end = main_start + 40
+
+    target_program = await root_resource.view_as(Program)
+    await target_program.get_code_region_for_vaddr(main_start)
+
+    cfg = SegmentInjectorModifierConfig(
+        (
+            (
+                Segment(".text", main_start, 0, False, main_end - main_start, MemoryPermissions.RX),
+                b"\x00" * (main_end - main_start),
+            ),
+        )
+    )
+
+    await root_resource.run(FastSegmentInjectorModifier, cfg)
+
+
+def test_segment_injector_modifier_config_from_fems(tmp_path):
+    """
+    Tests that SegmentInjectorModifierConfig.from_fems correctly merges segment data extracted
+    from multiple FEMs into a single config (REQ6.1).
+    """
+    exe_a_path = tmp_path / "exec_a"
+    exe_a_data = b"AAAABBBB"
+    exe_a_path.write_bytes(exe_a_data)
+
+    exe_b_path = tmp_path / "exec_b"
+    exe_b_data = b"CCCCDDDDEEEE"
+    exe_b_path.write_bytes(exe_b_data)
+
+    segment_a = Segment(".text", 0x1000, 0, False, 4, MemoryPermissions.RX)
+    segment_b_text = Segment(".text", 0x2000, 0, False, 4, MemoryPermissions.RX)
+    segment_b_data = Segment(".data", 0x3000, 4, False, 8, MemoryPermissions.RW)
+
+    fem_a = FEM(
+        name="patch_a",
+        executable=LinkedExecutable(
+            path=str(exe_a_path),
+            file_format=BinFileType.ELF,
+            segments=(segment_a,),
+            symbols={},
+            relocatable=False,
+        ),
+    )
+    fem_b = FEM(
+        name="patch_b",
+        executable=LinkedExecutable(
+            path=str(exe_b_path),
+            file_format=BinFileType.ELF,
+            segments=(segment_b_text, segment_b_data),
+            symbols={},
+            relocatable=False,
+        ),
+    )
+
+    cfg = SegmentInjectorModifierConfig.from_fems([fem_a, fem_b])
+
+    assert cfg.segments_and_data == (
+        (segment_a, b"AAAA"),
+        (segment_b_text, b"CCCC"),
+        (segment_b_data, b"DDDDEEEE"),
+    )

ofrak_core/version.py

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