From bb7a85ec4b7f2505210b1b2a474a5520966b0921 Mon Sep 17 00:00:00 2001
From: Tim Bradgate <timbradgate@hotmail.co.uk>
Date: Mon, 10 Jul 2023 21:41:03 +0100
Subject: [PATCH] Add error handling to WebSocket connection when sending
 message to closed client

---
 server/controllers/ws_controller.py | 17 +++++++++++++++--
 1 file changed, 15 insertions(+), 2 deletions(-)

diff --git a/server/controllers/ws_controller.py b/server/controllers/ws_controller.py
index 6c7737fb..6884e45c 100644
--- a/server/controllers/ws_controller.py
+++ b/server/controllers/ws_controller.py
@@ -1,10 +1,11 @@
 import json
-from typing import Optional, Awaitable, Union, TYPE_CHECKING
+from typing import Optional, Awaitable, Union, TYPE_CHECKING, Dict, Any
 
 from uuid import uuid4
 
 from tornado import gen
-from tornado.websocket import WebSocketHandler
+from tornado.concurrent import Future
+from tornado.websocket import WebSocketHandler, WebSocketClosedError
 from tornado_sqlalchemy import SessionMixin
 
 from digi_server.logger import get_logger
@@ -246,3 +247,15 @@ def on_ping(self, data: bytes) -> None:
         self.update_session()
         get_logger().trace(
             f'Ping from {self.request.remote_ip} : {data.hex()}')
+
+    @gen.coroutine
+    def write_message(self, message: Union[bytes, str, Dict[str, Any]],
+                      binary: bool = False) -> Future[None]:
+        try:
+            return super().write_message(message, binary)
+        except WebSocketClosedError:
+            get_logger().error(f'Trying to send message to closed websocket '
+                               f'{self.__getattribute__("internal_id")} at IP address '
+                               f'{self.request.remote_ip}, closing.')
+            self.on_close()
+            return None