From afb0a94219e9f689a999c15c55ec9f145b1bd6c8 Mon Sep 17 00:00:00 2001 From: Wyatt Date: Fri, 2 Feb 2024 17:38:44 -0500 Subject: [PATCH] Refactor FreeSpaceModifier, PartialFreeSpaceModifier to add stub - stub is optional bytes that displace the beginning of the free space - update tests to assert correct functionality of this --- ofrak_core/ofrak/core/free_space.py | 82 +++++++++++++------ .../test_ofrak/components/test_free_space.py | 61 ++++++++++++-- 2 files changed, 110 insertions(+), 33 deletions(-) diff --git a/ofrak_core/ofrak/core/free_space.py b/ofrak_core/ofrak/core/free_space.py index 44d01aabf..66a1b4922 100644 --- a/ofrak_core/ofrak/core/free_space.py +++ b/ofrak_core/ofrak/core/free_space.py @@ -450,14 +450,19 @@ async def _find_and_delete_overlapping_children(resource: Resource, freed_range: await overlapping_child.resource.save() -def _get_fill(freed_range: Range, fill: Optional[bytes]): - if not fill: - return b"\x00" * freed_range.length() - else: - diff_len = freed_range.length() - len(fill) - if diff_len < 0: - raise ValueError("config.fill value cannot be longer than the range to be freed.") - return fill + b"\x00" * diff_len +def _get_patch(freed_range: Range, stub: bytes, fill: bytes) -> bytes: + total_fill_length = freed_range.length() - len(stub) + remainder = total_fill_length % len(fill) + final = b"".join( + [ + stub, + fill * (total_fill_length // len(fill)), + fill[:remainder] + ] + ) + assert len(final) == freed_range.length() + return final + @dataclass @@ -465,19 +470,24 @@ class FreeSpaceModifierConfig(ComponentConfig): """ Configuration for modifier which marks some free space. - :var permissions: memory permissions to give the created free space. - :var fill: bytes to fill the free space with + :var permissions: Memory permissions to give the created free space. + :var stub: Bytes for a stub to be injected before the free space. The stub will not be marked as + [FreeSpace][ofrak.core.free_space.FreeSpace]. + :var fill: Pattern of bytes to fill the free space with. """ permissions: MemoryPermissions - fill: Optional[bytes] = None + stub: bytes = b"" + fill: bytes = b"\x00" class FreeSpaceModifier(Modifier[FreeSpaceModifierConfig]): """ Turn a [MemoryRegion][ofrak.core.memory_region.MemoryRegion] resource into allocatable free space by replacing its data with b'\x00' or optionally specified bytes. - [FreeSpace][ofrak.core.free_space.FreeSpace]. + + The modifier allows for an optional "stub", bytes to be injected at the beginning of the target resource. The stub + bytes are not marked as [FreeSpace][ofrak.core.free_space.FreeSpace]. """ targets = (MemoryRegion,) @@ -489,18 +499,34 @@ async def modify(self, resource: Resource, config: FreeSpaceModifierConfig): mem_region_view.virtual_address, mem_region_view.virtual_address + mem_region_view.size, ) - patch_data = _get_fill(freed_range, config.fill) + patch_data = _get_patch(freed_range, config.stub, config.fill) parent = await resource.get_parent() patch_offset = (await resource.get_data_range_within_parent()).start patch_range = freed_range.translate(patch_offset - freed_range.start) + # Grab tags, so they can be saved to the stub. + # At some point, it might be nice to save the attributes as well. + current_tags = resource.get_tags() + # One interesting side effect here is the Resource used to call this modifier no longer exists + # when this modifier returns. This can be confusing. Would an update work better in this case? await resource.delete() await resource.save() # Patch in the patch_data await parent.run(BinaryPatchModifier, BinaryPatchConfig(patch_offset, patch_data)) - free_offset = len(config.fill) if config.fill else 0 + free_offset = len(config.stub) + + if len(config.stub) > 0: + # Create the stub + await parent.create_child_from_view( + MemoryRegion( + mem_region_view.virtual_address + free_offset, + len(config.stub) + ), + data_range=Range.from_size(patch_range.start, len(config.stub)), + additional_tags=current_tags, + ) # Create the FreeSpace child await parent.create_child_from_view( @@ -509,7 +535,7 @@ async def modify(self, resource: Resource, config: FreeSpaceModifierConfig): mem_region_view.size - free_offset, config.permissions, ), - data_range=patch_range, + data_range=Range(patch_range.start + free_offset, patch_range.end) ) @@ -517,20 +543,26 @@ async def modify(self, resource: Resource, config: FreeSpaceModifierConfig): class PartialFreeSpaceModifierConfig(ComponentConfig): """ :var permissions: memory permissions to give the created free space. - :var range_to_remove: the ranges to consider as free space (remove) - :var fill: bytes to fill the free space with + :var range_to_remove: The ranges to consider as free space (remove). + :var stub: Bytes for a stub to be injected before the free space. If a stub is specified, then the FreeSpace created + will decrease in size. For example, with a stub of b"HA" and range_to_remove=Range(4,10), the final FreeSpace will + end up corresponding to Range(6,10). + :var fill: Pattern of bytes to fill the free space with. """ permissions: MemoryPermissions range_to_remove: Range - fill: Optional[bytes] = None + stub: bytes = b"" + fill: bytes = b"\x00" class PartialFreeSpaceModifier(Modifier[PartialFreeSpaceModifierConfig]): """ Turn part of a [MemoryRegion][ofrak.core.memory_region.MemoryRegion] resource into allocatable - free space by replacing a range of its data with b'\x00' or optionally specified fill bytes. - [FreeSpace][ofrak.core.free_space.FreeSpace] child resource at that range. + free space by replacing a range of its data with fill bytes (b'\x00' by default). + + The modifier supports optionally injecting a "stub", bytes at the beginning of the targeted range that will not be + marked as [FreeSpace][ofrak.core.free_space.FreeSpace]. """ targets = (MemoryRegion,) @@ -548,15 +580,17 @@ async def modify(self, resource: Resource, config: PartialFreeSpaceModifierConfi patch_offset = mem_region_view.get_offset_in_self(freed_range.start) patch_range = Range.from_size(patch_offset, freed_range.length()) - patch_data = _get_fill(freed_range, config.fill) + patch_data = _get_patch(freed_range, config.stub, config.fill) await mem_region_view.resource.run( BinaryPatchModifier, BinaryPatchConfig(patch_offset, patch_data) ) + + free_offset = len(config.stub) await mem_region_view.resource.create_child_from_view( FreeSpace( - freed_range.start, - freed_range.length(), + freed_range.start + free_offset, + freed_range.length() - free_offset, config.permissions, ), - data_range=patch_range, + data_range=Range(patch_range.start + free_offset, patch_range.end), ) diff --git a/ofrak_core/test_ofrak/components/test_free_space.py b/ofrak_core/test_ofrak/components/test_free_space.py index dad8e9321..76eb1e2b9 100644 --- a/ofrak_core/test_ofrak/components/test_free_space.py +++ b/ofrak_core/test_ofrak/components/test_free_space.py @@ -1,9 +1,9 @@ import pytest -from ofrak.core import FreeSpace +from ofrak.core import FreeSpace, FreeSpaceModifier, FreeSpaceModifierConfig from ofrak.component.modifier import ModifierError -from ofrak import OFRAKContext, Resource +from ofrak import OFRAKContext, Resource, ResourceFilter, ResourceAttributeValueFilter from ofrak.core import ( MemoryRegion, Program, @@ -19,11 +19,12 @@ async def resource_under_test(ofrak_context: OFRAKContext) -> Resource: resource = await ofrak_context.create_root_resource( "mock_memory_region", b"\xff" * 0x100, - (Program, MemoryRegion), + (Program,), ) - resource.add_view(MemoryRegion(0x0, 0x100)) + memory_region = await resource.create_child_from_view(MemoryRegion(0, 0x100), data_range=Range(0, 0x100)) await resource.save() - return resource + await memory_region.save() + return memory_region async def test_partial_free_modifier_out_of_bounds(resource_under_test: Resource): @@ -35,7 +36,8 @@ async def test_partial_free_modifier_out_of_bounds(resource_under_test: Resource config = PartialFreeSpaceModifierConfig( MemoryPermissions.RX, range_to_remove=Range.from_size(0, data_length + 4), - fill=b"\xfe\xed\xfa\xce", + stub=b"\xfe\xed\xfa\xce", + fill=b"\x00", ) with pytest.raises(ModifierError): await resource_under_test.run(PartialFreeSpaceModifier, config) @@ -45,13 +47,54 @@ async def test_partial_free_modifier(resource_under_test: Resource): """ Test that the PartialFreeSpaceModifier returns expected results. """ + partial_start_offset = 4 + partial_end_offset = 10 + parent = await resource_under_test.get_parent() data_length = await resource_under_test.get_data_length() + range_to_remove = Range.from_size(4, data_length - 4 - 10) config = PartialFreeSpaceModifierConfig( MemoryPermissions.RX, - range_to_remove=Range.from_size(0, data_length - 10), - fill=b"\xfe\xed\xfa\xce", + range_to_remove=range_to_remove, + stub=b"\xfe\xed\xfa\xce", + fill=b"\x00", ) await resource_under_test.run(PartialFreeSpaceModifier, config) + + # Assert free space is as required free_space = await resource_under_test.get_only_child_as_view(FreeSpace) free_space_data = await free_space.resource.get_data() - assert free_space_data == config.fill + (b"\x00" * (data_length - 10 - len(config.fill))) + assert free_space_data == (b"\x00" * (range_to_remove.length() - len(config.stub))) + + # Assert stub is injected + memory_region_data = await resource_under_test.get_data() + assert memory_region_data[partial_start_offset: partial_start_offset + len(config.stub)] == config.stub + + +async def test_free_space_modifier(resource_under_test: Resource): + """ + Test that the FreeSpaceModifier returns expected results + """ + data_length = await resource_under_test.get_data_length() + config = FreeSpaceModifierConfig( + MemoryPermissions.RX, + stub=b"\xfe\xed\xfa\xce", + fill=b"\x00", + ) + parent = await resource_under_test.get_parent() + await resource_under_test.run(FreeSpaceModifier, config) + + # Assert free space created as required + free_space = await parent.get_only_child_as_view(FreeSpace, r_filter=ResourceFilter.with_tags(FreeSpace)) + free_space_data = await free_space.resource.get_data() + # Free space should not include the stub + assert free_space_data == (config.fill * (data_length - len(config.stub))) + + # If stub exists, assert that it matches + child = await parent.get_only_child( + r_filter=ResourceFilter( + tags=(MemoryRegion,), + attribute_filters=(ResourceAttributeValueFilter(MemoryRegion.Size, len(config.stub)),), + ) + ) + child_data = await child.get_data() + assert child_data == config.stub