mirror of
https://github.com/Stone-Red-Code/RemoteExec.git
synced 2026-09-04 00:56:17 +02:00
Code improvements
This commit is contained in:
@@ -45,6 +45,11 @@ public class RemoteExecutor(string url)
|
||||
await _connection.StartAsync(cancellationToken);
|
||||
}
|
||||
|
||||
public async Task StopAsync(CancellationToken cancellationToken = default)
|
||||
{
|
||||
await _connection.StopAsync(cancellationToken);
|
||||
}
|
||||
|
||||
public bool TryExecute<TDelegate, TResult>(TDelegate del, out TResult? result, params object[] args) where TDelegate : Delegate
|
||||
{
|
||||
try
|
||||
|
||||
@@ -6,3 +6,5 @@
|
||||
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")]
|
||||
[assembly: SuppressMessage("Major Code Smell", "S2139:Exceptions should be either logged or rethrown but not both", Justification = "<Pending>", Scope = "member", Target = "~M:RemoteExec.Server.Hubs.RemoteExecutionHub.RequestAssemblyAsync(System.String)~System.Threading.Tasks.Task{System.Reflection.Assembly}")]
|
||||
[assembly: SuppressMessage("Major Code Smell", "S3011:Reflection should not be used to increase accessibility of classes, methods, or fields", Justification = "<Pending>", Scope = "member", Target = "~M:RemoteExec.Server.Hubs.RemoteExecutionHub.Execute(RemoteExec.Shared.RemoteExecutionRequest)~System.Threading.Tasks.Task{RemoteExec.Shared.RemoteExecutionResult}")]
|
||||
|
||||
@@ -4,7 +4,6 @@ using RemoteExec.Shared;
|
||||
|
||||
using System.Collections.Concurrent;
|
||||
using System.Reflection;
|
||||
using System.Runtime.Loader;
|
||||
using System.Text.Json;
|
||||
using System.Threading.Channels;
|
||||
|
||||
@@ -20,46 +19,19 @@ public class RemoteExecutionHub(ILogger<RemoteExecutionHub> logger) : Hub
|
||||
{
|
||||
RemoteJobAssemblyLoadContext assemblyLoadContext = new RemoteJobAssemblyLoadContext($"RemoteJob_{Guid.NewGuid()}");
|
||||
|
||||
assemblyLoadContext.Resolving += AssemblyLoadContext_Resolving;
|
||||
|
||||
_ = connections.TryAdd(Context.ConnectionId, assemblyLoadContext);
|
||||
|
||||
logger.LogInformation("Connection {ConnectionId} established", Context.ConnectionId);
|
||||
|
||||
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)
|
||||
{
|
||||
if (connections.TryRemove(Context.ConnectionId, out RemoteJobAssemblyLoadContext? assemblyLoadContext))
|
||||
{
|
||||
assemblyLoadContext.Unload();
|
||||
logger.LogInformation("Connection {ConnectionId} disconnected", Context.ConnectionId);
|
||||
}
|
||||
|
||||
return base.OnDisconnectedAsync(exception);
|
||||
@@ -71,6 +43,7 @@ public class RemoteExecutionHub(ILogger<RemoteExecutionHub> logger) : Hub
|
||||
{
|
||||
if (!connections.TryGetValue(Context.ConnectionId, out RemoteJobAssemblyLoadContext? assemblyLoadContext))
|
||||
{
|
||||
logger.LogError("Connection {ConnectionId} not found", Context.ConnectionId);
|
||||
throw new InvalidOperationException("Connection not found");
|
||||
}
|
||||
|
||||
@@ -104,6 +77,7 @@ public class RemoteExecutionHub(ILogger<RemoteExecutionHub> logger) : Hub
|
||||
|
||||
if (parameters.Length != req.Arguments.Length)
|
||||
{
|
||||
logger.LogError("Argument count mismatch for method {Method} in type {Type} for connection {ConnectionId}", req.MethodName, req.TypeName, Context.ConnectionId);
|
||||
throw new ArgumentException("Argument count mismatch");
|
||||
}
|
||||
|
||||
@@ -143,6 +117,8 @@ public class RemoteExecutionHub(ILogger<RemoteExecutionHub> logger) : Hub
|
||||
}
|
||||
catch (Exception ex)
|
||||
{
|
||||
logger.LogError(ex, "Error executing remote method {Method} in type {Type} for connection {ConnectionId}", req.MethodName, req.TypeName, Context.ConnectionId);
|
||||
|
||||
return new RemoteExecutionResult
|
||||
{
|
||||
Exception = ex.ToString()
|
||||
@@ -150,6 +126,26 @@ public class RemoteExecutionHub(ILogger<RemoteExecutionHub> logger) : Hub
|
||||
}
|
||||
}
|
||||
|
||||
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);
|
||||
}
|
||||
}
|
||||
|
||||
private async Task<Assembly> RequestAssemblyAsync(string assemblyName)
|
||||
{
|
||||
try
|
||||
@@ -159,17 +155,11 @@ public class RemoteExecutionHub(ILogger<RemoteExecutionHub> logger) : Hub
|
||||
|
||||
_ = 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);
|
||||
@@ -190,14 +180,10 @@ public class RemoteExecutionHub(ILogger<RemoteExecutionHub> logger) : Hub
|
||||
|
||||
_ = 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;
|
||||
}
|
||||
|
||||
@@ -225,9 +211,6 @@ public class RemoteExecutionHub(ILogger<RemoteExecutionHub> logger) : Hub
|
||||
}
|
||||
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);
|
||||
|
||||
|
||||
Reference in New Issue
Block a user