From b9afb74d13b801b7a6f0eb339a1d1307396f1ade Mon Sep 17 00:00:00 2001 From: Stone_Red <56473591+Stone-Red-Code@users.noreply.github.com> Date: Sat, 27 Dec 2025 00:08:15 +0100 Subject: [PATCH] Cleanup code --- RemoteExec.Server/Hubs/RemoteExecutionHub.cs | 2 - .../ApiKeyAuthenticationMiddleware.cs | 23 ++++----- ...AssemblyLoadContextExecutionEnvironment.cs | 12 ++--- .../DockerContainerExecutionEnvironment.cs | 48 +++++-------------- .../appsettings.Development.json | 2 +- RemoteExec.Worker/Program.cs | 7 +-- 6 files changed, 26 insertions(+), 68 deletions(-) diff --git a/RemoteExec.Server/Hubs/RemoteExecutionHub.cs b/RemoteExec.Server/Hubs/RemoteExecutionHub.cs index 47bf9bb..932e0e6 100644 --- a/RemoteExec.Server/Hubs/RemoteExecutionHub.cs +++ b/RemoteExec.Server/Hubs/RemoteExecutionHub.cs @@ -115,10 +115,8 @@ public class RemoteExecutionHub : Hub { await foreach (TaskItem taskItem in taskStream.WithCancellation(Context.ConnectionAborted)) { - // Wait for available slot before processing await taskSemaphore.WaitAsync(Context.ConnectionAborted); - // Process task asynchronously without blocking the stream _ = Task.Run(async () => { try diff --git a/RemoteExec.Server/Middleware/ApiKeyAuthenticationMiddleware.cs b/RemoteExec.Server/Middleware/ApiKeyAuthenticationMiddleware.cs index 2778c8d..192a5e2 100644 --- a/RemoteExec.Server/Middleware/ApiKeyAuthenticationMiddleware.cs +++ b/RemoteExec.Server/Middleware/ApiKeyAuthenticationMiddleware.cs @@ -14,7 +14,7 @@ public class ApiKeyAuthenticationMiddleware { private readonly RequestDelegate _next; private readonly ILogger _logger; - private readonly ConcurrentDictionary _apiKeys; + private readonly ConcurrentDictionary apiKeys; /// /// Initializes a new instance of the class. @@ -22,27 +22,20 @@ public class ApiKeyAuthenticationMiddleware /// The next middleware in the pipeline. /// The authentication configuration options. /// The logger instance. - public ApiKeyAuthenticationMiddleware( - RequestDelegate next, - IOptionsMonitor authOptions, - ILogger logger) + public ApiKeyAuthenticationMiddleware(RequestDelegate next, IOptionsMonitor authOptions, ILogger logger) { _next = next; _logger = logger; + apiKeys = new ConcurrentDictionary(); - // Build lookup dictionary from configuration - _apiKeys = new ConcurrentDictionary(); - - // Initial load LoadApiKeys(authOptions.CurrentValue); - // Watch for configuration changes _ = authOptions.OnChange(LoadApiKeys); } private void LoadApiKeys(AuthenticationConfiguration config) { - _apiKeys.Clear(); + apiKeys.Clear(); if (config.ApiKeys == null || config.ApiKeys.Count == 0) { @@ -58,7 +51,7 @@ public class ApiKeyAuthenticationMiddleware continue; } - if (_apiKeys.TryAdd(apiKey.Key, apiKey)) + if (apiKeys.TryAdd(apiKey.Key, apiKey)) { _logger.LogInformation( "Registered API key: {Description}", @@ -72,7 +65,7 @@ public class ApiKeyAuthenticationMiddleware } } - _logger.LogInformation("Loaded {Count} active API keys", _apiKeys.Count); + _logger.LogInformation("Loaded {Count} active API keys", apiKeys.Count); } /// @@ -89,7 +82,7 @@ public class ApiKeyAuthenticationMiddleware } // Check if any API keys are configured - if (_apiKeys.IsEmpty) + if (apiKeys.IsEmpty) { _logger.LogError("No API keys configured. Rejecting request to {Path}", context.Request.Path); @@ -113,7 +106,7 @@ public class ApiKeyAuthenticationMiddleware string providedKey = extractedApiKey.ToString(); // Validate API key - if (!_apiKeys.TryGetValue(providedKey, out ApiKeyConfiguration? apiKeyConfig)) + if (!apiKeys.TryGetValue(providedKey, out ApiKeyConfiguration? apiKeyConfig)) { _logger.LogWarning("Invalid API Key provided for request to {Path} from {RemoteIp}", context.Request.Path, diff --git a/RemoteExec.Server/Services/AssemblyLoadContextExecutionEnvironment.cs b/RemoteExec.Server/Services/AssemblyLoadContextExecutionEnvironment.cs index 2020885..5263b92 100644 --- a/RemoteExec.Server/Services/AssemblyLoadContextExecutionEnvironment.cs +++ b/RemoteExec.Server/Services/AssemblyLoadContextExecutionEnvironment.cs @@ -31,12 +31,7 @@ public class AssemblyLoadContextExecutionEnvironment : ExecutionEnvironment .Select(Type.GetType) .ToArray()!; - MethodInfo? method = type.GetMethod( - request.MethodName, - BindingFlags.Static | BindingFlags.Public | BindingFlags.NonPublic, - binder: null, - argTypes, - modifiers: null) ?? throw new MissingMethodException(request.TypeName, request.MethodName); + MethodInfo? method = type.GetMethod(request.MethodName, BindingFlags.Static | BindingFlags.Public | BindingFlags.NonPublic, binder: null, argTypes, modifiers: null) ?? throw new MissingMethodException(request.TypeName, request.MethodName); // Pre-load all referenced assemblies to avoid triggering Resolving event during Invoke await AssemblyUtilities.PreLoadReferencedAssembliesAsync(assemblyLoadContext, assembly, RequestAssemblyAsync); @@ -57,7 +52,6 @@ public class AssemblyLoadContextExecutionEnvironment : ExecutionEnvironment if (arg is JsonElement je) { - // Deserialize the JSON element into the expected CLR type invokeArgs[i] = JsonSerializer.Deserialize(je.GetRawText(), targetType); } else if (arg == null) @@ -66,7 +60,6 @@ public class AssemblyLoadContextExecutionEnvironment : ExecutionEnvironment } else if (!targetType.IsInstanceOfType(arg)) { - // Fallback for simple primitive conversions invokeArgs[i] = Convert.ChangeType(arg, targetType); } else @@ -80,7 +73,9 @@ public class AssemblyLoadContextExecutionEnvironment : ExecutionEnvironment if (result is Task taskResult) { await taskResult.ConfigureAwait(false); + Type returnType = method.ReturnType; + if (returnType.IsGenericType && returnType.GetGenericTypeDefinition() == typeof(Task<>)) { PropertyInfo resultProperty = returnType.GetProperty("Result")!; @@ -88,7 +83,6 @@ public class AssemblyLoadContextExecutionEnvironment : ExecutionEnvironment } else { - // For non-generic Task, result is null result = null; } } diff --git a/RemoteExec.Server/Services/DockerContainerExecutionEnvironment.cs b/RemoteExec.Server/Services/DockerContainerExecutionEnvironment.cs index 6803a60..fd4eeca 100644 --- a/RemoteExec.Server/Services/DockerContainerExecutionEnvironment.cs +++ b/RemoteExec.Server/Services/DockerContainerExecutionEnvironment.cs @@ -71,14 +71,12 @@ public class DockerContainerExecutionEnvironment : ExecutionEnvironment try { - // Get assembly bytes from cache if (!assemblyCache.TryGetValue(request.AssemblyName, out byte[]? assemblyBytes)) { assemblyBytes = await RequestAssemblyAsync(request.AssemblyName); assemblyCache[request.AssemblyName] = assemblyBytes; } - // Prepare execution request ContainerExecutionRequest containerRequest = new() { AssemblyBytes = Convert.ToBase64String(assemblyBytes), @@ -90,17 +88,13 @@ public class DockerContainerExecutionEnvironment : ExecutionEnvironment string requestJson = JsonSerializer.Serialize(containerRequest); - // Create and start container containerId = await CreateAndStartContainerAsync(requestJson, CancellationToken.None); - // Start monitoring logs for assembly requests in background using CancellationTokenSource timeoutCts = new(containerTimeout); Task logMonitorTask = MonitorContainerLogsAsync(containerId, timeoutCts.Token); - // Wait for container to complete with timeout ContainerWaitResponse waitResponse = await dockerClient.Containers.WaitContainerAsync(containerId, timeoutCts.Token); - // Cancel log monitoring await timeoutCts.CancelAsync(); try @@ -124,7 +118,6 @@ public class DockerContainerExecutionEnvironment : ExecutionEnvironment }; } - // Parse result from stdout - get last JSON line (filter out assembly protocol lines) string[] lines = stdout.Split('\n', StringSplitOptions.RemoveEmptyEntries); string? resultLine = lines.LastOrDefault(l => { @@ -181,16 +174,12 @@ public class DockerContainerExecutionEnvironment : ExecutionEnvironment { try { - MultiplexedStream logStream = await dockerClient.Containers.GetContainerLogsAsync( - containerId, - false, - new ContainerLogsParameters - { - ShowStdout = true, - ShowStderr = false, - Follow = true - }, - cancellationToken); + MultiplexedStream logStream = await dockerClient.Containers.GetContainerLogsAsync(containerId, false, new ContainerLogsParameters + { + ShowStdout = true, + ShowStderr = false, + Follow = true + }, cancellationToken); byte[] buffer = new byte[4096]; StringBuilder lineBuffer = new(); @@ -207,7 +196,6 @@ public class DockerContainerExecutionEnvironment : ExecutionEnvironment string text = Encoding.UTF8.GetString(buffer, 0, result.Count); _ = lineBuffer.Append(text); - // Process complete lines string bufferContent = lineBuffer.ToString(); int lastNewline = bufferContent.LastIndexOf('\n'); @@ -336,14 +324,11 @@ public class DockerContainerExecutionEnvironment : ExecutionEnvironment private async Task GetContainerLogsAsync(string containerId) { - MultiplexedStream logStream = await dockerClient.Containers.GetContainerLogsAsync( - containerId, - false, - new ContainerLogsParameters - { - ShowStdout = true, - ShowStderr = true - }); + MultiplexedStream logStream = await dockerClient.Containers.GetContainerLogsAsync(containerId, false, new ContainerLogsParameters + { + ShowStdout = true, + ShowStderr = true + }); StringBuilder output = new(); byte[] buffer = new byte[4096]; @@ -365,15 +350,8 @@ public class DockerContainerExecutionEnvironment : ExecutionEnvironment { try { - // Stop container if still running - _ = await dockerClient.Containers.StopContainerAsync( - containerId, - new ContainerStopParameters { WaitBeforeKillSeconds = 5 }); - - // Remove container - await dockerClient.Containers.RemoveContainerAsync( - containerId, - new ContainerRemoveParameters { Force = true, RemoveVolumes = true }); + _ = await dockerClient.Containers.StopContainerAsync(containerId, new ContainerStopParameters { WaitBeforeKillSeconds = 5 }); + await dockerClient.Containers.RemoveContainerAsync(containerId, new ContainerRemoveParameters { Force = true, RemoveVolumes = true }); _ = runningContainers.TryRemove(containerId, out _); } diff --git a/RemoteExec.Server/appsettings.Development.json b/RemoteExec.Server/appsettings.Development.json index 65b7966..a9a2c15 100644 --- a/RemoteExec.Server/appsettings.Development.json +++ b/RemoteExec.Server/appsettings.Development.json @@ -17,7 +17,7 @@ }, "AllowedHosts": "*", "Execution": { - "ExecutionEnvironment": "DockerContainer" + "ExecutionEnvironment": "AssemblyLoadContext" }, "DockerExecution": { "DockerHost": null, diff --git a/RemoteExec.Worker/Program.cs b/RemoteExec.Worker/Program.cs index 794617c..5ad13e8 100644 --- a/RemoteExec.Worker/Program.cs +++ b/RemoteExec.Worker/Program.cs @@ -45,12 +45,7 @@ public static class Program .Select(Type.GetType) .ToArray()!; - MethodInfo? method = type.GetMethod( - request.MethodName, - BindingFlags.Static | BindingFlags.Public | BindingFlags.NonPublic, - binder: null, - argTypes, - modifiers: null); + MethodInfo? method = type.GetMethod(request.MethodName, BindingFlags.Static | BindingFlags.Public | BindingFlags.NonPublic, binder: null, argTypes, modifiers: null); if (method == null) {