UnityMono: Resolve assemblies into memory using mono API

This fixes several issues with loading assemblies. In turn, this commit also removes old Location/CodeBase fixes since now assemblies are resolved automatically correctly.

Additionally, integrate the mono loader into preloading which fixes assembly location resolving

BepInEx.Core/Utility.cs

@@ -176,10 +176,10 @@ public static class Utility
     /// <param name="directory">Directory to search the assembly from.</param>
     /// <param name="assembly">The loaded assembly.</param>
     /// <returns>True, if the assembly was found and loaded. Otherwise, false.</returns>
-    private static bool TryResolveDllAssembly<T>(AssemblyName assemblyName,
-                                                 string directory,
-                                                 Func<string, T> loader,
-                                                 out T assembly) where T : class
+    public static bool TryResolveDllAssembly<T>(AssemblyName assemblyName,
+                                                string directory,
+                                                Func<string, T> loader,
+                                                out T assembly) where T : class
     {
         assembly = null;
 

BepInEx.IL2CPP/Preloader.cs

@@ -62,7 +62,7 @@ public static class Preloader
             InitializeUnityVersion();
             Il2CppInteropManager.Initialize(UnityVersion);
 
-            using (var assemblyPatcher = new AssemblyPatcher())
+            using (var assemblyPatcher = new AssemblyPatcher((data, _) => Assembly.Load(data)))
             {
                 assemblyPatcher.AddPatchersFromDirectory(Paths.PatcherPluginPath);
 
@@ -90,7 +90,7 @@ public static class Preloader
             throw;
         }
     }
-    
+
     private static IntPtr DllImportResolver(string libraryName, Assembly assembly, DllImportSearchPath? searchPath)
     {
         if (libraryName == "GameAssembly")

BepInEx.NetCore/NetCorePreloader.cs

@@ -1,5 +1,7 @@
 using System;
 using System.Diagnostics;
+using System.IO;
+using System.Reflection;
 using BepInEx.Bootstrap;
 using BepInEx.Logging;
 using BepInEx.NetLauncher.Common;
@@ -50,7 +52,7 @@ namespace BepInEx.NetCore
 
             Log.LogMessage("Preloader started");
 
-            using (var assemblyPatcher = new AssemblyPatcher())
+            using (var assemblyPatcher = new AssemblyPatcher((data, _) => Assembly.Load(data)))
             {
                 assemblyPatcher.AddPatchersFromDirectory(Paths.PatcherPluginPath);
 

BepInEx.NetLauncher/NetPreloader.cs

@@ -95,7 +95,8 @@ public static class NetPreloader
             Log.LogError($"Game executable is {(executableInfo.Is64Bit ? "64" : "32")}-bit while BepInEx has been compiled as {(launcherInfo.Is64Bit ? "64" : "32")}-bit. Expect crashes");
         }
 
-        if (executableInfo.NetFrameworkVersion != launcherInfo.NetFrameworkVersion || executableInfo.AssemblyFrameworkType != launcherInfo.AssemblyFrameworkType)
+        if (executableInfo.NetFrameworkVersion != launcherInfo.NetFrameworkVersion ||
+            executableInfo.AssemblyFrameworkType != launcherInfo.AssemblyFrameworkType)
         {
             Log.LogWarning($"Game executable has been compiled as {executableInfo}, while BepInEx has been compiled as {launcherInfo}. There may be issues within the game caused by this");
         }
@@ -107,7 +108,7 @@ public static class NetPreloader
 
         Assembly entrypointAssembly;
 
-        using (var assemblyPatcher = new AssemblyPatcher())
+        using (var assemblyPatcher = new AssemblyPatcher((data, _) => Assembly.Load(data)))
         {
             assemblyPatcher.AddPatchersFromDirectory(Paths.PatcherPluginPath);
 

BepInEx.Preloader.Core/Patching/AssemblyPatcher.cs

@@ -22,6 +22,13 @@ public class AssemblyPatcher : IDisposable
 {
     private static readonly string CurrentAssemblyName = Assembly.GetExecutingAssembly().GetName().Name;
 
+    private Func<byte[], string, Assembly> assemblyLoader;
+
+    public AssemblyPatcher(Func<byte[], string, Assembly> assemblyLoader)
+    {
+        this.assemblyLoader = assemblyLoader;
+    }
+
     /// <summary>
     ///     The context of this assembly patcher instance that is passed to all patcher plugins.
     /// </summary>
@@ -50,6 +57,8 @@ public class AssemblyPatcher : IDisposable
 
         PatcherContext.AvailableAssemblies.Clear();
 
+        PatcherContext.AvailableAssembliesPaths.Clear();
+
         // Clear to allow GC collection.
         PatcherContext.PatcherPlugins.Clear();
     }
@@ -251,7 +260,9 @@ public class AssemblyPatcher : IDisposable
                 continue;
             }
 
-            PatcherContext.AvailableAssemblies.Add(Path.GetFileName(assemblyPath), assembly);
+            var fileName = Path.GetFileName(assemblyPath);
+            PatcherContext.AvailableAssemblies.Add(fileName, assembly);
+            PatcherContext.AvailableAssembliesPaths.Add(fileName, assemblyPath);
 
             Logger.LogDebug($"Assembly loaded: {Path.GetFileName(assemblyPath)}");
 
@@ -491,7 +502,8 @@ public class AssemblyPatcher : IDisposable
                 {
                     using var assemblyStream = new MemoryStream();
                     assembly.Write(assemblyStream);
-                    loadedAssembly = Assembly.Load(assemblyStream.ToArray());
+                    loadedAssembly = assemblyLoader(assemblyStream.ToArray(),
+                                                    PatcherContext.AvailableAssembliesPaths[filename]);
                 }
 
                 PatcherContext.LoadedAssemblies.Add(filename, loadedAssembly);

BepInEx.Preloader.Core/Patching/PatcherContext.cs

@@ -70,6 +70,11 @@ public class PatcherContext
     /// </summary>
     public Dictionary<string, AssemblyDefinition> AvailableAssemblies { get; } = new();
 
+    /// <summary>
+    ///     <para>Contains a mapping of available assembly name to their original filenames.</para>
+    /// </summary>
+    public Dictionary<string, string> AvailableAssembliesPaths { get; } = new();
+
     /// <summary>
     ///     <para>Contains a dictionary of assemblies that have been loaded as part of executing this assembly patcher.</para>
     ///     <para>

BepInEx.Preloader.UnityMono/RuntimeFixes/UnityPatches.cs

@@ -1,6 +1,6 @@
-using System;
-using System.Collections.Generic;
+using System.IO;
 using System.Reflection;
+using BepInEx.Preloader.UnityMono.Utils;
 using HarmonyLib;
 
 namespace BepInEx.Preloader.UnityMono.RuntimeFixes;
@@ -9,9 +9,6 @@ internal static class UnityPatches
 {
     private static Harmony HarmonyInstance { get; set; }
 
-    public static Dictionary<string, string> AssemblyLocations { get; } =
-        new(StringComparer.InvariantCultureIgnoreCase);
-
     public static void Apply()
     {
         HarmonyInstance = Harmony.CreateAndPatchAll(typeof(UnityPatches));
@@ -23,19 +20,14 @@ internal static class UnityPatches
         catch { } //ignore everything, if it's thrown an exception, we're using an assembly that has already fixed this
     }
 
-    [HarmonyPostfix]
-    [HarmonyPatch(typeof(Assembly), nameof(Assembly.Location), MethodType.Getter)]
-    public static void GetLocation(ref string __result, Assembly __instance)
+    [HarmonyPrefix]
+    [HarmonyPatch(typeof(Assembly), nameof(Assembly.LoadFile), typeof(string))]
+    [HarmonyPatch(typeof(Assembly), nameof(Assembly.LoadFrom), typeof(string))]
+    public static bool LoadFilePrefix(string __0, ref Assembly __result)
     {
-        if (AssemblyLocations.TryGetValue(__instance.FullName, out var location))
-            __result = location;
-    }
-
-    [HarmonyPostfix]
-    [HarmonyPatch(typeof(Assembly), nameof(Assembly.CodeBase), MethodType.Getter)]
-    public static void GetCodeBase(ref string __result, Assembly __instance)
-    {
-        if (AssemblyLocations.TryGetValue(__instance.FullName, out var location))
-            __result = $"file://{location.Replace('\\', '/')}";
+        if (!File.Exists(__0))
+            throw new FileNotFoundException(__0);
+        __result = MonoAssemblyHelper.Load(__0);
+        return false;
     }
 }

BepInEx.Preloader.UnityMono/UnityPreloader.cs

@@ -10,6 +10,7 @@ using BepInEx.Preloader.Core.Logging;
 using BepInEx.Preloader.Core.Patching;
 using BepInEx.Preloader.RuntimeFixes;
 using BepInEx.Preloader.UnityMono.RuntimeFixes;
+using BepInEx.Preloader.UnityMono.Utils;
 using MonoMod.Utils;
 
 namespace BepInEx.Preloader.UnityMono;
@@ -77,7 +78,7 @@ internal static class UnityPreloader
 
             TypeLoader.SearchDirectories.UnionWith(Paths.DllSearchPaths);
 
-            using (var assemblyPatcher = new AssemblyPatcher())
+            using (var assemblyPatcher = new AssemblyPatcher(MonoAssemblyHelper.LoadFromMemory))
             {
                 assemblyPatcher.AddPatchersFromDirectory(Paths.BepInExAssemblyDirectory);
                 assemblyPatcher.AddPatchersFromDirectory(Paths.PatcherPluginPath);

BepInEx.Preloader.UnityMono/UnityPreloaderRunner.cs

@@ -5,6 +5,7 @@ using System.Reflection;
 using BepInEx.Preloader.Core;
 using BepInEx.Preloader.RuntimeFixes;
 using BepInEx.Preloader.UnityMono.RuntimeFixes;
+using BepInEx.Preloader.UnityMono.Utils;
 
 namespace BepInEx.Preloader.UnityMono;
 
@@ -29,7 +30,7 @@ internal static class UnityPreloaderRunner
         foreach (var criticalAssembly in CriticalAssemblies)
             try
             {
-                Assembly.LoadFile(Path.Combine(Paths.BepInExAssemblyDirectory, criticalAssembly));
+                MonoAssemblyHelper.Load(Path.Combine(Paths.BepInExAssemblyDirectory, criticalAssembly));
             }
             catch (Exception)
             {
@@ -94,9 +95,9 @@ internal static class UnityPreloaderRunner
         if (foundAssembly != null)
             return foundAssembly;
 
-        if (Utility.TryResolveDllAssembly(assemblyName, Paths.BepInExAssemblyDirectory, out foundAssembly)
-         || Utility.TryResolveDllAssembly(assemblyName, Paths.PatcherPluginPath, out foundAssembly)
-         || Utility.TryResolveDllAssembly(assemblyName, Paths.PluginPath, out foundAssembly))
+        if (MonoAssemblyHelper.TryResolveDllAssembly(assemblyName, Paths.BepInExAssemblyDirectory, out foundAssembly)
+         || MonoAssemblyHelper.TryResolveDllAssembly(assemblyName, Paths.PatcherPluginPath, out foundAssembly)
+         || MonoAssemblyHelper.TryResolveDllAssembly(assemblyName, Paths.PluginPath, out foundAssembly))
             return foundAssembly;
 
         return null;

BepInEx.Preloader.UnityMono/Utils/MonoAssemblyHelper.cs

@@ -0,0 +1,100 @@
+using System;
+using System.IO;
+using System.Linq;
+using System.Reflection;
+using System.Runtime.InteropServices;
+using Mono.Cecil;
+
+namespace BepInEx.Preloader.UnityMono.Utils;
+
+internal static class MonoAssemblyHelper
+{
+    [DllImport("__Internal", CallingConvention = CallingConvention.Cdecl)]
+    private static extern nint mono_image_open_from_data_with_name(nint data,
+                                                                   uint dataLength,
+                                                                   bool needCopy,
+                                                                   out MonoImageOpenStatus status,
+                                                                   bool refOnly,
+                                                                   [MarshalAs(UnmanagedType.LPStr)] string name);
+
+    [DllImport("__Internal", CallingConvention = CallingConvention.Cdecl)]
+    private static extern nint mono_assembly_load_from_full(nint image,
+                                                            [MarshalAs(UnmanagedType.LPStr)] string fileName,
+                                                            out MonoImageOpenStatus status,
+                                                            bool refOnly);
+
+    private static ReadAssemblyResult ReadAssemblyData(string filePath)
+    {
+        return ReadAssemblyData(File.ReadAllBytes(filePath));
+    }
+
+    private static ReadAssemblyResult ReadAssemblyData(byte[] assemblyData)
+    {
+        using var ms = new MemoryStream(assemblyData);
+        using var ad = AssemblyDefinition.ReadAssembly(ms);
+
+        return new ReadAssemblyResult
+        {
+            AssemblyName = ad.Name.Name,
+            AssemblyData = assemblyData
+        };
+    }
+
+    private static Assembly GetAssemblyByName(string assemblyName)
+    {
+        return AppDomain.CurrentDomain.GetAssemblies()
+                        .FirstOrDefault(a => Utility.TryParseAssemblyName(a.FullName, out var name) &&
+                                             name.Name == assemblyName);
+    }
+
+    public static bool TryResolveDllAssembly(AssemblyName assemblyName, string directory, out Assembly assembly) =>
+        Utility.TryResolveDllAssembly(assemblyName, directory, Load, out assembly);
+
+    public static Assembly LoadFromMemory(byte[] data, string filePath)
+    {
+        var fullPath = Path.GetFullPath(filePath);
+        return ReadAssemblyData(data).Load(fullPath);
+    }
+
+    public static Assembly Load(string filePath)
+    {
+        var fullPath = Path.GetFullPath(filePath);
+        return ReadAssemblyData(fullPath).Load(fullPath);
+    }
+
+    private enum MonoImageOpenStatus
+    {
+        MONO_IMAGE_OK,
+        MONO_IMAGE_ERROR_ERRNO,
+        MONO_IMAGE_MISSING_ASSEMBLYREF,
+        MONO_IMAGE_IMAGE_INVALID
+    }
+
+    private class ReadAssemblyResult
+    {
+        public byte[] AssemblyData;
+        public string AssemblyName;
+
+        public unsafe Assembly Load(string fullPath)
+        {
+            var assembly = GetAssemblyByName(AssemblyName);
+            File.AppendAllText("log.log", $"Load 1: {assembly?.FullName}\n");
+            if (assembly != null)
+                return assembly;
+
+            fixed (byte* data = &AssemblyData[0])
+            {
+                var image = mono_image_open_from_data_with_name((nint) data, (uint) AssemblyData.Length, true,
+                                                                out var status, false, fullPath);
+                if (status != MonoImageOpenStatus.MONO_IMAGE_OK)
+                    throw new BadImageFormatException($"Failed to load image {fullPath}: {status}");
+                mono_assembly_load_from_full(image, fullPath, out status, false);
+                if (status != MonoImageOpenStatus.MONO_IMAGE_OK)
+                    throw new BadImageFormatException($"Failed to load assembly {fullPath}: {status}");
+                var ass = GetAssemblyByName(AssemblyName);
+                File.AppendAllText("log.log", $"Load 2: {ass?.FullName}\n");
+                return ass;
+            }
+        }
+    }
+}