#895: separate out remove and disconnect functions

Badgerati · May 12, 2022 · 4a9573d · 4a9573d
1 parent 7d2f20f
commit 4a9573d
Show file tree

Hide file tree

Showing 3 changed files with 46 additions and 6 deletions.
diff --git a/src/Listener/PodeReceiver.cs b/src/Listener/PodeReceiver.cs
@@ -62,6 +62,19 @@ public void DisconnectWebSocket(string name)
                     return;
                 }
 
+                WebSockets[name].Dispose();
+            }
+        }
+
+        public void RemoveWebSocket(string name)
+        {
+            lock (WebSockets)
+            {
+                if (!WebSockets.ContainsKey(name))
+                {
+                    return;
+                }
+
                 WebSockets[name].Dispose();
                 WebSockets.Remove(name);
             }

diff --git a/src/Pode.psd1 b/src/Pode.psd1
@@ -311,6 +311,7 @@
         'Set-PodeWebSocketConcurrency',
         'Connect-PodeWebSocket',
         'Disconnect-PodeWebSocket',
+        'Remove-PodeWebSocket',
         'Send-PodeWebSocket',
         'Reset-PodeWebSocket',
         'Test-PodeWebSocket'

diff --git a/src/Public/WebSockets.ps1 b/src/Public/WebSockets.ps1
@@ -115,11 +115,28 @@ function Disconnect-PodeWebSocket
 
     if (Test-PodeWebSocket -Name $Name) {
         $PodeContext.Server.WebSockets.Receiver.DisconnectWebSocket($Name)
-        $PodeContext.Server.WebSockets.Connections.Remove($Name)
+    }
+}
+
+function Remove-PodeWebSocket
+{
+    [CmdletBinding()]
+    param(
+        [Parameter()]
+        [string]
+        $Name
+    )
+
+    if ([string]::IsNullOrWhiteSpace($Name) -and ($null -ne $WsEvent)) {
+        $Name = $WsEvent.Request.WebSocket.Name
+    }
+
+    if ([string]::IsNullOrWhiteSpace($Name)) {
+        throw "No Name for a WebSocket to remove supplied"
     }
 
-    #TODO: this should just disconnect, not disconnect and remove
-    # make a new Remove-WS to disconnect and remove
+    $PodeContext.Server.WebSockets.Receiver.RemoveWebSocket($Name)
+    $PodeContext.Server.WebSockets.Connections.Remove($Name)
 }
 
 function Send-PodeWebSocket
@@ -150,7 +167,7 @@ function Send-PodeWebSocket
     }
 
     if (Test-PodeWebSocket -Name $Name) {
-        $ws.Send($Message, $Type)
+        $PodeContext.Server.WebSockets.Receiver.GetWebSocket($Name).Send($Message, $Type)
     }
 }
 
@@ -177,7 +194,7 @@ function Reset-PodeWebSocket
     }
 
     if (Test-PodeWebSocket -Name $Name) {
-        $ws.Reconnect($Url)
+        $PodeContext.Server.WebSockets.Receiver.GetWebSocket($Name).Reconnect($Url)
     }
 }
 
@@ -190,5 +207,14 @@ function Test-PodeWebSocket
         $Name
     )
 
-    return ($null -ne $PodeContext.Server.WebSockets.Receiver.GetWebSocket($Name))
+    $found = ($null -ne $PodeContext.Server.WebSockets.Receiver.GetWebSocket($Name))
+    if ($found) {
+        return $true
+    }
+
+    if ($PodeContext.Server.WebSockets.Connections.ContainsKey($Name)) {
+        Remove-PodeWebSocket -Name $Name
+    }
+
+    return $false
 }