remote controller working!
This commit is contained in:
Nathan
@@ -45,50 +45,52 @@ $logStream = [System.IO.FileStream]::new(
|
||||
$logWriter = [System.IO.StreamWriter]::new($logStream, [System.Text.Encoding]::UTF8)
|
||||
$logWriter.AutoFlush = $true
|
||||
|
||||
if (-not ("UnifiedWorkers.WorkerLogSink" -as [type])) {
|
||||
Add-Type -Namespace UnifiedWorkers -Name WorkerLogSink -MemberDefinition @"
|
||||
# Create C# event handler class that doesn't require PowerShell runspace
|
||||
if (-not ("UnifiedWorkers.ProcessLogHandler" -as [type])) {
|
||||
$csharpCode = @'
|
||||
using System;
|
||||
using System.Diagnostics;
|
||||
using System.IO;
|
||||
|
||||
public sealed class WorkerLogSink
|
||||
namespace UnifiedWorkers
|
||||
{
|
||||
private readonly TextWriter _writer;
|
||||
private readonly string _prefix;
|
||||
private readonly object _sync = new object();
|
||||
|
||||
public WorkerLogSink(TextWriter writer, string prefix)
|
||||
public sealed class ProcessLogHandler
|
||||
{
|
||||
if (writer == null)
|
||||
private readonly TextWriter _writer;
|
||||
private readonly string _prefix;
|
||||
private readonly object _lock = new object();
|
||||
|
||||
public ProcessLogHandler(TextWriter writer, string prefix)
|
||||
{
|
||||
throw new ArgumentNullException(\"writer\");
|
||||
}
|
||||
if (prefix == null)
|
||||
{
|
||||
throw new ArgumentNullException(\"prefix\");
|
||||
_writer = writer ?? throw new ArgumentNullException("writer");
|
||||
_prefix = prefix ?? throw new ArgumentNullException("prefix");
|
||||
}
|
||||
|
||||
_writer = writer;
|
||||
_prefix = prefix;
|
||||
}
|
||||
|
||||
public void OnData(object sender, DataReceivedEventArgs e)
|
||||
{
|
||||
if (e == null || string.IsNullOrEmpty(e.Data))
|
||||
public void OnDataReceived(object sender, DataReceivedEventArgs e)
|
||||
{
|
||||
return;
|
||||
}
|
||||
if (e == null || string.IsNullOrEmpty(e.Data))
|
||||
{
|
||||
return;
|
||||
}
|
||||
|
||||
lock (_sync)
|
||||
{
|
||||
var timestamp = DateTime.UtcNow.ToString("u");
|
||||
var line = string.Format("[{0} {1}] {2}", _prefix, timestamp, e.Data);
|
||||
_writer.WriteLine(line);
|
||||
_writer.Flush();
|
||||
lock (_lock)
|
||||
{
|
||||
try
|
||||
{
|
||||
var timestamp = DateTime.UtcNow.ToString("u");
|
||||
_writer.WriteLine(string.Format("[{0} {1}] {2}", _prefix, timestamp, e.Data));
|
||||
_writer.Flush();
|
||||
}
|
||||
catch
|
||||
{
|
||||
// Ignore write errors to prevent cascading failures
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
"@
|
||||
'@
|
||||
Add-Type -TypeDefinition $csharpCode -ErrorAction Stop
|
||||
}
|
||||
|
||||
function Write-LogLine {
|
||||
@@ -99,20 +101,39 @@ function Write-LogLine {
|
||||
|
||||
if (-not $logWriter) { return }
|
||||
$timestamp = (Get-Date).ToString('u')
|
||||
$logWriter.WriteLine("[{0} {1}] {2}" -f $Prefix, $timestamp, $Message)
|
||||
$logWriter.WriteLine("[$Prefix $timestamp] $Message")
|
||||
}
|
||||
|
||||
function Write-ControllerLog {
|
||||
param([string]$Message)
|
||||
Write-LogLine -Prefix 'CTRL' -Message $Message
|
||||
}
|
||||
|
||||
function Write-FatalLog {
|
||||
param([string]$Message)
|
||||
|
||||
try {
|
||||
Write-ControllerLog $Message
|
||||
}
|
||||
catch {
|
||||
$timestamp = (Get-Date).ToString('u')
|
||||
$fallback = "[CTRL $timestamp] $Message"
|
||||
try {
|
||||
[System.IO.File]::AppendAllText($logPath, $fallback + [Environment]::NewLine, [System.Text.Encoding]::UTF8)
|
||||
}
|
||||
catch {
|
||||
# last resort: write to host
|
||||
Write-Error $fallback
|
||||
}
|
||||
}
|
||||
}
|
||||
# endregion
|
||||
|
||||
# region Helpers
|
||||
|
||||
function Resolve-PayloadBase64 {
|
||||
if ($PayloadBase64) {
|
||||
return $PayloadBase64
|
||||
return $PayloadBase64.Trim()
|
||||
}
|
||||
|
||||
if ($PayloadBase64Path) {
|
||||
@@ -120,7 +141,12 @@ function Resolve-PayloadBase64 {
|
||||
throw "Payload file '$PayloadBase64Path' not found."
|
||||
}
|
||||
|
||||
return (Get-Content -Path $PayloadBase64Path -Raw)
|
||||
$content = Get-Content -Path $PayloadBase64Path -Raw
|
||||
if ([string]::IsNullOrWhiteSpace($content)) {
|
||||
throw "Payload file '$PayloadBase64Path' is empty."
|
||||
}
|
||||
|
||||
return $content.Trim()
|
||||
}
|
||||
|
||||
throw "No payload data provided to controller."
|
||||
@@ -168,117 +194,137 @@ function Get-PendingCommands {
|
||||
}
|
||||
# endregion
|
||||
|
||||
# record initial state before launching worker
|
||||
Write-Metadata -Status 'initializing' -WorkerPid $null -ControllerPid $PID -Restarts 0
|
||||
|
||||
$resolvedPayloadBase64 = Resolve-PayloadBase64
|
||||
$PayloadBase64 = $resolvedPayloadBase64
|
||||
|
||||
try {
|
||||
# Write payload script to disk
|
||||
$payloadBytes = [Convert]::FromBase64String($resolvedPayloadBase64)
|
||||
[IO.File]::WriteAllBytes($payloadPath, $payloadBytes)
|
||||
Write-ControllerLog "Payload written to $payloadPath"
|
||||
}
|
||||
catch {
|
||||
Write-Error "Unable to write payload: $($_.Exception.Message)"
|
||||
exit 1
|
||||
}
|
||||
# record initial state before launching worker
|
||||
Write-Metadata -Status 'initializing' -WorkerPid $null -ControllerPid $PID -Restarts 0
|
||||
|
||||
$restartCount = 0
|
||||
$controllerPid = $PID
|
||||
$resolvedPayloadBase64 = Resolve-PayloadBase64
|
||||
$PayloadBase64 = $resolvedPayloadBase64
|
||||
|
||||
while ($restartCount -le $MaxRestarts) {
|
||||
try {
|
||||
# Initialize worker process
|
||||
$psi = [System.Diagnostics.ProcessStartInfo]::new()
|
||||
$pwsh = Get-Command pwsh -ErrorAction SilentlyContinue
|
||||
if ($pwsh) {
|
||||
$psi.FileName = $pwsh.Source
|
||||
}
|
||||
else {
|
||||
$psi.FileName = (Get-Command powershell -ErrorAction Stop).Source
|
||||
}
|
||||
|
||||
$psi.Arguments = "-NoLogo -NoProfile -ExecutionPolicy Bypass -File `"$payloadPath`""
|
||||
$psi.UseShellExecute = $false
|
||||
$psi.RedirectStandardInput = $true
|
||||
$psi.RedirectStandardOutput = $true
|
||||
$psi.RedirectStandardError = $true
|
||||
$psi.CreateNoWindow = $true
|
||||
|
||||
$workerProcess = New-Object System.Diagnostics.Process
|
||||
$workerProcess.StartInfo = $psi
|
||||
|
||||
if (-not $workerProcess.Start()) {
|
||||
throw "Failed to start worker process."
|
||||
}
|
||||
|
||||
Write-ControllerLog "Worker process started with PID $($workerProcess.Id)"
|
||||
Write-Metadata -Status 'running' -WorkerPid $workerProcess.Id -ControllerPid $controllerPid -Restarts $restartCount
|
||||
|
||||
$stdoutSink = [UnifiedWorkers.WorkerLogSink]::new($logWriter, 'OUT')
|
||||
$stderrSink = [UnifiedWorkers.WorkerLogSink]::new($logWriter, 'ERR')
|
||||
$outputHandler = [System.Diagnostics.DataReceivedEventHandler]$stdoutSink.OnData
|
||||
$errorHandler = [System.Diagnostics.DataReceivedEventHandler]$stderrSink.OnData
|
||||
|
||||
$workerProcess.add_OutputDataReceived($outputHandler)
|
||||
$workerProcess.add_ErrorDataReceived($errorHandler)
|
||||
$workerProcess.BeginOutputReadLine()
|
||||
$workerProcess.BeginErrorReadLine()
|
||||
|
||||
while (-not $workerProcess.HasExited) {
|
||||
$commands = Get-PendingCommands
|
||||
foreach ($command in $commands) {
|
||||
$trimmed = $command.Trim()
|
||||
if (-not $trimmed) { continue }
|
||||
|
||||
Write-ControllerLog "Received command '$trimmed'"
|
||||
try {
|
||||
$workerProcess.StandardInput.WriteLine($trimmed)
|
||||
$workerProcess.StandardInput.Flush()
|
||||
}
|
||||
catch {
|
||||
Write-ControllerLog "Failed to forward command '$trimmed': $($_.Exception.Message)"
|
||||
}
|
||||
|
||||
if ($trimmed -ieq 'quit') {
|
||||
Write-ControllerLog "Quit command issued. Waiting for worker to exit."
|
||||
}
|
||||
}
|
||||
|
||||
Start-Sleep -Milliseconds 500
|
||||
}
|
||||
|
||||
$exitCode = $workerProcess.ExitCode
|
||||
Write-ControllerLog "Worker exited with code $exitCode"
|
||||
Write-Metadata -Status 'stopped' -WorkerPid $null -ControllerPid $controllerPid -Restarts $restartCount -LastExitCode $exitCode
|
||||
|
||||
if ($exitCode -eq 0) {
|
||||
break
|
||||
}
|
||||
# Write payload script to disk
|
||||
# The payload is base64-encoded UTF-16 (Unicode), so decode it properly
|
||||
Write-ControllerLog "Decoding payload base64 (length: $($resolvedPayloadBase64.Length))"
|
||||
$payloadBytes = [Convert]::FromBase64String($resolvedPayloadBase64)
|
||||
Write-ControllerLog "Decoded payload to $($payloadBytes.Length) bytes"
|
||||
|
||||
# Convert UTF-16 bytes back to string, then write as UTF-8 (PowerShell's preferred encoding)
|
||||
$payloadText = [Text.Encoding]::Unicode.GetString($payloadBytes)
|
||||
[IO.File]::WriteAllText($payloadPath, $payloadText, [Text.Encoding]::UTF8)
|
||||
Write-ControllerLog "Payload written to $payloadPath ($($payloadText.Length) characters)"
|
||||
}
|
||||
catch {
|
||||
Write-ControllerLog "Controller error: $($_.Exception.Message)"
|
||||
Write-FatalLog "Unable to write payload: $($_.Exception.Message)"
|
||||
if ($_.Exception.InnerException) {
|
||||
Write-FatalLog "Inner exception: $($_.Exception.InnerException.Message)"
|
||||
}
|
||||
throw
|
||||
}
|
||||
|
||||
$restartCount++
|
||||
if ($restartCount -gt $MaxRestarts) {
|
||||
Write-ControllerLog "Maximum restart attempts reached. Controller stopping."
|
||||
break
|
||||
$restartCount = 0
|
||||
$controllerPid = $PID
|
||||
|
||||
while ($restartCount -le $MaxRestarts) {
|
||||
try {
|
||||
# Initialize worker process
|
||||
$psi = [System.Diagnostics.ProcessStartInfo]::new()
|
||||
$pwsh = Get-Command pwsh -ErrorAction SilentlyContinue
|
||||
if ($pwsh) {
|
||||
$psi.FileName = $pwsh.Source
|
||||
}
|
||||
else {
|
||||
$psi.FileName = (Get-Command powershell -ErrorAction Stop).Source
|
||||
}
|
||||
|
||||
$psi.Arguments = "-NoLogo -NoProfile -ExecutionPolicy Bypass -File `"$payloadPath`""
|
||||
$psi.UseShellExecute = $false
|
||||
$psi.RedirectStandardInput = $true
|
||||
$psi.RedirectStandardOutput = $true
|
||||
$psi.RedirectStandardError = $true
|
||||
$psi.CreateNoWindow = $true
|
||||
|
||||
$workerProcess = New-Object System.Diagnostics.Process
|
||||
$workerProcess.StartInfo = $psi
|
||||
|
||||
if (-not $workerProcess.Start()) {
|
||||
throw "Failed to start worker process."
|
||||
}
|
||||
|
||||
Write-ControllerLog "Worker process started with PID $($workerProcess.Id)"
|
||||
Write-Metadata -Status 'running' -WorkerPid $workerProcess.Id -ControllerPid $controllerPid -Restarts $restartCount
|
||||
|
||||
$stdoutHandler = [UnifiedWorkers.ProcessLogHandler]::new($logWriter, 'OUT')
|
||||
$stderrHandler = [UnifiedWorkers.ProcessLogHandler]::new($logWriter, 'ERR')
|
||||
|
||||
$outputHandler = [System.Diagnostics.DataReceivedEventHandler]$stdoutHandler.OnDataReceived
|
||||
$errorHandler = [System.Diagnostics.DataReceivedEventHandler]$stderrHandler.OnDataReceived
|
||||
|
||||
$workerProcess.add_OutputDataReceived($outputHandler)
|
||||
$workerProcess.add_ErrorDataReceived($errorHandler)
|
||||
$workerProcess.BeginOutputReadLine()
|
||||
$workerProcess.BeginErrorReadLine()
|
||||
|
||||
while (-not $workerProcess.HasExited) {
|
||||
$commands = Get-PendingCommands
|
||||
foreach ($command in $commands) {
|
||||
$trimmed = $command.Trim()
|
||||
if (-not $trimmed) { continue }
|
||||
|
||||
Write-ControllerLog "Received command '$trimmed'"
|
||||
try {
|
||||
$workerProcess.StandardInput.WriteLine($trimmed)
|
||||
$workerProcess.StandardInput.Flush()
|
||||
}
|
||||
catch {
|
||||
Write-ControllerLog "Failed to forward command '$trimmed': $($_.Exception.Message)"
|
||||
}
|
||||
|
||||
if ($trimmed -ieq 'quit') {
|
||||
Write-ControllerLog "Quit command issued. Waiting for worker to exit."
|
||||
}
|
||||
}
|
||||
|
||||
Start-Sleep -Milliseconds 500
|
||||
}
|
||||
|
||||
$exitCode = $workerProcess.ExitCode
|
||||
Write-ControllerLog "Worker exited with code $exitCode"
|
||||
Write-Metadata -Status 'stopped' -WorkerPid $null -ControllerPid $controllerPid -Restarts $restartCount -LastExitCode $exitCode
|
||||
|
||||
if ($exitCode -eq 0) {
|
||||
break
|
||||
}
|
||||
}
|
||||
catch {
|
||||
Write-ControllerLog "Controller error: $($_.Exception.Message)"
|
||||
}
|
||||
|
||||
$restartCount++
|
||||
if ($restartCount -gt $MaxRestarts) {
|
||||
Write-ControllerLog "Maximum restart attempts reached. Controller stopping."
|
||||
break
|
||||
}
|
||||
|
||||
Write-ControllerLog "Restarting worker in $RestartDelaySeconds seconds (attempt $restartCount of $MaxRestarts)."
|
||||
Start-Sleep -Seconds $RestartDelaySeconds
|
||||
}
|
||||
|
||||
Write-ControllerLog "Restarting worker in $RestartDelaySeconds seconds (attempt $restartCount of $MaxRestarts)."
|
||||
Start-Sleep -Seconds $RestartDelaySeconds
|
||||
}
|
||||
|
||||
Write-Metadata -Status 'inactive' -WorkerPid $null -ControllerPid $controllerPid -Restarts $restartCount
|
||||
Write-ControllerLog "Controller exiting."
|
||||
|
||||
if ($logWriter) {
|
||||
$logWriter.Dispose()
|
||||
}
|
||||
if ($logStream) {
|
||||
$logStream.Dispose()
|
||||
Write-Metadata -Status 'inactive' -WorkerPid $null -ControllerPid $controllerPid -Restarts $restartCount
|
||||
Write-ControllerLog "Controller exiting."
|
||||
}
|
||||
catch {
|
||||
Write-FatalLog "Fatal controller error: $($_.Exception.Message)"
|
||||
if ($_.ScriptStackTrace) {
|
||||
Write-FatalLog "Stack: $($_.ScriptStackTrace)"
|
||||
}
|
||||
throw
|
||||
}
|
||||
finally {
|
||||
if ($logWriter) {
|
||||
$logWriter.Dispose()
|
||||
}
|
||||
if ($logStream) {
|
||||
$logStream.Dispose()
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user