XFESystem

using System.Buffers.Binary;
using System.Reflection.Metadata;
using System.Reflection.Metadata.Ecma335;
using System.Reflection.PortableExecutable;
using System.Text;

namespace XFESystem.ManagedRuntime.Tests;

/// <summary>Mutates copies of real compiled CLI PE images without loading any assembly into the CLR.</summary>
internal sealed class ManagedImage : IDisposable
{
    private readonly MemoryStream _stream;
    private readonly PEReader _pe;
    public byte[] Bytes { get; }
    public MetadataReader Metadata { get; }
    public MethodDefinition Entry { get; }
    public int CodeOffset { get; }
    public int CodeLength { get; }
    public int MaximumStack { get; }

    public ManagedImage(string name = "ManagedFixture.dll")
    {
        Bytes = File.ReadAllBytes(Path.Combine(AppContext.BaseDirectory, "Fixtures", name));
        _stream = new MemoryStream(Bytes, false);
        _pe = new PEReader(_stream);
        Metadata = _pe.GetMetadataReader();
        var token = _pe.PEHeaders.CorHeader!.EntryPointTokenOrRelativeVirtualAddress;
        Entry = Metadata.GetMethodDefinition((MethodDefinitionHandle)MetadataTokens.EntityHandle(token));
        var body = _pe.GetMethodBody(Entry.RelativeVirtualAddress);
        CodeLength = body.GetILBytes()!.Length;
        MaximumStack = body.MaxStack;
        var header = FileOffset(Entry.RelativeVirtualAddress);
        CodeOffset = header + ((Bytes[header] & 3) == 2 ? 1 : (BinaryPrimitives.ReadUInt16LittleEndian(Bytes.AsSpan(header, 2)) >> 12) * 4);
    }

    public byte[] WithCode(params byte[] instructions)
    {
        if (instructions.Length > CodeLength) throw new ArgumentException("The mutation exceeds the original method body.");
        var copy = Bytes.ToArray();
        copy.AsSpan(CodeOffset, CodeLength).Clear();
        instructions.CopyTo(copy, CodeOffset);
        return copy;
    }

    public byte[] WithTokenInstruction(byte opcode, int token, params byte[] tail)
    {
        var code = new byte[5 + tail.Length];
        code[0] = opcode;
        BinaryPrimitives.WriteInt32LittleEndian(code.AsSpan(1), token);
        tail.CopyTo(code, 5);
        return WithCode(code);
    }

    public int Member(string owner, string name)
    {
        foreach (var handle in Metadata.MemberReferences)
        {
            var member = Metadata.GetMemberReference(handle);
            if (Metadata.GetString(member.Name) != name || member.Parent.Kind != HandleKind.TypeReference) continue;
            var type = Metadata.GetTypeReference((TypeReferenceHandle)member.Parent);
            if (Metadata.GetString(type.Namespace) + "." + Metadata.GetString(type.Name) == owner)
                return MetadataTokens.GetToken(handle);
        }
        throw new InvalidOperationException($"No member token: {owner}.{name}");
    }

    public int UserString(string value)
    {
        // User-string heap entries carry their own compressed length. Read handles from actual ldstr instructions.
        foreach (var methodHandle in Metadata.MethodDefinitions)
        {
            var method = Metadata.GetMethodDefinition(methodHandle);
            if (method.RelativeVirtualAddress == 0) continue;
            var code = _pe.GetMethodBody(method.RelativeVirtualAddress).GetILBytes()!;
            for (int index = 0; index + 4 < code.Length; index++)
                if (code[index] == 0x72)
                {
                    var token = BinaryPrimitives.ReadInt32LittleEndian(code.AsSpan(index + 1, 4));
                    if ((token & unchecked((int)0xFF000000)) == 0x70000000 &&
                        Metadata.GetUserString(MetadataTokens.UserStringHandle(token & 0xFFFFFF)) == value) return token;
                }
        }
        throw new InvalidOperationException("The compiled fixture has no requested user string.");
    }

    public int SignaturePayloadOffset(BlobHandle signature)
    {
        var metadataStart = FileOffset(_pe.PEHeaders.CorHeader!.MetadataDirectory.RelativeVirtualAddress);
        int versionLength = BinaryPrimitives.ReadInt32LittleEndian(Bytes.AsSpan(metadataStart + 12, 4));
        int at = metadataStart + 16 + ((versionLength + 3) & ~3);
        int streams = BinaryPrimitives.ReadUInt16LittleEndian(Bytes.AsSpan(at + 2, 2));
        at += 4;
        for (int index = 0; index < streams; index++)
        {
            int offset = BinaryPrimitives.ReadInt32LittleEndian(Bytes.AsSpan(at, 4));
            int nameAt = at + 8, end = nameAt;
            while (Bytes[end] != 0) end++;
            string name = Encoding.ASCII.GetString(Bytes.AsSpan(nameAt, end - nameAt));
            at = nameAt + ((end - nameAt + 4) & ~3);
            if (name != "#Blob") continue;
            int blob = metadataStart + offset + MetadataTokens.GetHeapOffset(signature);
            return blob + (Bytes[blob] < 0x80 ? 1 : Bytes[blob] < 0xC0 ? 2 : 4);
        }
        throw new InvalidOperationException("Missing metadata blob heap.");
    }

    public int CorHeaderOffset => FileOffset(_pe.PEHeaders.PEHeader!.CorHeaderTableDirectory.RelativeVirtualAddress);

    private int FileOffset(int rva)
    {
        foreach (var section in _pe.PEHeaders.SectionHeaders)
            if (rva >= section.VirtualAddress && rva - section.VirtualAddress < Math.Max(section.VirtualSize, section.SizeOfRawData))
                return section.PointerToRawData + rva - section.VirtualAddress;
        throw new BadImageFormatException("Fixture RVA is not in a PE section.");
    }

    public void Dispose() { _pe.Dispose(); _stream.Dispose(); }
}