XFESystem

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

/// <summary>The NativeAOT update ABI supplies end addresses, not byte counts.</summary>
internal static class FrozenRangeRuntimePatch
{
    public const string ValidateName = "XfeValidateFrozenEndAddressesV1";
    public const string FatalName = "XfeFailFrozenEndAddressesV1";

    public static bool Apply(AssemblyDefinition assembly)
    {
        var module = assembly.MainModule;
        var gc = module.Types.Single(type => type.FullName == "Cosmos.Kernel.Core.Memory.GarbageCollector.GarbageCollector");
        var update = gc.Methods.Single(method => method.Name == "UpdateFrozenSegment");
        var info = gc.NestedTypes.Single(type => type.Name == "FrozenSegmentInfo");
        var start = info.Fields.Single(field => field.Name == "Start");
        var reserved = info.Fields.Single(field => field.Name == "ReservedSize");
        var allocated = info.Fields.Single(field => field.Name == "AllocSize");
        var committed = info.Fields.Single(field => field.Name == "CommitSize");
        if (gc.Methods.Any(method => method.Name == ValidateName)) { Verify(); return true; }
        if (update.Body.CodeSize != 45 || update.Body.Variables.Count != 1 || update.Body.ExceptionHandlers.Count != 0 ||
            update.Body.Instructions.Count(instruction => instruction.OpCode == OpCodes.Stfld) != 2 ||
            update.Body.Instructions.Single(instruction => instruction.Offset == 0x13).Operand is not FieldReference commitStore || commitStore.FullName != committed.FullName ||
            update.Body.Instructions.Single(instruction => instruction.Offset == 0x1a).Operand is not FieldReference allocStore || allocStore.FullName != allocated.FullName)
            throw new InvalidOperationException("Pinned frozen-segment update shape changed.");

        var validate = new MethodDefinition(ValidateName, MethodAttributes.Assembly | MethodAttributes.Static | MethodAttributes.HideBySig, module.TypeSystem.Int32);
        foreach (string name in new[] { "start", "reservedSize", "allocatedEnd", "committedEnd" })
            validate.Parameters.Add(new(name, ParameterAttributes.None, module.TypeSystem.UInt64));
        gc.Methods.Add(validate);
        var vi = validate.Body.GetILProcessor(); var invalid = vi.Create(OpCodes.Ldc_I4_0);
        vi.Append(vi.Create(OpCodes.Ldarg_0)); vi.Append(vi.Create(OpCodes.Brfalse, invalid));
        vi.Append(vi.Create(OpCodes.Ldarg_2)); vi.Append(vi.Create(OpCodes.Ldarg_0)); vi.Append(vi.Create(OpCodes.Blt_Un, invalid));
        vi.Append(vi.Create(OpCodes.Ldarg_3)); vi.Append(vi.Create(OpCodes.Ldarg_2)); vi.Append(vi.Create(OpCodes.Blt_Un, invalid));
        vi.Append(vi.Create(OpCodes.Ldarg_3)); vi.Append(vi.Create(OpCodes.Ldarg_0)); vi.Append(vi.Create(OpCodes.Sub));
        vi.Append(vi.Create(OpCodes.Ldarg_1)); vi.Append(vi.Create(OpCodes.Bgt_Un, invalid));
        vi.Append(vi.Create(OpCodes.Ldc_I4_1)); vi.Append(vi.Create(OpCodes.Ret)); vi.Append(invalid); vi.Append(vi.Create(OpCodes.Ret)); validate.Body.MaxStackSize = 2;

        var fatal = new MethodDefinition(FatalName, MethodAttributes.Private | MethodAttributes.Static | MethodAttributes.HideBySig, module.TypeSystem.Void);
        gc.Methods.Add(fatal); var fi = fatal.Body.GetILProcessor();
        fi.Append(fi.Create(OpCodes.Call, module.Types.Single(type => type.FullName == "Cosmos.Kernel.Core.CPU.InternalCpu").Methods.Single(method => method.Name == "DisableInterruptsScope" && method.Parameters.Count == 0)));
        fi.Append(fi.Create(OpCodes.Pop));
        fi.Append(fi.Create(OpCodes.Ldstr, "\nXFE_RUNTIME_FROZEN_RANGE_INVALID\n[FAILFAST] Invalid frozen-segment end address; guest halted.\n"));
        fi.Append(fi.Create(OpCodes.Call, module.Types.Single(type => type.FullName == "Cosmos.Kernel.Core.IO.Serial").Methods.Single(method => method.Name == "WriteString" && method.Parameters.Count == 1 && method.Parameters[0].ParameterType.MetadataType == MetadataType.String)));
        var stop = fi.Create(OpCodes.Br, fi.Create(OpCodes.Nop)); stop.Operand = stop; fi.Append(stop); fatal.Body.MaxStackSize = 1;

        var ui = update.Body.GetILProcessor();
        var commitValue = update.Body.Instructions.Single(instruction => instruction.Offset == 0x12);
        var allocValue = update.Body.Instructions.Single(instruction => instruction.Offset == 0x19);
        var matched = commitValue.Previous;
        if (commitValue.OpCode != OpCodes.Ldarg_2 || allocValue.OpCode != OpCodes.Ldarg_1 || matched.OpCode != OpCodes.Ldloc_0 ||
            update.Body.Instructions.Any(instruction => ReferenceEquals(instruction.Operand, matched)))
            throw new InvalidOperationException("Pinned frozen-segment stores can bypass normalization.");
        foreach (var instruction in new[]
        {
            ui.Create(OpCodes.Ldarg_0), ui.Create(OpCodes.Conv_U8), ui.Create(OpCodes.Ldloc_0), ui.Create(OpCodes.Ldfld, reserved), ui.Create(OpCodes.Conv_U8),
            ui.Create(OpCodes.Ldarg_1), ui.Create(OpCodes.Conv_U8), ui.Create(OpCodes.Ldarg_2), ui.Create(OpCodes.Conv_U8), ui.Create(OpCodes.Call, validate),
            ui.Create(OpCodes.Brtrue, matched), ui.Create(OpCodes.Call, fatal)
        }) ui.InsertBefore(matched, instruction);
        // Fatal never returns; retain an explicit fail-closed edge for its caller.
        var terminal = ui.Create(OpCodes.Br, matched); terminal.Operand = terminal; ui.InsertBefore(matched, terminal);
        foreach (var value in new[] { commitValue, allocValue })
        {
            var origin = ui.Create(OpCodes.Ldarg_0); var subtract = ui.Create(OpCodes.Sub); var convert = ui.Create(OpCodes.Conv_U);
            ui.InsertAfter(value, origin); ui.InsertAfter(origin, subtract); ui.InsertAfter(subtract, convert);
        }
        foreach (var instruction in update.Body.Instructions)
            instruction.OpCode = instruction.OpCode.Code switch { Code.Br_S => OpCodes.Br, Code.Bne_Un_S => OpCodes.Bne_Un, _ => instruction.OpCode };
        update.Body.MaxStackSize = 4;

        var export = Export("XfeFrozenRangeValidationExportV1", "XfeRuntimeValidateFrozenRange");
        foreach (var parameter in validate.Parameters) export.Parameters.Add(new(parameter.Name, ParameterAttributes.None, module.TypeSystem.UInt64));
        var ei = export.Body.GetILProcessor();
        ei.Append(ei.Create(OpCodes.Ldarg_0)); ei.Append(ei.Create(OpCodes.Ldarg_1)); ei.Append(ei.Create(OpCodes.Ldarg_2)); ei.Append(ei.Create(OpCodes.Ldarg_3));
        ei.Append(ei.Create(OpCodes.Call, validate)); ei.Append(ei.Create(OpCodes.Ret)); export.Body.MaxStackSize = 4;

        var registered = Export("XfeRegisteredFrozenRangesValidV1", "XfeRuntimeRegisteredFrozenRangesValid");
        var head = gc.Fields.Single(field => field.Name == "s_frozenSegments");
        var next = info.Fields.Single(field => field.Name == "Next");
        var cursor = new VariableDefinition(head.FieldType); registered.Body.Variables.Add(cursor); registered.Body.InitLocals = true;
        var ri = registered.Body.GetILProcessor(); var body = ri.Create(OpCodes.Ldloc, cursor); var test = ri.Create(OpCodes.Ldloc, cursor);
        var bad = ri.Create(OpCodes.Ldc_I4_0);
        ri.Append(ri.Create(OpCodes.Ldsfld, head)); ri.Append(ri.Create(OpCodes.Stloc, cursor)); ri.Append(ri.Create(OpCodes.Br, test));
        ri.Append(body); ri.Append(ri.Create(OpCodes.Ldfld, start)); ri.Append(ri.Create(OpCodes.Brfalse, bad));
        ri.Append(ri.Create(OpCodes.Ldloc, cursor)); ri.Append(ri.Create(OpCodes.Ldfld, allocated));
        ri.Append(ri.Create(OpCodes.Ldloc, cursor)); ri.Append(ri.Create(OpCodes.Ldfld, committed)); ri.Append(ri.Create(OpCodes.Bgt_Un, bad));
        ri.Append(ri.Create(OpCodes.Ldloc, cursor)); ri.Append(ri.Create(OpCodes.Ldfld, committed));
        ri.Append(ri.Create(OpCodes.Ldloc, cursor)); ri.Append(ri.Create(OpCodes.Ldfld, reserved)); ri.Append(ri.Create(OpCodes.Bgt_Un, bad));
        ri.Append(ri.Create(OpCodes.Ldloc, cursor)); ri.Append(ri.Create(OpCodes.Ldfld, next)); ri.Append(ri.Create(OpCodes.Stloc, cursor));
        ri.Append(test); ri.Append(ri.Create(OpCodes.Brtrue, body)); ri.Append(ri.Create(OpCodes.Ldc_I4_1)); ri.Append(ri.Create(OpCodes.Ret));
        ri.Append(bad); ri.Append(ri.Create(OpCodes.Ret)); registered.Body.MaxStackSize = 2;
        Verify(); return false;

        MethodDefinition Export(string name, string entryPoint)
        {
            var method = new MethodDefinition(name, MethodAttributes.Assembly | MethodAttributes.Static | MethodAttributes.HideBySig, module.TypeSystem.Int32);
            var attribute = new CustomAttribute(module.ImportReference(typeof(System.Runtime.InteropServices.UnmanagedCallersOnlyAttribute).GetConstructor(Type.EmptyTypes)!));
            attribute.Fields.Add(new("EntryPoint", new CustomAttributeArgument(module.TypeSystem.String, entryPoint)));
            method.CustomAttributes.Add(attribute); gc.Methods.Add(method); return method;
        }

        void Verify()
        {
            var validate = gc.Methods.Single(method => method.Name == ValidateName);
            if (validate.Body.Instructions.Count != 17 || validate.Body.Instructions.Count(instruction => instruction.OpCode == OpCodes.Blt_Un) != 2 ||
                validate.Body.Instructions.Count(instruction => instruction.OpCode == OpCodes.Bgt_Un) != 1 || validate.Body.Instructions.Count(instruction => instruction.OpCode == OpCodes.Sub) != 1 ||
                update.Body.Instructions.Count(instruction => instruction.Operand is MethodReference call && call.FullName == validate.FullName) != 1 ||
                update.Body.Instructions.Count(instruction => instruction.OpCode == OpCodes.Sub) != 2 ||
                gc.Methods.Count(method => method.Name is "XfeFrozenRangeValidationExportV1" or "XfeRegisteredFrozenRangesValidV1" or FatalName) != 3 ||
                !gc.Methods.Single(method => method.Name == FatalName).Body.Instructions.Any(instruction => instruction.Operand is MethodReference call && call.Name == "DisableInterruptsScope"))
                throw new InvalidOperationException("Unexpected frozen-segment normalization patch shape.");
            foreach (var field in new[] { allocated, committed })
            {
                var store = update.Body.Instructions.Single(instruction => instruction.OpCode == OpCodes.Stfld && instruction.Operand is FieldReference target && target.FullName == field.FullName);
                if (store.Previous.OpCode != OpCodes.Conv_U || store.Previous.Previous.OpCode != OpCodes.Sub || store.Previous.Previous.Previous.OpCode != OpCodes.Ldarg_0)
                    throw new InvalidOperationException("A frozen size store is missing its start-address subtraction.");
            }
        }
    }
}