XFESystem

using Mono.Cecil;
using Mono.Cecil.Cil;

internal static class AllocationRuntimePatch
{
    public const string GuardName = "XfeRequireManagedAllocationV1";
    public const string FatalName = "XfeFailManagedAllocationV1";

    public static bool Apply(AssemblyDefinition assembly)
    {
        var module = assembly.MainModule;
        var memory = module.Types.Single(type => type.FullName == "Cosmos.Kernel.Core.Runtime.Memory");
        var gc = module.Types.Single(type => type.FullName == "Cosmos.Kernel.Core.Memory.GarbageCollector.GarbageCollector");
        var managed = gc.Methods.Single(method => method.Name == "AllocObject" && method.Parameters.Count == 4 && method.HasBody);
        var runtime = memory.Methods.Single(method => method.Name == "AllocObject" && method.Parameters.Count == 4 && method.HasBody);
        var existing = memory.Methods.Where(method => method.Name == GuardName).ToArray();
        if (existing.Length != 0)
        {
            VerifyExisting(memory, managed, runtime, existing);
            return true;
        }

        // This is the shared GC exit AFTER the interrupt-scope finally. Throwing inside
        // AllocObjectSlow would skip that finally on Cosmos and could leave IRQs disabled.
        if (managed.Body.CodeSize != 67 || managed.Body.Variables.Count != 3 || managed.Body.ExceptionHandlers.Count != 1 ||
            managed.Body.ExceptionHandlers[0].HandlerType != ExceptionHandlerType.Finally ||
            managed.Body.Instructions.Count(instruction => instruction.OpCode == OpCodes.Ret) != 1)
            throw new InvalidOperationException("Pinned managed allocation boundary changed.");
        var result = managed.Body.Instructions.Single(instruction => instruction.Offset == 0x41);
        var scope = managed.Body.ExceptionHandlers[0];
        if (result.OpCode != OpCodes.Ldloc_2 || result.Next.OpCode != OpCodes.Ret || !ReferenceEquals(scope.HandlerEnd, result) ||
            managed.Body.Instructions.Single(instruction => instruction.Offset == 0x3b).Operand is not MethodReference dispose || dispose.Name != "Dispose" ||
            managed.Body.Instructions.Count(instruction => instruction.Operand is MethodReference call && call.Name == "AllocObjectCore") != 1 ||
            managed.Body.Instructions.Count(instruction => ReferenceEquals(instruction.Operand, result)) != 1)
            throw new InvalidOperationException("Pinned allocation exit is not the interrupt-scope return.");
        if (runtime.Body.CodeSize != 57 || runtime.Body.Variables.Count != 1 || runtime.Body.ExceptionHandlers.Count != 0)
            throw new InvalidOperationException("Pinned pre-GC allocation boundary changed.");
        var raw = runtime.Body.Instructions.Single(instruction => instruction.Offset == 0x13);
        if (raw.OpCode != OpCodes.Call || raw.Operand is not MethodReference allocator || allocator.Name != "Alloc" ||
            allocator.DeclaringType.FullName != "Cosmos.Kernel.Core.Memory.MemoryOp" || raw.Next.OpCode != OpCodes.Stloc_0 ||
            runtime.Body.Instructions.Single(instruction => instruction.Offset == 0x1c).Operand is not MethodReference clear || clear.Name != "MemSet" ||
            runtime.Body.Instructions.Count(instruction => instruction.Operand is MethodReference call && call.Name == "AllocObject" && call.DeclaringType.FullName == gc.FullName) != 1 ||
            runtime.Body.Instructions.Any(instruction => ReferenceEquals(instruction.Operand, raw.Next)))
            throw new InvalidOperationException("Pinned pre-GC allocator/zeroing path changed.");

        var fatal = CreateFatal(module, memory);
        var guard = new MethodDefinition(GuardName, MethodAttributes.Assembly | MethodAttributes.Static | MethodAttributes.HideBySig, managed.ReturnType)
        {
            ImplAttributes = MethodImplAttributes.AggressiveInlining
        };
        guard.Parameters.Add(new("result", ParameterAttributes.None, managed.ReturnType));
        guard.Parameters.Add(new("requestedBytes", ParameterAttributes.None, module.TypeSystem.UInt64));
        memory.Methods.Add(guard);
        var guardIl = guard.Body.GetILProcessor();
        var success = guardIl.Create(OpCodes.Ldarg_0);
        var stop = guardIl.Create(OpCodes.Br, success);
        stop.Operand = stop;
        guardIl.Append(guardIl.Create(OpCodes.Ldarg_0)); guardIl.Append(guardIl.Create(OpCodes.Brtrue, success));
        guardIl.Append(guardIl.Create(OpCodes.Ldarg_1)); guardIl.Append(guardIl.Create(OpCodes.Call, fatal));
        guardIl.Append(stop); guardIl.Append(success); guardIl.Append(guardIl.Create(OpCodes.Ret));
        guard.Body.MaxStackSize = 1;

        InsertCheck(managed, result, guard);
        InsertCheck(runtime, raw, guard);
        VerifyExisting(memory, managed, runtime, [guard]);
        return false;
    }

    private static void InsertCheck(MethodDefinition owner, Instruction predecessor, MethodDefinition guard)
    {
        var il = owner.Body.GetILProcessor();
        var size = il.Create(OpCodes.Ldarg_0);
        var convert = il.Create(OpCodes.Conv_U8);
        var call = il.Create(OpCodes.Call, guard);
        il.InsertAfter(predecessor, size); il.InsertAfter(size, convert); il.InsertAfter(convert, call);
        owner.Body.MaxStackSize = Math.Max(owner.Body.MaxStackSize, 3);
    }

    private static MethodDefinition CreateFatal(ModuleDefinition module, TypeDefinition memory)
    {
        var serial = module.Types.Single(type => type.FullName == "Cosmos.Kernel.Core.IO.Serial");
        var write = serial.Methods.Single(method => method.Name == "WriteString" && method.Parameters.Count == 1 && method.Parameters[0].ParameterType.MetadataType == MetadataType.String);
        var number = serial.Methods.Single(method => method.Name == "WriteNumber" && method.Parameters.Count == 2 && method.Parameters[0].ParameterType.MetadataType == MetadataType.UInt64);
        var fatal = new MethodDefinition(FatalName, MethodAttributes.Private | MethodAttributes.Static | MethodAttributes.HideBySig, module.TypeSystem.Void)
        {
            ImplAttributes = MethodImplAttributes.NoInlining
        };
        fatal.Parameters.Add(new("requestedBytes", ParameterAttributes.None, module.TypeSystem.UInt64));
        memory.Methods.Add(fatal);
        var il = fatal.Body.GetILProcessor();
        // Terminal single-CPU path: retain the IRQ mask from the first byte
        // through the permanent stop. Other guest threads cannot keep running
        // or interleave diagnostics after a whole-guest allocation failure.
        il.Append(il.Create(OpCodes.Call, module.Types.Single(type => type.FullName == "Cosmos.Kernel.Core.CPU.InternalCpu").Methods.Single(method => method.Name == "DisableInterruptsScope" && method.Parameters.Count == 0)));
        il.Append(il.Create(OpCodes.Pop));
        Text("\nXFE_RUNTIME_OOM|RequestedBytes|");
        il.Append(il.Create(OpCodes.Ldarg_0)); il.Append(il.Create(OpCodes.Ldc_I4_0)); il.Append(il.Create(OpCodes.Call, number));
        Text("|Action|GuestHalt\n");
        var pages = module.Types.Single(type => type.FullName == "Cosmos.Kernel.Core.Memory.PageAllocator");
        var largest = module.Types.Single(type => type.FullName == "Cosmos.Kernel.Core.Scheduler.SchedulerThread").Methods.Single(method => method.Name == ThreadStackRuntimePatch.LargestName);
        Text("XFE_RUNTIME_OOM_PHYSICAL|RawFreePages|");
        il.Append(il.Create(OpCodes.Call, pages.Methods.Single(method => method.Name == "get_FreePageCount"))); Number();
        Text("|RawTotalPages|"); il.Append(il.Create(OpCodes.Ldsfld, pages.Fields.Single(field => field.Name == "TotalPageCount"))); Number();
        Text("|RawLargestFreeRunPages|"); il.Append(il.Create(OpCodes.Call, largest)); Number();
        Text("\n");
        Text("[FAILFAST] Managed allocation failed; guest halted.\n");
        // Do not call ExceptionHelper.FailFast: its first static initialization allocates
        // tracking arrays. This path must also work during pre-GC heap exhaustion.
        var stop = il.Create(OpCodes.Br, il.Create(OpCodes.Nop)); stop.Operand = stop; il.Append(stop);
        fatal.Body.MaxStackSize = 2;
        return fatal;
        void Text(string value) { il.Append(il.Create(OpCodes.Ldstr, value)); il.Append(il.Create(OpCodes.Call, write)); }
        void Number() { il.Append(il.Create(OpCodes.Ldc_I4_0)); il.Append(il.Create(OpCodes.Call, number)); }
    }

    private static void VerifyExisting(TypeDefinition memory, MethodDefinition managed, MethodDefinition runtime, MethodDefinition[] guards)
    {
        var fatals = memory.Methods.Where(method => method.Name == FatalName).ToArray();
        if (guards.Length != 1 || fatals.Length != 1 || !guards[0].HasBody || !fatals[0].HasBody)
            throw new InvalidOperationException("An allocation guard has an unexpected identity.");
        var guard = guards[0]; var fatal = fatals[0]; var instructions = guard.Body.Instructions;
        if (guard.Parameters.Count != 2 || guard.ReturnType.FullName != managed.ReturnType.FullName || guard.Parameters[0].ParameterType.FullName != managed.ReturnType.FullName ||
            guard.Parameters[1].ParameterType.MetadataType != MetadataType.UInt64 || instructions.Count != 7 ||
            instructions[0].OpCode != OpCodes.Ldarg_0 || instructions[1].OpCode != OpCodes.Brtrue || !ReferenceEquals(instructions[1].Operand, instructions[5]) ||
            instructions[2].OpCode != OpCodes.Ldarg_1 || instructions[3].OpCode != OpCodes.Call || instructions[3].Operand is not MethodReference failure || failure.FullName != fatal.FullName ||
            instructions[4].OpCode != OpCodes.Br || !ReferenceEquals(instructions[4].Operand, instructions[4]) || instructions[5].OpCode != OpCodes.Ldarg_0 || instructions[6].OpCode != OpCodes.Ret ||
            fatal.Body.Instructions.Count != 29 || fatal.Body.Instructions[0].Operand is not MethodReference disable || disable.Name != "DisableInterruptsScope" || fatal.Body.Instructions[1].OpCode != OpCodes.Pop ||
            fatal.Body.Instructions[^1].OpCode != OpCodes.Br || !ReferenceEquals(fatal.Body.Instructions[^1].Operand, fatal.Body.Instructions[^1]) ||
            fatal.Body.Instructions.Any(instruction => instruction.Operand is MethodReference call && call.DeclaringType.FullName != "Cosmos.Kernel.Core.IO.Serial" &&
                call.FullName != "System.UInt64 Cosmos.Kernel.Core.Memory.PageAllocator::get_FreePageCount()" && call.Name != ThreadStackRuntimePatch.LargestName && call.Name != "DisableInterruptsScope") ||
            fatal.Body.Instructions.Any(instruction => instruction.OpCode == OpCodes.Newobj || instruction.OpCode == OpCodes.Newarr || instruction.OpCode == OpCodes.Ret))
            throw new InvalidOperationException("The nonallocating whole-guest OOM halt shape changed.");
        foreach (var owner in new[] { managed, runtime })
        {
            var calls = owner.Body.Instructions.Where(instruction => instruction.OpCode == OpCodes.Call && instruction.Operand is MethodReference call && call.FullName == guard.FullName).ToArray();
            if (calls.Length != 1 || calls[0].Previous.OpCode != OpCodes.Conv_U8 || calls[0].Previous.Previous.OpCode != OpCodes.Ldarg_0 ||
                owner.Body.Instructions.Any(instruction => ReferenceEquals(instruction.Operand, calls[0].Next)))
                throw new InvalidOperationException("The allocation boundary guard can be bypassed.");
        }
        var managedCall = managed.Body.Instructions.Single(instruction => instruction.Operand is MethodReference call && call.FullName == guard.FullName);
        var runtimeCall = runtime.Body.Instructions.Single(instruction => instruction.Operand is MethodReference call && call.FullName == guard.FullName);
        if (managedCall.Next.OpCode != OpCodes.Ret || managedCall.Previous.Previous.Previous.OpCode != OpCodes.Ldloc_2 ||
            managed.Body.ExceptionHandlers.Count != 1 || !ReferenceEquals(managed.Body.ExceptionHandlers[0].HandlerEnd, managedCall.Previous.Previous.Previous) ||
            runtimeCall.Next.OpCode != OpCodes.Stloc_0 || runtimeCall.Previous.Previous.Previous.Operand is not MethodReference raw || raw.Name != "Alloc")
            throw new InvalidOperationException("The allocation guard is not after IRQ restoration/before raw zeroing.");
    }
}