@@ -484,142 +484,205 @@ function Get-WebSocket {
484484 break
485485 }
486486
487+ # If we're authenticating, and haven't yet said hello
487488 if ($Authenticate -and -not $SaidHello ) {
488- # a number of websockets require some handshaking to authenticate
489- $authenticationMessage =
489+ # then we should say hello.
490+ # Determine the authentication message
491+ $authenticationMessage =
492+ # If the authentication message is a scriptblock,
490493 if ($Authenticate -is [ScriptBlock ]) {
491- & $Authenticate $MessageObject
494+ & $Authenticate # run it
492495 } else {
493- $authenticate
496+ $authenticate # otherwise, use it as-is.
494497 }
495498
499+ # If we have an authentication message
496500 if ($authenticationMessage ) {
501+ # and it's not a string
497502 if ($authenticationMessage -isnot [string ]) {
503+ # then we should send it as JSON and mark that we've said hello.
498504 $saidHello = $ws.SendAsync ([ArraySegment [byte ]]::new(
499505 $OutputEncoding.GetBytes ((ConvertTo-Json - InputObject $authenticationMessage - Depth 10 ))
500506 ), ' Text' , $true , $CT )
501507 }
502508 }
503509 }
504510
511+ # Ok, let's get the next message.
505512 $Buf = [byte []]::new($BufferSize )
506513 $Seg = [ArraySegment [byte ]]::new($Buf )
507514 $receivingWebSocket = $ws.ReceiveAsync ($Seg , $CT )
515+ # use this tight loop to let us cancel the await if we need to.
508516 while (-not ($receivingWebSocket.IsCompleted -or $receivingWebSocket.IsFaulted -or $receivingWebSocket.IsCanceled )) {
509517
518+ }
519+ # If we had a problem, write an error.
520+ if ($receivingWebSocket.Exception ) {
521+ Write-Error - Exception $receivingWebSocket.Exception - Category ProtocolError
522+ continue
510523 }
511524 $MessageCount ++
512525
513526 try {
527+ # If we have a handshake and we haven't yet shaken hands
514528 if ($Handshake -and -not $shookHands ) {
515- # a number of websockets require some handshaking
516- $handShakeMessage =
529+ # then we should shake hands.
530+ # Get the message string
531+ $messageString = $OutputEncoding.GetString ($Buf , 0 , $Buf.Count )
532+ # and try to convert it from JSON.
533+ $messageObject = ConvertFrom-Json - InputObject $messageString * > & 1
534+ # Determine the handshake message
535+ $handShakeMessage =
536+ # If the handshake message is a scriptblock,
517537 if ($Handshake -is [ScriptBlock ]) {
518- & $Handshake $MessageObject
538+ & $Handshake $MessageObject # run it and pass the message
519539 } else {
520- $Handshake
540+ $Handshake # otherwise, use it as-is.
521541 }
522542
543+ # If we have a handshake message
523544 if ($handShakeMessage ) {
545+ # and it's not a string
524546 if ($handShakeMessage -isnot [string ]) {
547+ # then we should send it as JSON and mark that we've shaken hands.
525548 $saidHello = $ws.SendAsync ([ArraySegment [byte ]]::new(
526549 $OutputEncoding.GetBytes ((ConvertTo-Json - InputObject $handShakeMessage - Depth 10 ))
527550 ), ' Text' , $true , $CT )
528551 }
529552 }
530553 }
531554
555+ # Get the message from the websocket
532556 $webSocketMessage =
533- if ($Binary ) {
534- $Buf -gt 0
557+ if ($Binary ) { # If we wanted binary
558+ $Buf -gt 0 -as [ byte []] # then return non-null bytes
535559 } else {
560+ # otherwise, get the message as a string
536561 $messageString = $OutputEncoding.GetString ($Buf , 0 , $Buf.Count )
562+ # if we have any filters
537563 if ($Filter ) {
564+ # then we see if we can apply them now.
538565 foreach ($fil in $filter ) {
566+ # Wilcard filters can be applied to the raw text
539567 if ($fil -is [string ] -and $messageString -like " *$fil *" ) {
540568 $FilteredCount ++
541569 continue WebSocketMessageLoop
542570 }
571+ # and so can regex filters.
543572 if ($fil -is [regex ] -and $fil.IsMatch ($messageString )) {
544573 $FilteredCount ++
545574 continue WebSocketMessageLoop
546575 }
547576 }
548577 }
578+ # If we have asked for -RawText
549579 if ($RawText ) {
550- $messageString
580+ $messageString # then return the raw text
551581 } else {
552- $MessageObject = ConvertFrom-Json - InputObject $messageString
582+ # Otherwise, try to convert the message from JSON.
583+ $MessageObject = ConvertFrom-Json - InputObject $messageString
584+
585+ # Now we can run any filters that are scriptblocks or commands.
553586 if ($filter ) {
554587 foreach ($fil in $Filter ) {
555- if ($fil -is [ScriptBlock ] -or
588+ if ($fil -is [ScriptBlock ] -or
556589 $fil -is [Management.Automation.CommandInfo ]
557590 ) {
558- if (& $fil $MessageObject ) {
559- $FilteredCount ++
591+ # Capture the output of the filter
592+ $filterOutput = $MessageObject | & $fil $MessageObject
593+ # if the output was falsy,
594+ if (-not $filterOutput ) {
595+ $FilteredCount ++ # filter out the message.
560596 continue WebSocketMessageLoop
561597 }
562598 }
563599 }
564- }
600+ }
565601
602+ # If -Skip was provided and we haven't skipped enough messages
566603 if ($Skip -and ($SkipCount -le $Skip )) {
604+ # then skip this message.
567605 $SkipCount ++
568606 continue WebSocketMessageLoop
569607 }
570-
571- $MessageObject
572-
573- # If we are forwarding events
574- if ($ForwardEvent -and $MainRunspace.Events.GenerateEvent ) {
575- # generate an event in the main runspace
576- $null = $MainRunspace.Events.GenerateEvent (
577- " $SocketUrl " ,
578- $ws ,
579- @ ($MessageObject ),
580- $MessageObject
581- )
582- }
608+
609+ # Now, emit the message object.
610+ # (expressions that are not assigned will be outputted)
611+ $MessageObject
583612
613+ # If we have a -First parameter, and we have not yet reached the maximum
614+ # (after accounting for skips and filters)
584615 if ($First -and ($MessageCount - $FilteredCount - $SkipCount ) -ge $First ) {
616+ # then set the maximum to first (which will cancel this after the next loop)
585617 $Maximum = $first
586618 }
587619 }
588620 }
621+
622+ # If we want to decorate the output
589623 if ($PSTypeName ) {
624+ # clear it's typenames
590625 $webSocketMessage.pstypenames.clear ()
591- [Array ]::Reverse($PSTypeName )
592- foreach ($psType in $psTypeName ) {
593- $webSocketMessage.pstypenames.add ($psType )
594- }
626+ for ($typeNameIndex = $PSTypeName.Length - 1 ; $typeNameIndex -ge 0 ; $typeNameIndex -- ) {
627+ # and add each type name in reverse order
628+ $webSocketMessage.pstypenames.add ($PSTypeName [$typeNameIndex ])
629+ }
630+ }
631+
632+ # If we are forwarding events
633+ if ($ForwardEvent -and $MainRunspace.Events.GenerateEvent ) {
634+ # generate an event in the main runspace
635+ $null = $MainRunspace.Events.GenerateEvent (
636+ " $SocketUrl " ,
637+ $ws ,
638+ @ ($webSocketMessage ),
639+ $webSocketMessage
640+ )
595641 }
596- if ($handler ) {
597- $psCmd =
642+
643+ # If we have an output handler, try to run it and get the output
644+ $handledResponse = if ($handler ) {
645+ # We may need to run the handler in a `[PowerShell]` command.
646+ $psCmd =
647+ # This is true if we want `NoLanguage` mode.
598648 if ($runspace.LanguageMode -eq ' NoLanguage' -or
599649 $runspacePool.InitialSessionState.LanguageMode -eq ' NoLanguage' ) {
650+ # (in which case we'll call .GetPowerShell())
600651 $handler.GetPowerShell ()
601- } elseif ($Runspace -or $RunspacePool ) {
602- [PowerShell ]::Create().AddScript($handler )
652+ } elseif (
653+ # or if we have a runspace or runspace pool
654+ $Runspace -or $RunspacePool
655+ ) {
656+ # (in which case we'll `.Create()` and `.AddScript()`)
657+ [PowerShell ]::Create().AddScript($handler , $true )
603658 }
604659 if ($psCmd ) {
660+ # If we have a runspace, we'll use that.
605661 if ($Runspace ) {
606662 $psCmd.Runspace = $Runspace
607663 } elseif ($RunspacePool ) {
664+ # or, alternatively, we can use a runspace pool.
608665 $psCmd.RunspacePool = $RunspacePool
609666 }
667+ # Now, we can invoke the command.
668+ $psCmd.Invoke (@ ($webSocketMessage ))
610669 } else {
670+ # Otherwise, we'll just run the handler.
611671 $webSocketMessage | . $handler
612672 }
613-
614673 } else {
615674 $webSocketMessage
616- }
675+ }
617676 } catch {
618677 Write-Error $_
619678 }
620679 }
621680
681+ # Now that the socket is closed,
682+ # check for a status description.
683+ # If there is one,
622684 if ($ws.CloseStatusDescription ) {
685+ # write an error.
623686 Write-Error $ws.CloseStatusDescription - TargetObject $ws
624687 }
625688 }
0 commit comments