Bugfix: more selective criteria for resources to delete when injecting (

#200) * more selective criteria for resources to delete when injecting * update changelog * add test for segment injector delete behavior, and include patched resources in components' "modified" set * update changelog * add some clarifications to address PR comments --------- Co-authored-by: edward <[email protected]>
redballoonsecurity · Feb 3, 2023 · b602045 · b602045
1 parent c2227f6
commit b602045
Show file tree

Hide file tree

Showing 7 changed files with 108 additions and 11 deletions.
diff --git a/ofrak_core/CHANGELOG.md b/ofrak_core/CHANGELOG.md
@@ -7,6 +7,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/)
 ### Fixed
 - Fix bug in data service that can cause mangled internal state [#197](https://github.com/redballoonsecurity/ofrak/pull/197)
 - Fix long-broken `OFRAK.set_id_service` [#198](https://github.com/redballoonsecurity/ofrak/pull/198)
+- Fix bug in `SegmentInjectorModifier` that resulted in deleting more resources than necessary [#200](https://github.com/redballoonsecurity/ofrak/pull/200)
 
 ### Added
 - Replace unofficial p7zip with official 7zip package. 
@@ -15,6 +16,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/)
 
 ### Changed
 - GUI is much faster, especially for resources with hundreds of thousands of children [#191](https://github.com/redballoonsecurity/ofrak/pull/191)
+- Resources whose data gets modified are now listed in the `resource_modified` field of component results [#200](https://github.com/redballoonsecurity/ofrak/pull/200)
 
 ## [2.1.1](https://github.com/redballoonsecurity/ofrak/compare/ofrak-v2.1.0...ofrak-v2.1.1) - 2023-01-25
 ### Fixed

diff --git a/ofrak_core/ofrak/component/abstract.py b/ofrak_core/ofrak/component/abstract.py
@@ -149,6 +149,17 @@ async def run(
                     f"the resource context"
                 )
 
+        # Include resources modified by data patches in `modified_resource_ids`
+        data_ids_to_models = await dependency_handler.map_data_ids_to_resources(
+            patch_result.data_id for patch_result in patch_results
+        )
+        for m in data_ids_to_models.values():
+            modified_resource_ids.add(m.id)
+
+        # Exclude deleted resources from `modified_resource_ids`
+        # (deleting and modifying are handled separately)
+        modified_resource_ids.difference_update(component_context.resources_deleted)
+
         # Save modified resources
         await self._save_resources(
             job_id,

diff --git a/ofrak_core/ofrak/core/patch_maker/modifiers.py b/ofrak_core/ofrak/core/patch_maker/modifiers.py
@@ -209,15 +209,22 @@ async def modify(self, resource: Resource, config: SegmentInjectorModifierConfig
 
             injection_tasks.append((region.resource, BinaryInjectorModifierConfig(patches)))
 
-        all_decendants = await resource.get_descendants()
         for injected_resource, injection_config in injection_tasks:
             result = await injected_resource.run(BinaryInjectorModifier, injection_config)
-            parents_to_delete = list(
-                filter(lambda r: r.get_id() in result.resources_modified, all_decendants)
+            # 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 = []
-            for parent in parents_to_delete:
-                to_delete.extend(list(await parent.get_descendants()))
+            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))
 
 

diff --git a/ofrak_core/ofrak/resource.py b/ofrak_core/ofrak/resource.py
@@ -16,6 +16,7 @@
     Awaitable,
     Sequence,
     Callable,
+    Set,
 )
 
 from ofrak.component.interface import ComponentInterface
@@ -292,7 +293,8 @@ async def _fetch(self, resource: MutableResourceModel):
         try:
             fetched_resource = await self._resource_service.get_by_id(resource.id)
         except NotFoundError:
-            del self._resource_context.resource_models[resource.id]
+            if resource.id in self._component_context.modification_trackers:
+                del self._resource_context.resource_models[resource.id]
             return
 
         resource.reset(fetched_resource)
@@ -305,10 +307,12 @@ async def _fetch_resources(self, resource_ids: Iterable[bytes]):
                 tasks.append(self._fetch(context_resource))
         await asyncio.gather(*tasks)
 
-    async def _update_views(self, modified: Iterable[bytes], deleted: Iterable[bytes]):
+    async def _update_views(self, modified: Set[bytes], deleted: Set[bytes]):
         for resource_id in modified:
             views_in_context = self._resource_view_context.views_by_resource[resource_id]
             for view in views_in_context.values():
+                if resource_id not in self._resource_context.resource_models:
+                    await self._fetch(view.resource.get_model())  # type: ignore
                 updated_model = self._resource_context.resource_models[resource_id]
                 fresh_view = view.create(updated_model)
                 for field in dataclasses.fields(fresh_view):
@@ -344,6 +348,9 @@ async def run(
             ),
             job_context,
         )
+        for deleted_id in component_result.resources_deleted:
+            if deleted_id in self._component_context.modification_trackers:
+                del self._component_context.modification_trackers[deleted_id]
         await self._fetch_resources(component_result.resources_modified)
         await self._update_views(
             component_result.resources_modified, component_result.resources_deleted
@@ -386,6 +393,9 @@ async def auto_run(
                 all_packers=all_packers,
             )
         )
+        for deleted_id in components_result.resources_deleted:
+            if deleted_id in self._component_context.modification_trackers:
+                del self._component_context.modification_trackers[deleted_id]
         await self._fetch_resources(components_result.resources_modified)
         await self._update_views(
             components_result.resources_modified, components_result.resources_deleted

diff --git a/ofrak_core/ofrak/service/dependency_handler.py b/ofrak_core/ofrak/service/dependency_handler.py
@@ -1,3 +1,4 @@
+import functools
 import logging
 from typing import Set, List, Iterable, Dict, cast
 
@@ -33,7 +34,8 @@ def __init__(
         self._component_context = component_context
         self._resource_context = resource_context
 
-    async def _map_data_ids_to_resources(
+    @functools.lru_cache(None)
+    async def map_data_ids_to_resources(
         self, data_ids: Iterable[bytes]
     ) -> Dict[bytes, MutableResourceModel]:
         resources_by_data_id = dict()
@@ -59,7 +61,7 @@ async def _map_data_ids_to_resources(
 
     async def handle_post_patch_dependencies(self, patch_results: List[DataPatchesResult]):
         # Create look up maps for resources and dependencies
-        resources_by_data_id = await self._map_data_ids_to_resources(
+        resources_by_data_id = await self.map_data_ids_to_resources(
             patch_result.data_id for patch_result in patch_results
         )
 

diff --git a/ofrak_core/ofrak/service/resource_service.py b/ofrak_core/ofrak/service/resource_service.py
@@ -966,6 +966,7 @@ async def delete_resource(self, resource_id: bytes):
             former_parent_resource_node.remove_child(resource_node)
 
         deleted_models = self._delete_resource_helper(resource_node)
+        LOGGER.debug(f"Deleted {resource_id.hex()}")
         return deleted_models
 
     async def delete_resources(self, resource_ids: Iterable[bytes]):
@@ -981,6 +982,8 @@ async def delete_resources(self, resource_ids: Iterable[bytes]):
                 former_parent_resource_node.remove_child(resource_node)
 
             deleted_models.extend(self._delete_resource_helper(resource_node))
+        if resource_ids:
+            LOGGER.debug(f"Deleted {', '.join(resource_id.hex() for resource_id in resource_ids)}")
         return deleted_models
 
     def _delete_resource_helper(self, _resource_node: ResourceNode):

diff --git a/ofrak_core/test_ofrak/components/test_patch_maker_component.py b/ofrak_core/test_ofrak/components/test_patch_maker_component.py
@@ -6,14 +6,16 @@
 import re
 import subprocess
 
+from ofrak.core import MemoryRegion
 from ofrak_patch_maker.toolchain.gnu_x64 import GNU_X86_64_LINUX_EABI_10_3_0_Toolchain
 
 from ofrak_patch_maker.toolchain.gnu_arm import GNU_ARM_NONE_EABI_10_2_1_Toolchain
 from ofrak_patch_maker.toolchain.llvm_12 import LLVM_12_0_1_Toolchain
 from ofrak_patch_maker.toolchain.abstract import Toolchain
 
-from ofrak import OFRAKContext
+from ofrak import OFRAKContext, ResourceFilter, ResourceAttributeRangeFilter
 from ofrak.core.architecture import ProgramAttributes
+from ofrak_type import MemoryPermissions
 from ofrak_type.architecture import (
     InstructionSet,
     SubInstructionSet,
@@ -27,11 +29,14 @@
 from ofrak.core.patch_maker.modifiers import (
     FunctionReplacementModifierConfig,
     FunctionReplacementModifier,
+    SegmentInjectorModifierConfig,
+    SegmentInjectorModifier,
 )
 from ofrak_patch_maker.toolchain.model import (
     CompilerOptimizationLevel,
     BinFileType,
     ToolchainConfig,
+    Segment,
 )
 from ofrak_patch_maker.toolchain.utils import get_repository_config
 from ofrak_type.bit_width import BitWidth
@@ -236,3 +241,60 @@ async def test_function_replacement_modifier(ofrak_context: OFRAKContext, config
     expected_objdump_output_str = "\n".join(config.expected_objdump_output)
 
     assert normalize_assembly(expected_objdump_output_str) in normalize_assembly(readobj_output)
+
+
+async def test_segment_injector_deletes_patched_descendants(ofrak_context: OFRAKContext):
+    # unpack_recursively an ELF
+    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
+
+    function_cb = ComplexBlock(
+        virtual_address=main_start,
+        size=main_end - main_start,
+        name="main",
+    )
+
+    target_program = await root_resource.view_as(Program)
+
+    function_cb_parent_code_region = await target_program.get_code_region_for_vaddr(main_start)
+
+    function_cb.resource = await function_cb_parent_code_region.create_child_region(function_cb)
+
+    # Create a dummy basic block in the complex block, so its `get_mode` method won't fail.
+    dummy_bb = BasicBlock(0x8068, 8, InstructionSetMode.NONE, False, None)
+    await function_cb.create_child_region(dummy_bb)
+
+    # get IDs of resources in a vaddr range
+    expected_deleted_ids = set()
+    for r in await root_resource.get_descendants(
+        r_filter=ResourceFilter(
+            attribute_filters=(
+                ResourceAttributeRangeFilter(MemoryRegion.VirtualAddress, min=main_start),
+                ResourceAttributeRangeFilter(MemoryRegion.EndVaddr, max=main_end + 1),
+            )
+        )
+    ):
+        expected_deleted_ids.add(r.get_id())
+        for r in await r.get_descendants():
+            expected_deleted_ids.add(r.get_id())
+
+    assert len(expected_deleted_ids) > 0
+
+    # create a SegmentInjectorModifierConfig
+    cfg = SegmentInjectorModifierConfig(
+        (
+            (
+                Segment(".text", main_start, 0, False, main_end - main_start, MemoryPermissions.RX),
+                b"\x00" * (main_end - main_start),
+            ),
+        )
+    )
+
+    # run SegmentInjectorModifier
+    results = await root_resource.run(SegmentInjectorModifier, cfg)
+
+    # check that resources have been deleted
+    assert results.resources_deleted == expected_deleted_ids