feat: register local agent project workspaces
This commit is contained in:
Hamza Ayed
@@ -0,0 +1,218 @@
|
||||
param(
|
||||
[switch]$DirectElevated,
|
||||
[string]$ReportPath
|
||||
)
|
||||
|
||||
$ErrorActionPreference = 'Stop'
|
||||
|
||||
function Test-Administrator {
|
||||
$identity = [Security.Principal.WindowsIdentity]::GetCurrent()
|
||||
$principal = [Security.Principal.WindowsPrincipal]::new($identity)
|
||||
return $principal.IsInRole([Security.Principal.WindowsBuiltInRole]::Administrator)
|
||||
}
|
||||
|
||||
if (-not $DirectElevated) {
|
||||
if (Test-Administrator) {
|
||||
$DirectElevated = $true
|
||||
} else {
|
||||
$ReportPath = Join-Path $env:TEMP ("SovereignAI-VHDQuota-{0}.json" -f [guid]::NewGuid().ToString('N'))
|
||||
$powershell = Join-Path $env:WINDIR 'System32\WindowsPowerShell\v1.0\powershell.exe'
|
||||
$scriptPath = $MyInvocation.MyCommand.Path
|
||||
$arguments = @(
|
||||
'-NoProfile',
|
||||
'-ExecutionPolicy', 'Bypass',
|
||||
'-File', ('"{0}"' -f $scriptPath),
|
||||
'-DirectElevated',
|
||||
'-ReportPath', ('"{0}"' -f $ReportPath)
|
||||
) -join ' '
|
||||
try {
|
||||
$child = Start-Process -FilePath $powershell -ArgumentList $arguments `
|
||||
-Verb RunAs -WindowStyle Hidden -PassThru -Wait
|
||||
} catch {
|
||||
Write-Error "Windows did not grant the required administrator token: $($_.Exception.Message)"
|
||||
exit 1
|
||||
}
|
||||
if (Test-Path -LiteralPath $ReportPath -PathType Leaf) {
|
||||
Get-Content -LiteralPath $ReportPath -Raw
|
||||
Remove-Item -LiteralPath $ReportPath -Force
|
||||
} else {
|
||||
Write-Error "The elevated quota probe exited $($child.ExitCode) without producing its report."
|
||||
}
|
||||
exit $child.ExitCode
|
||||
}
|
||||
}
|
||||
|
||||
if (-not (Test-Administrator)) {
|
||||
throw 'The direct probe requires an elevated administrator token.'
|
||||
}
|
||||
|
||||
$probeParent = Join-Path $env:LOCALAPPDATA 'SovereignAI\execution-sandbox-probes'
|
||||
$probeDirectory = Join-Path $probeParent ("vhd-quota-{0}" -f [guid]::NewGuid().ToString('N'))
|
||||
$vhdPath = Join-Path $probeDirectory 'quota-probe.vhdx'
|
||||
$mountPath = Join-Path $probeDirectory 'volume'
|
||||
$diskpartScript = Join-Path $probeDirectory 'diskpart.txt'
|
||||
$detachScript = Join-Path $probeDirectory 'detach.txt'
|
||||
$testFile = Join-Path $mountPath 'capacity-probe.bin'
|
||||
$maximumMegabytes = 64
|
||||
$report = $null
|
||||
$resultCode = 0
|
||||
$cleanupError = $null
|
||||
$mounted = $false
|
||||
|
||||
try {
|
||||
$systemDrive = [System.IO.Path]::GetPathRoot($env:LOCALAPPDATA)
|
||||
$availableBytes = [System.IO.DriveInfo]::new($systemDrive).AvailableFreeSpace
|
||||
if ($availableBytes -lt 256MB) {
|
||||
throw 'Less than 256 MiB is free on the system drive; refusing to create the 64 MiB probe disk.'
|
||||
}
|
||||
|
||||
New-Item -ItemType Directory -Path $probeDirectory -Force | Out-Null
|
||||
New-Item -ItemType Directory -Path $mountPath -Force | Out-Null
|
||||
$diskpartCommands = @(
|
||||
"create vdisk file=`"$vhdPath`" maximum=$maximumMegabytes type=expandable",
|
||||
"select vdisk file=`"$vhdPath`"",
|
||||
'attach vdisk',
|
||||
'create partition primary',
|
||||
'format fs=ntfs quick label=SOVEREIGNAI_PROBE',
|
||||
"assign mount=`"$mountPath`"",
|
||||
'exit'
|
||||
)
|
||||
Set-Content -LiteralPath $diskpartScript -Value $diskpartCommands -Encoding ascii
|
||||
$diskpartOutput = & "$env:WINDIR\System32\diskpart.exe" /s $diskpartScript 2>&1 | Out-String
|
||||
$diskpartExitCode = $LASTEXITCODE
|
||||
if ($diskpartExitCode -ne 0) {
|
||||
throw "DiskPart failed with exit code $diskpartExitCode. $diskpartOutput"
|
||||
}
|
||||
|
||||
$deadline = (Get-Date).AddSeconds(15)
|
||||
$volumes = @(Get-Volume -FileSystemLabel 'SOVEREIGNAI_PROBE' -ErrorAction SilentlyContinue)
|
||||
while ($volumes.Count -eq 0 -and (Get-Date) -lt $deadline) {
|
||||
Start-Sleep -Milliseconds 250
|
||||
$volumes = @(Get-Volume -FileSystemLabel 'SOVEREIGNAI_PROBE' -ErrorAction SilentlyContinue)
|
||||
}
|
||||
if ($volumes.Count -ne 1) {
|
||||
$mountvolOutput = & "$env:WINDIR\System32\mountvol.exe" $mountPath /L 2>&1 | Out-String
|
||||
throw "Expected exactly one formatted probe volume. mountvol: $mountvolOutput DiskPart: $diskpartOutput"
|
||||
}
|
||||
$volume = $volumes[0]
|
||||
$mountvolOutput = (& "$env:WINDIR\System32\mountvol.exe" $mountPath /L 2>&1 | Out-String).Trim()
|
||||
if ($LASTEXITCODE -ne 0 -or
|
||||
-not $mountvolOutput.Trim().Equals($volume.Path.Trim(), [StringComparison]::OrdinalIgnoreCase)) {
|
||||
throw "The test mount path does not resolve to the formatted probe VHDX. mountvol='$mountvolOutput'; volume='$($volume.Path)'. DiskPart: $diskpartOutput"
|
||||
}
|
||||
if ($volume.Size -gt (($maximumMegabytes + 2) * 1MB) -or $volume.Size -lt (48MB)) {
|
||||
$mountvolOutput = & "$env:WINDIR\System32\mountvol.exe" $mountPath /L 2>&1 | Out-String
|
||||
throw "Unexpected probe volume capacity: $($volume.Size) bytes. mountvol: $mountvolOutput DiskPart: $diskpartOutput"
|
||||
}
|
||||
$mounted = $true
|
||||
|
||||
$buffer = [byte[]]::new(1MB)
|
||||
for ($index = 0; $index -lt $buffer.Length; $index++) {
|
||||
$buffer[$index] = [byte](($index * 31 + 97) % 251)
|
||||
}
|
||||
$writtenBytes = [long]0
|
||||
$diskFullException = $null
|
||||
$stream = [System.IO.FileStream]::new(
|
||||
$testFile,
|
||||
[System.IO.FileMode]::CreateNew,
|
||||
[System.IO.FileAccess]::Write,
|
||||
[System.IO.FileShare]::None,
|
||||
$buffer.Length,
|
||||
[System.IO.FileOptions]::SequentialScan
|
||||
)
|
||||
try {
|
||||
while ($true) {
|
||||
$stream.Write($buffer, 0, $buffer.Length)
|
||||
$writtenBytes += $buffer.Length
|
||||
}
|
||||
} catch [System.IO.IOException] {
|
||||
$diskFullException = $_.Exception
|
||||
} finally {
|
||||
$stream.Dispose()
|
||||
}
|
||||
|
||||
$errorCode = if ($diskFullException) { $diskFullException.HResult -band 0xffff } else { $null }
|
||||
$volume = Get-Volume -UniqueId $volume.UniqueId
|
||||
$freeAtLimit = $volume.SizeRemaining
|
||||
if (-not $diskFullException -or $errorCode -ne 112) {
|
||||
throw "The write did not stop with ERROR_DISK_FULL (112). HResult=$($diskFullException.HResult); message=$($diskFullException.Message)"
|
||||
}
|
||||
if ($freeAtLimit -gt 1MB) {
|
||||
throw "The write stopped with $freeAtLimit bytes still free; the volume-capacity boundary was not confirmed."
|
||||
}
|
||||
|
||||
$vhdBytes = (Get-Item -LiteralPath $vhdPath).Length
|
||||
$report = [pscustomobject]@{
|
||||
passed = $true
|
||||
administrator_token = $true
|
||||
diskpart_exit_code = $diskpartExitCode
|
||||
volume_label = $volume.FileSystemLabel
|
||||
virtual_volume_bytes = $volume.Size
|
||||
bytes_written_before_error = $writtenBytes
|
||||
error_code = $errorCode
|
||||
error_message = $diskFullException.Message
|
||||
free_bytes_at_limit = $freeAtLimit
|
||||
dynamic_vhdx_bytes_at_limit = $vhdBytes
|
||||
probe_vhdx_max_megabytes = $maximumMegabytes
|
||||
isolation = 'temporary NTFS filesystem inside a dynamically expanding VHDX with a fixed virtual capacity'
|
||||
}
|
||||
} catch {
|
||||
$resultCode = 1
|
||||
$report = [pscustomobject]@{
|
||||
passed = $false
|
||||
administrator_token = (Test-Administrator)
|
||||
error = $_.Exception.Message
|
||||
probe_vhdx_max_megabytes = $maximumMegabytes
|
||||
}
|
||||
} finally {
|
||||
try {
|
||||
if ($stream) { $stream.Dispose() }
|
||||
if (Test-Path -LiteralPath $testFile -PathType Leaf) {
|
||||
Remove-Item -LiteralPath $testFile -Force
|
||||
}
|
||||
if (Test-Path -LiteralPath $vhdPath -PathType Leaf) {
|
||||
$mountItem = Get-Item -LiteralPath $mountPath -Force -ErrorAction SilentlyContinue
|
||||
if ($mountItem -and ($mountItem.Attributes -band [System.IO.FileAttributes]::ReparsePoint)) {
|
||||
& "$env:WINDIR\System32\mountvol.exe" $mountPath /D | Out-Null
|
||||
if ($LASTEXITCODE -ne 0) { throw "MountVol could not remove the exact temporary mount point ($LASTEXITCODE)." }
|
||||
}
|
||||
Set-Content -LiteralPath $detachScript -Value @(
|
||||
"select vdisk file=`"$vhdPath`"",
|
||||
'detach vdisk',
|
||||
'exit'
|
||||
) -Encoding ascii
|
||||
& "$env:WINDIR\System32\diskpart.exe" /s $detachScript 2>&1 | Out-Null
|
||||
if ($LASTEXITCODE -ne 0) { throw "DiskPart could not detach the temporary VHDX ($LASTEXITCODE)." }
|
||||
Remove-Item -LiteralPath $vhdPath -Force
|
||||
}
|
||||
foreach ($file in @($diskpartScript, $detachScript)) {
|
||||
if (Test-Path -LiteralPath $file -PathType Leaf) { Remove-Item -LiteralPath $file -Force }
|
||||
}
|
||||
if (Test-Path -LiteralPath $mountPath -PathType Container) {
|
||||
$mountItem = Get-Item -LiteralPath $mountPath -Force
|
||||
if (($mountItem.Attributes -band [System.IO.FileAttributes]::ReparsePoint) -eq 0 -and
|
||||
-not (Get-ChildItem -LiteralPath $mountPath -Force)) {
|
||||
Remove-Item -LiteralPath $mountPath -Force
|
||||
}
|
||||
}
|
||||
if ((Test-Path -LiteralPath $probeDirectory -PathType Container) -and
|
||||
-not (Get-ChildItem -LiteralPath $probeDirectory -Force)) {
|
||||
Remove-Item -LiteralPath $probeDirectory -Force
|
||||
}
|
||||
if ((Test-Path -LiteralPath $probeParent -PathType Container) -and
|
||||
-not (Get-ChildItem -LiteralPath $probeParent -Force)) {
|
||||
Remove-Item -LiteralPath $probeParent -Force
|
||||
}
|
||||
} catch {
|
||||
$cleanupError = $_.Exception.Message
|
||||
$resultCode = 1
|
||||
}
|
||||
if ($cleanupError) { $report | Add-Member -NotePropertyName cleanup_error -NotePropertyValue $cleanupError -Force }
|
||||
if ($ReportPath) {
|
||||
$report | ConvertTo-Json -Depth 5 | Set-Content -LiteralPath $ReportPath -Encoding utf8
|
||||
} else {
|
||||
$report | ConvertTo-Json -Depth 5
|
||||
}
|
||||
}
|
||||
|
||||
exit $resultCode
|
||||
Reference in New Issue
Block a user