Add dependency resolution

This commit is contained in:
Stone_Red
2025-12-19 21:48:37 +01:00
parent e647916f78
commit fad12f7271
8 changed files with 196 additions and 54 deletions
@@ -1,5 +1,5 @@
namespace RemoteExec.Client;
[Serializable]
internal class RemoteExecutionException(string message) : Exception(message) public class RemoteExecutionException(string message) : Exception(message)
{ {
} }
+26 -14
View File
@@ -5,6 +5,7 @@ using RemoteExec.Shared;
using System.Diagnostics.CodeAnalysis; using System.Diagnostics.CodeAnalysis;
using System.Reflection; using System.Reflection;
using System.Text.Json; using System.Text.Json;
using System.Threading.Channels;
namespace RemoteExec.Client; namespace RemoteExec.Client;
@@ -18,6 +19,29 @@ public class RemoteExecutor(string url)
public async Task StartAsync(CancellationToken cancellationToken = default) public async Task StartAsync(CancellationToken cancellationToken = default)
{ {
_ = _connection.On($"RequestAssembly", async (string assemblyName, Guid requestId) =>
{
Assembly? assembly = AppDomain.CurrentDomain.GetAssemblies().FirstOrDefault(a => a.GetName().FullName == assemblyName);
if (assembly == null)
{
assembly = Assembly.Load(new AssemblyName(assemblyName));
}
byte[] dllBytes = await File.ReadAllBytesAsync(assembly.Location!);
Channel<byte> channel = Channel.CreateUnbounded<byte>();
foreach (byte b in dllBytes)
{
await channel.Writer.WriteAsync(b);
}
channel.Writer.Complete();
await _connection.InvokeAsync("ProvideAssembly", requestId, channel.Reader);
});
await _connection.StartAsync(cancellationToken); await _connection.StartAsync(cancellationToken);
} }
@@ -73,28 +97,16 @@ public class RemoteExecutor(string url)
{ {
MethodInfo method = del.Method; MethodInfo method = del.Method;
Type declaringType = method.DeclaringType!; Type declaringType = method.DeclaringType!;
Assembly asm = declaringType.Assembly; Assembly assembly = declaringType.Assembly;
if (!method.IsStatic) if (!method.IsStatic)
{ {
throw new InvalidOperationException("Only static methods supported"); throw new InvalidOperationException("Only static methods supported");
} }
if (asm.IsDynamic)
{
throw new InvalidOperationException("Dynamic assemblies are not supported");
}
if (string.IsNullOrEmpty(asm.Location))
{
throw new InvalidOperationException("Assembly location is not available");
}
byte[] dllBytes = File.ReadAllBytes(asm.Location);
RemoteExecutionRequest request = new RemoteExecutionRequest RemoteExecutionRequest request = new RemoteExecutionRequest
{ {
AssemblyBytes = dllBytes, AssemblyName = assembly.GetName().FullName,
TypeName = declaringType.FullName!, TypeName = declaringType.FullName!,
MethodName = method.Name, MethodName = method.Name,
ArgumentTypes = [.. method.GetParameters().Select(p => p.ParameterType.AssemblyQualifiedName!)], ArgumentTypes = [.. method.GetParameters().Select(p => p.ParameterType.AssemblyQualifiedName!)],
+8
View File
@@ -0,0 +1,8 @@
// This file is used by Code Analysis to maintain SuppressMessage
// attributes that are applied to this project.
// Project-level suppressions either have no target or are given
// a specific target and scoped to a namespace, type, member, etc.
using System.Diagnostics.CodeAnalysis;
[assembly: SuppressMessage("Minor Code Smell", "S2325:Methods and properties that don't access instance data should be static", Justification = "<Pending>", Scope = "type", Target = "~T:RemoteExec.Server.Hubs.RemoteExecutionHub")]
+151 -20
View File
@@ -2,30 +2,64 @@
using RemoteExec.Shared; using RemoteExec.Shared;
using System.Collections.Concurrent;
using System.Reflection; using System.Reflection;
using System.Runtime.Loader; using System.Runtime.Loader;
using System.Text.Json; using System.Text.Json;
using System.Threading.Channels;
namespace RemoteExec.Server.Hubs; namespace RemoteExec.Server.Hubs;
public class RemoteExecutionHub : Hub public class RemoteExecutionHub(ILogger<RemoteExecutionHub> logger) : Hub
{ {
private readonly Dictionary<string, AssemblyLoadContext> connections = []; private static readonly ConcurrentDictionary<string, RemoteJobAssemblyLoadContext> connections = new();
private static readonly ConcurrentDictionary<Guid, TaskCompletionSource<byte[]>> pendingAssemblyRequests = new();
public override Task OnConnectedAsync() public override Task OnConnectedAsync()
{ {
connections.Add(Context.ConnectionId, new RemoteJobAssemblyLoadContext($"RemoteJob_{Guid.NewGuid()}", true)); RemoteJobAssemblyLoadContext assemblyLoadContext = new RemoteJobAssemblyLoadContext($"RemoteJob_{Guid.NewGuid()}");
assemblyLoadContext.Resolving += AssemblyLoadContext_Resolving;
_ = connections.TryAdd(Context.ConnectionId, assemblyLoadContext);
return base.OnConnectedAsync(); return base.OnConnectedAsync();
} }
private Assembly? AssemblyLoadContext_Resolving(AssemblyLoadContext assemblyLoadContext, AssemblyName assemblyName)
{
logger.LogWarning("Assembly resolution triggered synchronously for {AssemblyName}. This should have been pre-loaded.", assemblyName.FullName);
// Return null to let other resolution mechanisms try
return null;
}
public async Task ProvideAssembly(Guid requestId, ChannelReader<byte> stream)
{
using MemoryStream ms = new MemoryStream();
while (await stream.WaitToReadAsync())
{
while (stream.TryRead(out byte item))
{
ms.WriteByte(item);
}
}
byte[] assemblyBytes = ms.ToArray();
if (pendingAssemblyRequests.TryRemove(requestId, out TaskCompletionSource<byte[]>? tcs))
{
tcs.SetResult(assemblyBytes);
}
}
public override Task OnDisconnectedAsync(Exception? exception) public override Task OnDisconnectedAsync(Exception? exception)
{ {
if (connections.TryGetValue(Context.ConnectionId, out RemoteJobAssemblyLoadContext assemblyLoadContext)) if (connections.TryRemove(Context.ConnectionId, out RemoteJobAssemblyLoadContext? assemblyLoadContext))
{ {
assemblyLoadContext.Unload(); assemblyLoadContext.Unload();
_ = connections.Remove(Context.ConnectionId);
} }
return base.OnDisconnectedAsync(exception); return base.OnDisconnectedAsync(exception);
@@ -35,16 +69,24 @@ public class RemoteExecutionHub : Hub
{ {
try try
{ {
AssemblyLoadContext alc = new AssemblyLoadContext( if (!connections.TryGetValue(Context.ConnectionId, out RemoteJobAssemblyLoadContext? assemblyLoadContext))
name: $"RemoteJob_{Guid.NewGuid()}", {
isCollectible: true); throw new InvalidOperationException("Connection not found");
}
using MemoryStream ms = new MemoryStream(req.AssemblyBytes); Assembly? assembly = assemblyLoadContext.Assemblies.FirstOrDefault(a => a.GetName().FullName == req.AssemblyName);
Assembly asm = alc.LoadFromStream(ms);
Type type = asm.GetType(req.TypeName, throwOnError: true)!; assembly ??= await RequestAssemblyAsync(req.AssemblyName);
Type?[] argTypes = req.ArgumentTypes if (!assemblyLoadContext.Assemblies.Contains(assembly))
{
using MemoryStream ms = new MemoryStream(await GetAssemblyBytesAsync(assembly));
assembly = assemblyLoadContext.LoadFromStream(ms);
}
Type type = assembly.GetType(req.TypeName, throwOnError: true)!;
Type[] argTypes = req.ArgumentTypes
.Select(Type.GetType) .Select(Type.GetType)
.ToArray()!; .ToArray()!;
@@ -53,14 +95,13 @@ public class RemoteExecutionHub : Hub
BindingFlags.Static | BindingFlags.Public | BindingFlags.NonPublic, BindingFlags.Static | BindingFlags.Public | BindingFlags.NonPublic,
binder: null, binder: null,
argTypes, argTypes,
modifiers: null); modifiers: null) ?? throw new MissingMethodException(req.TypeName, req.MethodName);
if (method == null) // Pre-load all referenced assemblies to avoid triggering Resolving event during Invoke
{ await PreLoadReferencedAssembliesAsync(assemblyLoadContext, assembly);
throw new MissingMethodException(req.TypeName, req.MethodName);
}
ParameterInfo[] parameters = method.GetParameters(); ParameterInfo[] parameters = method.GetParameters();
if (parameters.Length != req.Arguments.Length) if (parameters.Length != req.Arguments.Length)
{ {
throw new ArgumentException("Argument count mismatch"); throw new ArgumentException("Argument count mismatch");
@@ -95,8 +136,6 @@ public class RemoteExecutionHub : Hub
object? result = method.Invoke(null, invokeArgs); object? result = method.Invoke(null, invokeArgs);
alc.Unload();
return new RemoteExecutionResult return new RemoteExecutionResult
{ {
Result = result Result = result
@@ -110,4 +149,96 @@ public class RemoteExecutionHub : Hub
}; };
} }
} }
private async Task<Assembly> RequestAssemblyAsync(string assemblyName)
{
try
{
Guid guid = Guid.NewGuid();
TaskCompletionSource<byte[]> tcs = new TaskCompletionSource<byte[]>();
_ = pendingAssemblyRequests.TryAdd(guid, tcs);
logger.LogInformation("Requesting assembly {Assembly} with request ID {RequestId}", assemblyName, guid);
await Clients.Caller.SendAsync("RequestAssembly", assemblyName, guid);
logger.LogInformation("Waiting for assembly {Assembly} with request ID {RequestId}", assemblyName, guid);
// Wait for the assembly with a timeout
byte[] assemblyBytes = await tcs.Task.WaitAsync(TimeSpan.FromSeconds(30));
logger.LogInformation("Received assembly {Assembly} with request ID {RequestId}", assemblyName, guid);
// Return a temporary assembly just for metadata inspection
using MemoryStream ms = new MemoryStream(assemblyBytes);
return Assembly.Load(assemblyBytes);
}
catch (Exception ex)
{
logger.LogError(ex, "Error requesting assembly {Assembly}", assemblyName);
throw;
}
}
private async Task<byte[]> GetAssemblyBytesAsync(Assembly assembly)
{
string assemblyName = assembly.GetName().FullName!;
Guid guid = Guid.NewGuid();
TaskCompletionSource<byte[]> tcs = new TaskCompletionSource<byte[]>();
_ = pendingAssemblyRequests.TryAdd(guid, tcs);
logger.LogInformation("Requesting assembly bytes for {Assembly} with request ID {RequestId}", assemblyName, guid);
await Clients.Caller.SendAsync("RequestAssembly", assemblyName, guid);
byte[] assemblyBytes = await tcs.Task.WaitAsync(TimeSpan.FromSeconds(30));
logger.LogInformation("Received assembly bytes for {Assembly} with request ID {RequestId}", assemblyName, guid);
return assemblyBytes;
}
private async Task PreLoadReferencedAssembliesAsync(RemoteJobAssemblyLoadContext assemblyLoadContext, Assembly assembly)
{
AssemblyName[] referencedAssemblies = assembly.GetReferencedAssemblies();
foreach (AssemblyName referencedAssembly in referencedAssemblies)
{
try
{
// Try to load from the assembly load context first
Assembly? loadedAssembly = assemblyLoadContext.Assemblies.FirstOrDefault(a => a.GetName().FullName == referencedAssembly.FullName);
if (loadedAssembly != null)
{
continue; // Already loaded in the context
}
// Try to load from default context (BCL assemblies)
try
{
_ = assemblyLoadContext.LoadFromAssemblyName(referencedAssembly);
continue; // Successfully loaded from default context
}
catch
{
// If not in default context, request from client
logger.LogInformation("Pre-loading referenced assembly {Assembly}", referencedAssembly.FullName);
Assembly tempAssembly = await RequestAssemblyAsync(referencedAssembly.FullName!);
byte[] assemblyBytes = await GetAssemblyBytesAsync(tempAssembly);
using MemoryStream ms = new MemoryStream(assemblyBytes);
_ = assemblyLoadContext.LoadFromStream(ms);
}
}
catch (Exception ex)
{
logger.LogWarning(ex, "Could not pre-load referenced assembly {Assembly}", referencedAssembly.FullName);
}
}
}
} }
+2 -2
View File
@@ -5,7 +5,7 @@ WebApplicationBuilder builder = WebApplication.CreateBuilder(args);
// Add services to the container. // Add services to the container.
builder.Services.AddControllers(); builder.Services.AddControllers();
builder.Services.AddSignalR(options => options.MaximumReceiveMessageSize = null); builder.Services.AddSignalR();
builder.Services.AddOpenApi(); builder.Services.AddOpenApi();
WebApplication app = builder.Build(); WebApplication app = builder.Build();
@@ -24,4 +24,4 @@ app.UseAuthorization();
app.MapControllers(); app.MapControllers();
app.Run(); await app.RunAsync();
@@ -1,18 +1,7 @@
using System.Reflection; using System.Runtime.Loader;
using System.Runtime.Loader;
namespace RemoteExec.Server; namespace RemoteExec.Server;
public class RemoteJobAssemblyLoadContext(string name) : AssemblyLoadContext(name, true) public class RemoteJobAssemblyLoadContext(string name) : AssemblyLoadContext(name, true)
{ {
public event EventHandler<RequestAssemblyEventArgs>? RequestAssembly;
protected override Assembly? Load(AssemblyName assemblyName)
{
RequestAssemblyEventArgs requestAssemblyEventArgs = new RequestAssemblyEventArgs(assemblyName);
RequestAssembly?.Invoke(this, requestAssemblyEventArgs);
return requestAssemblyEventArgs.GetAssemblyAsync().ConfigureAwait(false).GetAwaiter().GetResult();
}
} }
+1 -1
View File
@@ -2,7 +2,7 @@
public sealed class RemoteExecutionRequest public sealed class RemoteExecutionRequest
{ {
public required byte[] AssemblyBytes { get; set; } public required string AssemblyName { get; set; }
public required string TypeName { get; set; } public required string TypeName { get; set; }
public required string MethodName { get; set; } public required string MethodName { get; set; }
public required string[] ArgumentTypes { get; set; } public required string[] ArgumentTypes { get; set; }
+4 -2
View File
@@ -1,4 +1,6 @@
using RemoteExec.Client; using CuteUtils.FluentMath.TypeExtensions;
using RemoteExec.Client;
RemoteExecutor remoteExecutor = new RemoteExecutor("https://localhost:7109/remote"); RemoteExecutor remoteExecutor = new RemoteExecutor("https://localhost:7109/remote");
await remoteExecutor.StartAsync(); await remoteExecutor.StartAsync();
@@ -13,5 +15,5 @@ Console.WriteLine($"Result: {result}");
static int Multiply(int x, int y) static int Multiply(int x, int y)
{ {
return x * y; return x.Multiply(y);
} }