MercuryToolbox/scripts/cargo-flamegraph-windows.ps1

[CmdletBinding()]
param(
    [Parameter(Mandatory = $true)]
    [string]$Package,
    [string]$Binary,
    [ValidateSet('Debug', 'Release', 'ReleaseFast', 'ReleaseSize')]
    [string]$BuildProfile = 'Release',
    [string]$Output,
    [string]$LogPath,
    [string]$WorkspaceRoot = (Split-Path -Parent $PSScriptRoot),
    [switch]$NoElevate,
    [switch]$Describe,
    [string]$TargetArgumentJson,
    [Parameter(ValueFromRemainingArguments = $true)]
    [string[]]$TargetArgument = @()
)

$ErrorActionPreference = 'Stop'
Set-StrictMode -Version Latest

function Convert-ToSingleQuotedLiteral {
    param(
        [Parameter(Mandatory = $true)]
        [AllowEmptyString()]
        [string]$Value
    )

    return "'{0}'" -f $Value.Replace("'", "''")
}

function Convert-ToEncodedCommand {
    param(
        [Parameter(Mandatory = $true)]
        [AllowEmptyString()]
        [string]$Command
    )

    return [Convert]::ToBase64String([Text.Encoding]::Unicode.GetBytes($Command))
}

function Test-IsAdministrator {
    $currentIdentity = [Security.Principal.WindowsIdentity]::GetCurrent()
    $currentPrincipal = [Security.Principal.WindowsPrincipal]::new($currentIdentity)
    return $currentPrincipal.IsInRole([Security.Principal.WindowsBuiltInRole]::Administrator)
}

function Write-RelayLog {
    param(
        [Parameter(Mandatory = $true)]
        [string]$Message
    )

    if ([string]::IsNullOrWhiteSpace($LogPath)) {
        return
    }

    $logDirectory = Split-Path -Parent $LogPath
    if (-not [string]::IsNullOrWhiteSpace($logDirectory)) {
        New-Item -ItemType Directory -Force -Path $logDirectory | Out-Null
    }

    $timestamp = Get-Date -Format 'O'
    $line = "[$timestamp] $Message{0}" -f [Environment]::NewLine
    $utf8NoBom = [Text.UTF8Encoding]::new($false)
    [System.IO.File]::AppendAllText($LogPath, $line, $utf8NoBom)
}

function Get-ElevationHostPath {
    $currentProcess = Get-Process -Id $PID -ErrorAction SilentlyContinue
    if ($null -ne $currentProcess -and -not [string]::IsNullOrWhiteSpace($currentProcess.Path)) {
        return $currentProcess.Path
    }

    $pwshCommand = Get-Command -Name 'pwsh.exe' -ErrorAction SilentlyContinue
    if ($null -ne $pwshCommand) {
        return $pwshCommand.Source
    }

    return 'powershell.exe'
}

function Get-TranscriptPath {
    if ([string]::IsNullOrWhiteSpace($LogPath)) {
        return $null
    }

    return '{0}.transcript.txt' -f $LogPath
}

function Get-DTracePath {
    if (-not [string]::IsNullOrWhiteSpace($env:DTRACE)) {
        if (Test-Path -LiteralPath $env:DTRACE) {
            return (Resolve-Path -LiteralPath $env:DTRACE).Path
        }

        $resolvedOverride = Get-Command -Name $env:DTRACE -ErrorAction SilentlyContinue
        if ($null -ne $resolvedOverride) {
            return $resolvedOverride.Source
        }
    }

    $dtraceCommand = Get-Command -Name 'dtrace' -ErrorAction SilentlyContinue
    if ($null -eq $dtraceCommand) {
        return $null
    }

    return $dtraceCommand.Source
}

function Get-ProfileArgumentList {
    param(
        [Parameter(Mandatory = $true)]
        [string]$Profile
    )

    switch ($Profile) {
        'Debug' { return @('--dev') }
        'Release' { return @('--release') }
        'ReleaseFast' { return @('--profile', 'release-fast') }
        'ReleaseSize' { return @('--profile', 'release-size') }
        default { throw "Unsupported build profile: $Profile" }
    }
}

function New-FlamegraphArgumentList {
    param(
        [Parameter(Mandatory = $true)]
        [string]$ResolvedBinary,
        [Parameter(Mandatory = $true)]
        [AllowEmptyCollection()]
        [string[]]$ResolvedTargetArgument
    )

    $arguments = @('flamegraph') + (Get-ProfileArgumentList -Profile $BuildProfile) + @(
        '-p',
        $Package,
        '--bin',
        $ResolvedBinary
    )

    if (-not [string]::IsNullOrWhiteSpace($Output)) {
        $arguments += @('-o', $Output)
    }

    if ($ResolvedTargetArgument.Count -gt 0) {
        $arguments += '--'
        $arguments += $ResolvedTargetArgument
    }

    return $arguments
}

function New-RelayCommand {
    param(
        [Parameter(Mandatory = $true)]
        [string]$ResolvedBinary,
        [Parameter(Mandatory = $true)]
        [AllowEmptyCollection()]
        [string[]]$ResolvedTargetArgument
    )

    $relayTokens = @(
        '&',
        (Convert-ToSingleQuotedLiteral -Value $PSCommandPath),
        '-Package',
        (Convert-ToSingleQuotedLiteral -Value $Package),
        '-Binary',
        (Convert-ToSingleQuotedLiteral -Value $ResolvedBinary),
        '-BuildProfile',
        (Convert-ToSingleQuotedLiteral -Value $BuildProfile),
        '-WorkspaceRoot',
        (Convert-ToSingleQuotedLiteral -Value $WorkspaceRoot),
        '-NoElevate'
    )

    if (-not [string]::IsNullOrWhiteSpace($Output)) {
        $relayTokens += @('-Output', (Convert-ToSingleQuotedLiteral -Value $Output))
    }

    if (-not [string]::IsNullOrWhiteSpace($LogPath)) {
        $relayTokens += @('-LogPath', (Convert-ToSingleQuotedLiteral -Value $LogPath))
    }

    if ($Describe) {
        $relayTokens += '-Describe'
    }

    if ($ResolvedTargetArgument.Count -gt 0) {
        $targetArgumentJson = $ResolvedTargetArgument | ConvertTo-Json -Compress
        $relayTokens += @('-TargetArgumentJson', (Convert-ToSingleQuotedLiteral -Value $targetArgumentJson))
    }

    $commandParts = @(
        "Set-Location -LiteralPath $(Convert-ToSingleQuotedLiteral -Value $WorkspaceRoot)",
        ($relayTokens -join ' ')
    )

    return ($commandParts -join '; ')
}

$resolvedWorkspaceRoot = [System.IO.Path]::GetFullPath($WorkspaceRoot)
$resolvedBinary = if ([string]::IsNullOrWhiteSpace($Binary)) { $Package } else { $Binary }
$dtracePath = Get-DTracePath
$transcriptPath = Get-TranscriptPath
$elevationHostPath = Get-ElevationHostPath
$requiresElevation = [string]::IsNullOrWhiteSpace($dtracePath) -and -not (Test-IsAdministrator)
$resolvedTargetArgument = @()
if (-not [string]::IsNullOrWhiteSpace($TargetArgumentJson)) {
    if ($TargetArgument.Count -gt 0) {
        throw 'Pass either -TargetArgumentJson or trailing -TargetArgument values, not both.'
    }

    $decodedArgument = ConvertFrom-Json -InputObject $TargetArgumentJson
    foreach ($argument in @($decodedArgument)) {
        $resolvedTargetArgument += [string]$argument
    }
}
else {
    $resolvedTargetArgument = $TargetArgument
}

$flamegraphArguments = New-FlamegraphArgumentList -ResolvedBinary $resolvedBinary -ResolvedTargetArgument $resolvedTargetArgument
$relayCommand = New-RelayCommand -ResolvedBinary $resolvedBinary -ResolvedTargetArgument $resolvedTargetArgument

if ($Describe) {
    [pscustomobject]@{
        workspace_root = $resolvedWorkspaceRoot
        package = $Package
        binary = $resolvedBinary
        build_profile = $BuildProfile
        output = $Output
        log_path = $LogPath
        transcript_path = $transcriptPath
        dtrace_path = $dtracePath
        elevation_host_path = $elevationHostPath
        requires_elevation = $requiresElevation
        target_arguments = $resolvedTargetArgument
        relay_command = $relayCommand
        cargo_arguments = $flamegraphArguments
    } | ConvertTo-Json -Depth 4
    exit 0
}

if ($requiresElevation) {
    if ($NoElevate) {
        Write-RelayLog 'Refusing to continue without elevation because no DTrace executable was detected.'
        Write-Error 'cargo flamegraph on Windows needs an elevated session here because no DTrace executable was detected and blondie would fail with NotAnAdmin.'
    }

    $encodedRelayCommand = Convert-ToEncodedCommand -Command $relayCommand
    Write-RelayLog "Requesting elevation with encoded relay command for package '$Package'."
    try {
        Start-Process -FilePath $elevationHostPath -Verb RunAs -WorkingDirectory $resolvedWorkspaceRoot -ArgumentList @(
            '-NoLogo',
            '-NoProfile',
            '-ExecutionPolicy',
            'Bypass',
            '-EncodedCommand',
            $encodedRelayCommand
        ) | Out-Null
        Write-RelayLog 'Elevation request handed off to a new PowerShell window.'
        Write-Host 'Elevation requested in a new PowerShell window. After UAC approval, cargo flamegraph will run there.'
    }
    catch {
        Write-RelayLog "Elevation request failed: $($_.Exception.Message)"
        throw
    }
    exit 0
}

Write-RelayLog "Running cargo flamegraph in elevated/session-ready mode from '$resolvedWorkspaceRoot'."
$transcriptStarted = $false
if (-not [string]::IsNullOrWhiteSpace($LogPath)) {
    $transcriptPath = Get-TranscriptPath
    $logDirectory = Split-Path -Parent $transcriptPath
    if (-not [string]::IsNullOrWhiteSpace($logDirectory)) {
        New-Item -ItemType Directory -Force -Path $logDirectory | Out-Null
    }
    Start-Transcript -LiteralPath $transcriptPath -Append | Out-Null
    $transcriptStarted = $true
}

Push-Location $resolvedWorkspaceRoot
try {
    Write-RelayLog "Invoking cargo $($flamegraphArguments -join ' ')"
    cargo @flamegraphArguments
    Write-RelayLog "cargo flamegraph finished with exit code $LASTEXITCODE"
}
finally {
    Pop-Location
    if ($transcriptStarted) {
        Stop-Transcript | Out-Null
    }
}