XFESystem

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

/// <summary>Reject a null raw stack before the native ThreadContext writes to it.</summary>
internal static class ThreadStackRuntimePatch
{
    public const string GuardName = "XfeRequireThreadStackV1";
    public const string FatalName = "XfeFailThreadStackV1";
    public const string AllocateName = "XfeAllocateThreadStackV1";
    public const string LargestName = "XfeRawLargestFreeRunPagesV1";
    public static bool Apply(AssemblyDefinition assembly)
    {
        var module = assembly.MainModule;
        var thread = module.Types.Single(type => type.FullName == "Cosmos.Kernel.Core.Scheduler.SchedulerThread");
        var initialize = thread.Methods.Single(method => method.Name == "InitializeStack");
        var existing = thread.Methods.SingleOrDefault(method => method.Name == GuardName);
        if (existing is not null)
        {
            if (existing.Body.Instructions.Count != 7 || initialize.Body.Instructions.Count(instruction => instruction.Operand is MethodReference method && method.FullName == existing.FullName) != 1 ||
                thread.Methods.Count(method => method.Name is FatalName or AllocateName or LargestName or "XfeRawLargestFreeRunPagesExportV1" or "XfeRawFreePageCountV1" or "XfeRawTotalPageCountV1") != 6)
                throw new InvalidOperationException("Unexpected raw thread stack patch shape.");
            return true;
        }
        if (initialize.Body.CodeSize != 84 || initialize.Body.Variables.Count != 3 || initialize.Body.ExceptionHandlers.Count != 0 || initialize.Body.Instructions[0].OpCode != OpCodes.Ldarg_0 ||
            initialize.Body.Instructions.Single(instruction => instruction.Offset == 0x0c) is not { OpCode.Code: Code.Call, Operand: MethodReference allocator } raw ||
            allocator.Name != "Alloc" || allocator.DeclaringType.FullName != "Cosmos.Kernel.Core.Memory.MemoryOp" ||
            raw.Next.Operand is not MethodReference store || store.Name != "set_StackBase")
            throw new InvalidOperationException("Pinned raw thread stack allocation changed; refusing to guess.");

        var serial = module.Types.Single(type => type.FullName == "Cosmos.Kernel.Core.IO.Serial");
        var pageAllocator = module.Types.Single(type => type.FullName == "Cosmos.Kernel.Core.Memory.PageAllocator");
        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 free = pageAllocator.Methods.Single(method => method.Name == "get_FreePageCount");
        var total = pageAllocator.Fields.Single(field => field.Name == "TotalPageCount");
        var largest = CreateLargestRun(module, thread, pageAllocator, total);
        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));
        thread.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));
        Text("\nXFE_RUNTIME_STACK_OOM|RequestedBytes|"); fi.Append(fi.Create(OpCodes.Ldarg_0)); Number();
        Text("|RawFreePages|"); fi.Append(fi.Create(OpCodes.Call, free)); Number();
        Text("|RawTotalPages|"); fi.Append(fi.Create(OpCodes.Ldsfld, total)); Number();
        Text("|RawLargestFreeRunPages|"); fi.Append(fi.Create(OpCodes.Call, largest)); Number();
        Text("|Action|GuestHalt\n[FAILFAST] Raw thread stack allocation failed; guest halted.\n");
        var stop = fi.Create(OpCodes.Br, fi.Create(OpCodes.Nop)); stop.Operand = stop; fi.Append(stop); fatal.Body.MaxStackSize = 2;

        var guard = new MethodDefinition(GuardName, MethodAttributes.Assembly | MethodAttributes.Static | MethodAttributes.HideBySig, allocator.ReturnType)
        { ImplAttributes = MethodImplAttributes.AggressiveInlining };
        guard.Parameters.Add(new("result", ParameterAttributes.None, allocator.ReturnType));
        guard.Parameters.Add(new("requestedBytes", ParameterAttributes.None, module.TypeSystem.UInt64));
        thread.Methods.Add(guard);
        var gi = guard.Body.GetILProcessor(); var valid = gi.Create(OpCodes.Ldarg_0); var halted = gi.Create(OpCodes.Br, valid); halted.Operand = halted;
        gi.Append(gi.Create(OpCodes.Ldarg_0)); gi.Append(gi.Create(OpCodes.Brtrue, valid));
        gi.Append(gi.Create(OpCodes.Ldarg_1)); gi.Append(gi.Create(OpCodes.Call, fatal)); gi.Append(halted); gi.Append(valid); gi.Append(gi.Create(OpCodes.Ret)); guard.Body.MaxStackSize = 1;
        var ii = initialize.Body.GetILProcessor(); var size = ii.Create(OpCodes.Ldarg, initialize.Parameters[3]); var convert = ii.Create(OpCodes.Conv_U8); var check = ii.Create(OpCodes.Call, guard);
        ii.InsertAfter(raw, size); ii.InsertAfter(size, convert); ii.InsertAfter(convert, check);
        initialize.Body.MaxStackSize = Math.Max(3, initialize.Body.MaxStackSize);
        var allocate = new MethodDefinition(AllocateName, MethodAttributes.Private | MethodAttributes.Static | MethodAttributes.HideBySig, allocator.ReturnType);
        allocate.Parameters.Add(new("size", ParameterAttributes.None, module.TypeSystem.UInt32));
        allocate.Body.InitLocals = true;
        var result = new VariableDefinition(allocator.ReturnType); allocate.Body.Variables.Add(result); thread.Methods.Add(allocate);
        var ai = allocate.Body.GetILProcessor(); var success = ai.Create(OpCodes.Ldloc, result);
        ai.Append(ai.Create(OpCodes.Ldarg_0)); ai.Append(ai.Create(OpCodes.Call, allocator)); ai.Append(ai.Create(OpCodes.Stloc, result));
        ai.Append(ai.Create(OpCodes.Ldloc, result)); ai.Append(ai.Create(OpCodes.Brtrue, success));
        var gc = module.Types.Single(type => type.FullName == "Cosmos.Kernel.Core.Memory.GarbageCollector.GarbageCollector");
        ai.Append(ai.Create(OpCodes.Call, gc.Methods.Single(method => method.Name == "Collect" && method.Parameters.Count == 0))); ai.Append(ai.Create(OpCodes.Pop));
        ai.Append(ai.Create(OpCodes.Ldarg_0)); ai.Append(ai.Create(OpCodes.Call, allocator)); ai.Append(ai.Create(OpCodes.Stloc, result)); ai.Append(success); ai.Append(ai.Create(OpCodes.Ret)); allocate.Body.MaxStackSize = 1;
        raw.Operand = allocate;
        Export("XfeRawFreePageCountV1", "XfeRuntimeRawFreePageCount", OpCodes.Call, free);
        Export("XfeRawTotalPageCountV1", "XfeRuntimeRawTotalPageCount", OpCodes.Ldsfld, total);
        Export("XfeRawLargestFreeRunPagesExportV1", "XfeRuntimeRawLargestFreeRunPages", OpCodes.Call, largest);
        return false;

        void Text(string text) { fi.Append(fi.Create(OpCodes.Ldstr, text)); fi.Append(fi.Create(OpCodes.Call, write)); }
        void Number() { fi.Append(fi.Create(OpCodes.Ldc_I4_0)); fi.Append(fi.Create(OpCodes.Call, number)); }
        void Export(string name, string entryPoint, OpCode load, MemberReference source)
        {
            var method = new MethodDefinition(name, MethodAttributes.Assembly | MethodAttributes.Static | MethodAttributes.HideBySig, module.TypeSystem.UInt64);
            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); thread.Methods.Add(method);
            var il = method.Body.GetILProcessor();
            il.Append(source is MethodReference call ? il.Create(load, call) : il.Create(load, (FieldReference)source)); il.Append(il.Create(OpCodes.Ret)); method.Body.MaxStackSize = 1;
        }
    }

    private static MethodDefinition CreateLargestRun(ModuleDefinition module, TypeDefinition owner, TypeDefinition allocator, FieldDefinition total)
    {
        var method = new MethodDefinition(LargestName, MethodAttributes.Assembly | MethodAttributes.Static | MethodAttributes.HideBySig, module.TypeSystem.UInt64);
        owner.Methods.Add(method);
        var cpu = module.Types.Single(type => type.FullName == "Cosmos.Kernel.Core.CPU.InternalCpu");
        var disable = cpu.Methods.Single(candidate => candidate.Name == "DisableInterruptsScope" && candidate.Parameters.Count == 0);
        var scopeType = disable.ReturnType.Resolve(); var dispose = scopeType.Methods.Single(candidate => candidate.Name == "Dispose" && candidate.Parameters.Count == 0);
        var ratField = allocator.Fields.Single(field => field.Name == "s_mRAT");
        var scope = new VariableDefinition(disable.ReturnType); var rat = new VariableDefinition(ratField.FieldType);
        var count = new VariableDefinition(module.TypeSystem.UInt64); var index = new VariableDefinition(module.TypeSystem.UInt64);
        var run = new VariableDefinition(module.TypeSystem.UInt64); var best = new VariableDefinition(module.TypeSystem.UInt64);
        foreach (var local in new[] { scope, rat, count, index, run, best }) method.Body.Variables.Add(local);
        method.Body.InitLocals = true;
        var il = method.Body.GetILProcessor(); var finish = il.Create(OpCodes.Ldloca, scope); var test = il.Create(OpCodes.Ldloc, index);
        var body = il.Create(OpCodes.Ldloc, rat); var used = il.Create(OpCodes.Ldc_I4_0); var advance = il.Create(OpCodes.Ldloc, index);
        il.Append(il.Create(OpCodes.Call, disable)); il.Append(il.Create(OpCodes.Stloc, scope));
        il.Append(il.Create(OpCodes.Ldsfld, ratField)); il.Append(il.Create(OpCodes.Stloc, rat));
        il.Append(il.Create(OpCodes.Ldsfld, total)); il.Append(il.Create(OpCodes.Stloc, count));
        il.Append(il.Create(OpCodes.Ldloc, rat)); il.Append(il.Create(OpCodes.Brfalse, finish)); il.Append(il.Create(OpCodes.Br, test));
        il.Append(body); il.Append(il.Create(OpCodes.Ldloc, index)); il.Append(il.Create(OpCodes.Conv_U)); il.Append(il.Create(OpCodes.Add)); il.Append(il.Create(OpCodes.Ldind_U1)); il.Append(il.Create(OpCodes.Brtrue, used));
        il.Append(il.Create(OpCodes.Ldloc, run)); il.Append(il.Create(OpCodes.Ldc_I4_1)); il.Append(il.Create(OpCodes.Conv_U8)); il.Append(il.Create(OpCodes.Add)); il.Append(il.Create(OpCodes.Stloc, run));
        il.Append(il.Create(OpCodes.Ldloc, run)); il.Append(il.Create(OpCodes.Ldloc, best)); il.Append(il.Create(OpCodes.Ble_Un, advance));
        il.Append(il.Create(OpCodes.Ldloc, run)); il.Append(il.Create(OpCodes.Stloc, best)); il.Append(il.Create(OpCodes.Br, advance));
        il.Append(used); il.Append(il.Create(OpCodes.Conv_U8)); il.Append(il.Create(OpCodes.Stloc, run));
        il.Append(advance); il.Append(il.Create(OpCodes.Ldc_I4_1)); il.Append(il.Create(OpCodes.Conv_U8)); il.Append(il.Create(OpCodes.Add)); il.Append(il.Create(OpCodes.Stloc, index));
        il.Append(test); il.Append(il.Create(OpCodes.Ldloc, count)); il.Append(il.Create(OpCodes.Blt_Un, body));
        il.Append(finish); il.Append(il.Create(OpCodes.Call, dispose)); il.Append(il.Create(OpCodes.Ldloc, best)); il.Append(il.Create(OpCodes.Ret)); method.Body.MaxStackSize = 2;
        return method;
    }
}