From 65f8a97499a5625ebfbe813a001ac5c4b3509acc Mon Sep 17 00:00:00 2001 From: num0005 Date: Fri, 5 Jul 2024 19:26:41 +0100 Subject: [PATCH] [h2] Patch tool in memory to fix lightmap overwrite bug. --- Launcher/HashHelpers.cs | 15 +- Launcher/NativeMethods.txt | 2 + Launcher/ToolkitInterface/H2Toolkit.cs | 266 +++++++++++++----- Launcher/ToolkitInterface/ToolkitBase.cs | 11 +- Launcher/ToolkitLauncher.csproj | 4 + Launcher/Utility/H2ToolLightmapFixInjector.cs | 100 +++++++ Launcher/Utility/IProcessInjector.cs | 27 ++ Launcher/Utility/Process.Windows.cs | 28 +- Launcher/Utility/Process.cs | 30 +- 9 files changed, 398 insertions(+), 85 deletions(-) create mode 100644 Launcher/NativeMethods.txt create mode 100644 Launcher/Utility/H2ToolLightmapFixInjector.cs create mode 100644 Launcher/Utility/IProcessInjector.cs diff --git a/Launcher/HashHelpers.cs b/Launcher/HashHelpers.cs index 85916c3..da36e74 100644 --- a/Launcher/HashHelpers.cs +++ b/Launcher/HashHelpers.cs @@ -50,14 +50,21 @@ private static IEnumerable GetExecutableNames(string directory) List<(string, string)> result = new(); foreach (string fileName in GetExecutableNames(directory)) { - using var md5 = System.Security.Cryptography.MD5.Create(); - using var stream = File.OpenRead(fileName); - string hash = BitConverter.ToString(md5.ComputeHash(stream)).Replace("-", ""); + string hash = GetMD5Hash(fileName); - result.Add((fileName, hash)); + result.Add((fileName, hash)); } return result; } + + public static string GetMD5Hash(string fileName) + { + using var md5 = System.Security.Cryptography.MD5.Create(); + using var stream = File.OpenRead(fileName); + string hash = BitConverter.ToString(md5.ComputeHash(stream)).Replace("-", ""); + + return hash; + } } } diff --git a/Launcher/NativeMethods.txt b/Launcher/NativeMethods.txt new file mode 100644 index 0000000..307614d --- /dev/null +++ b/Launcher/NativeMethods.txt @@ -0,0 +1,2 @@ +OpenProcess +WriteProcessMemory \ No newline at end of file diff --git a/Launcher/ToolkitInterface/H2Toolkit.cs b/Launcher/ToolkitInterface/H2Toolkit.cs index f0d5f73..5756ed3 100644 --- a/Launcher/ToolkitInterface/H2Toolkit.cs +++ b/Launcher/ToolkitInterface/H2Toolkit.cs @@ -1,8 +1,11 @@ using System.Collections.Generic; using System.Diagnostics; using System.IO; +using System.Linq; using System.Threading; using System.Threading.Tasks; +using ToolkitLauncher; +using ToolkitLauncher.Utility; using static ToolkitLauncher.ToolkitProfiles; namespace ToolkitLauncher.ToolkitInterface @@ -54,79 +57,210 @@ private static string GetLightmapQuality(LightmapArgs lightmapArgs) return lightmapArgs.level_combobox.ToLower(); } - public override async Task BuildLightmap(string scenario, string bsp, LightmapArgs args, ICancellableProgress? progress) + private record NopFillFormat(uint BaseAddress, List CallsToPatch); + + + readonly static Dictionary _calls_to_patch_md5 = new() { - LogFolder = $"lightmaps_{Path.GetFileNameWithoutExtension(scenario)}"; - try + // tool regular, latest MCC build + { "C2011BB9B07A7325492D7A804BD939EB", + new NopFillFormat(0x400000, new() {0x4F833F, 0x4F867B}) }, // tag_save lightmap_tag, tag_save scenario_editable + // tool_fast, latest MCC build + { "3A889D370A7BE537AF47FF8035ACD201", + new NopFillFormat(0x400000, new() {0x4ADD50, 0x4ADFF5}) } // tag_save lightmap_tag, tag_save scenario_editable + }; + + private H2ToolLightmapFixInjector? GetInjector(LightmapArgs args) + { + if (!Profile.IsMCC) + return null; + + ToolType tool = args.NoAssert ? ToolType.ToolFast: ToolType.Tool; + + + string tool_Path = GetToolExecutable(tool); + if (!Path.IsPathRooted(tool_Path)) { - string quality = GetLightmapQuality(args); + tool_Path = Path.Join(BaseDirectory, tool_Path); + } - if (args.instanceCount > 1 && (Profile.IsMCC || Profile.CommunityTools)) // multi instance? - { - if (progress is not null) - progress.MaxValue += 1 + args.instanceCount; + string tool_hash = HashHelpers.GetMD5Hash(tool_Path).ToUpper(); - async Task RunInstance(int index) - { - if (index == 0 && !Profile.IsH2Codez()) // not needed for H2Codez - { - if (progress is not null) - progress.Status = "Delaying launch of zeroth instance"; - await Task.Delay(1000 * 70, progress.GetCancellationToken()); - } - Utility.Process.Result result = await RunLightmapWorker( - scenario, - bsp, - quality, - args.instanceCount, - index, - args.NoAssert, - progress.GetCancellationToken(), - args.outputSetting - ); - if (result is not null && result.HasErrorOccured) - progress.Cancel($"Tool worker {index} has failed - exit code {result.ReturnCode}"); - if (progress is not null) - progress.Report(1); - } - - var instances = new List(); - for (int i = args.instanceCount - 1; i >= 0; i--) - { - instances.Add(RunInstance(i)); - } - if (progress is not null) - progress.Status = $"Running {args.instanceCount} instances"; - await Task.WhenAll(instances); - if (progress is not null) - progress.Status = "Merging output"; - - if (progress.IsCancelled) - return; - - await RunMergeLightmap(scenario, bsp, args.instanceCount, args.NoAssert); - if (progress is not null) - progress.Report(1); - } - else - { - Debug.Assert(args.instanceCount == 1); // should be one, otherwise we got bad args - if (progress is not null) - { - progress.DisableCancellation(); - progress.MaxValue += 1; - } - await RunTool((args.NoAssert && Profile.IsMCC) ? ToolType.ToolFast : ToolType.Tool, new() { "lightmaps", scenario, bsp, quality }); - if (progress is not null) - progress.Report(1); - } - } finally + if (_calls_to_patch_md5.ContainsKey(tool_hash)) + { + NopFillFormat config = _calls_to_patch_md5[tool_hash]; + + IEnumerable nopFills = config.CallsToPatch.Select(offset => new H2ToolLightmapFixInjector.NopFill(offset, 5)); + + return new H2ToolLightmapFixInjector(config.BaseAddress, nopFills); + + } + else { - LogFolder = null; + return null; } - } + } + + public override async Task BuildLightmap(string scenario, string bsp, LightmapArgs args, ICancellableProgress? progress) + { + LogFolder = $"lightmaps_{Path.GetFileNameWithoutExtension(scenario)}"; + try + { + string quality = GetLightmapQuality(args); + + if (args.instanceCount > 1 && (Profile.IsMCC || Profile.CommunityTools)) // multi instance? + { + if (progress is not null) + progress.Status = $"Running {args.instanceCount} instances"; + + if (progress is not null) + progress.MaxValue += 1 + args.instanceCount; + + + Utility.H2ToolLightmapFixInjector? injector = null; + Dictionary injectionState = new(); + if (Profile.IsMCC) + injector = GetInjector(args); + + async Task RunInstance(int index) + { + bool delayZerothInstance = true; + if (index == 0 && !Profile.IsH2Codez()) // not needed for H2Codez + { + Trace.WriteLine("Launcher worker zero, checking patch success, etc"); + if (injector is not null) + { + for (int i = 0; i < 20; i++) + { + await Task.Delay(200); + if (injectionState.Values.Any(c => c.Success)) + { + Trace.WriteLine("fix injection succeeded for all processes!"); + delayZerothInstance = false; + break; + } + } + + if (delayZerothInstance) + { + Trace.WriteLine("Failed to inject the fix into some processes"); + Trace.Indent(); + foreach (var entry in injectionState) + { + if (!entry.Value.Success) + Trace.WriteLine($"{entry.Key} worker injection failed"); + } + Trace.Unindent(); + } + + } + + if (delayZerothInstance) + { + Trace.WriteLine("Unable to patch workers, worker zero will be delayed to compensate"); + if (progress is not null) + progress.Status = "Delaying launch of zeroth instance"; + await Task.Delay(1000 * 70, progress.GetCancellationToken()); + progress.Status = $"Running {args.instanceCount} instances"; + } + } + + Utility.Process.Result result; + + LogFileSuffix = $"-{index}"; + if (Profile.IsMCC) + { + bool wereWeExperts = Profile.ElevatedToExpert; + Profile.ElevatedToExpert = true; + + Utility.Process.InjectionConfig? config = null; + if (injector is not null && index != 0) + { + Trace.WriteLine($"Configuring injector for worker {index}"); + config = new(injector); + injectionState[index] = config; + + } + + try + { + result = await RunTool(args.NoAssert ? ToolType.ToolFast : ToolType.Tool, + new List(){ + "lightmaps-farm-worker", + scenario, + bsp, + quality, + index.ToString(), + args.instanceCount.ToString(), + }, + outputMode: args.outputSetting, + lowPriority: index == 0 && delayZerothInstance, + injectionOptions: config, + cancellationToken: progress.GetCancellationToken()); + } + finally + { + Profile.ElevatedToExpert = wereWeExperts; + } + } + else + { + // todo: Remove this code + result = await RunTool(ToolType.Tool, + new List(){ + "lightmaps-slave",// the long legacy of h2codez + scenario, + bsp, + quality, + args.instanceCount.ToString(), + index.ToString() + }, + outputMode: args.outputSetting, + cancellationToken: progress.GetCancellationToken()); + } + + if (result is not null && result.HasErrorOccured) + progress.Cancel($"Tool worker {index} has failed - exit code {result.ReturnCode}"); + if (progress is not null) + progress.Report(1); + } + + var instances = new List(); + for (int i = args.instanceCount - 1; i >= 0; i--) + { + instances.Add(RunInstance(i)); + } + await Task.WhenAll(instances); + if (progress is not null) + progress.Status = "Merging output"; + + if (progress.IsCancelled) + return; + + await RunMergeLightmap(scenario, bsp, args.instanceCount, args.NoAssert); + if (progress is not null) + progress.Report(1); + } + else + { + Debug.Assert(args.instanceCount == 1); // should be one, otherwise we got bad args + if (progress is not null) + { + progress.DisableCancellation(); + progress.MaxValue += 1; + } + await RunTool((args.NoAssert && Profile.IsMCC) ? ToolType.ToolFast : ToolType.Tool, new() { "lightmaps", scenario, bsp, quality }); + if (progress is not null) + progress.Report(1); + } + } + finally + { + LogFolder = null; + } + } - private async Task RunMergeLightmap(string scenario, string bsp, int workerCount, bool useFast) + private async Task RunMergeLightmap(string scenario, string bsp, int workerCount, bool useFast) { if (Profile.IsMCC) diff --git a/Launcher/ToolkitInterface/ToolkitBase.cs b/Launcher/ToolkitInterface/ToolkitBase.cs index ba84cf4..084ccfb 100644 --- a/Launcher/ToolkitInterface/ToolkitBase.cs +++ b/Launcher/ToolkitInterface/ToolkitBase.cs @@ -8,6 +8,7 @@ using System.Threading.Tasks; using System.Xml; using static ToolkitLauncher.ToolkitProfiles; +using static ToolkitLauncher.Utility.Process; #nullable enable @@ -555,9 +556,9 @@ private string GetLogFileName(List? args) /// Lower priority if possible /// Kill the tool before it exits /// Results of running the tool if possible - public async Task RunTool(ToolType tool, List? args = null, OutputMode? outputMode = null, bool lowPriority = false, CancellationToken cancellationToken = default) + public async Task RunTool(ToolType tool, List? args = null, OutputMode? outputMode = null, bool lowPriority = false, InjectionConfig? injectionOptions = null, CancellationToken cancellationToken = default) { - Utility.Process.Result? result = await RunToolInternal(tool, args, outputMode, lowPriority, cancellationToken); + Utility.Process.Result? result = await RunToolInternal(tool, args, outputMode, lowPriority, injectionOptions, cancellationToken); if (result is not null && result.ReturnCode != 0 && ToolFailure is not null) ToolFailure(result); return result; @@ -566,7 +567,7 @@ private string GetLogFileName(List? args) /// /// Implementation of RunTool /// - private async Task RunToolInternal(ToolType tool, List? args, OutputMode? outputMode, bool lowPriority, CancellationToken cancellationToken) + private async Task RunToolInternal(ToolType tool, List? args, OutputMode? outputMode, bool lowPriority, InjectionConfig? injectionOptions, CancellationToken cancellationToken) { bool has_window = outputMode != OutputMode.slient && outputMode != OutputMode.logToDisk; bool enabled_log = outputMode == OutputMode.logToDisk; @@ -588,9 +589,9 @@ private string GetLogFileName(List? args) } if (outputMode == OutputMode.keepOpen) - return await Utility.Process.StartProcessWithShell(BaseDirectory, tool_path, full_args, lowPriority, cancellationToken); + return await Utility.Process.StartProcessWithShell(BaseDirectory, tool_path, full_args, lowPriority, injectionOptions, cancellationToken: cancellationToken); else - return await Utility.Process.StartProcess(BaseDirectory, executable: tool_path, args: full_args, lowPriority: lowPriority, logFileName: log_path, noWindow: !has_window, cancellationToken: cancellationToken); + return await Utility.Process.StartProcess(BaseDirectory, executable: tool_path, args: full_args, lowPriority: lowPriority, logFileName: log_path, noWindow: !has_window, injectionOptions: injectionOptions, cancellationToken: cancellationToken); } /// diff --git a/Launcher/ToolkitLauncher.csproj b/Launcher/ToolkitLauncher.csproj index 82dbf28..d445c3a 100644 --- a/Launcher/ToolkitLauncher.csproj +++ b/Launcher/ToolkitLauncher.csproj @@ -81,6 +81,10 @@ + + all + runtime; build; native; contentfiles; analyzers; buildtransitive + all runtime; build; native; contentfiles; analyzers; buildtransitive diff --git a/Launcher/Utility/H2ToolLightmapFixInjector.cs b/Launcher/Utility/H2ToolLightmapFixInjector.cs new file mode 100644 index 0000000..b13e0cc --- /dev/null +++ b/Launcher/Utility/H2ToolLightmapFixInjector.cs @@ -0,0 +1,100 @@ +using System; +using System.Collections.Generic; +using System.Diagnostics; +using System.Runtime.InteropServices; +using System.Runtime.Versioning; +using System.Threading.Tasks; +using Windows.Win32; +using Windows.Win32.Foundation; + +namespace ToolkitLauncher.Utility +{ + internal class H2ToolLightmapFixInjector : IProcessInjector + { + private static readonly Guid _uuid = Guid.Parse("1B6E9E6C-DBA1-47DA-A7B4-7B940808FCB2"); + + public struct NopFill + { + public NopFill(uint offset, uint length) + { + Offset = offset; + Length = length; + } + + public uint Offset; + public uint Length; + } + + private readonly uint _baseAddress; + private readonly IEnumerable _nopfills; + + public H2ToolLightmapFixInjector(uint baseAddress, IEnumerable nopFills) + { + _baseAddress = baseAddress; + _nopfills = nopFills; + } + + public Guid SetupEnviroment(ProcessStartInfo startInfo) + { + return _uuid; + } + + [SupportedOSPlatform("windows5.1.2600")] + [DllImport("ntdll.dll", PreserveSig = false)] + public static extern void NtSuspendProcess(IntPtr processHandle); + + [SupportedOSPlatform("windows5.1.2600")] + [DllImport("ntdll.dll", PreserveSig = false)] + public static extern void NtResumeProcess(IntPtr processHandle); + + [SupportedOSPlatform("windows5.1.2600")] + public Task Inject(Guid id, System.Diagnostics.Process process) + { + Debug.Assert(id == _uuid); + + // use try-finally to ensure the process is always resumed no matter whatever the patching was sucessful or not + try + { + NtSuspendProcess(process.Handle); + Trace.WriteLine($"[H2 LM Patcher] Target process suspended, ready for patching (last error = {Marshal.GetLastWin32Error()})"); + + ProcessModule mainModule = process.MainModule; + Trace.Assert(mainModule is not null); + + Trace.WriteLine($"[H2 LM Patcher] Patch target: {mainModule.ModuleName} base offset: {mainModule.BaseAddress:X}"); + + foreach (NopFill fill in _nopfills) + { + IntPtr translatedAddress = IntPtr.Add(mainModule.BaseAddress, (int)(fill.Offset - _baseAddress)); + Trace.WriteLine($"Nopfilling {fill.Length} bytes at {fill.Offset:X} (translated address => {translatedAddress:X})"); + + byte[] nopValues = new byte[fill.Length]; + Array.Fill(nopValues, 0x90); + + bool success; + + unsafe + { + fixed (byte* nopVals = nopValues) + success = PInvoke.WriteProcessMemory((HANDLE)process.Handle, translatedAddress.ToPointer(), nopVals, (nuint)nopValues.Length, null); + } + + if (!success) + { + Trace.WriteLine("Failed to write patch to memory!"); + return Task.FromResult(false); + } + } + + Trace.WriteLine($"[H2 LM Patcher] Done patching"); + + return Task.FromResult(true); + } finally + { + NtResumeProcess(process.Handle); + Trace.WriteLine($"[H2 LM Patcher] Process resumed, all done (last error = {Marshal.GetLastWin32Error()})"); + } + + } + } +} diff --git a/Launcher/Utility/IProcessInjector.cs b/Launcher/Utility/IProcessInjector.cs new file mode 100644 index 0000000..553d2d2 --- /dev/null +++ b/Launcher/Utility/IProcessInjector.cs @@ -0,0 +1,27 @@ +using System; +using System.Collections.Generic; +using System.Diagnostics; +using System.Linq; +using System.Text; +using System.Threading.Tasks; + +namespace ToolkitLauncher.Utility +{ + public interface IProcessInjector + { + /// + /// Modify process startup options relating to enviroment. + /// + /// + /// ID of the process, may not be unique if the injector does not need to distinguish between processes + public Guid SetupEnviroment(ProcessStartInfo startInfo); + + /// + /// Inject our changes into a process. + /// + /// UUID returned by SetupEnviroment + /// Process object + /// Was the process sucessfully modified? + public Task Inject(Guid id, System.Diagnostics.Process process); + } +} diff --git a/Launcher/Utility/Process.Windows.cs b/Launcher/Utility/Process.Windows.cs index e94c10b..903447f 100644 --- a/Launcher/Utility/Process.Windows.cs +++ b/Launcher/Utility/Process.Windows.cs @@ -34,7 +34,7 @@ private static ProcessPriorityClass LowerPriority(ProcessPriorityClass old) return ProcessPriorityClass.Idle; } } - static public async Task StartProcess(string directory, string executable, List args, bool lowPriority, bool admin, bool noWindow, string? logFileName, CancellationToken cancellationToken) + static public async Task StartProcess(string directory, string executable, List args, bool lowPriority, bool admin, bool noWindow, string? logFileName, InjectionConfig? injectionOptions, CancellationToken cancellationToken) { try { @@ -43,6 +43,11 @@ static public async Task StartProcess(string directory, string executabl info.WorkingDirectory = directory; info.CreateNoWindow = noWindow; + Guid injector_id = Guid.Empty; + if (injectionOptions is not null) + { + injector_id = injectionOptions.Injector.SetupEnviroment(info); + } bool loggingToDisk = false; if (!String.IsNullOrWhiteSpace(logFileName)) @@ -66,6 +71,12 @@ static public async Task StartProcess(string directory, string executabl } System.Diagnostics.Process proc = System.Diagnostics.Process.Start(info); + + if (injectionOptions is not null) + { + injectionOptions.Success = await injectionOptions.Injector.Inject(injector_id, proc); + } + if (lowPriority) { try @@ -153,7 +164,7 @@ static public async Task StartProcess(string directory, string executabl } } - static public async Task StartProcessWithShell(string directory, string executable, string args, bool lowPriority, CancellationToken cancellationToken) + static public async Task StartProcessWithShell(string directory, string executable, string args, bool lowPriority, InjectionConfig? injectionOptions, CancellationToken cancellationToken) { // build command line string commnad_line = "/c \"" + escape_arg(executable) + " " + args + " & pause\""; @@ -161,6 +172,13 @@ static public async Task StartProcess(string directory, string executabl // run shell process ProcessStartInfo info = new("cmd", commnad_line); info.WorkingDirectory = directory; + + Guid injector_id = Guid.Empty; + if (injectionOptions is not null) + { + injector_id = injectionOptions.Injector.SetupEnviroment(info); + } + System.Diagnostics.Process proc = System.Diagnostics.Process.Start(info); // TODO: find a way to do this without System.Management or P/invoke @@ -177,6 +195,12 @@ async Task HandleProcess(System.Diagnostics.Process process) if (lowPriority) process.PriorityClass = LowerPriority(process.PriorityClass); Trace.WriteLine($"final priority: {process.PriorityClass}"); + + if (injectionOptions is not null) + { + injectionOptions.Success = await injectionOptions.Injector.Inject(injector_id, process); + } + await process.WaitForExitAsync(cancellationToken); } catch (OperationCanceledException) { }; diff --git a/Launcher/Utility/Process.cs b/Launcher/Utility/Process.cs index fd22342..f963db5 100644 --- a/Launcher/Utility/Process.cs +++ b/Launcher/Utility/Process.cs @@ -24,6 +24,18 @@ public bool Success get => ReturnCode == 0; } } + + public record InjectionConfig + { + public InjectionConfig(IProcessInjector injector) + { + Injector = injector; + } + + public IProcessInjector Injector { get; set; } + public bool Success { get; set; } = false; + } + /// /// Run an executable and wait for it to exit /// @@ -33,11 +45,11 @@ public bool Success /// Cancellation token for canceling the process before it exists /// Lower priority if possible /// A task that will complete when the executable exits - static public Task StartProcess(string directory, string executable, List args, bool lowPriority = false, bool admin = false, bool noWindow = false, string? logFileName = null, CancellationToken cancellationToken = default) + static public Task StartProcess(string directory, string executable, List args, bool lowPriority = false, bool admin = false, bool noWindow = false, string? logFileName = null, InjectionConfig? injectionOptions = null, CancellationToken cancellationToken = default) { Trace.WriteLine($"starting(): directory: {directory}, executable:{executable}, args:{args}, admin: {admin}, low priority {lowPriority}, noWindow {noWindow} log {logFileName}"); if (OperatingSystem.IsWindows()) - return Windows.StartProcess(directory, executable, args, lowPriority, admin, noWindow, logFileName, cancellationToken); + return Windows.StartProcess(directory, executable, args, lowPriority, admin, noWindow, logFileName, injectionOptions, cancellationToken); throw new PlatformNotSupportedException(); } @@ -48,13 +60,14 @@ static public Task StartProcess(string directory, string executable, Lis /// unescaped name /// escaped arguments string /// Lower priority if possible - /// Cancellation token for canceling the process before it exists + /// /// A task that will complete when the executable exits - static public Task StartProcessWithShell(string directory, string executable, string args, bool lowPriority = false, CancellationToken cancellationToken = default) + /// Cancellation token for canceling the process before it exists + static public Task StartProcessWithShell(string directory, string executable, string args, bool lowPriority = false, InjectionConfig? injectionOptions = null, CancellationToken cancellationToken = default) { Trace.WriteLine($"starting_with_shell(): directory: {directory}, executable:{executable}, args:{args}"); if (OperatingSystem.IsWindows()) - return Windows.StartProcessWithShell(directory, executable, args, lowPriority, cancellationToken); + return Windows.StartProcessWithShell(directory, executable, args, lowPriority, injectionOptions, cancellationToken); throw new PlatformNotSupportedException(); } @@ -65,11 +78,12 @@ static public Task StartProcess(string directory, string executable, Lis /// unescaped name /// unescaped arguments /// Lower priority if possible - /// Cancellation token for canceling the process before it exists + /// /// A task that will complete when the executable exits - static public async Task StartProcessWithShell(string directory, string executable, List args, bool lowPriority = false, CancellationToken cancellationToken = default) + /// Cancellation token for canceling the process before it exists + static public async Task StartProcessWithShell(string directory, string executable, List args, bool lowPriority = false, InjectionConfig? injectionOptions = null, CancellationToken cancellationToken = default) { - return await StartProcessWithShell(directory, executable, EscapeArgList(args), lowPriority, cancellationToken); + return await StartProcessWithShell(directory, executable, EscapeArgList(args), lowPriority, injectionOptions, cancellationToken: cancellationToken); } ///