OnePAM

1.0.2

Private/WebSocketHelper.ps1

function New-OpWebSocket {
    [OutputType([System.Net.WebSockets.ClientWebSocket])]
    param(
        [Parameter(Mandatory)]
        [string]$Uri,

        [hashtable]$Headers,

        [int]$TimeoutSeconds = 15
    )

    $ws = [System.Net.WebSockets.ClientWebSocket]::new()
    $ws.Options.SetRequestHeader('User-Agent', "onepam-powershell/$script:ModuleVersion")

    if ($Headers) {
        foreach ($key in $Headers.Keys) {
            $ws.Options.SetRequestHeader($key, $Headers[$key])
        }
    }

    $cts = [System.Threading.CancellationTokenSource]::new([TimeSpan]::FromSeconds($TimeoutSeconds))

    try {
        $ws.ConnectAsync([Uri]$Uri, $cts.Token).GetAwaiter().GetResult()
    }
    catch {
        $ws.Dispose()
        throw "WebSocket connection failed: $_"
    }
    finally {
        $cts.Dispose()
    }

    $ws
}

function Send-OpWsMessage {
    param(
        [Parameter(Mandatory)]
        [System.Net.WebSockets.ClientWebSocket]$WebSocket,

        [Parameter(Mandatory)]
        [string]$Message,

        [System.Net.WebSockets.WebSocketMessageType]$MessageType = [System.Net.WebSockets.WebSocketMessageType]::Text,

        [int]$TimeoutSeconds = 30
    )

    $bytes = [System.Text.Encoding]::UTF8.GetBytes($Message)
    $segment = [ArraySegment[byte]]::new($bytes)
    $cts = [System.Threading.CancellationTokenSource]::new([TimeSpan]::FromSeconds($TimeoutSeconds))

    try {
        $WebSocket.SendAsync($segment, $MessageType, $true, $cts.Token).GetAwaiter().GetResult()
    }
    finally {
        $cts.Dispose()
    }
}

function Send-OpWsBinaryMessage {
    param(
        [Parameter(Mandatory)]
        [System.Net.WebSockets.ClientWebSocket]$WebSocket,

        [Parameter(Mandatory)]
        [byte[]]$Data,

        [int]$TimeoutSeconds = 30
    )

    $segment = [ArraySegment[byte]]::new($Data)
    $cts = [System.Threading.CancellationTokenSource]::new([TimeSpan]::FromSeconds($TimeoutSeconds))

    try {
        $WebSocket.SendAsync($segment, [System.Net.WebSockets.WebSocketMessageType]::Binary, $true, $cts.Token).GetAwaiter().GetResult()
    }
    finally {
        $cts.Dispose()
    }
}

$script:MaxWsMessageSize = 100 * 1024 * 1024  # 100 MB

function Receive-OpWsMessage {
    [OutputType([PSCustomObject])]
    param(
        [Parameter(Mandatory)]
        [System.Net.WebSockets.ClientWebSocket]$WebSocket,

        [int]$BufferSize = 65536,

        [int]$TimeoutSeconds = 30
    )

    $buffer = [byte[]]::new($BufferSize)
    $segment = [ArraySegment[byte]]::new($buffer)
    $cts = [System.Threading.CancellationTokenSource]::new([TimeSpan]::FromSeconds($TimeoutSeconds))
    $ms = [System.IO.MemoryStream]::new()

    try {
        do {
            $result = $WebSocket.ReceiveAsync($segment, $cts.Token).GetAwaiter().GetResult()

            if ($result.MessageType -eq [System.Net.WebSockets.WebSocketMessageType]::Close) {
                return [PSCustomObject]@{
                    MessageType = 'Close'
                    Data        = $null
                    Text        = $null
                }
            }

            $ms.Write($buffer, 0, $result.Count)

            if ($ms.Length -gt $script:MaxWsMessageSize) {
                throw "WebSocket message exceeds maximum size of $($script:MaxWsMessageSize / 1MB) MB."
            }
        } while (-not $result.EndOfMessage)

        $allBytes = $ms.ToArray()

        [PSCustomObject]@{
            MessageType = $result.MessageType.ToString()
            Data        = $allBytes
            Text        = if ($result.MessageType -eq [System.Net.WebSockets.WebSocketMessageType]::Text) {
                [System.Text.Encoding]::UTF8.GetString($allBytes)
            } else { $null }
        }
    }
    finally {
        $ms.Dispose()
        $cts.Dispose()
    }
}

function Close-OpWebSocket {
    param(
        [Parameter(Mandatory)]
        [System.Net.WebSockets.ClientWebSocket]$WebSocket
    )

    if ($WebSocket.State -eq [System.Net.WebSockets.WebSocketState]::Open) {
        $cts = [System.Threading.CancellationTokenSource]::new([TimeSpan]::FromSeconds(5))
        try {
            $WebSocket.CloseAsync(
                [System.Net.WebSockets.WebSocketCloseStatus]::NormalClosure,
                'Closing',
                $cts.Token
            ).GetAwaiter().GetResult()
        }
        catch { }
        finally {
            $cts.Dispose()
        }
    }
    $WebSocket.Dispose()
}