From 54af55651cf9f7d43a5d74d10404801d73398e5a Mon Sep 17 00:00:00 2001
From: Marc Scholten <marc@digitallyinduced.com>
Date: Sat, 5 Feb 2022 13:33:49 +0100
Subject: [PATCH] Implemented expected timeout handling in websockets

The 'withPingThread'  function of the websockets package is not dealing with missing pong messages at all. This means that websocket connection might never be cleaned up when the connection is not closed correctly.

This change implements a manual ping and pong handling that closes a connection after not receiving a pong message within 10 seconds after sending a ping to the client.
---
 IHP/WebSocket.hs | 60 +++++++++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 59 insertions(+), 1 deletion(-)

diff --git a/IHP/WebSocket.hs b/IHP/WebSocket.hs
index 91662fafd..1489eebec 100644
--- a/IHP/WebSocket.hs
+++ b/IHP/WebSocket.hs
@@ -29,6 +29,12 @@ import qualified Data.Aeson as Aeson
 import qualified IHP.Log.Types as Log
 import qualified IHP.Log as Log
 
+import Control.Concurrent.Chan
+import Control.Concurrent
+import System.Timeout
+import Data.Function (fix)
+import qualified Network.WebSockets.Connection as WebSocket
+
 class WSApp state where
     initialState :: state
 
@@ -47,7 +53,15 @@ startWSApp connection = do
     let ?state = state
     let ?connection = connection
 
-    result <- Exception.try ((Websocket.withPingThread connection 30 (onPing @state) (run @state)) `Exception.finally` onClose @state)
+    let runWithPongChan pongChan = do
+        let connectionOnPong = writeChan pongChan ()
+        let ?connection = connection
+                { WebSocket.connectionOptions = (get #connectionOptions connection) { WebSocket.connectionOnPong } 
+                }
+            in
+                run @state
+
+    result <- Exception.try ((withPinger connection runWithPongChan) `Exception.finally` onClose @state)
     case result of
         Left (e@Exception.SomeException{}) ->
             case Exception.fromException e of
@@ -90,3 +104,47 @@ instance Websocket.WebSocketsData UUID where
     fromDataMessage (Websocket.Binary byteString) = UUID.fromLazyASCIIBytes byteString |> Maybe.fromJust
     fromLazyByteString byteString = UUID.fromLazyASCIIBytes byteString |> Maybe.fromJust
     toLazyByteString = UUID.toLazyASCIIBytes
+
+data PongTimeout
+    = PongTimeout
+    deriving (Show)
+
+instance Exception PongTimeout
+
+pingWaitTime :: Int
+pingWaitTime = 30
+
+
+-- | Pings the client every 30 seconds and expects a pong response within 10 secons. If no pong response
+-- is received within 10 seconds, it will kill the connection.
+--
+-- We cannot use the withPingThread of the websockets package as this doesn't deal with pong messages. So
+-- open connection will stay around forever.
+--
+-- This implementation is based on https://github.com/jaspervdj/websockets/issues/159#issuecomment-552776502
+withPinger conn action = do
+    pongChan <- newChan
+    mainAsync <- async $ action pongChan
+    pingerAsync <- async $ runPinger conn pongChan
+
+    waitEitherCatch mainAsync pingerAsync >>= \case
+        -- If the application async died for any reason, kill the pinger async
+        Left result -> do
+            cancel pingerAsync
+            case result of
+                Left exception -> throw exception
+                Right result -> pure ()
+        -- The pinger thread should never throw an exception. If it does, kill the app thread
+        Right (Left exception) -> do
+            cancel mainAsync
+            throw exception
+        -- The pinger thread exited due to a pong timeout. Tell the app thread about it.
+        Right (Right ()) -> cancelWith mainAsync PongTimeout
+
+runPinger conn pongChan = fix $ \loop -> do
+    Websocket.sendPing conn (mempty :: ByteString)
+    threadDelay pingWaitTime
+    -- See if we got a pong in that time
+    timeout 1000000 (readChan pongChan) >>= \case
+        Just () -> loop
+        Nothing -> return ()
\ No newline at end of file