AssemblyPatcher.cs 6.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161
  1. using System;
  2. using System.Collections.Generic;
  3. using System.IO;
  4. using System.Reflection;
  5. using BepInEx.Harmony;
  6. using Harmony;
  7. using Mono.Cecil;
  8. namespace BepInEx.Bootstrap
  9. {
  10. /// <summary>
  11. /// Delegate used in patching assemblies.
  12. /// </summary>
  13. /// <param name="assembly">The assembly that is being patched.</param>
  14. public delegate void AssemblyPatcherDelegate(ref AssemblyDefinition assembly);
  15. /// <summary>
  16. /// Worker class which is used for loading and patching entire folders of assemblies, or alternatively patching and loading assemblies one at a time.
  17. /// </summary>
  18. public static class AssemblyPatcher
  19. {
  20. /// <summary>
  21. /// Configuration value of whether assembly dumping is enabled or not.
  22. /// </summary>
  23. private static bool DumpingEnabled => Utility.SafeParseBool(Config.GetEntry("dump-assemblies", "false", "Preloader"));
  24. /// <summary>
  25. /// Patches and loads an entire directory of assemblies.
  26. /// </summary>
  27. /// <param name="directory">The directory to load assemblies from.</param>
  28. /// <param name="patcherMethodDictionary">The dictionary of patchers and their targeted assembly filenames which they are patching.</param>
  29. /// <param name="initializers">List of initializers to run before any patching starts</param>
  30. /// <param name="finalizers">List of finalizers to run before returning</param>
  31. public static void PatchAll(string directory, IDictionary<AssemblyPatcherDelegate, IEnumerable<string>> patcherMethodDictionary, IEnumerable<Action> initializers = null, IEnumerable<Action> finalizers = null)
  32. {
  33. //run all initializers
  34. if (initializers != null)
  35. foreach (Action init in initializers)
  36. init.Invoke();
  37. //load all the requested assemblies
  38. Dictionary<string, AssemblyDefinition> assemblies = new Dictionary<string, AssemblyDefinition>();
  39. foreach (string assemblyPath in Directory.GetFiles(directory, "*.dll"))
  40. {
  41. var assembly = AssemblyDefinition.ReadAssembly(assemblyPath);
  42. //NOTE: this is special cased here because the dependency handling for System.dll is a bit wonky
  43. //System has an assembly reference to itself, and it also has a reference to Mono.Security causing a circular dependency
  44. //It's also generally dangerous to change system.dll since so many things rely on it,
  45. // and it's already loaded into the appdomain since this loader references it, so we might as well skip it
  46. if (assembly.Name.Name == "System"
  47. || assembly.Name.Name == "mscorlib") //mscorlib is already loaded into the appdomain so it can't be patched
  48. {
  49. assembly.Dispose();
  50. continue;
  51. }
  52. assemblies.Add(Path.GetFileName(assemblyPath), assembly);
  53. PatchedAssemblyResolver.AssemblyLocations.Add(assembly.FullName, Path.GetFullPath(assemblyPath));
  54. }
  55. HashSet<string> patchedAssemblies = new HashSet<string>();
  56. //call the patchers on the assemblies
  57. foreach (var patcherMethod in patcherMethodDictionary)
  58. {
  59. foreach (string assemblyFilename in patcherMethod.Value)
  60. {
  61. if (assemblies.TryGetValue(assemblyFilename, out var assembly))
  62. {
  63. Patch(ref assembly, patcherMethod.Key);
  64. assemblies[assemblyFilename] = assembly;
  65. patchedAssemblies.Add(assemblyFilename);
  66. }
  67. }
  68. }
  69. // Finally, load all assemblies into memory
  70. foreach (var kv in assemblies)
  71. {
  72. string filename = kv.Key;
  73. var assembly = kv.Value;
  74. if (DumpingEnabled && patchedAssemblies.Contains(filename))
  75. {
  76. using (MemoryStream mem = new MemoryStream())
  77. {
  78. string dirPath = Path.Combine(Paths.PluginPath, "DumpedAssemblies");
  79. if (!Directory.Exists(dirPath))
  80. Directory.CreateDirectory(dirPath);
  81. assembly.Write(mem);
  82. File.WriteAllBytes(Path.Combine(dirPath, filename), mem.ToArray());
  83. }
  84. }
  85. Load(assembly);
  86. assembly.Dispose();
  87. }
  88. // Patch Assembly.Location and Assembly.CodeBase only if the assemblies were loaded from memory
  89. PatchedAssemblyResolver.ApplyPatch();
  90. //run all finalizers
  91. if (finalizers != null)
  92. foreach (Action finalizer in finalizers)
  93. finalizer.Invoke();
  94. }
  95. /// <summary>
  96. /// Patches an individual assembly, without loading it.
  97. /// </summary>
  98. /// <param name="assembly">The assembly definition to apply the patch to.</param>
  99. /// <param name="patcherMethod">The patcher to use to patch the assembly definition.</param>
  100. public static void Patch(ref AssemblyDefinition assembly, AssemblyPatcherDelegate patcherMethod)
  101. {
  102. patcherMethod.Invoke(ref assembly);
  103. }
  104. /// <summary>
  105. /// Loads an individual assembly defintion into the CLR.
  106. /// </summary>
  107. /// <param name="assembly">The assembly to load.</param>
  108. public static void Load(AssemblyDefinition assembly)
  109. {
  110. using (MemoryStream assemblyStream = new MemoryStream())
  111. {
  112. assembly.Write(assemblyStream);
  113. Assembly.Load(assemblyStream.ToArray());
  114. }
  115. }
  116. }
  117. internal static class PatchedAssemblyResolver
  118. {
  119. public static HarmonyInstance HarmonyInstance { get; } = HarmonyInstance.Create("com.bepis.bepinex.asmlocationfix");
  120. public static Dictionary<string, string> AssemblyLocations { get; } = new Dictionary<string, string>(StringComparer.InvariantCultureIgnoreCase);
  121. public static void ApplyPatch()
  122. {
  123. HarmonyWrapper.PatchAll(typeof(PatchedAssemblyResolver), HarmonyInstance);
  124. }
  125. [HarmonyPostfix, HarmonyPatch(typeof(Assembly), nameof(Assembly.Location), MethodType.Getter)]
  126. public static void GetLocation(ref string __result, Assembly __instance)
  127. {
  128. if (AssemblyLocations.TryGetValue(__instance.FullName, out string location))
  129. __result = location;
  130. }
  131. [HarmonyPostfix, HarmonyPatch(typeof(Assembly), nameof(Assembly.CodeBase), MethodType.Getter)]
  132. public static void GetCodeBase(ref string __result, Assembly __instance)
  133. {
  134. if (AssemblyLocations.TryGetValue(__instance.FullName, out string location))
  135. __result = $"file://{location.Replace('\\', '/')}";
  136. }
  137. }
  138. }