From f9a0fd732a4562adf6a38f1b1df12c162b9f19ac Mon Sep 17 00:00:00 2001
From: "Frank A. Krueger" <fak@praeclarum.org>
Date: Sun, 15 Apr 2018 17:40:48 -0700
Subject: [PATCH] Better error handling with web sockets

---
 Ooui.AspNetCore/ElementResult.cs    | 11 +++++++----
 Ooui.AspNetCore/WebSocketHandler.cs | 29 +++++++++++++++++++++++------
 2 files changed, 30 insertions(+), 10 deletions(-)

diff --git a/Ooui.AspNetCore/ElementResult.cs b/Ooui.AspNetCore/ElementResult.cs
index 4c1f0fe..c7f64b6 100644
--- a/Ooui.AspNetCore/ElementResult.cs
+++ b/Ooui.AspNetCore/ElementResult.cs
@@ -3,6 +3,7 @@ using System.Text;
 using System.Threading.Tasks;
 using Microsoft.AspNetCore.Http;
 using Microsoft.AspNetCore.Mvc;
+using Microsoft.Extensions.Logging;
 
 namespace Ooui.AspNetCore
 {
@@ -10,10 +11,12 @@ namespace Ooui.AspNetCore
     {
         readonly Element element;
         readonly string title;
-        readonly bool disposeWhenDone;
-
-        public ElementResult (Element element, string title = "", bool disposeWhenDone = true)
-        {
+        readonly bool disposeWhenDone;
+        readonly ILogger logger;
+
+        public ElementResult (Element element, string title = "", bool disposeWhenDone = true, ILogger logger = null)
+        {
+            this.logger = logger;
             this.element = element;
             this.title = title;
             this.disposeWhenDone = disposeWhenDone;
diff --git a/Ooui.AspNetCore/WebSocketHandler.cs b/Ooui.AspNetCore/WebSocketHandler.cs
index f43ae73..7b479ea 100644
--- a/Ooui.AspNetCore/WebSocketHandler.cs
+++ b/Ooui.AspNetCore/WebSocketHandler.cs
@@ -4,6 +4,8 @@ using System.Linq;
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.AspNetCore.Http;
+using System.IO;
+using Microsoft.Extensions.Logging;
 
 namespace Ooui.AspNetCore
 {
@@ -33,9 +35,7 @@ namespace Ooui.AspNetCore
             return id;
         }
 
-
-
-        public static async Task HandleWebSocketRequestAsync (HttpContext context)
+        public static async Task HandleWebSocketRequestAsync (HttpContext context, ILogger logger = null)
         {
             void BadRequest (string message)
             {
@@ -97,9 +97,26 @@ namespace Ooui.AspNetCore
             // OK, Run
             //
             var token = CancellationToken.None;
-            var webSocket = await context.WebSockets.AcceptWebSocketAsync ("ooui");
-            var session = new Ooui.WebSocketSession (webSocket, activeSession.Element, activeSession.DisposeElementWhenDone, w, h, token);
-            await session.RunAsync ().ConfigureAwait (false);
+            System.Net.WebSockets.WebSocket webSocket = null;
+
+            //
+            // Create a new session and let it handle everything from here
+            //
+            try {
+                webSocket = await context.WebSockets.AcceptWebSocketAsync ("ooui").ConfigureAwait (false);
+                var session = new Ooui.WebSocketSession (webSocket, activeSession.Element, activeSession.DisposeElementWhenDone, w, h, token);
+                await session.RunAsync ().ConfigureAwait (false);
+            }
+            catch (System.Net.WebSockets.WebSocketException ex) when (ex.WebSocketErrorCode == System.Net.WebSockets.WebSocketError.ConnectionClosedPrematurely) {
+                // The remote party closed the WebSocket connection without completing the close handshake.
+            }
+            catch (Exception ex) {
+                context.Abort ();
+                Console.WriteLine (ex);
+            }
+            finally {
+                webSocket?.Dispose ();
+            }
         }
 
         class PendingSession