mirror of
https://github.com/Stone-Red-Code/RemoteExec.git
synced 2026-09-04 00:56:17 +02:00
Cleanup code
This commit is contained in:
@@ -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
|
||||
|
||||
@@ -14,7 +14,7 @@ public class ApiKeyAuthenticationMiddleware
|
||||
{
|
||||
private readonly RequestDelegate _next;
|
||||
private readonly ILogger<ApiKeyAuthenticationMiddleware> _logger;
|
||||
private readonly ConcurrentDictionary<string, ApiKeyConfiguration> _apiKeys;
|
||||
private readonly ConcurrentDictionary<string, ApiKeyConfiguration> apiKeys;
|
||||
|
||||
/// <summary>
|
||||
/// Initializes a new instance of the <see cref="ApiKeyAuthenticationMiddleware"/> class.
|
||||
@@ -22,27 +22,20 @@ public class ApiKeyAuthenticationMiddleware
|
||||
/// <param name="next">The next middleware in the pipeline.</param>
|
||||
/// <param name="authOptions">The authentication configuration options.</param>
|
||||
/// <param name="logger">The logger instance.</param>
|
||||
public ApiKeyAuthenticationMiddleware(
|
||||
RequestDelegate next,
|
||||
IOptionsMonitor<AuthenticationConfiguration> authOptions,
|
||||
ILogger<ApiKeyAuthenticationMiddleware> logger)
|
||||
public ApiKeyAuthenticationMiddleware(RequestDelegate next, IOptionsMonitor<AuthenticationConfiguration> authOptions, ILogger<ApiKeyAuthenticationMiddleware> logger)
|
||||
{
|
||||
_next = next;
|
||||
_logger = logger;
|
||||
apiKeys = new ConcurrentDictionary<string, ApiKeyConfiguration>();
|
||||
|
||||
// Build lookup dictionary from configuration
|
||||
_apiKeys = new ConcurrentDictionary<string, ApiKeyConfiguration>();
|
||||
|
||||
// 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);
|
||||
}
|
||||
|
||||
/// <summary>
|
||||
@@ -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,
|
||||
|
||||
@@ -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;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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<string> 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 _);
|
||||
}
|
||||
|
||||
@@ -17,7 +17,7 @@
|
||||
},
|
||||
"AllowedHosts": "*",
|
||||
"Execution": {
|
||||
"ExecutionEnvironment": "DockerContainer"
|
||||
"ExecutionEnvironment": "AssemblyLoadContext"
|
||||
},
|
||||
"DockerExecution": {
|
||||
"DockerHost": null,
|
||||
|
||||
@@ -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)
|
||||
{
|
||||
|
||||
Reference in New Issue
Block a user