Files
Torch/Torch/Plugins/AssemblyRewriter.cs
zznty b6e88b359f
All checks were successful
Release / Get Version (push) Successful in 11s
Release / Build and Publish Nuget (push) Successful in 2m14s
Release / Build and Publish Package (push) Successful in 19m29s
fix assembly rewriter stripping assembly resources
2024-10-14 00:15:05 +07:00

123 lines
3.9 KiB
C#

using System;
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 ZipResolver _zipResolver;
private static readonly DefaultAssemblyResolver _defaultResolver;
static AssemblyRewriter()
{
_defaultResolver = new();
_zipResolver = new(_defaultResolver);
_defaultResolver.AddSearchDirectory(AppContext.BaseDirectory);
_defaultResolver.AddSearchDirectory(Path.Combine(Directory.GetCurrentDirectory(), "DedicatedServer64"));
#if !DEBUG
_defaultResolver.AddSearchDirectory(Path.Join(AppContext.BaseDirectory, "torch64"));
#endif
}
public static Assembly ProcessWeavers(this Stream stream, ZipArchive archive)
{
_zipResolver.Archive = archive;
using var assStream = new MemoryStream();
stream.CopyTo(assStream);
assStream.Position = 0;
try
{
var ass = ProcessInternal(assStream, _zipResolver);
return ass;
}
catch (Exception e)
{
Log.Error(e, "Unable to process assembly, skipping");
return Assembly.Load(assStream.ToArray());
}
finally
{
_zipResolver.Archive = null;
}
}
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 assembly = AssemblyDefinition.ReadAssembly(inputStream, new()
{
AssemblyResolver = resolver
});
foreach (var fieldDefinition in FindAllToRewrite(assembly.MainModule))
{
fieldDefinition.IsInitOnly = false;
}
using var memStream = new MemoryStream();
assembly.Write(memStream);
return Assembly.Load(memStream.ToArray());
}
private static IEnumerable<FieldDefinition> FindAllToRewrite(ModuleDefinition definition)
{
return definition.Types.SelectMany(b => b.Fields.Where(HasValidAttributes));
}
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);
}
}
}