Stone_Red/RemoteExec
1
0
Fork 0
mirror of https://github.com/Stone-Red-Code/RemoteExec.git synced 2026-09-04 09:06:20 +02:00
Code Issues Packages Projects Releases Wiki Activity

Code improvements

Browse Source
This commit is contained in:
Stone_Red
2025-12-19 23:06:45 +01:00
parent fad12f7271
commit 7c3ab3a154
3 changed files with 34 additions and 44 deletions
Download Patch File Download Diff File Expand all files Collapse all files
RemoteExec.Server/Hubs/RemoteExecutionHub.cs
+27 -44
View File
@@ -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);