Workspace/Assets/Modules/ShrinkModFramework/Runtime/Loading/ShrinkExternalModAssemblyLoader.cs

using System;
using System.Collections.Generic;
using System.IO;
using System.Linq;
using System.Reflection;
using System.Security.Cryptography;
using UnityEngine;

namespace ShrinkModFramework
{
    internal static class ShrinkExternalModAssemblyLoader
    {
        internal sealed class ExternalAssemblyRevision
        {
            public string Path;
            public string Revision;
            public Assembly Assembly;
        }

        private static readonly Dictionary<string, ExternalAssemblyRevision> CurrentAssemblyRevisions =
            new(StringComparer.OrdinalIgnoreCase);
        private static readonly HashSet<Assembly> ExternalAssemblyHistory = new();
        private static readonly Dictionary<string, string> KnownAssemblyFiles = new(StringComparer.OrdinalIgnoreCase);
        private static readonly object ResolveLock = new();
        private static bool _resolveRegistered;

        public static IReadOnlyList<Assembly> LoadExternalAssemblies(ShrinkModFrameworkSettings settings, bool verboseLogging)
        {
            return ScanExternalAssemblyRevisions(settings, verboseLogging)
                .Select(revision => revision.Assembly)
                .ToArray();
        }

        internal static IReadOnlyList<ExternalAssemblyRevision> ScanExternalAssemblyRevisions(
            ShrinkModFrameworkSettings settings, bool verboseLogging)
        {
            if (settings != null && !settings.enableExternalDllMods)
            {
                CurrentAssemblyRevisions.Clear();
                KnownAssemblyFiles.Clear();
                return Array.Empty<ExternalAssemblyRevision>();
            }

#if ENABLE_IL2CPP && !UNITY_EDITOR
            Debug.LogWarning("[ShrinkModFramework] IL2CPP 运行时不支持外部 DLL 热加载，已跳过外部模组扫描。");
            return Array.Empty<ExternalAssemblyRevision>();
#else
            var modsDirectory = GetExternalModsDirectory(settings);
            if (settings == null || settings.autoCreateExternalModsDirectory)
                Directory.CreateDirectory(modsDirectory);

            RegisterAssemblyResolve();

            var dllPaths = Directory.GetFiles(modsDirectory, "*.dll", SearchOption.AllDirectories)
                .Select(Path.GetFullPath)
                .OrderBy(path => path, StringComparer.OrdinalIgnoreCase)
                .ToArray();
            var presentPaths = new HashSet<string>(dllPaths, StringComparer.OrdinalIgnoreCase);

            foreach (var dllPath in dllPaths
                         .OrderBy(path => path, StringComparer.OrdinalIgnoreCase))
            {
                var assemblyName = Path.GetFileNameWithoutExtension(dllPath);
                KnownAssemblyFiles[assemblyName] = dllPath;

                try
                {
                    var bytes = File.ReadAllBytes(dllPath);
                    var revision = ComputeSha256(bytes);
                    if (CurrentAssemblyRevisions.TryGetValue(dllPath, out var current) &&
                        string.Equals(current.Revision, revision, StringComparison.Ordinal))
                    {
                        continue;
                    }

                    var pdbPath = Path.ChangeExtension(dllPath, ".pdb");
                    var assembly = File.Exists(pdbPath)
                        ? Assembly.Load(bytes, File.ReadAllBytes(pdbPath))
                        : Assembly.Load(bytes);

                    var loadedRevision = new ExternalAssemblyRevision
                    {
                        Path = dllPath,
                        Revision = revision,
                        Assembly = assembly
                    };
                    CurrentAssemblyRevisions[dllPath] = loadedRevision;
                    ExternalAssemblyHistory.Add(assembly);

                    if (verboseLogging)
                        Debug.Log($"[ShrinkModFramework] 已加载外部模组程序集：{assembly.GetName().Name} ({revision[..12]})");
                }
                catch (Exception ex)
                {
                    Debug.LogError($"[ShrinkModFramework] 加载外部 DLL 失败：{dllPath}\n{ex.Message}");
                }
            }

            foreach (var missingPath in CurrentAssemblyRevisions.Keys
                         .Where(path => !presentPaths.Contains(path))
                         .ToArray())
            {
                CurrentAssemblyRevisions.Remove(missingPath);
            }

            return CurrentAssemblyRevisions.Values
                .OrderBy(revision => revision.Path, StringComparer.OrdinalIgnoreCase)
                .ToArray();
#endif
        }

        internal static bool IsExternalAssembly(Assembly assembly) =>
            assembly != null && ExternalAssemblyHistory.Contains(assembly);

        internal static bool TryGetCurrentRevision(Assembly assembly, out string revision)
        {
            foreach (var current in CurrentAssemblyRevisions.Values)
            {
                if (ReferenceEquals(current.Assembly, assembly))
                {
                    revision = current.Revision;
                    return true;
                }
            }

            revision = null;
            return false;
        }

        public static string GetExternalModsDirectory(ShrinkModFrameworkSettings settings)
        {
            var folderName = settings != null && !string.IsNullOrWhiteSpace(settings.externalModsFolderName)
                ? settings.externalModsFolderName.Trim()
                : "Mods";
            return Path.Combine(Application.persistentDataPath, folderName);
        }

        internal static void ResetForTesting()
        {
            CurrentAssemblyRevisions.Clear();
            ExternalAssemblyHistory.Clear();
            KnownAssemblyFiles.Clear();
            if (_resolveRegistered)
            {
                AppDomain.CurrentDomain.AssemblyResolve -= OnAssemblyResolve;
                _resolveRegistered = false;
            }
        }

        private static string ComputeSha256(byte[] bytes)
        {
            using var sha256 = SHA256.Create();
            return BitConverter.ToString(sha256.ComputeHash(bytes)).Replace("-", string.Empty);
        }

        private static void RegisterAssemblyResolve()
        {
            lock (ResolveLock)
            {
                if (_resolveRegistered)
                    return;

                AppDomain.CurrentDomain.AssemblyResolve += OnAssemblyResolve;
                _resolveRegistered = true;
            }
        }

        private static Assembly OnAssemblyResolve(object sender, ResolveEventArgs args)
        {
            var requestedName = new AssemblyName(args.Name).Name;
            if (string.IsNullOrWhiteSpace(requestedName))
                return null;

            if (!KnownAssemblyFiles.TryGetValue(requestedName, out var path) || !File.Exists(path))
                return null;

            try
            {
                return Assembly.Load(File.ReadAllBytes(path));
            }
            catch
            {
                return null;
            }
        }
    }
}