XFESystem

using System.Security.Cryptography;
using System.Text.Json;
using Mono.Cecil;
using Mono.Cecil.Cil;

namespace XFESystem.Tools.JintCleanupPatch;

public static partial class Program
{
    private const string Marker = "XFESystem.NativeJintCleanupPatchV2";
    public static int Main(string[] arguments)
    {
        if (arguments.Length == 2 && arguments[1] == "--bind-only") return PatchBoundTargets(Path.GetFullPath(arguments[0]), hostOnly: true);
        if (arguments.Length == 2 && arguments[1] == "--bind-and-cleanup")
        {
            var nativePath = Path.GetFullPath(arguments[0]); VerifyPrivateCopy(nativePath);
            ValidateBoundInput(nativePath, hostOnly: false);
            Main([nativePath]);
            PatchBoundTargets(nativePath, hostOnly: false);
            var receiptPath = nativePath + ".XfeCleanup.json";
            var combinedReceipt = System.Text.Json.Nodes.JsonNode.Parse(File.ReadAllText(receiptPath))!;
            combinedReceipt["AfterSha256"] = Hash(nativePath); combinedReceipt["BoundCallablePatchVersion"] = 1;
            combinedReceipt["WeakCollectionsPatchVersion"] = 1;
            combinedReceipt["EngineWeakCollectionsPatchVersion"] = 1;
            File.WriteAllText(receiptPath, combinedReceipt.ToJsonString(new JsonSerializerOptions { WriteIndented = true }));
            return 0;
        }
        if (arguments.Length is < 1 or > 2 || arguments.Length == 2 && arguments[1] != "--inspect")
            throw new ArgumentException("JintCleanupPatch <project-private obj Jint.dll> [--inspect|--bind-only|--bind-and-cleanup]");
        var path = Path.GetFullPath(arguments[0]); var inspect = arguments.Length == 2;
        if (!inspect) VerifyPrivateCopy(path);
        using var assembly = AssemblyDefinition.ReadAssembly(path, new ReaderParameters { InMemory = true });
        if (assembly.Name.Name != "Jint" || assembly.Name.Version != new Version(4, 11, 0, 0))
            throw new InvalidOperationException("The cleanup patch only supports the pinned Jint 4.11.0 assembly.");
        var before = Hash(path);
        if (!inspect && assembly.MainModule.Types.Any(type => type.FullName.StartsWith("XFESystem.NativeJintCleanupPatch", StringComparison.Ordinal) && type.FullName != Marker))
            throw new InvalidOperationException("An older cleanup patch is present; refresh this private copy from the pinned original Jint assembly.");
        if (assembly.MainModule.Types.Any(type => type.FullName == Marker))
        {
            var receiptPath = path + ".XfeCleanup.json";
            if (!File.Exists(receiptPath)) throw new InvalidOperationException("A cleanup marker without its verification receipt is not trusted.");
            using var previousReceipt = JsonDocument.Parse(File.ReadAllText(receiptPath));
            if (previousReceipt.RootElement.GetProperty("AfterSha256").GetString() != before)
                throw new InvalidOperationException("The marked Jint assembly differs from its cleanup verification receipt.");
            Console.WriteLine("XFE_JINT_CLEANUP_PATCH_ALREADY_READY|" + before); return 0;
        }
        var reports = new List<MethodReport>();
        var exception = assembly.MainModule.GetTypeReferences().First(type => type.FullName == "System.Exception");
        foreach (var method in Types(assembly.MainModule.Types).SelectMany(type => type.Methods).Where(method => method.HasBody))
        {
            var clauses = method.Body.ExceptionHandlers.Where(handler => handler.HandlerType == ExceptionHandlerType.Finally).ToArray();
            if (clauses.Length == 0) continue;
            var patched = 0; var unsupported = new List<string>();
            // Inner protected regions first. A newly injected inner catch then remains inside
            // the surrounding original try; each outer scope performs its own cleanup once.
            foreach (var handler in clauses.OrderBy(handler => Span(method, handler.TryStart, handler.TryEnd)))
            {
                var reason = ValidateCleanup(method, handler);
                if (reason is not null) { unsupported.Add(reason); continue; }
                if (!inspect) Patch(method, handler, exception);
                ++patched;
            }
            reports.Add(new(method.FullName, clauses.Length, patched, unsupported.ToArray()));
        }
        if (inspect) { Console.WriteLine(JsonSerializer.Serialize(reports, new JsonSerializerOptions { WriteIndented = true })); return 0; }
        var criticalUnsupported = reports.Where(report => IsCritical(report.Method) && report.Unsupported.Length != 0).ToArray();
        if (criticalUnsupported.Length != 0)
            throw new InvalidOperationException("A required interpreter scope has an unsupported cleanup shape: " + string.Join("; ", criticalUnsupported.Select(report => report.Method)));
        assembly.MainModule.Types.Add(new TypeDefinition("XFESystem", "NativeJintCleanupPatchV2", TypeAttributes.NotPublic | TypeAttributes.Abstract | TypeAttributes.Sealed, assembly.MainModule.TypeSystem.Object));
        var temporary = path + ".XfeCleanup.tmp";
        assembly.Write(temporary); File.Move(temporary, path, true);
        var receipt = new
        {
            Kind = "ProjectLocalJintExplicitExceptionalCleanup", Version = 2, JintVersion = "4.11.0", Path = path,
            BeforeSha256 = before, AfterSha256 = Hash(path), GlobalPackageCacheModified = false,
            OriginalFinallyNormalPathPreserved = true, ExceptionalCleanupUsesCatchAndRethrow = true,
            Methods = reports, UnsupportedClauses = reports.Sum(report => report.Unsupported.Length)
        };
        File.WriteAllText(path + ".XfeCleanup.json", JsonSerializer.Serialize(receipt, new JsonSerializerOptions { WriteIndented = true }));
        Console.WriteLine("XFE_JINT_CLEANUP_PATCH_READY|" + reports.Sum(report => report.Patched) + "|Unsupported|" + receipt.UnsupportedClauses + "|" + receipt.AfterSha256);
        return 0;
    }
    public sealed record MethodReport(string Method, int FinallyClauses, int Patched, string[] Unsupported);
    private static bool IsCritical(string method) => method.Contains("Jint.Engine::", StringComparison.Ordinal) || method.Contains("Jint.Native.Function.ScriptFunction::Call", StringComparison.Ordinal) || method.Contains("Jint.Runtime.Interpreter.Statements.JintTryStatement::", StringComparison.Ordinal);
    private static IEnumerable<TypeDefinition> Types(IEnumerable<TypeDefinition> types) { foreach (var type in types) { yield return type; foreach (var nested in Types(type.NestedTypes)) yield return nested; } }
    private static int Span(MethodDefinition method, Instruction? start, Instruction? end) => (end is null ? method.Body.Instructions.Count : method.Body.Instructions.IndexOf(end)) - (start is null ? 0 : method.Body.Instructions.IndexOf(start));
    private static string? ValidateCleanup(MethodDefinition method, ExceptionHandler handler)
    {
        if (handler.TryEnd != handler.HandlerStart || handler.HandlerEnd is null) return "Noncontiguous try/finally boundaries";
        var instructions = method.Body.Instructions; var start = instructions.IndexOf(handler.HandlerStart); var end = instructions.IndexOf(handler.HandlerEnd);
        var cleanup = instructions.Skip(start).Take(end - start).ToHashSet();
        if (!cleanup.Any(instruction => instruction.OpCode == OpCodes.Endfinally)) return "No endfinally in cleanup";
        foreach (var other in method.Body.ExceptionHandlers)
        {
            if (other == handler) continue;
            if (cleanup.Contains(other.TryStart) || cleanup.Contains(other.HandlerStart) || other.FilterStart is not null && cleanup.Contains(other.FilterStart))
                return "Nested exception region inside cleanup";
        }
        foreach (var instruction in cleanup)
        {
            if (instruction.Operand is Instruction target && !cleanup.Contains(target)) return "Cleanup branch leaves handler";
            if (instruction.Operand is Instruction[] targets && targets.Any(target => !cleanup.Contains(target))) return "Cleanup switch leaves handler";
            if (instruction.OpCode == OpCodes.Ret || instruction.OpCode == OpCodes.Jmp || instruction.OpCode == OpCodes.Rethrow) return "Unsupported cleanup control transfer";
        }
        return null;
    }
    private static void Patch(MethodDefinition method, ExceptionHandler handler, TypeReference exception)
    {
        var body = method.Body; ExpandBranches(body); body.InitLocals = true;
        var done = new VariableDefinition(method.Module.TypeSystem.Boolean); body.Variables.Add(done);
        // A try region can be re-entered by a loop after an outer catch handled its
        // exception. Reset the guard on each entry, not merely at method initialization.
        var oldTryStart = handler.TryStart; var reset = Instruction.Create(OpCodes.Ldc_I4_0);
        var resetStore = Instruction.Create(OpCodes.Stloc, done); var processor = body.GetILProcessor();
        foreach (var instruction in body.Instructions)
        {
            if (instruction.Operand == oldTryStart) instruction.Operand = reset;
            else if (instruction.Operand is Instruction[] targets) instruction.Operand = targets.Select(target => target == oldTryStart ? reset : target).ToArray();
        }
        foreach (var region in body.ExceptionHandlers)
        {
            if (region.TryStart == oldTryStart) region.TryStart = reset;
            if (region.TryEnd == oldTryStart) region.TryEnd = reset;
            if (region.HandlerStart == oldTryStart) region.HandlerStart = reset;
            if (region.HandlerEnd == oldTryStart) region.HandlerEnd = reset;
            if (region.FilterStart == oldTryStart) region.FilterStart = reset;
        }
        processor.InsertBefore(oldTryStart, reset); processor.InsertBefore(oldTryStart, resetStore);
        var originalStart = handler.HandlerStart; var start = body.Instructions.IndexOf(originalStart); var end = body.Instructions.IndexOf(handler.HandlerEnd);
        var cleanup = body.Instructions.Skip(start).Take(end - start).ToArray();
        var il = body.GetILProcessor(); var catchStart = il.Create(OpCodes.Pop); var rethrow = il.Create(OpCodes.Rethrow);
        var clones = cleanup.ToDictionary(instruction => instruction, _ => il.Create(OpCodes.Nop));
        foreach (var instruction in cleanup)
        {
            var clone = clones[instruction]; clone.OpCode = instruction.OpCode; clone.Operand = instruction.Operand;
            if (instruction.OpCode == OpCodes.Endfinally) { clone.OpCode = OpCodes.Br; clone.Operand = rethrow; }
            else if (instruction.Operand is Instruction target) clone.Operand = clones[target];
            else if (instruction.Operand is Instruction[] targets) clone.Operand = targets.Select(target => clones[target]).ToArray();
        }
        // Inserting before the old boundary must not extend an existing nested handler into
        // this new catch. End boundaries are exclusive and retain their original extent.
        foreach (var other in body.ExceptionHandlers)
        {
            if (other == handler) continue;
            if (other.TryEnd == originalStart) other.TryEnd = catchStart;
            if (other.HandlerEnd == originalStart) other.HandlerEnd = catchStart;
        }
        il.InsertBefore(originalStart, catchStart); il.InsertBefore(originalStart, il.Create(OpCodes.Ldc_I4_1)); il.InsertBefore(originalStart, il.Create(OpCodes.Stloc, done));
        foreach (var instruction in cleanup) il.InsertBefore(originalStart, clones[instruction]);
        il.InsertBefore(originalStart, rethrow);
        var guard = il.Create(OpCodes.Ldloc, done); il.InsertBefore(originalStart, guard);
        il.InsertBefore(originalStart, il.Create(OpCodes.Brtrue, cleanup.First(instruction => instruction.OpCode == OpCodes.Endfinally)));
        var catchHandler = new ExceptionHandler(ExceptionHandlerType.Catch)
        { CatchType = exception, TryStart = handler.TryStart, TryEnd = catchStart, HandlerStart = catchStart, HandlerEnd = guard };
        body.ExceptionHandlers.Insert(body.ExceptionHandlers.IndexOf(handler), catchHandler);
        handler.TryEnd = guard; handler.HandlerStart = guard;
        body.MaxStackSize = Math.Max(8, body.MaxStackSize + 2);
    }
    private static void ExpandBranches(MethodBody body)
    {
        foreach (var instruction in body.Instructions)
            instruction.OpCode = instruction.OpCode.Code switch
            {
                Code.Br_S => OpCodes.Br, Code.Brfalse_S => OpCodes.Brfalse, Code.Brtrue_S => OpCodes.Brtrue, Code.Beq_S => OpCodes.Beq,
                Code.Bge_S => OpCodes.Bge, Code.Bge_Un_S => OpCodes.Bge_Un, Code.Bgt_S => OpCodes.Bgt, Code.Bgt_Un_S => OpCodes.Bgt_Un,
                Code.Ble_S => OpCodes.Ble, Code.Ble_Un_S => OpCodes.Ble_Un, Code.Blt_S => OpCodes.Blt, Code.Blt_Un_S => OpCodes.Blt_Un,
                Code.Bne_Un_S => OpCodes.Bne_Un, Code.Leave_S => OpCodes.Leave, _ => instruction.OpCode
            };
    }
    private static string Hash(string path) => Convert.ToHexStringLower(SHA256.HashData(File.ReadAllBytes(path)));
    private static void VerifyPrivateCopy(string path)
    {
        string? project = Path.GetDirectoryName(path);
        while (project is not null && !File.Exists(Path.Combine(project, "XFESystem.Kernel.csproj"))) project = Path.GetDirectoryName(project);
        if (project is null || Path.GetFileName(path) != "Jint.dll" || !path.StartsWith(Path.Combine(project, "obj") + Path.DirectorySeparatorChar, StringComparison.OrdinalIgnoreCase))
            throw new ArgumentException("Only the kernel project's private obj Jint.dll copy may be patched.");
        for (string? candidate = path; candidate is not null && !candidate.Equals(project, StringComparison.OrdinalIgnoreCase); candidate = Path.GetDirectoryName(candidate))
            if ((File.GetAttributes(candidate) & FileAttributes.ReparsePoint) != 0) throw new ArgumentException("A private Jint copy cannot be reached through a symbolic link or junction.");
    }
}