Fluent Patching API

[Washi1337]

Problem Description

Some use-cases of AsmResolver depend on changing only very small pieces in the input binary (e.g., just a couple bytes or words) while preserving the structure of everything else as-is.

Currently, AsmResolver's API promotes to always fully rebuild the target binary, or manually calculating offsets and patching the file after ISegment::Write was called. The first option is very likely to destroy the original structure of the segment, as many segments are reconstructed completely. The second option does preserve the structure, however depending on the type of patch that should be applied, it often requires a lot of ugly and/or non-trivial code to be written by the end-user. Furthermore, if these patches depend on the location of symbols (e.g., an address in the IAT) or result in new base-relocations to be computed, this becomes very error-prone very quickly.

An end-user should have an easier and more reliable way to quickly patch segments (and by extension entire files) without having to worry about the nitty-gritty details on how this should be implemented.

Proposal

TL;DR
Add functionality that allows for constructions like the following:

var section = peFile.Sections.First(s => s.Name == ".text");

var someSymbol = peImage
   .Imports.First(m => m.Name == "ucrtbased.dll")
   .Symbols.First(s => s.Name == "puts");
var newSubSegment = ...

section.Contents = section.Contents.AsPatchedSegment()                      // Create patched segment.
   .Patch(offset: 0x10, data: new byte[] {1, 2, 3, 4})                      // Apply literal bytes patch
   .Patch(offset: 0x20, AddressFixupType.Absolute64BitAddress, someSymbol); // Apply address fixup patch.
   .Patch(offset: 0x30, data: newSubSegment)                                // Replace entire sub segment.

Details:

This could be implemented in the following manner. Add a decorator class PatchedSegment that wraps an ISegment instance and attaches additional metadata describing how to patch the segment after calling ISegment::Write.

The basic contract would be as follows:

public class PatchedSegment : ISegment
{
    public PatchedSegment(ISegment contents)
    {
        Contents = contents;
    }

    public ISegment Contents
    {
        get;
    }

    public IList<IPatch> Patches
    {
        get;
    } = new List<IPatch>();

    public void Write(IBinaryStreamWriter writer)
    {
        Contents.Write(writer);
        ApplyPatches(writer);
    }

    /* ... */ 
}

public interface IPatch
{
    void Apply(/* ... */);
}

In their most basic form, they could be added to the Patches list. For instance:

// Create a patched version of the .text section.
var section = file.Sections.First(s => s.Name == ".text");
var segment = new PatchedSegment(section.Contents);

// Apply a literal bytes patch.
segment.Patches.Add(new BytesPatch(offset: 0x10, data: new byte[] {1, 2, 3, 4});

// Apply an address fixup patch (similar to how is done in CodeSegment and NativeMethodBody)):
ISymbol someSymbol = peImage
   .Imports.First(m => m.Name == "ucrtbased.dll")
   .Symbols.First(s => s.Name == "puts");
segment.Patches.Add(new AddressFixupPatch(offset: 0x20, AddressFixupType.Absolute64BitAddress, someSymbol));

// Replace entire sub segment.
ISegment newSubSegment = ...
segment.Patches.Add(new SegmentPatch(offset: 0x30, data: newSubSegment));

// Persist patches in the section:
section.Contents = segment;

Extension methods could be provided to allow for a more fluent syntax. For example, the following could be equivalent to the previous code snippet:

var section = file.Sections.First(s => s.Name == ".text");

var someSymbol = peImage
   .Imports.First(m => m.Name == "ucrtbased.dll")
   .Symbols.First(s => s.Name == "puts");
var newSubSegment = ...

section.Contents = section.Contents.AsPatchedSegment()                      // Create patched segment.
   .Patch(offset: 0x10, data: new byte[] {1, 2, 3, 4})                      // Apply literal bytes patch
   .Patch(offset: 0x20, AddressFixupType.Absolute64BitAddress, someSymbol); // Apply address fixup patch.
   .Patch(offset: 0x30, data: newSubSegment)                                // Replace entire sub segment.

Alternatives

A current alternative would be to manually write the segment / file to be patched to a temporary MemoryStream, and the patches would be applied to that. This would require a lot of manual work that isn't always trivial to implement (especially when it comes to finding the right patch addresses as well as the right values to write into the file).

Another alternative that could be considered is extending ISegment itself such that we do not need the PatchedSegment decorator class. This has the advantage that every segment can be patched regardless of where it is used and/or how it is implemented. However, this would increase the memory consumption unnecessarily for everything that is not meant to be patched.

Additional Context

This concept is similar to how CodeSegment is implemented, but is more generalized as PatchedSegment takes any ISegment as contents as opposed to hard-wiring it to a byte[]. Furthermore, more complex patches can be applied by deriving from IPatch. This means it can be used not only for raw code/data streams but also more higher-level objects and structures in the PE model. As such, this will likely obsolete CodeSegment.

This is going to be an important feature that will lay the foundation of full native and/or mixed-mode PE image rebuilding that try to preserve as much data as possible.