Skip to content

Commit

Permalink
[h2] Patch tool in memory to fix lightmap overwrite bug.
Browse files Browse the repository at this point in the history
  • Loading branch information
num0005 committed Jul 5, 2024
1 parent 9fb4ee6 commit 65f8a97
Show file tree
Hide file tree
Showing 9 changed files with 398 additions and 85 deletions.
15 changes: 11 additions & 4 deletions Launcher/HashHelpers.cs
Original file line number Diff line number Diff line change
Expand Up @@ -50,14 +50,21 @@ private static IEnumerable<string> 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;
}
}
}
2 changes: 2 additions & 0 deletions Launcher/NativeMethods.txt
Original file line number Diff line number Diff line change
@@ -0,0 +1,2 @@
OpenProcess
WriteProcessMemory
266 changes: 200 additions & 66 deletions Launcher/ToolkitInterface/H2Toolkit.cs
Original file line number Diff line number Diff line change
@@ -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
Expand Down Expand Up @@ -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<int>? progress)
private record NopFillFormat(uint BaseAddress, List<uint> CallsToPatch);


readonly static Dictionary<string, NopFillFormat> _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<Task>();
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<H2ToolLightmapFixInjector.NopFill> 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<int>? 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<int, Utility.Process.InjectionConfig> 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<string>(){
"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<string>(){
"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<Task>();
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)
Expand Down
11 changes: 6 additions & 5 deletions Launcher/ToolkitInterface/ToolkitBase.cs
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@
using System.Threading.Tasks;
using System.Xml;
using static ToolkitLauncher.ToolkitProfiles;
using static ToolkitLauncher.Utility.Process;

#nullable enable

Expand Down Expand Up @@ -555,9 +556,9 @@ private string GetLogFileName(List<string>? args)
/// <param name="lowPriority">Lower priority if possible</param>
/// <param name="cancellationToken">Kill the tool before it exits</param>
/// <returns>Results of running the tool if possible</returns>
public async Task<Utility.Process.Result?> RunTool(ToolType tool, List<string>? args = null, OutputMode? outputMode = null, bool lowPriority = false, CancellationToken cancellationToken = default)
public async Task<Utility.Process.Result?> RunTool(ToolType tool, List<string>? 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;
Expand All @@ -566,7 +567,7 @@ private string GetLogFileName(List<string>? args)
/// <summary>
/// Implementation of <c>RunTool</c>
/// </summary>
private async Task<Utility.Process.Result?> RunToolInternal(ToolType tool, List<string>? args, OutputMode? outputMode, bool lowPriority, CancellationToken cancellationToken)
private async Task<Utility.Process.Result?> RunToolInternal(ToolType tool, List<string>? args, OutputMode? outputMode, bool lowPriority, InjectionConfig? injectionOptions, CancellationToken cancellationToken)
{
bool has_window = outputMode != OutputMode.slient && outputMode != OutputMode.logToDisk;
bool enabled_log = outputMode == OutputMode.logToDisk;
Expand All @@ -588,9 +589,9 @@ private string GetLogFileName(List<string>? 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);
}

/// <summary>
Expand Down
4 changes: 4 additions & 0 deletions Launcher/ToolkitLauncher.csproj
Original file line number Diff line number Diff line change
Expand Up @@ -81,6 +81,10 @@
<ItemGroup>
<PackageReference Include="Microsoft.CSharp" Version="4.7.0" />
<PackageReference Include="Microsoft.VisualBasic" Version="10.3.0" />
<PackageReference Include="Microsoft.Windows.CsWin32" Version="0.3.106">
<PrivateAssets>all</PrivateAssets>
<IncludeAssets>runtime; build; native; contentfiles; analyzers; buildtransitive</IncludeAssets>
</PackageReference>
<PackageReference Include="Nerdbank.GitVersioning" Version="3.5.119">
<PrivateAssets>all</PrivateAssets>
<IncludeAssets>runtime; build; native; contentfiles; analyzers; buildtransitive</IncludeAssets>
Expand Down
Loading

0 comments on commit 65f8a97

Please sign in to comment.