diff --git a/.github/workflows/studio-windows-inference-smoke.yml b/.github/workflows/studio-windows-inference-smoke.yml index 50b733673c..db95ebbee5 100644 --- a/.github/workflows/studio-windows-inference-smoke.yml +++ b/.github/workflows/studio-windows-inference-smoke.yml @@ -131,6 +131,8 @@ jobs: if ($LASTEXITCODE) { exit $LASTEXITCODE } pwsh -NoProfile -File tests/studio/test_path_probe_access_denied.ps1 if ($LASTEXITCODE) { exit $LASTEXITCODE } + pwsh -NoProfile -File tests/studio/test_process_image_path_map.ps1 + if ($LASTEXITCODE) { exit $LASTEXITCODE } # uninstall.ps1: native uninstall must keep the shared unsloth.ico while a # WSL shortcut still references it (dual install), else that shortcut blanks. diff --git a/install.ps1 b/install.ps1 index e3e5a03538..6cf3075dc6 100644 --- a/install.ps1 +++ b/install.ps1 @@ -440,22 +440,6 @@ public static class UnslothStudioFinalPathV2 uint pathLength, uint flags); - [DllImport("kernel32.dll", SetLastError = true)] - private static extern IntPtr OpenProcess( - uint desiredAccess, - bool inheritHandle, - int processId); - - [DllImport("kernel32.dll", CharSet = CharSet.Unicode, SetLastError = true)] - private static extern bool QueryFullProcessImageNameW( - IntPtr process, - uint flags, - StringBuilder path, - ref uint pathLength); - - [DllImport("kernel32.dll", SetLastError = true)] - private static extern bool CloseHandle(IntPtr handle); - public static string Resolve(string path) { using (SafeFileHandle handle = CreateFileW( @@ -488,28 +472,6 @@ public static class UnslothStudioFinalPathV2 return buffer.ToString(); } } - - public static string GetProcessImagePath(int processId) - { - const uint ProcessQueryLimitedInformation = 0x1000; - IntPtr process = OpenProcess(ProcessQueryLimitedInformation, false, processId); - if (process == IntPtr.Zero) - { - return null; - } - try - { - StringBuilder path = new StringBuilder(32768); - uint pathLength = (uint)path.Capacity; - return QueryFullProcessImageNameW(process, 0, path, ref pathLength) - ? path.ToString() - : null; - } - finally - { - CloseHandle(process); - } - } } '@ } @@ -1906,6 +1868,23 @@ exit 0 } } + # PID -> image path for every process this user can see, in ONE query. This used to open a + # handle to every running PID through the inline C# type above, which is a shape AV + # heuristics score hard and bought nothing here: Win32_Process reports ExecutablePath for + # exactly the processes that could be opened, and answers for all of them at once. + function Get-StudioProcessImagePathMap { + $map = @{} + try { + foreach ($entry in @(Get-CimInstance Win32_Process -ErrorAction Stop)) { + if ($entry.ExecutablePath) { $map[[int]$entry.ProcessId] = $entry.ExecutablePath } + } + } catch { + # A degraded WMI repository leaves the map empty; the per-process .Path fallback + # below still resolves the accessible ones. + } + return $map + } + function Get-RunningStudioVenvProcesses { param( [Parameter(Mandatory = $true)][string]$VenvPath, @@ -1917,11 +1896,18 @@ exit 0 throw "Could not resolve managed Studio process path '$VenvPath': $($_.Exception.Message)" } + $imagePaths = Get-StudioProcessImagePathMap + # Block only confirmed executable identities: a command line or working # directory that merely mentions the path is not proof of an open file. foreach ($process in @(Get-Process -ErrorAction SilentlyContinue)) { $executable = $null - try { $executable = [UnslothStudioFinalPathV2]::GetProcessImagePath($process.Id) } catch { continue } + if ($imagePaths.ContainsKey([int]$process.Id)) { $executable = $imagePaths[[int]$process.Id] } + # .Path reads MainModule, which needs rights Win32_Process does not: only a fallback, + # and a protected or cross-bitness process throws here exactly as it did before. + if (-not $executable) { + try { $executable = $process.Path } catch { continue } + } if (-not $executable) { continue } try { $executable = Get-StudioFinalPath -Path $executable } catch { continue } if (Test-StudioProtectedPathMatch -Candidate $executable -ProtectedPath $resolvedPath -Exact:$Exact) { diff --git a/tests/studio/test_process_image_path_map.ps1 b/tests/studio/test_process_image_path_map.ps1 new file mode 100644 index 0000000000..9951242a39 --- /dev/null +++ b/tests/studio/test_process_image_path_map.ps1 @@ -0,0 +1,58 @@ +# Regression tests for install.ps1's venv-holder probe. +# +# Resolving a PID to its image path used to open a handle to every running process from inline +# C# compiled at runtime -- a shape AV heuristics score hard, on top of the csc.exe compile the +# type already costs. Win32_Process answers the same question for the same set of processes in +# one query, so the pair was removed. These tests pin both halves: the P/Invoke must stay gone, +# and the replacement must actually resolve a real process. +$ErrorActionPreference = "Stop" +$script:failures = 0 +function Check($name, $cond) { + if ($cond) { Write-Host " PASS $name" } + else { Write-Host " FAIL $name" -ForegroundColor Red; $script:failures++ } +} + +$repoRoot = (Resolve-Path ([System.IO.Path]::Combine($PSScriptRoot, "..", ".."))).Path +$installPath = [System.IO.Path]::Combine($repoRoot, "install.ps1") +$installText = Get-Content -Raw -LiteralPath $installPath + +# ── Source contract ── +Check "install.ps1 no longer opens a handle per PID from inline C#" ( + $installText -notmatch 'private static extern IntPtr OpenProcess') +Check "install.ps1 no longer imports the image-path query" ( + $installText -notmatch 'QueryFullProcessImageNameW') +Check "the inline type still resolves final paths" ( + $installText -match 'GetFinalPathNameByHandleW') +Check "the venv-holder probe reads the prefetched map" ( + $installText -match '\$imagePaths = Get-StudioProcessImagePathMap') +Check "one Win32_Process query, not one per PID" ( + ([regex]::Matches($installText, 'Get-CimInstance Win32_Process')).Count -eq 1) + +# ── Behaviour ── +. ([System.IO.Path]::Combine($repoRoot, "tests", "studio_setup_ps1", "Get-FunctionSource.ps1")) +$src = Get-FunctionSource -Path $installPath -Name "Get-StudioProcessImagePathMap" +Check "install.ps1 defines Get-StudioProcessImagePathMap" ($null -ne $src) + +if ($src -and $IsWindows -ne $false) { + . ([scriptblock]::Create($src)) + $map = Get-StudioProcessImagePathMap + Check "the map is a hashtable" ($map -is [hashtable]) + # This process is running an interpreter off disk, so it must resolve. + $own = $PID + Check "the map resolves this process to an image path" ( + $map.ContainsKey([int]$own) -and $map[[int]$own]) + if ($map.ContainsKey([int]$own)) { + Check "the resolved image path exists on disk" (Test-Path -LiteralPath $map[[int]$own]) + } + # Every value must be a path, never a bare process name: the venv match compares full paths. + $bad = @($map.Values | Where-Object { $_ -and -not ([System.IO.Path]::IsPathRooted($_)) }) + Check "every resolved image path is rooted" ($bad.Count -eq 0) +} else { + Write-Host " SKIP runtime map checks (Win32_Process is Windows-only)" +} + +if ($script:failures -gt 0) { + Write-Host "$($script:failures) check(s) failed" -ForegroundColor Red + exit 1 +} +Write-Host "All checks passed"