AssemblyPatcher.cs 6.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189
  1. using System;
  2. using System.Collections.Generic;
  3. using System.IO;
  4. using System.Reflection;
  5. using BepInEx.Configuration;
  6. using BepInEx.Logging;
  7. using BepInEx.Preloader.RuntimeFixes;
  8. using Mono.Cecil;
  9. namespace BepInEx.Preloader.Patching
  10. {
  11. /// <summary>
  12. /// Delegate used in patching assemblies.
  13. /// </summary>
  14. /// <param name="assembly">The assembly that is being patched.</param>
  15. internal delegate void AssemblyPatcherDelegate(ref AssemblyDefinition assembly);
  16. /// <summary>
  17. /// Worker class which is used for loading and patching entire folders of assemblies, or alternatively patching and
  18. /// loading assemblies one at a time.
  19. /// </summary>
  20. internal static class AssemblyPatcher
  21. {
  22. public static List<PatcherPlugin> PatcherPlugins { get; } = new List<PatcherPlugin>();
  23. /// <summary>
  24. /// Adds a single assembly patcher to the pool of applicable patches.
  25. /// </summary>
  26. /// <param name="patcher">Patcher to apply.</param>
  27. public static void AddPatcher(PatcherPlugin patcher)
  28. {
  29. PatcherPlugins.Add(patcher);
  30. }
  31. /// <summary>
  32. /// Adds all patchers from all managed assemblies specified in a directory.
  33. /// </summary>
  34. /// <param name="directory">Directory to search patcher DLLs from.</param>
  35. /// <param name="patcherLocator">A function that locates assembly patchers in a given managed assembly.</param>
  36. public static void AddPatchersFromDirectory(string directory,
  37. Func<Assembly, List<PatcherPlugin>> patcherLocator)
  38. {
  39. if (!Directory.Exists(directory))
  40. return;
  41. var sortedPatchers = new SortedDictionary<string, PatcherPlugin>();
  42. foreach (string assemblyPath in Directory.GetFiles(directory, "*.dll", SearchOption.AllDirectories))
  43. try
  44. {
  45. var assembly = Assembly.LoadFrom(assemblyPath);
  46. foreach (var patcher in patcherLocator(assembly))
  47. sortedPatchers.Add(patcher.Name, patcher);
  48. }
  49. catch (BadImageFormatException) { } //unmanaged DLL
  50. catch (ReflectionTypeLoadException) { } //invalid references
  51. foreach (KeyValuePair<string, PatcherPlugin> patcher in sortedPatchers)
  52. AddPatcher(patcher.Value);
  53. }
  54. private static void InitializePatchers()
  55. {
  56. foreach (var assemblyPatcher in PatcherPlugins)
  57. assemblyPatcher.Initializer?.Invoke();
  58. }
  59. private static void FinalizePatching()
  60. {
  61. foreach (var assemblyPatcher in PatcherPlugins)
  62. assemblyPatcher.Finalizer?.Invoke();
  63. }
  64. /// <summary>
  65. /// Releases all patchers to let them be collected by GC.
  66. /// </summary>
  67. public static void DisposePatchers()
  68. {
  69. PatcherPlugins.Clear();
  70. }
  71. /// <summary>
  72. /// Applies patchers to all assemblies in the given directory and loads patched assemblies into memory.
  73. /// </summary>
  74. /// <param name="directory">Directory to load CLR assemblies from.</param>
  75. public static void PatchAndLoad(string directory)
  76. {
  77. // First, load patchable assemblies into Cecil
  78. var assemblies = new Dictionary<string, AssemblyDefinition>();
  79. foreach (string assemblyPath in Directory.GetFiles(directory, "*.dll"))
  80. {
  81. var assembly = AssemblyDefinition.ReadAssembly(assemblyPath);
  82. //NOTE: this is special cased here because the dependency handling for System.dll is a bit wonky
  83. //System has an assembly reference to itself, and it also has a reference to Mono.Security causing a circular dependency
  84. //It's also generally dangerous to change system.dll since so many things rely on it,
  85. // and it's already loaded into the appdomain since this loader references it, so we might as well skip it
  86. if (assembly.Name.Name == "System"
  87. || assembly.Name.Name == "mscorlib"
  88. ) //mscorlib is already loaded into the appdomain so it can't be patched
  89. {
  90. assembly.Dispose();
  91. continue;
  92. }
  93. if (UnityPatches.AssemblyLocations.ContainsKey(assembly.FullName))
  94. {
  95. Logger.LogWarning($"Tried to load duplicate assembly {Path.GetFileName(assemblyPath)} from Managed folder! Skipping...");
  96. continue;
  97. }
  98. assemblies.Add(Path.GetFileName(assemblyPath), assembly);
  99. UnityPatches.AssemblyLocations.Add(assembly.FullName, Path.GetFullPath(assemblyPath));
  100. }
  101. // Next, initialize all the patchers
  102. InitializePatchers();
  103. // Then, perform the actual patching
  104. var patchedAssemblies = new HashSet<string>();
  105. foreach (var assemblyPatcher in PatcherPlugins)
  106. foreach (string targetDll in assemblyPatcher.TargetDLLs)
  107. if (assemblies.TryGetValue(targetDll, out var assembly))
  108. {
  109. Logger.LogInfo($"Patching [{assembly.Name.Name}] with [{assemblyPatcher.Name}]");
  110. assemblyPatcher.Patcher?.Invoke(ref assembly);
  111. assemblies[targetDll] = assembly;
  112. patchedAssemblies.Add(targetDll);
  113. }
  114. // Finally, load patched assemblies into memory
  115. foreach (KeyValuePair<string, AssemblyDefinition> kv in assemblies)
  116. {
  117. string filename = kv.Key;
  118. var assembly = kv.Value;
  119. if (ConfigDumpAssemblies.Value && patchedAssemblies.Contains(filename))
  120. using (var mem = new MemoryStream())
  121. {
  122. string dirPath = Path.Combine(Paths.BepInExRootPath, "DumpedAssemblies");
  123. if (!Directory.Exists(dirPath))
  124. Directory.CreateDirectory(dirPath);
  125. assembly.Write(mem);
  126. File.WriteAllBytes(Path.Combine(dirPath, filename), mem.ToArray());
  127. }
  128. // Note that since we only *load* assemblies, they shouldn't trigger dependency loading
  129. // Not loading all assemblies is very important not only because of memory reasons,
  130. // but because some games *rely* on that because of messed up internal dependencies.
  131. if (patchedAssemblies.Contains(filename))
  132. Load(assembly);
  133. // Though we have to dispose of all assemblies regardless of them being patched or not
  134. assembly.Dispose();
  135. }
  136. //run all finalizers
  137. FinalizePatching();
  138. }
  139. /// <summary>
  140. /// Loads an individual assembly definition into the CLR.
  141. /// </summary>
  142. /// <param name="assembly">The assembly to load.</param>
  143. public static void Load(AssemblyDefinition assembly)
  144. {
  145. using (var assemblyStream = new MemoryStream())
  146. {
  147. assembly.Write(assemblyStream);
  148. Assembly.Load(assemblyStream.ToArray());
  149. }
  150. }
  151. #region Config
  152. private static readonly ConfigWrapper<bool> ConfigDumpAssemblies = ConfigFile.CoreConfig.Wrap(
  153. "Preloader",
  154. "DumpAssemblies",
  155. "If enabled, BepInEx will save patched assemblies into BepInEx/DumpedAssemblies.\nThis can be used by developers to inspect and debug preloader patchers.",
  156. false);
  157. #endregion
  158. }
  159. }