From 3bd7a6972422d4594708c0fd19e01b6272ee8501 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Domen=20Ko=C5=BEar?= <domen@dev.si>
Date: Tue, 9 Jan 2024 05:48:27 +0000
Subject: [PATCH] server: add handleException

---
 src/Network/WebSockets/Simple/Server.hs | 10 ++++++----
 websockets-simple.cabal                 |  2 +-
 2 files changed, 7 insertions(+), 5 deletions(-)

diff --git a/src/Network/WebSockets/Simple/Server.hs b/src/Network/WebSockets/Simple/Server.hs
index bb49755..daa1184 100644
--- a/src/Network/WebSockets/Simple/Server.hs
+++ b/src/Network/WebSockets/Simple/Server.hs
@@ -9,7 +9,7 @@ module Network.WebSockets.Simple.Server
   )
 where
 
-import Control.Exception (throwIO)
+import Control.Exception.Safe (SomeException, handle, throwIO)
 import Control.Monad (when)
 import Data.ByteString (ByteString)
 import Data.ByteString.Char8 (unpack)
@@ -22,7 +22,8 @@ import Network.WebSockets.Simple.Utils qualified as Utils
 data Options a = Options
   { handlePendingConnection :: (ClientConnection a) => WS.PendingConnection -> IO (Maybe a),
     pingPongOptions :: WS.PingPongOptions -> IO WS.PingPongOptions,
-    messageLimit :: Int
+    messageLimit :: Int,
+    handleException :: WS.PendingConnection -> SomeException -> IO ()
   }
 
 --
@@ -31,7 +32,8 @@ defaultOptions =
   Options
     { handlePendingConnection = (fmap Just) . WS.acceptRequest,
       pingPongOptions = return,
-      messageLimit = 10000
+      messageLimit = 10000,
+      handleException = \_ _ -> return ()
     }
 
 class ClientConnection a where
@@ -53,7 +55,7 @@ run uriBS options app receiveApp = do
   WS.runServerWithOptions serverOptions (application pingpongOpts)
   where
     application :: PingPong.PingPongOptions -> WS.ServerApp
-    application pingpongOpts pendingConnection = do
+    application pingpongOpts pendingConnection = handle (handleException options pendingConnection) $ do
       maybeClient <- handlePendingConnection options pendingConnection
       for_ maybeClient $ \client -> PingPong.withPingPong pingpongOpts (getConnection client) $ \_ ->
         Session.run (messageLimit options) (getConnection client) (app client) (receiveApp client)
diff --git a/websockets-simple.cabal b/websockets-simple.cabal
index aa132c0..39b9763 100644
--- a/websockets-simple.cabal
+++ b/websockets-simple.cabal
@@ -26,7 +26,7 @@ common common
                    wuss >= 2.0.1.7,
                    unliftio-core,
                    bytestring,
-                   exceptions,
+                   safe-exceptions,
                    stamina
                    
     default-extensions: OverloadedStrings