Compare commits
2 Commits
18683b4ad0
...
428d9bac6b
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
428d9bac6b | ||
|
|
e568e02cdf |
File diff suppressed because one or more lines are too long
@@ -88,7 +88,16 @@ try {
|
||||
continue
|
||||
}
|
||||
|
||||
if ($input.Trim().ToLower() -eq 'detach') {
|
||||
$normalized = $input.Trim()
|
||||
if ($normalized.StartsWith(':')) {
|
||||
$normalized = $normalized.TrimStart(':').Trim()
|
||||
}
|
||||
|
||||
if ($normalized.Length -eq 0) {
|
||||
continue
|
||||
}
|
||||
|
||||
if ($normalized.ToLowerInvariant() -eq 'detach') {
|
||||
break
|
||||
}
|
||||
|
||||
@@ -97,9 +106,9 @@ try {
|
||||
}
|
||||
finally {
|
||||
if ($logJob) {
|
||||
Stop-Job -Job $logJob -Force -ErrorAction SilentlyContinue
|
||||
Stop-Job -Job $logJob -ErrorAction SilentlyContinue
|
||||
Receive-Job -Job $logJob -ErrorAction SilentlyContinue | Out-Null
|
||||
Remove-Job -Job $logJob -Force -ErrorAction SilentlyContinue
|
||||
Remove-Job -Job $logJob -ErrorAction SilentlyContinue
|
||||
}
|
||||
|
||||
Write-Host "Detached from worker $WorkerName." -ForegroundColor Cyan
|
||||
|
||||
@@ -5,9 +5,10 @@ param(
|
||||
[Parameter(Mandatory = $true)]
|
||||
[string]$WorkerType,
|
||||
|
||||
[Parameter(Mandatory = $true)]
|
||||
[string]$PayloadBase64,
|
||||
|
||||
[string]$PayloadBase64Path,
|
||||
|
||||
[string]$DataRoot = (Join-Path ([Environment]::GetFolderPath('LocalApplicationData')) 'UnifiedWorkers'),
|
||||
|
||||
[int]$MaxRestarts = 5,
|
||||
@@ -16,6 +17,7 @@ param(
|
||||
)
|
||||
|
||||
$ErrorActionPreference = 'Stop'
|
||||
[System.Management.Automation.Runspaces.Runspace]::DefaultRunspace = $Host.Runspace
|
||||
|
||||
# region Path setup
|
||||
$workerRoot = Join-Path -Path $DataRoot -ChildPath $WorkerType
|
||||
@@ -33,11 +35,121 @@ $commandPath = Join-Path -Path $stateRoot -ChildPath 'commands.txt'
|
||||
$payloadPath = Join-Path -Path $stateRoot -ChildPath "payload.ps1"
|
||||
# endregion
|
||||
|
||||
# region Helpers
|
||||
# region Logging
|
||||
$logStream = [System.IO.FileStream]::new(
|
||||
$logPath,
|
||||
[System.IO.FileMode]::Append,
|
||||
[System.IO.FileAccess]::Write,
|
||||
[System.IO.FileShare]::ReadWrite
|
||||
)
|
||||
$logWriter = [System.IO.StreamWriter]::new($logStream, [System.Text.Encoding]::UTF8)
|
||||
$logWriter.AutoFlush = $true
|
||||
|
||||
# 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;
|
||||
|
||||
namespace UnifiedWorkers
|
||||
{
|
||||
public sealed class ProcessLogHandler
|
||||
{
|
||||
private readonly TextWriter _writer;
|
||||
private readonly string _prefix;
|
||||
private readonly object _lock = new object();
|
||||
|
||||
public ProcessLogHandler(TextWriter writer, string prefix)
|
||||
{
|
||||
_writer = writer ?? throw new ArgumentNullException("writer");
|
||||
_prefix = prefix ?? throw new ArgumentNullException("prefix");
|
||||
}
|
||||
|
||||
public void OnDataReceived(object sender, DataReceivedEventArgs e)
|
||||
{
|
||||
if (e == null || string.IsNullOrEmpty(e.Data))
|
||||
{
|
||||
return;
|
||||
}
|
||||
|
||||
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 {
|
||||
param(
|
||||
[string]$Prefix,
|
||||
[string]$Message
|
||||
)
|
||||
|
||||
if (-not $logWriter) { return }
|
||||
$timestamp = (Get-Date).ToString('u')
|
||||
$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')
|
||||
Add-Content -Path $logPath -Value "[CTRL $timestamp] $Message" -Encoding UTF8
|
||||
$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.Trim()
|
||||
}
|
||||
|
||||
if ($PayloadBase64Path) {
|
||||
if (-not (Test-Path $PayloadBase64Path)) {
|
||||
throw "Payload file '$PayloadBase64Path' not found."
|
||||
}
|
||||
|
||||
$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."
|
||||
}
|
||||
|
||||
function Write-Metadata {
|
||||
@@ -82,18 +194,31 @@ function Get-PendingCommands {
|
||||
}
|
||||
# endregion
|
||||
|
||||
try {
|
||||
# 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($PayloadBase64)
|
||||
[IO.File]::WriteAllBytes($payloadPath, $payloadBytes)
|
||||
Write-ControllerLog "Payload written to $payloadPath"
|
||||
# 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-Error "Unable to write payload: $($_.Exception.Message)"
|
||||
exit 1
|
||||
Write-FatalLog "Unable to write payload: $($_.Exception.Message)"
|
||||
if ($_.Exception.InnerException) {
|
||||
Write-FatalLog "Inner exception: $($_.Exception.InnerException.Message)"
|
||||
}
|
||||
throw
|
||||
}
|
||||
|
||||
$restartCount = 0
|
||||
@@ -128,23 +253,11 @@ while ($restartCount -le $MaxRestarts) {
|
||||
Write-ControllerLog "Worker process started with PID $($workerProcess.Id)"
|
||||
Write-Metadata -Status 'running' -WorkerPid $workerProcess.Id -ControllerPid $controllerPid -Restarts $restartCount
|
||||
|
||||
$logWriter = [System.IO.StreamWriter]::new($logPath, $true, [System.Text.Encoding]::UTF8)
|
||||
$logWriter.AutoFlush = $true
|
||||
$stdoutHandler = [UnifiedWorkers.ProcessLogHandler]::new($logWriter, 'OUT')
|
||||
$stderrHandler = [UnifiedWorkers.ProcessLogHandler]::new($logWriter, 'ERR')
|
||||
|
||||
$outputHandler = [System.Diagnostics.DataReceivedEventHandler]{
|
||||
param($s, $e)
|
||||
if ($e.Data) {
|
||||
$ts = (Get-Date).ToString('u')
|
||||
$logWriter.WriteLine("[OUT $ts] $($e.Data)")
|
||||
}
|
||||
}
|
||||
$errorHandler = [System.Diagnostics.DataReceivedEventHandler]{
|
||||
param($s, $e)
|
||||
if ($e.Data) {
|
||||
$ts = (Get-Date).ToString('u')
|
||||
$logWriter.WriteLine("[ERR $ts] $($e.Data)")
|
||||
}
|
||||
}
|
||||
$outputHandler = [System.Diagnostics.DataReceivedEventHandler]$stdoutHandler.OnDataReceived
|
||||
$errorHandler = [System.Diagnostics.DataReceivedEventHandler]$stderrHandler.OnDataReceived
|
||||
|
||||
$workerProcess.add_OutputDataReceived($outputHandler)
|
||||
$workerProcess.add_ErrorDataReceived($errorHandler)
|
||||
@@ -198,4 +311,20 @@ while ($restartCount -le $MaxRestarts) {
|
||||
|
||||
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()
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -402,8 +402,11 @@ New-Item -ItemType Directory -Path `$logsRoot -Force | Out-Null
|
||||
New-Item -ItemType Directory -Path `$stateRoot -Force | Out-Null
|
||||
`$logPath = Join-Path `$logsRoot 'worker.log'
|
||||
`$commandPath = Join-Path `$stateRoot 'commands.txt'
|
||||
`$payloadPath = Join-Path `$stateRoot 'payload.ps1'
|
||||
`$payloadBase64Path = Join-Path `$stateRoot 'payload.b64'
|
||||
if (-not (Test-Path `$logPath)) { New-Item -Path `$logPath -ItemType File -Force | Out-Null }
|
||||
if (-not (Test-Path `$commandPath)) { New-Item -Path `$commandPath -ItemType File -Force | Out-Null }
|
||||
[IO.File]::WriteAllText(`$payloadBase64Path, `$payloadBase64, [System.Text.Encoding]::UTF8)
|
||||
`$metaPath = Join-Path `$instanceRoot 'state\worker-info.json'
|
||||
`$controllerPath = Join-Path `$dataRoot 'controller.ps1'
|
||||
|
||||
@@ -437,7 +440,7 @@ if (`$shouldStart) {
|
||||
LastExitCode = `$null
|
||||
LogPath = `$logPath
|
||||
CommandPath = `$commandPath
|
||||
PayloadPath = `$null
|
||||
PayloadPath = `$payloadPath
|
||||
UpdatedAtUtc = (Get-Date).ToUniversalTime()
|
||||
} | ConvertTo-Json -Depth 5
|
||||
`$initialMeta | Set-Content -Path `$metaPath -Encoding UTF8
|
||||
@@ -455,7 +458,7 @@ if (`$shouldStart) {
|
||||
'-File',"`$controllerPath",
|
||||
'-WorkerName',"`$workerName",
|
||||
'-WorkerType',"`$workerType",
|
||||
'-PayloadBase64',"`$payloadBase64"
|
||||
'-PayloadBase64Path',"`$payloadBase64Path"
|
||||
)
|
||||
|
||||
Start-Process -FilePath `$psExe -ArgumentList `$controllerArgs -WindowStyle Hidden | Out-Null
|
||||
@@ -767,7 +770,11 @@ while ($true) {
|
||||
Write-Host "5. Quit all workers" -ForegroundColor Yellow
|
||||
Write-Host "6. Exit" -ForegroundColor Yellow
|
||||
|
||||
$choice = Read-Host "Select option (1-3)"
|
||||
$choice = Read-Host "Select option (1-6)"
|
||||
|
||||
if ($choice -eq '6') {
|
||||
break
|
||||
}
|
||||
|
||||
switch ($choice) {
|
||||
'1' { Select-SheepItWorker }
|
||||
@@ -775,7 +782,6 @@ while ($true) {
|
||||
'3' { Send-SheepItCommandAll -CommandText 'pause' }
|
||||
'4' { Send-SheepItCommandAll -CommandText 'resume' }
|
||||
'5' { Send-SheepItCommandAll -CommandText 'quit' }
|
||||
'6' { break }
|
||||
default {
|
||||
Write-Host "Invalid selection." -ForegroundColor Red
|
||||
Start-Sleep -Seconds 1
|
||||
|
||||
Reference in New Issue
Block a user