AssemblyPatcher.cs 16 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418
  1. using System;
  2. using System.Collections.Generic;
  3. using System.Diagnostics;
  4. using System.IO;
  5. using System.Linq;
  6. using System.Reflection;
  7. using System.Text;
  8. using BepInEx.Bootstrap;
  9. using BepInEx.Configuration;
  10. using BepInEx.Logging;
  11. using Mono.Cecil;
  12. namespace BepInEx.Preloader.Core
  13. {
  14. /// <summary>
  15. /// Delegate used in patching assemblies.
  16. /// </summary>
  17. /// <param name="assembly">The assembly that is being patched.</param>
  18. public delegate void AssemblyPatcherDelegate(ref AssemblyDefinition assembly);
  19. /// <summary>
  20. /// Worker class which is used for loading and patching entire folders of assemblies, or alternatively patching and
  21. /// loading assemblies one at a time.
  22. /// </summary>
  23. public class AssemblyPatcher : IDisposable
  24. {
  25. private const BindingFlags ALL = BindingFlags.Public | BindingFlags.NonPublic | BindingFlags.Static | BindingFlags.IgnoreCase;
  26. /// <summary>
  27. /// A list of plugins that will be initialized and executed, in the order of the list.
  28. /// </summary>
  29. public List<PatcherPlugin> PatcherPlugins { get; } = new List<PatcherPlugin>();
  30. /// <summary>
  31. /// A cloned version of <see cref="PatcherPlugins"/> to ensure that any foreach loops do not break when the collection gets modified.
  32. /// </summary>
  33. private IEnumerable<PatcherPlugin> PatcherPluginsSafe => PatcherPlugins.ToList();
  34. /// <summary>
  35. /// <para>Contains a list of assemblies that will be patched and loaded into the runtime.</para>
  36. /// <para>The dictionary has the name of the file, without any directories. These are used by the dumping functionality, and as such, these are also required to be unique. They do not have to be exactly the same as the real filename, however they have to be mapped deterministically.</para>
  37. /// <para>Order is not respected, as it will be sorted by dependencies.</para>
  38. /// </summary>
  39. public Dictionary<string, AssemblyDefinition> AssembliesToPatch { get; } = new Dictionary<string, AssemblyDefinition>();
  40. /// <summary>
  41. /// <para>Contains a dictionary of assemblies that have been loaded as part of executing this assembly patcher..</para>
  42. /// <para>The key is the same key as used in <see cref="LoadedAssemblies"/>, while the value is the actual assembly itself.</para>
  43. /// </summary>
  44. public Dictionary<string, Assembly> LoadedAssemblies { get; } = new Dictionary<string, Assembly>();
  45. /// <summary>
  46. /// The directory location as to where patched assemblies will be saved to and loaded from disk, for debugging purposes. Defaults to BepInEx/DumpedAssemblies
  47. /// </summary>
  48. public string DumpedAssembliesPath { get; set; } = Path.Combine(Paths.BepInExRootPath, "DumpedAssemblies");
  49. public ManualLogSource Logger { get; } = BepInEx.Logging.Logger.CreateLogSource("AssemblyPatcher");
  50. private static T CreateDelegate<T>(MethodInfo method) where T : class => method != null ? Delegate.CreateDelegate(typeof(T), method) as T : null;
  51. private static PatcherPlugin ToPatcherPlugin(TypeDefinition type, string assemblyPath)
  52. {
  53. if (type.IsInterface || type.IsAbstract && !type.IsSealed)
  54. return null;
  55. var targetDlls = type.Methods.FirstOrDefault(m => m.Name.Equals("get_TargetDLLs", StringComparison.InvariantCultureIgnoreCase) &&
  56. m.IsPublic &&
  57. m.IsStatic);
  58. if (targetDlls == null ||
  59. targetDlls.ReturnType.FullName != "System.Collections.Generic.IEnumerable`1<System.String>")
  60. return null;
  61. var patch = type.Methods.FirstOrDefault(m => m.Name.Equals("Patch") &&
  62. m.IsPublic &&
  63. m.IsStatic &&
  64. m.ReturnType.FullName == "System.Void" &&
  65. m.Parameters.Count == 1 &&
  66. (m.Parameters[0].ParameterType.FullName == "Mono.Cecil.AssemblyDefinition&" ||
  67. m.Parameters[0].ParameterType.FullName == "Mono.Cecil.AssemblyDefinition"));
  68. if (patch == null)
  69. return null;
  70. return new PatcherPlugin
  71. {
  72. TypeName = type.FullName
  73. };
  74. }
  75. /// <summary>
  76. /// Adds all patchers from all managed assemblies specified in a directory.
  77. /// </summary>
  78. /// <param name="directory">Directory to search patcher DLLs from.</param>
  79. public void AddPatchersFromDirectory(string directory)
  80. {
  81. if (!Directory.Exists(directory))
  82. return;
  83. var sortedPatchers = new SortedDictionary<string, PatcherPlugin>();
  84. var patchers = TypeLoader.FindPluginTypes(directory, ToPatcherPlugin);
  85. foreach (var keyValuePair in patchers)
  86. {
  87. var assemblyPath = keyValuePair.Key;
  88. var patcherCollection = keyValuePair.Value;
  89. if(patcherCollection.Count == 0)
  90. continue;
  91. var ass = Assembly.LoadFile(assemblyPath);
  92. foreach (var patcherPlugin in patcherCollection)
  93. {
  94. try
  95. {
  96. var type = ass.GetType(patcherPlugin.TypeName);
  97. var methods = type.GetMethods(ALL);
  98. patcherPlugin.Initializer = CreateDelegate<Action>(methods.FirstOrDefault(m => m.Name.Equals("Initialize", StringComparison.InvariantCultureIgnoreCase) &&
  99. m.GetParameters().Length == 0 &&
  100. m.ReturnType == typeof(void)));
  101. patcherPlugin.Finalizer = CreateDelegate<Action>(methods.FirstOrDefault(m => m.Name.Equals("Finish", StringComparison.InvariantCultureIgnoreCase) &&
  102. m.GetParameters().Length == 0 &&
  103. m.ReturnType == typeof(void)));
  104. patcherPlugin.TargetDLLs = CreateDelegate<Func<IEnumerable<string>>>(type.GetProperty("TargetDLLs", ALL).GetGetMethod());
  105. var patcher = methods.FirstOrDefault(m => m.Name.Equals("Patch", StringComparison.CurrentCultureIgnoreCase) &&
  106. m.ReturnType == typeof(void) &&
  107. m.GetParameters().Length == 1 &&
  108. (m.GetParameters()[0].ParameterType == typeof(AssemblyDefinition) ||
  109. m.GetParameters()[0].ParameterType == typeof(AssemblyDefinition).MakeByRefType()));
  110. patcherPlugin.Patcher = (ref AssemblyDefinition pAss) =>
  111. {
  112. //we do the array fuckery here to get the ref result out
  113. object[] args = { pAss };
  114. patcher.Invoke(null, args);
  115. pAss = (AssemblyDefinition)args[0];
  116. };
  117. sortedPatchers.Add($"{ass.GetName().Name}/{type.FullName}", patcherPlugin);
  118. }
  119. catch (Exception e)
  120. {
  121. Logger.LogError($"Failed to load patcher [{patcherPlugin.TypeName}]: {e.Message}");
  122. if (e is ReflectionTypeLoadException re)
  123. Logger.LogDebug(TypeLoader.TypeLoadExceptionToString(re));
  124. else
  125. Logger.LogDebug(e.ToString());
  126. }
  127. }
  128. var assName = ass.GetName();
  129. Logger.Log(patcherCollection.Any() ? LogLevel.Info : LogLevel.Debug,
  130. $"Loaded {patcherCollection.Count} patcher method{(patcherCollection.Count == 1 ? "" : "s")} from [{assName.Name} {assName.Version}]");
  131. }
  132. foreach (KeyValuePair<string, PatcherPlugin> patcher in sortedPatchers)
  133. PatcherPlugins.Add(patcher.Value);
  134. }
  135. /// <summary>
  136. /// Adds all .dll assemblies in a directory to be patched and loaded by this patcher instance. Non-managed assemblies are skipped.
  137. /// </summary>
  138. /// <param name="directory">The directory to search.</param>
  139. public void LoadAssemblyDirectory(string directory)
  140. {
  141. LoadAssemblyDirectory(directory, "dll");
  142. }
  143. /// <summary>
  144. /// Adds all assemblies in a directory to be patched and loaded by this patcher instance. Non-managed assemblies are skipped.
  145. /// </summary>
  146. /// <param name="directory">The directory to search.</param>
  147. /// <param name="assemblyExtensions">The file extensions to attempt to load.</param>
  148. public void LoadAssemblyDirectory(string directory, params string[] assemblyExtensions)
  149. {
  150. var filesToSearch = assemblyExtensions
  151. .SelectMany(ext => Directory.GetFiles(directory, "*." + ext, SearchOption.TopDirectoryOnly));
  152. foreach (string assemblyPath in filesToSearch)
  153. {
  154. if (!TryLoadAssembly(assemblyPath, out var assembly))
  155. continue;
  156. // NOTE: this is special cased here because the dependency handling for System.dll is a bit wonky
  157. // System has an assembly reference to itself, and it also has a reference to Mono.Security causing a circular dependency
  158. // It's also generally dangerous to change system.dll since so many things rely on it,
  159. // and it's already loaded into the appdomain since this loader references it, so we might as well skip it
  160. if (assembly.Name.Name == "System" || assembly.Name.Name == "mscorlib") //mscorlib is already loaded into the appdomain so it can't be patched
  161. {
  162. assembly.Dispose();
  163. continue;
  164. }
  165. AssembliesToPatch.Add(Path.GetFileName(assemblyPath), assembly);
  166. Logger.LogDebug($"Assembly loaded: {Path.GetFileName(assemblyPath)}");
  167. //if (UnityPatches.AssemblyLocations.ContainsKey(assembly.FullName))
  168. //{
  169. // Logger.LogWarning($"Tried to load duplicate assembly {Path.GetFileName(assemblyPath)} from Managed folder! Skipping...");
  170. // continue;
  171. //}
  172. //assemblies.Add(Path.GetFileName(assemblyPath), assembly);
  173. //UnityPatches.AssemblyLocations.Add(assembly.FullName, Path.GetFullPath(assemblyPath));
  174. }
  175. }
  176. /// <summary>
  177. /// Attempts to load a managed assembly as an <see cref="AssemblyDefinition"/>. Returns true if successful.
  178. /// </summary>
  179. /// <param name="path">The path of the assembly.</param>
  180. /// <param name="assembly">The loaded assembly. Null if not successful in loading.</param>
  181. public static bool TryLoadAssembly(string path, out AssemblyDefinition assembly)
  182. {
  183. try
  184. {
  185. assembly = AssemblyDefinition.ReadAssembly(path);
  186. return true;
  187. }
  188. catch (BadImageFormatException)
  189. {
  190. // Not a managed assembly
  191. assembly = null;
  192. return false;
  193. }
  194. }
  195. /// <summary>
  196. /// Performs work to dispose collection objects.
  197. /// </summary>
  198. public void Dispose()
  199. {
  200. foreach (var assembly in AssembliesToPatch)
  201. assembly.Value.Dispose();
  202. AssembliesToPatch.Clear();
  203. // Clear to allow GC collection.
  204. PatcherPlugins.Clear();
  205. }
  206. /// <summary>
  207. /// Applies patchers to all assemblies in the given directory and loads patched assemblies into memory.
  208. /// </summary>
  209. /// <param name="directory">Directory to load CLR assemblies from.</param>
  210. public void PatchAndLoad()
  211. {
  212. // First, create a copy of the assembly dictionary as the initializer can change them
  213. var assemblies = new Dictionary<string, AssemblyDefinition>(AssembliesToPatch, StringComparer.InvariantCultureIgnoreCase);
  214. // Next, initialize all the patchers
  215. foreach (var assemblyPatcher in PatcherPluginsSafe)
  216. {
  217. try
  218. {
  219. assemblyPatcher.Initializer?.Invoke();
  220. }
  221. catch (Exception ex)
  222. {
  223. Logger.LogError($"Failed to run initializer of {assemblyPatcher.TypeName}: {ex}");
  224. }
  225. }
  226. // Then, perform the actual patching
  227. var patchedAssemblies = new HashSet<string>(StringComparer.InvariantCultureIgnoreCase);
  228. var resolvedAssemblies = new Dictionary<string, string>();
  229. // TODO: Maybe instead reload the assembly and repatch with other valid patchers?
  230. var invalidAssemblies = new HashSet<string>(StringComparer.InvariantCultureIgnoreCase);
  231. foreach (var assemblyPatcher in PatcherPluginsSafe)
  232. foreach (string targetDll in assemblyPatcher.TargetDLLs())
  233. if (AssembliesToPatch.TryGetValue(targetDll, out var assembly) && !invalidAssemblies.Contains(targetDll))
  234. {
  235. Logger.LogInfo($"Patching [{assembly.Name.Name}] with [{assemblyPatcher.TypeName}]");
  236. try
  237. {
  238. assemblyPatcher.Patcher?.Invoke(ref assembly);
  239. }
  240. catch (Exception e)
  241. {
  242. Logger.LogError($"Failed to run [{assemblyPatcher.TypeName}] when patching [{assembly.Name.Name}]. This assembly will not be patched. Error: {e}");
  243. patchedAssemblies.Remove(targetDll);
  244. invalidAssemblies.Add(targetDll);
  245. continue;
  246. }
  247. AssembliesToPatch[targetDll] = assembly;
  248. patchedAssemblies.Add(targetDll);
  249. foreach (var resolvedAss in AppDomain.CurrentDomain.GetAssemblies())
  250. {
  251. var name = Utility.TryParseAssemblyName(resolvedAss.FullName, out var assName) ? assName.Name : resolvedAss.FullName;
  252. // Report only the first type that caused the assembly to load, because any subsequent ones can be false positives
  253. if (!resolvedAssemblies.ContainsKey(name))
  254. resolvedAssemblies[name] = assemblyPatcher.TypeName;
  255. }
  256. }
  257. // Check if any patched assemblies have been already resolved by the CLR
  258. // If there are any, they cannot be loaded by the preloader
  259. var patchedAssemblyNames = new HashSet<string>(assemblies.Where(kv => patchedAssemblies.Contains(kv.Key)).Select(kv => kv.Value.Name.Name), StringComparer.InvariantCultureIgnoreCase);
  260. var earlyLoadAssemblies = resolvedAssemblies.Where(kv => patchedAssemblyNames.Contains(kv.Key)).ToList();
  261. if (earlyLoadAssemblies.Count != 0)
  262. {
  263. Logger.LogWarning(new StringBuilder()
  264. .AppendLine("The following assemblies have been loaded too early and will not be patched by preloader:")
  265. .AppendLine(string.Join(Environment.NewLine, earlyLoadAssemblies.Select(kv => $"* [{kv.Key}] (first loaded by [{kv.Value}])").ToArray()))
  266. .AppendLine("Expect unexpected behavior and issues with plugins and patchers not being loaded.")
  267. .ToString());
  268. }
  269. // Finally, load patched assemblies into memory
  270. if (ConfigDumpAssemblies.Value || ConfigLoadDumpedAssemblies.Value)
  271. {
  272. if (!Directory.Exists(DumpedAssembliesPath))
  273. Directory.CreateDirectory(DumpedAssembliesPath);
  274. foreach (KeyValuePair<string, AssemblyDefinition> kv in assemblies)
  275. {
  276. string filename = kv.Key;
  277. var assembly = kv.Value;
  278. if (patchedAssemblies.Contains(filename))
  279. assembly.Write(Path.Combine(DumpedAssembliesPath, filename));
  280. }
  281. }
  282. if (ConfigBreakBeforeLoadAssemblies.Value)
  283. {
  284. Logger.LogInfo($"BepInEx is about load the following assemblies:\n{String.Join("\n", patchedAssemblies.ToArray())}");
  285. Logger.LogInfo($"The assemblies were dumped into {DumpedAssembliesPath}");
  286. Logger.LogInfo("Load any assemblies into the debugger, set breakpoints and continue execution.");
  287. Debugger.Break();
  288. }
  289. foreach (var kv in assemblies)
  290. {
  291. string filename = kv.Key;
  292. var assembly = kv.Value;
  293. // Note that since we only *load* assemblies, they shouldn't trigger dependency loading
  294. // Not loading all assemblies is very important not only because of memory reasons,
  295. // but because some games *rely* on that because of messed up internal dependencies.
  296. if (patchedAssemblies.Contains(filename))
  297. {
  298. Assembly loadedAssembly;
  299. if (ConfigLoadDumpedAssemblies.Value)
  300. loadedAssembly = Assembly.LoadFile(Path.Combine(DumpedAssembliesPath, filename));
  301. else
  302. {
  303. using (var assemblyStream = new MemoryStream())
  304. {
  305. assembly.Write(assemblyStream);
  306. loadedAssembly =Assembly.Load(assemblyStream.ToArray());
  307. }
  308. }
  309. LoadedAssemblies.Add(filename, loadedAssembly);
  310. Logger.LogDebug($"Loaded '{assembly.FullName}' into memory");
  311. }
  312. // Though we have to dispose of all assemblies regardless of them being patched or not
  313. assembly.Dispose();
  314. }
  315. // Finally, run all finalizers
  316. foreach (var assemblyPatcher in PatcherPluginsSafe)
  317. {
  318. try
  319. {
  320. assemblyPatcher.Finalizer?.Invoke();
  321. }
  322. catch (Exception ex)
  323. {
  324. Logger.LogError($"Failed to run finalizer of {assemblyPatcher.TypeName}: {ex}");
  325. }
  326. }
  327. }
  328. #region Config
  329. private static readonly ConfigEntry<bool> ConfigDumpAssemblies = ConfigFile.CoreConfig.Bind(
  330. "Preloader", "DumpAssemblies",
  331. false,
  332. "If enabled, BepInEx will save patched assemblies into BepInEx/DumpedAssemblies.\nThis can be used by developers to inspect and debug preloader patchers.");
  333. private static readonly ConfigEntry<bool> ConfigLoadDumpedAssemblies = ConfigFile.CoreConfig.Bind(
  334. "Preloader", "LoadDumpedAssemblies",
  335. false,
  336. "If enabled, BepInEx will load patched assemblies from BepInEx/DumpedAssemblies instead of memory.\nThis can be used to be able to load patched assemblies into debuggers like dnSpy.\nIf set to true, will override DumpAssemblies.");
  337. private static readonly ConfigEntry<bool> ConfigBreakBeforeLoadAssemblies = ConfigFile.CoreConfig.Bind(
  338. "Preloader", "BreakBeforeLoadAssemblies",
  339. false,
  340. "If enabled, BepInEx will call Debugger.Break() once before loading patched assemblies.\nThis can be used with debuggers like dnSpy to install breakpoints into patched assemblies before they are loaded.");
  341. #endregion
  342. }
  343. }