CaiWuXiTong/scripts/keep-upstream-tunnel.ps1

#Requires -Version 5.1

[CmdletBinding(DefaultParameterSetName = 'Explicit')]
param(
    [Parameter(Mandatory = $true, ParameterSetName = 'Config')]
    [string] $ConfigPath,

    [Parameter(Mandatory = $true, ParameterSetName = 'Explicit')]
    [ValidatePattern('^[A-Za-z0-9._:-]+$')]
    [string] $ServerHost,

    [ValidateRange(1, 65535)]
    [int] $ServerPort = 22,

    [ValidatePattern('^[A-Za-z0-9._-]+$')]
    [string] $RemoteUser = 'caiwu-tunnel',

    [Parameter(Mandatory = $true, ParameterSetName = 'Explicit')]
    [string] $PrivateKeyPath,

    [Parameter(Mandatory = $true, ParameterSetName = 'Explicit')]
    [ValidatePattern('^[A-Za-z0-9._:-]+$')]
    [string] $RemoteBindAddress,

    [Parameter(Mandatory = $true, ParameterSetName = 'Explicit')]
    [ValidateRange(1, 65535)]
    [int] $RemoteBindPort,

    [Parameter(Mandatory = $true, ParameterSetName = 'Explicit')]
    [ValidatePattern('^[A-Za-z0-9._:-]+$')]
    [string] $UpstreamHost,

    [ValidateRange(1, 65535)]
    [int] $UpstreamPort = 1433,

    [ValidateRange(10, 300)]
    [int] $RetrySeconds = 30,

    [switch] $Once
)

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

if ($PSCmdlet.ParameterSetName -eq 'Config') {
    $resolvedConfigPath = (Resolve-Path -LiteralPath $ConfigPath -ErrorAction Stop).Path
    $statePath = $resolvedConfigPath + '.pid'
    $settings = Get-Content -LiteralPath $resolvedConfigPath -Raw -Encoding UTF8 | ConvertFrom-Json

    $ServerHost = [string] $settings.ServerHost
    $ServerPort = [int] $settings.ServerPort
    $RemoteUser = [string] $settings.RemoteUser
    $PrivateKeyPath = [string] $settings.PrivateKeyPath
    $RemoteBindAddress = [string] $settings.RemoteBindAddress
    $RemoteBindPort = [int] $settings.RemoteBindPort
    $UpstreamHost = [string] $settings.UpstreamHost
    $UpstreamPort = [int] $settings.UpstreamPort
    $RetrySeconds = [int] $settings.RetrySeconds
}

$endpointPattern = '^[A-Za-z0-9._:-]+$'
foreach ($endpoint in @($ServerHost, $RemoteBindAddress, $UpstreamHost)) {
    if ([string]::IsNullOrWhiteSpace($endpoint) -or $endpoint -notmatch $endpointPattern) {
        throw 'The tunnel configuration contains an invalid endpoint.'
    }
}

if ($RemoteUser -notmatch '^[A-Za-z0-9._-]+$') {
    throw 'The tunnel configuration contains an invalid SSH user.'
}

foreach ($port in @($ServerPort, $RemoteBindPort, $UpstreamPort)) {
    if ($port -lt 1 -or $port -gt 65535) {
        throw 'The tunnel configuration contains an invalid port.'
    }
}

if ($RetrySeconds -lt 10 -or $RetrySeconds -gt 300) {
    throw 'The tunnel retry interval must be between 10 and 300 seconds.'
}

$sshCommand = Get-Command 'ssh.exe' -ErrorAction Stop
$resolvedKeyPath = (Resolve-Path -LiteralPath $PrivateKeyPath -ErrorAction Stop).Path

$identity = '{0}|{1}|{2}|{3}' -f $ServerHost, $ServerPort, $RemoteBindAddress, $RemoteBindPort
$sha256 = [System.Security.Cryptography.SHA256]::Create()
try {
    $identityBytes = [System.Text.Encoding]::UTF8.GetBytes($identity)
    $identityHash = [System.BitConverter]::ToString($sha256.ComputeHash($identityBytes)).Replace('-', '')
}
finally {
    $sha256.Dispose()
}

$mutex = New-Object System.Threading.Mutex($false, ('Local\CaiWuXiTong-UpstreamTunnel-{0}' -f $identityHash.Substring(0, 20)))
$ownsMutex = $false

try {
    try {
        $ownsMutex = $mutex.WaitOne(0)
    }
    catch [System.Threading.AbandonedMutexException] {
        $ownsMutex = $true
    }

    if (-not $ownsMutex) {
        Write-Output 'The read-only upstream tunnel is already managed by another local process.'
        exit 0
    }

    if ($null -ne $statePath) {
        [System.IO.File]::WriteAllText($statePath, [string] $PID, [System.Text.Encoding]::ASCII)
        $currentSid = [System.Security.Principal.WindowsIdentity]::GetCurrent().User.Value
        & icacls.exe $statePath '/inheritance:r' '/grant:r' "*$currentSid`:(F)" '*S-1-5-18:(F)' | Out-Null
        if ($LASTEXITCODE -ne 0) {
            throw 'The tunnel PID file permissions could not be restricted.'
        }
    }

    $forward = '{0}:{1}:{2}:{3}' -f $RemoteBindAddress, $RemoteBindPort, $UpstreamHost, $UpstreamPort
    $destination = '{0}@{1}' -f $RemoteUser, $ServerHost
    $sshArguments = @(
        '-N',
        '-T',
        '-i', $resolvedKeyPath,
        '-p', $ServerPort,
        '-o', 'BatchMode=yes',
        '-o', 'ConnectTimeout=10',
        '-o', 'ExitOnForwardFailure=yes',
        '-o', 'ServerAliveInterval=20',
        '-o', 'ServerAliveCountMax=3',
        '-o', 'StrictHostKeyChecking=accept-new',
        '-o', 'TCPKeepAlive=yes',
        '-o', 'LogLevel=ERROR',
        '-R', $forward,
        $destination
    )

    do {
        & $sshCommand.Source @sshArguments
        $exitCode = $LASTEXITCODE

        if ($Once) {
            exit $exitCode
        }

        Write-Warning ('The read-only upstream tunnel disconnected (SSH exit code {0}); retrying shortly.' -f $exitCode)
        $retryDelay = $RetrySeconds + (Get-Random -Minimum 0 -Maximum 6)
        Start-Sleep -Seconds $retryDelay
    }
    while ($true)
}
finally {
    if ($null -ne $statePath -and (Test-Path -LiteralPath $statePath -PathType Leaf)) {
        $recordedPid = Get-Content -LiteralPath $statePath -Raw -ErrorAction SilentlyContinue
        if ([string] $recordedPid -eq [string] $PID) {
            Remove-Item -LiteralPath $statePath -Force -ErrorAction SilentlyContinue
        }
    }

    if ($ownsMutex) {
        $mutex.ReleaseMutex()
    }

    $mutex.Dispose()
}