Skip to content

Commit 87b29a5

Browse files
author
James Brundage
committed
feat: Get-WebSocket returns socket or listener jobs ( Fixes #68 )
Also, adding some aliases
1 parent 8a20eca commit 87b29a5

File tree

1 file changed

+245
-27
lines changed

1 file changed

+245
-27
lines changed

‎Commands/Get-WebSocket.ps1‎

Lines changed: 245 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -111,16 +111,30 @@ function Get-WebSocket {
111111
param(
112112
# The WebSocket Uri.
113113
[Parameter(Position=0,ValueFromPipelineByPropertyName)]
114-
[Alias('Url','Uri')]
115-
[uri]$WebSocketUri,
114+
[Alias('Url','Uri','WebSocketUrl','WebSocketUri')]
115+
[uri]
116+
$WebSocketUri,
116117

117118
# One or more root urls.
118119
# If these are provided, a WebSocket server will be created with these listener prefixes.
119120
[Parameter(Position=1,ValueFromPipelineByPropertyName)]
120-
[Alias('HostHeader','Host','ServerURL','ListenerPrefix','ListenerPrefixes','ListenerUrl')]
121+
[Alias('HostHeader','Host','CNAME','ServerURL','ListenerPrefix','ListenerPrefixes','ListenerUrl')]
121122
[string[]]
122123
$RootUrl,
123124

125+
# A route table for all requests.
126+
[Parameter(ValueFromPipelineByPropertyName)]
127+
[Alias('Routes','RouteTable','WebHook','WebHooks')]
128+
[Collections.IDictionary]
129+
$Route,
130+
131+
# The Default HTML.
132+
# This will be displayed when visiting the root url.
133+
[Parameter(ValueFromPipelineByPropertyName)]
134+
[Alias('DefaultHTML','Home','Index','IndexHTML','DefaultPage')]
135+
[string]
136+
$HTML,
137+
124138
# A collection of query parameters.
125139
# These will be appended onto the `-WebSocketUri`.
126140
[Collections.IDictionary]
@@ -421,7 +435,33 @@ function Get-WebSocket {
421435
# Take every every `-Variable` passed in and define it within the job
422436
foreach ($keyValue in $variable.GetEnumerator()) {
423437
$ExecutionContext.SessionState.PSVariable.Set($keyValue.Key, $keyValue.Value)
424-
}
438+
}
439+
440+
# If we have routes, we will cache all of their possible parameters now
441+
if ($route.Count) {
442+
# We want to keep the parameter sets
443+
$routeParameterSets = [Ordered]@{}
444+
# and the metadata about parameters.
445+
$routeParameters = [Ordered]@{}
446+
447+
# For each key and value in the route table, we will try to get the command info for the value.
448+
foreach ($routePair in $route.GetEnumerator()) {
449+
$routeToCmd =
450+
# If the value is a scriptblock
451+
if ($routePair.Value -is [ScriptBlock]) {
452+
# we have to create a temporary function
453+
$function:TempFunction = $routePair.Value
454+
# and get that function.
455+
$ExecutionContext.SessionState.InvokeCommand.GetCommand('TempFunction', 'Function')
456+
} elseif ($routePair.Value -is [Management.Automation.CommandInfo]) {
457+
$routePair.Value
458+
}
459+
if ($routeToCmd) {
460+
$routeParameterSets[$routePair.Name] = $routeToCmd.ParametersSets
461+
$routeParameters[$routePair.Name] = $routeToCmd.Parameters
462+
}
463+
}
464+
}
425465

426466
# If there's no listener, create one.
427467
if (-not $httpListener) {
@@ -454,7 +494,12 @@ function Get-WebSocket {
454494
# Get the context async result.
455495
# The context is basically the next request and response in the queue.
456496
$context = $(try { $contextAsync.Result } catch { $_ })
457-
$RequestedUrl = $context.Request.Url
497+
498+
# yield the context immediately, in case anything is watching the output of this job
499+
$context
500+
501+
$Request, $response = $context.Request, $context.Response
502+
$RequestedUrl = $Request.Url
458503
# Favicons are literally outdated, but they're still requested.
459504
if ($RequestedUrl -match '/favicon.ico$') {
460505
# by returning a 404 for them, we can make the browser stop asking.
@@ -464,16 +509,18 @@ function Get-WebSocket {
464509
}
465510
# Now, for the fun part.
466511
# We turn request into a PowerShell events.
467-
# Each event will have the source identifier of the request scheme
468-
$eventIdentifier = "$($context.Request.Url.Scheme)://"
512+
# The protocol is the scheme of the request url.
513+
$Protocol = $RequestedUrl.Scheme
514+
# Each event will have the source identifier of the protocol, followed by ://
515+
$eventIdentifier = "$($Protocol)://"
469516
# and by default it will pass a message containing the context.
470-
$messageData = [Ordered]@{Url = $context.Request.Url;Context = $context}
517+
$messageData = [Ordered]@{Protocol = $protocol; Url = $context.Request.Url;Context = $context}
471518

472519
# HttpListeners are quite nice, especially when it comes to websocket upgrades.
473520
# If the request is a websocket request
474-
if ($context.Request.IsWebSocketRequest) {
521+
if ($Request.IsWebSocketRequest) {
475522
# we will change the event identifier to a websocket scheme.
476-
$eventIdentifier = $eventIdentifier -replace '^http', 'ws'
523+
$eventIdentifier = $eventIdentifier -replace '^http', 'ws'
477524
# and call the `AcceptWebSocketAsync` method to upgrade the connection.
478525
$acceptWebSocket = $context.AcceptWebSocketAsync('json')
479526
# Once again, we'll use a tight loop to wait for the upgrade to complete or fail.
@@ -486,30 +533,196 @@ function Get-WebSocket {
486533
}
487534
# If it succeeds, capture the result.
488535
$webSocketResult = try { $acceptWebSocket.Result } catch { $_ }
489-
# and add it to the SocketRequests lookup table, using the request trace identifier as the key.
490-
$httpListener.SocketRequests[$context.Request.RequestTraceIdentifier] = $webSocketResult
491-
# and add the websocketcontext result to the message data.
492-
$messageData["WebSocketContext"] = $webSocketResult
493-
# also add the websocket result to the message data, since many might not exactly know what a "WebSocketContext" is.
494-
$messageData["WebSocket"] = $webSocketResult.WebSocket
536+
537+
# If the websocket is open
538+
if ($webSocketResult.WebSocket.State -eq 'open') {
539+
# we have switched protocols!
540+
$Protocol = $requestedUrl.Scheme -replace '^http', 'ws'
541+
542+
# Now add the result it to the SocketRequests lookup table, using the request trace identifier as the key.
543+
$httpListener.SocketRequests[$context.Request.RequestTraceIdentifier] = $webSocketResult
544+
# and add the websocketcontext result to the message data.
545+
$messageData["WebSocketContext"] = $webSocketResult
546+
# also add the websocket result to the message data,
547+
# since many might not exactly know what a "WebSocketContext" is.
548+
$messageData["WebSocket"] = $webSocketResult.WebSocket
549+
}
495550
}
496551

497552
# Now, we generate the event.
498553
$generateEventArguments = @(
499-
$eventIdentifier,
500-
$httpListener,
554+
$eventIdentifier,
555+
$httpListener,
501556
@($context)
502557
$messageData
503558
)
504559
# Get a pointer to the GenerateEvent method (we'll want this later)
505560
if ($MainRunspace.Events.GenerateEvent) {
506561
$MainRunspace.Events.GenerateEvent.Invoke($generateEventArguments)
507562
}
563+
564+
# Everything below this point is for HTTP requests.
565+
if ($protocol -notmatch '^http') {
566+
continue # so if we're already a websocket, we will skip the rest of this code.
567+
}
568+
569+
$routedTo = $null
570+
$routeKey = $null
571+
# If we have routes, we will try to find a route that matches the request.
572+
if ($route.Count) {
573+
$routeTable = $route
574+
$potentialRouteKeys = @(
575+
$request.Url.AbsolutePath,
576+
($request.Url.AbsolutePath -replace '/$'),
577+
"$($request.HttpMethod) $($request.Url.AbsolutePath)",
578+
"$($request.HttpMethod) $($request.Url.AbsolutePath -replace '/$')"
579+
"$($request.HttpMethod) $($request.Url.LocalPath)",
580+
"$($request.HttpMethod) $($request.Url.LocalPath -replace '/$')"
581+
)
582+
$routedTo = foreach ($potentialKey in $potentialRouteKeys) {
583+
if ($routeTable[$potentialKey]) {
584+
$routeTable[$potentialKey]
585+
$routeKey = $potentialKey
586+
break
587+
}
588+
}
589+
}
590+
591+
if (-not $routedTo -and $html) {
592+
$routedTo =
593+
# If the content is already html, we will use it as is.
594+
if ($html -match '\<html') {
595+
$html
596+
} else {
597+
# Otherwise, we will wrap it in an html tag.
598+
@(
599+
"<html>"
600+
"<head>"
601+
# and apply the site header.
602+
$SiteHeader
603+
"</head>"
604+
"<body>"
605+
$html
606+
"</body>"
607+
"</html>"
608+
) -join [Environment]::NewLine
609+
}
610+
}
611+
612+
# If we routed to a string, we will close the response with the string.
613+
if ($routedTo -is [string]) {
614+
$response.Close($OutputEncoding.GetBytes($routedTo), $true)
615+
continue
616+
}
617+
618+
# If we've routed to is a byte array, we will close the response with the byte array.
619+
if ($routedTo -is [byte[]]) {
620+
$response.Close($routedTo, $true)
621+
continue
622+
}
623+
624+
# If we routed to a script block or command, we will try to execute it.
625+
if ($routedTo -is [ScriptBlock] -or
626+
$routedTo -is [Management.Automation.CommandInfo]) {
627+
$routeSplat = [Ordered]@{}
628+
629+
# If the command had a `-Request` parameter, we will pass the request object.
630+
if ($routeParameters -and $routeParameters[$routeKey].Request) {
631+
$routeSplat['Request'] = $request
632+
}
633+
# If the command had a `-Response` parameter, we will pass the response object.
634+
if ($routeParameters -and $routeParameters[$routeKey].Response) {
635+
$routeSplat['Response'] = $response
636+
}
637+
638+
# If the request has a query string, we will parse it and pass the values to the command.
639+
if ($request.Url.QueryString) {
640+
$parsedQuery = [Web.HttpUtility]::ParseQueryString($request.Url.QueryString)
641+
foreach ($parsedQueryKey in $parsedQuery.Keys) {
642+
if ($routeParameters[$routeKey][$parsedQueryKey]) {
643+
$routeSplat[$parsedQueryKey] = $parsedQuery[$parsedQueryKey]
644+
}
645+
}
646+
}
647+
# If the request has a content type of json, we will parse the json and pass the values to the command.
648+
if ($request.ContentType -match '^(?>application|text)/json') {
649+
$streamReader = [IO.StreamReader]::new($request.InputStream)
650+
$json = $streamReader.ReadToEnd()
651+
$jsonHashtable = ConvertFrom-Json -InputObject $json -AsHashtable
652+
foreach ($keyValuePair in $jsonHashtable.GetEnumerator()) {
653+
if ($routeParameters[$routeKey][$keyValuePair.Key]) {
654+
$routeSplat[$keyValuePair.Key] = $keyValuePair.Value
655+
}
656+
}
657+
$streamReader.Close()
658+
$streamReader.Dispose()
659+
}
660+
661+
# If the request has a content type of form-urlencoded, we will parse the form and pass the values to the command.
662+
if ($request.ContentType -eq 'application/x-www-form-urlencoded') {
663+
$streamReader = [IO.StreamReader]::new($request.InputStream)
664+
$formData = [Web.HttpUtility]::ParseQueryString($streamReader.ReadToEnd())
665+
foreach ($formKey in $formData.Keys) {
666+
if ($routeParameters[$routeKey][$formKey]) {
667+
$routeSplat[$formKey] = $form[$formKey]
668+
}
669+
}
670+
$streamReader.Close()
671+
$streamReader.Dispose()
672+
}
673+
674+
# We will execute the command and get the output.
675+
$routeOutput = . $routedTo @routeSplat
676+
677+
# If the output is a string, we will close the response with the string.
678+
if ($routeOutput -is [string])
679+
{
680+
$response.Close($OutputEncoding.GetBytes($routeOutput), $true)
681+
continue
682+
}
683+
# If the output is a byte array, we will close the response with the byte array.
684+
elseif ($routeOutput -is [byte[]])
685+
{
686+
$response.Close($routeOutput, $true)
687+
continue
688+
}
689+
# If the response is an array, write the responses out one at a time.
690+
# (note: this will likely be changed in the future)
691+
elseif ($routeOutput -is [object[]]) {
692+
foreach ($routeOut in $routeOutput) {
693+
if ($routeOut -is [string]) {
694+
$routeOut = $OutputEncoding.GetBytes($routeOut)
695+
}
696+
if ($routeOut -is [byte[]]) {
697+
$response.OutputStream.Write($routeOut, 0, $routeOut.Length)
698+
}
699+
}
700+
$response.Close()
701+
}
702+
else {
703+
# If the response was an object, we will convert it to json and close the response with the json.
704+
$responseJson = ConvertTo-Json -InputObject $routeOutput -Depth 3
705+
$response.ContentType = 'application/json'
706+
$response.Close($OutputEncoding.GetBytes($responseJson), $true)
707+
}
708+
}
508709
}
509710
}
510711
}
511712

512-
process {
713+
process {
714+
if ((-not $WebSocketUri) -and (-not $RootUrl)) {
715+
$socketAndListenerJobs =
716+
foreach ($job in Get-Job) {
717+
if (
718+
$Job.WebSocket -is [Net.WebSockets.ClientWebSocket] -or
719+
$Job.HttpListener -is [Net.HttpListener]
720+
) {
721+
$job
722+
}
723+
}
724+
$socketAndListenerJobs
725+
}
513726
# First, let's pack all of the parameters into a dictionary of variables.
514727
foreach ($keyValuePair in $PSBoundParameters.GetEnumerator()) {
515728
$Variable[$keyValuePair.Key] = $keyValuePair.Value
@@ -520,17 +733,22 @@ function Get-WebSocket {
520733
# If we're going to be listening for HTTP requests, run a thread job for the server.
521734
if ($RootUrl) {
522735

523-
$httpListener = $variable['HttpListener'] = [Net.HttpListener]::new()
524-
foreach ($rootUrl in $RootUrl) {
525-
if ($rootUrl -match '^https?://') {
526-
$httpListener.Prefixes.Add($rootUrl)
736+
$variable['HttpListener'] = $httpListener = [Net.HttpListener]::new()
737+
foreach ($potentialPrefix in $RootUrl) {
738+
if ($potentialPrefix -match '^https?://') {
739+
$httpListener.Prefixes.Add($potentialPrefix)
527740
} else {
528-
$httpListener.Prefixes.Add("http://$rootUrl/")
529-
$httpListener.Prefixes.Add("https://$rootUrl/")
741+
$httpListener.Prefixes.Add("http://$potentialPrefix/")
742+
$httpListener.Prefixes.Add("https://$potentialPrefix/")
530743
}
531744
}
532745
$httpListener.Start()
533-
$httpListenerJob = Start-ThreadJob -ScriptBlock $SocketServerJob -Name "$RootUrl" -InitializationScript $InitializationScript -ArgumentList $Variable
746+
747+
if ($DebugPreference -notin 'SilentlyContinue','Ignore') {
748+
. $SocketServerJob -Variable $Variable
749+
} else {
750+
$httpListenerJob = Start-ThreadJob -ScriptBlock $SocketServerJob -Name "$RootUrl" -InitializationScript $InitializationScript -ArgumentList $Variable
751+
}
534752

535753
if ($httpListenerJob) {
536754
foreach ($keyValuePair in $Variable.GetEnumerator()) {
@@ -593,7 +811,7 @@ function Get-WebSocket {
593811
$variable['EventSubscriptions'] = $eventSubscriptions
594812
}
595813

596-
if ($webSocketJob) {
814+
if ($webSocketJob -and -not $webSocketJob.WebSocket) {
597815
$webSocketConnectTimeout = [DateTime]::Now + $ConnectionTimeout
598816
while (-not $variable['WebSocket'] -and
599817
([DateTime]::Now -lt $webSocketConnectTimeout)) {

0 commit comments

Comments
 (0)