From b60204515d87576a209f56ad9cf160d6ffd22541 Mon Sep 17 00:00:00 2001 From: Edward Larson Date: Fri, 3 Feb 2023 15:55:07 -0500 Subject: [PATCH] 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 --- ofrak_core/CHANGELOG.md | 2 + ofrak_core/ofrak/component/abstract.py | 11 ++++ .../ofrak/core/patch_maker/modifiers.py | 19 ++++-- ofrak_core/ofrak/resource.py | 14 +++- .../ofrak/service/dependency_handler.py | 6 +- ofrak_core/ofrak/service/resource_service.py | 3 + .../components/test_patch_maker_component.py | 64 ++++++++++++++++++- 7 files changed, 108 insertions(+), 11 deletions(-) diff --git a/ofrak_core/CHANGELOG.md b/ofrak_core/CHANGELOG.md index 63616d8f3..4fc9bc552 100644 --- 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 index 90700b12a..19a1c0c8b 100644 --- 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 index 6545f5300..7e347473d 100644 --- 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 index 30c2b5fce..a25d0acbf 100644 --- 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 index edaf7a9b2..102655872 100644 --- 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 index fb36fae67..a7b405520 100644 --- 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 index c51528122..ef26668ac 100644 --- 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