PoshAnywhereServer.ps1

#Exposes the Powershell Remoting Protocol via a TCP Port
using namespace System.Management.Automation
using namespace System.Management.Automation.Runspaces
using namespace System.Net.Sockets
using namespace System.Net.WebSockets
using namespace System.Collections.Generic
using namespace System.Diagnostics.CodeAnalysis
using namespace System.Diagnostics
using namespace System.Runtime.InteropServices
using namespace System.Threading
using namespace System.Threading.Tasks
using namespace System.IO.Pipes
using namespace System.IO
using namespace System.Net
using namespace System.Text

param (
  #Allow non-localhost connections. This may require additional permissions
  [Parameter(ParameterSetName = 'Server')][Switch]$AllowRemoteConnections,
  #How large to make the buffer for websockets. A larger buffer will have better performance but use more memory.
  [int]$BufferSize = 8192,
  #Which port to listen on, defaults to 7073
  [int]$Port = 7073,
  #Listen to a specific PowerShell process instead of the currently running one.
  [int]$PowerShellProcessId,
  #How long in milliseconds to wait for the named pipe to connect. This should almost always be pretty much immediately so the default is 500ms
  [int]$NamedPipeConnectTimeout,
  #Use a raw tcp listener instead of a websocket
  [Switch]$TCP,
  #Dont start the server but just load the functions. This is primarily for developer use
  [Switch]$NoStart,
  #Dont start a Cloudflare quick tunnel to reach this websocket over the internet
  [Switch]$NoCloudFlare
)

$ErrorActionPreference = 'Stop'

if ($AllowRemoteConnections) { $ListenAddress = '0.0.0.0' }

function Start-WebSocketNamedPipeServer ([ValidateNotNullOrEmpty()][int]$Port = 7073, [int]$ProcessId = $PID) {
  <#
  .SYNOPSIS
  Starts a websocket server that hosts psremoting for the named process
  #>
  try {
    $server = Start-WebSocketServer -Port $Port
    while ($true) {
      $pipeStream = Connect-PSRemotingNamedPipe
      Write-Verbose 'WEBSOCKET: Listening for new connection.'
      Receive-PSRPStreamSession -server $server -Stream $pipeStream
      $pipeStream.Dispose()
      Write-Verbose 'WEBSOCKET: Session Ended.'
    }
  } catch {
    throw
  } finally {
    if ($server) {
      Write-Verbose 'WEBSOCKET: Stopping Server...'
      $server.Stop()
    }

    if ($pipeStream) {
      Write-Verbose 'Closing and disposing named pipe connection'
      $pipeStream.Close()
      $pipeStream.Dispose()
    }
    Write-Verbose 'WEBSOCKET: Server Stopped'
  }
}

function Get-PSRemotingNamedPipe() {
  <#
  .SYNOPSIS
  Gets the PSRemoting autogenerated named pipe of the specified process
  .NOTES
  Reimplementation of internal powershell method NamedPipeUtils.CreateProcessPipeName
  #>
  param(
    [int]$ID = $PID,
    [switch]$PipeNameOnly
  )
  $process = Get-Process -Id $ID -ErrorAction Stop
  $processStartTime = $process.StartTime.ToFileTime()

  if ($PSEdition -eq 'Desktop' -or $isWindows) {
    $processStartId = $processStartTime
    $appDomain = 'DefaultAppDomain'
  } else {
    $processStartId = $processStartTime.ToString('X8').Substring(1, 8)
    $appDomain = 'None'
  }

  $pipeName = (@(
      'PSHost'
      $processStartID
      $process.Id
      $appDomain
      $process.ProcessName
    ) -join '.').ToString([cultureinfo]::InvariantCulture)
  if ($PipeNameOnly) { return $pipeName }

  if (-not ($PSEdition -eq 'Desktop' -or $isWindows)) {
    return $pipeName = "/tmp/CoreFxPipe_$pipeName"
  } else {
    return $pipeName = "\\.\pipe\$pipeName"
  }
}

function Connect-PSRemotingNamedPipe([String]$Name = (Get-PSRemotingNamedPipe -PipeNameOnly), [String]$ComputerName = '.') {
  #Connect the debug pipe
  $pipeclient = [NamedPipeClientStream]::new(
    $ComputerName, #string serverName
    $Name, #string pipeName
    [PipeDirection]::InOut, #PipeDirection direction
    [PipeOptions]::Asynchronous #PipeOptions options
  )
  Write-Verbose "Connecting to PowerShell Named Pipe $Name on ($ComputerName)..."
  try {
    $pipeclient.Connect($NamedPipeConnectTimeout)
  } catch [TimeoutException] {
    $pipeclient.Dispose()
    $PSCmdlet.ThrowTerminatingError(
      [ErrorRecord]::new(
        [TimeoutException]::new("Timed out connecting to named pipe $Name on ($ComputerName). The named pipe is probably busy, disconnected, or no longer exists"),
        'NamedPipeClientStream.TimeoutException',
        [ErrorCategory]::OperationTimeout,
        $pipeclient
      )
    )
  }
  Write-Verbose "CONNECTED to Powershell Named Pipe $Name"
  return $pipeClient
}

function Start-PSRemotingTCPListener ([Int]$Port = 7073, [ipaddress]$ListenAddress = '127.0.0.1') {
  #Connect the tcp listener. 7073 = ps in hex :)
  $tcpListener = [TcpListener]::new($ListenAddress, $port)
  $tcpListener.Start()
  Write-Verbose "Listening for TCP connection to ${listenAddress}:$port..."
  $tcpClient = $tcpListener.AcceptTcpClient()
  Write-Verbose "Accepted TCP connection to ${listenAddress}:$port..."
  return $tcpClient
}

function Join-Stream {
  param (
    #Provide an array of two streams to join together
    [ValidateCount(2, 2)][Stream[]]$Stream,
    #Wait for one of the streams to complete or disconnect. The script will return the first stream to disconnect
    [Switch]$Wait
  )
  $copyStreamTasks = @(
    $Stream[0].CopyToAsync($Stream[1])
    $Stream[1].CopyToAsync($Stream[0])
  )
  if ($Wait) {
    $completedTaskIndex = [Task]::WaitAny($copyStreamTasks)
    return $copyStreamTasks[$completedTaskIndex]
  } else {
    $copyStreamTasks
  }
}

#Region WebSocketServer

# Binds a websocket to a stream. You should provide a connected, async inout pipe stream with autoflush enabled.
function Start-WebSocketServer ([ValidateNotNullOrEmpty()][int]$Port = 7073) {
  $prefix = "http://localhost:$Port/psrp/"
  $server = [HttpListener]::new()
  $server.Prefixes.Add($Prefix)
  $server.Start()
  Write-Verbose "Websocket Server listening on $prefix"
  return $server
}

function Receive-PSRPStreamSession($server, [Stream]$Stream) {

  $context = $server.GetContextAsync() | Wait-Task

  if (-not $context.Request.IsWebSocketRequest) {
    Write-Verbose 'WEBSOCKET: non-websocket request received. Disconnecting'
    $context.Response.StatusCode = 400
    $context.Response.Close()
    continue
  }

  Write-Verbose "WEBSOCKET: Connection from $($context.Request.RemoteEndPoint)"
  try {
    # Powershell will cast null to an empty string, in this case we need to force an actual null
    $websocketContext = $context.AcceptWebSocketAsync([nullstring]::Value) | Wait-Task
    $websocket = $websocketContext.WebSocket
    Join-WebsocketToStream -Websocket $websocket -Stream $PipeStream
  } catch {
    Write-Verbose "WEBSOCKET: Error accepting websocket request: $($_.Exception.Message)"
    $context.Response.StatusCode = 500
    $context.Response.Close()
    continue
  } finally {
    Write-Verbose "WEBSOCKET: Connection CLOSE: $($context.Request.RemoteEndPoint)"
    if ($websocket) {
      $websocket.Dispose()
    }
  }
}

function Join-WebsocketToStream ([WebSocket]$Websocket, [Stream]$Stream) {
  <#
  .SYNOPSIS
  Joins a websocket to a stream and handles messaging between the two. You should provide a connected, async inout pipe stream
  #>

  $fromPipeTask = $null
  $fromWebSocketTask = $null
  $pendingWriteTask = $null
  [List[object]]$activeTasks = @()

  try {
    $pipeReader = [StreamReader]::new($Stream)
    [ArraySegment[byte]]$receiveBuffer = [ArraySegment[byte]]::new([byte[]]::new($BufferSize))

    #This outer loop uses Tasks to keep from blocking on either the websocket or the pipe, and acts on them as messages come in. This also ensures the websocket is never concurrently written to/read from, which is not allowed.
    while ($Websocket.State -eq [WebSocketState]::Open) {
      #This is the core of the loop, we wait until either we recieve data from the named pipe or the websocket. If the timeout is reached, check that the connection is still open. We use 500ms to allow for Ctrl-C to work in a reasonable timeframe.
      do {
        if (-not $fromPipeTask) {
          #Fetch a PSXML message from the named pipe. While we only care about the bytes, we use streamreader Readline() as a simple way to look for newlines as message delimiters, and that returns a string which we need to deconstruct back down into bytes to pass along. A faster solution would be to look for the newline ourselves within the byte array, but that's not enough of a perf gain to justify the complexity.
          $fromPipeTask = $pipeReader.ReadLineAsync()
          $activeTasks.Add($fromPipeTask)
        }
        if (-not $fromWebSocketTask) {
          $fromWebSocketTask = $websocket.ReceiveAsync($receiveBuffer, [CancellationToken]::None)
          $activeTasks.Add($fromWebSocketTask)
        }

        [int]$TIMEOUT_REACHED = -1
        [int]$completedTaskIndex = [Task]::WaitAny($activeTasks, 500)

        if (
          $completedTaskIndex -eq $TIMEOUT_REACHED -and
          $websocket.State -ne [WebSocketState]::Open
        ) {
          Write-Verbose "WEBSOCKET: Websocket is no longer open. Current Status: $websocket.State"
          return
        }
      } until ($completedTaskIndex -ne $TIMEOUT_REACHED)

      $completedTask = $activeTasks[$completedTaskIndex]
      $activeTasks.RemoveAt($completedTaskIndex)

      #If the websocket task completed, it means we received a message from the websocket and we need to send the data to the named pipe
      if ($completedTask -eq $fromWebSocketTask) {
        #Buffers can be tricky, so we instead capture everything into a memorystream and then extract the string out of that.
        do {
          [WebSocketReceiveResult]$result = $fromWebSocketTask | Wait-Task
          Write-Debug "WEBSOCKET SERVER RECV: Message of length $($result.count)"
          $Stream.Write($receiveBuffer, 0, $result.Count)

          if ($result.MessageType -eq [WebSocketMessageType]::Close) {
            Write-Verbose 'WEBSOCKET: Received Close Request from Client. Responding with NormalClosure'
            $websocket.CloseOutputAsync(
              [WebSocketCloseStatus]::NormalClosure,
              [String]::Empty,
              [CancellationToken]::None
            ) | Wait-Task | Out-Null
            return
          }

          if (-not $result.EndOfMessage) {
            # We don't need to reset $receiveBuffer here, it will be overwritten
            $result = $websocket.ReceiveAsync($receiveBuffer, [CancellationToken]::None) | Wait-Task
          }
        } until ($result.EndOfMessage)

        #Newline indicates the end of a message to the named pipe
        $newLineByte = [encoding]::UTF8.GetBytes([Environment]::NewLine)
        $Stream.Write($newLineByte)

        #Reset the websocket task variable which will get populated with a new listener when the loop restarts
        $fromWebSocketTask = $null
      }

      #If the pipe task completed, it means a message came from the PSRP pipe and we need to send the data to the websocket
      if ($completedTask -eq $fromPipeTask) {

        $pipeMessage = $fromPipeTask | Wait-Task

        # Write-Host "PIPE RECEIVED: $pipeMessage"

        #If an existing pendingWriteTask exists, wait for it to finish. We are not allowed to have simultaneous writes to a websocket.
        if ($pendingWriteTask) {
          $pendingWriteTask = $pendingWriteTask | Wait-Task | Out-Null
        }

        $messageBytes = [Encoding]::UTF8.GetBytes($pipeMessage)
        $payload = [ArraySegment[byte]]::new($messageBytes)
        $pendingWriteTask = $websocket.SendAsync($payload, [WebSocketMessageType]::Text, $true, [CancellationToken]::None)
        Write-Debug "WEBSOCKET SERVER SEND: Message of length $($payload.count)"

        #Reset the pipe task variable which will get populated with a new listener when the loop restarts
        $fromPipeTask = $null
      }
    }
  } catch [WebSocketException] {
    $PSCmdlet.WriteError($PSItem)
    return
  } catch {
    $PSCmdlet.ThrowTerminatingError($PSItem)
    $webSocket.CloseAsync([WebSocketCloseStatus]::InternalServerError, 'There was an error on the server side. Check the server side logs for details', [CancellationToken]::None) | Wait-Task | Out-Null
  }
}

filter Wait-Task ([int]$Timeout = 500) {
  $task = $PSItem

  #This makes the wait cancellable by Ctrl-C
  try {
    while (-not $task.Wait($Timeout)) {}
  } catch [AggregateException] {
    throw $task.Exception.InnerException
  }
  return $task.Result
}

#region Cloudflared

function Set-IsWindows {
  if ($null -eq [RuntimeInformation]::ProcessArchitecture) {
    #ProcessArchitecture is null on PowerShell 5.1 so we know we are on that.
    $GLOBAL:IsWindows = $true
  }
}

function Get-ProcessArchitecture {
  #Polyfill IsWindows to PowerShell 5.1
  $OSArchitecture = $null

  # Detect ARM architecture using [Environment]
  if ($null -eq [RuntimeInformation]::ProcessArchitecture) {
    #ProcessArchitecture is null on PowerShell 5.1 so we know we are on that.

    [SuppressMessageAttribute('PSAvoidAssignmentToAutomaticVariable', '', Justification = 'Polyfill')]
    $GLOBAL:IsWindows = $true
    if ([Environment]::Is64BitOperatingSystem) {
      $OSArchitecture = 'X64'
    } else {
      $OSArchitecture = 'X86'
    }
  } else {
    $OSArchitecture = [RuntimeInformation]::ProcessArchitecture
  }
  return $OSArchitecture
}

function Get-CloudflareAssetName {
  param (
    [string]$Release
  )

  $os = Get-CloudFlareAssetOS

  $arch = ConvertTo-CloudFlareArchitecture

  [string]$baseAsset = "cloudflared-$os-$arch"

  if ($os -eq 'darwin') {
    $baseAsset += '.tgz'
  }

  if ($os -eq 'windows') {
    $baseAsset += '.exe'
  }

  return $baseAsset
}

function Get-CloudFlareAssetOS {
  switch ($true) {
    $IsWindows { 'windows' }
    $IsLinux { 'linux' }
    $IsMacOS { 'darwin' }
    default { throw 'Could not detect operating system' }
  }
}

function ConvertTo-CloudFlareArchitecture ($architecture = (Get-ProcessArchitecture)) {
  switch ($architecture) {
    'X64' { 'amd64' }
    'X86' { '386' }
    'Arm' { 'arm' }
    'Arm64' { 'arm64' }
    default { throw 'Could not detect architecture' }
  }
}

function Save-CfGithubRelease {
  param (
    #The specific release to save. If not specified, the latest release will be downloaded
    [string]$Release,
    #Where to save the file, you should provide a directory.
    [string]$Destination,
    #Download even if the file exists
    [switch]$Force
  )

  # Use download/latest uri
  $baseGitHubUri = 'https://github.com/cloudflare/cloudflared/releases/'

  if ($Release) {
    $baseGitHubUri += "download/$Release/"
  } else {
    $baseGitHubUri += 'latest/download/'
  }
  $assetName = Get-CloudflareAssetName

  $gitHubReleaseUri = $baseGitHubUri + $assetName

  if (-not $Destination) {
    $Destination = [Path]::GetTempPath()
  }

  $AbsoluteDestination = Resolve-Path $Destination

  $OutFilePath = Join-Path $AbsoluteDestination $assetName

  if ((Test-Path $OutFilePath) -and -not $Force) {
    Write-Verbose "Cloudflared already found in $OutFilePath. Use -Force to download again"
  } else {
    Invoke-WebRequest -Uri $gitHubReleaseUri -OutFile $OutFilePath -ErrorAction stop
  }

  return $OutFilePath
}

function Start-Cloudflared {
  param(
    [string]$CloudflaredPath = $(Save-CfGithubRelease),
    [int]$Port
  )

  # Run cloudflared as a separate process and wait for the connection message
  $startInfo = [ProcessStartInfo]::new(
    $CloudflaredPath,
    "tunnel --http-host-header localhost --url localhost:$Port"
  )
  $startInfo.RedirectStandardOutput = $true
  $startInfo.RedirectStandardError = $true
  $startInfo.UseShellExecute = $false

  $process = [Process]::new()
  $process.StartInfo = $startInfo
  if ($isLinux -and -not (((Get-Item $CloudflaredPath).UnixFileMode -band 'UserExecute') -eq 'UserExecute')) {
    chmod +x $CloudflaredPath
  }
  if (-not $process.Start()) { throw 'Cloudflare process failed to start.' }

  #Read stdout until we get the connection message
  [string]$cfTunnelHostname = $null

  do {
    $cfOut = $process.StandardError
    $line = $cfOut.ReadLine()
    if ($line -match 'Your quick Tunnel has been created!') {
      if ($cfOut.ReadLine() -notmatch 'https\:\/\/(.+?).trycloudflare.com') {
        throw 'Tunnel was created but could not find the URL. The cloudflared output format may have changed'
      }
      $cfTunnelHostname = $Matches[1]
    }
  } until ($cfTunnelHostname)

  return @{
    Process  = $process
    Hostname = $cfTunnelHostname
  }
}

#endregion Cloudflared

#region main
if (-not $NoStart) {
  try {
    if ($TCP) {
      $pipeClient = Connect-PSRemotingNamedPipe
      $tcpClient = Start-PSRemotingTCPListener
      $firstClosedStream = Join-Stream -Wait $pipeClient, $tcpClient.GetStream()
      if (-not $firstClosedStream.IsCompletedSuccessfully) {
        throw $firstClosedStream.Exception
      }
      return
    }

    if (-not $NoCloudFlare) {
      #Download and start cloudflare
      $tunnelInfo = Start-Cloudflared -Port $Port
      $tunnelName = $tunnelInfo.Hostname

      Write-Host -Fore Green "Your cloudflare tunnel name is $tunnelName. Connect to it using: "
      Write-Host -Fore Cyan "New-WebsocketSession -Hostname $tunnelName.trycloudflare.com"
    }

    #Default Websocket Implementation
    Start-WebSocketNamedPipeServer -Verbose -Debug
  } catch {
    Write-Error $PSItem
  } finally {
    if ($tunnelInfo.Process) {
      $tunnelInfo.Process.Kill()
      $tunnelInfo.Process.Dispose()
    }

    if ($TCP) {
      $pipeClient.Close()
      $pipeClient.Dispose()
      $tcpClient.Close()
      $tcpClient.Dispose()
      if ($tcpListener) { $tcpListener.Stop() }
    }
  }
}
#endregion main