diff --git a/Torch/Plugins/AssemblyRewriter.cs b/Torch/Plugins/AssemblyRewriter.cs index 6721ec4..e2a006d 100644 --- a/Torch/Plugins/AssemblyRewriter.cs +++ b/Torch/Plugins/AssemblyRewriter.cs @@ -1,34 +1,52 @@ using System.Collections.Generic; using System.IO; +using System.IO.Compression; using System.Linq; using System.Reflection; using Mono.Cecil; -using NLog; namespace Torch.Plugins; internal static class AssemblyRewriter { - private static readonly ILogger Log = LogManager.GetCurrentClassLogger(); - - private static readonly IAssemblyResolver Resolver; + private static readonly ZipResolver _zipResolver; + private static readonly DefaultAssemblyResolver _defaultResolver; static AssemblyRewriter() { - var resolver = new DefaultAssemblyResolver(); - Resolver = resolver; - resolver.AddSearchDirectory(Directory.GetCurrentDirectory()); - resolver.AddSearchDirectory(Path.Combine(Directory.GetCurrentDirectory(), "DedicatedServer64")); + _defaultResolver = new(); + _zipResolver = new(_defaultResolver); + _defaultResolver.AddSearchDirectory(Directory.GetCurrentDirectory()); + _defaultResolver.AddSearchDirectory(Path.Combine(Directory.GetCurrentDirectory(), "DedicatedServer64")); } - public static Assembly ProcessWeavers(this Stream stream) + public static Assembly ProcessWeavers(this Stream stream, ZipArchive archive) { + _zipResolver.Archive = archive; using var assStream = new MemoryStream(); stream.CopyTo(assStream); assStream.Position = 0; - using var module = ModuleDefinition.ReadModule(assStream, new() + var ass = ProcessInternal(assStream, _zipResolver); + _zipResolver.Archive = null; + return ass; + } + + public static Assembly ProcessWeavers(this Stream stream, string path) + { + _defaultResolver.AddSearchDirectory(path); + using var assStream = new MemoryStream(); + stream.CopyTo(assStream); + assStream.Position = 0; + var ass = ProcessInternal(assStream, _defaultResolver); + _defaultResolver.RemoveSearchDirectory(path); + return ass; + } + + private static Assembly ProcessInternal(Stream inputStream, IAssemblyResolver resolver) + { + using var module = ModuleDefinition.ReadModule(inputStream, new() { - AssemblyResolver = Resolver + AssemblyResolver = _zipResolver }); foreach (var fieldDefinition in FindAllToRewrite(module)) { @@ -47,4 +65,40 @@ internal static class AssemblyRewriter private static bool HasValidAttributes(FieldDefinition definition) => definition.CustomAttributes.Any(b => b.AttributeType.Name.Contains("Reflected") || b.AttributeType.Name == "DependencyAttribute"); + + private class ZipResolver : IAssemblyResolver + { + private readonly IAssemblyResolver _fallbackResolver; + public ZipArchive Archive { get; set; } + + public ZipResolver(IAssemblyResolver fallbackResolver) + { + _fallbackResolver = fallbackResolver; + } + + public void Dispose() + { + _fallbackResolver.Dispose(); + } + + public AssemblyDefinition Resolve(AssemblyNameReference name) + { + return Resolve(name, new()); + } + + public AssemblyDefinition Resolve(AssemblyNameReference name, ReaderParameters parameters) + { + var fileName = $"{name.Name}.dll"; + + if (Archive.Entries.FirstOrDefault(entry => entry.Name == fileName) is not { } archiveEntry) + return _fallbackResolver.Resolve(name, parameters); + + using var stream = archiveEntry.Open(); + using var memStream = new MemoryStream(); + stream.CopyTo(memStream); + memStream.Position = 0; + + return AssemblyDefinition.ReadAssembly(memStream, parameters); + } + } } \ No newline at end of file diff --git a/Torch/Plugins/PluginManager.cs b/Torch/Plugins/PluginManager.cs index cbb45a3..8544e77 100644 --- a/Torch/Plugins/PluginManager.cs +++ b/Torch/Plugins/PluginManager.cs @@ -360,7 +360,7 @@ namespace Torch.Managers using var stream = entry.Open(); - assemblies.Add(stream.ProcessWeavers()); + assemblies.Add(stream.ProcessWeavers(zipFile)); } } else @@ -378,7 +378,7 @@ namespace Torch.Managers // continue; using var stream = File.OpenRead(file); - assemblies.Add(stream.ProcessWeavers()); + assemblies.Add(stream.ProcessWeavers(item.Path)); }