From 499e7597f70f74239cca0c2da17d938bc6fc87a6 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Domen=20Ko=C5=BEar?= <domen@dev.si>
Date: Fri, 22 Dec 2023 13:29:45 +0000
Subject: [PATCH] Client: default to 30s for connection timeout

---
 src/Network/WebSockets/Client.hs             | 11 +++++++----
 src/Network/WebSockets/Connection/Options.hs |  4 ++++
 src/Network/WebSockets/Http.hs               |  2 ++
 3 files changed, 13 insertions(+), 4 deletions(-)

diff --git a/src/Network/WebSockets/Client.hs b/src/Network/WebSockets/Client.hs
index 28880c3..3d41e55 100644
--- a/src/Network/WebSockets/Client.hs
+++ b/src/Network/WebSockets/Client.hs
@@ -25,6 +25,7 @@ import           Data.IORef                    (newIORef)
 import qualified Data.Text                     as T
 import qualified Data.Text.Encoding            as T
 import qualified Network.Socket                as S
+import           System.Timeout                (timeout)
 
 
 --------------------------------------------------------------------------------
@@ -74,10 +75,12 @@ runClientWith host port path0 opts customHeaders app = do
     S.setSocketOption sock S.NoDelay 1
 
     -- Connect WebSocket and run client
-    res <- finally
-        (S.connect sock (S.addrAddress addr) >>
-         runClientWithSocket sock fullHost path opts customHeaders app)
-        (S.close sock)
+    res <- bracket
+        (timeout (connectionTimeout opts * 1000 * 1000) $ S.connect sock (S.addrAddress addr))
+        (const $ S.close sock) $ \maybeConnected -> case maybeConnected of
+            Nothing -> throwIO $ ConnectionTimeout
+            Just () -> runClientWithSocket sock fullHost path opts customHeaders app
+
 
     -- Clean up
     return res
diff --git a/src/Network/WebSockets/Connection/Options.hs b/src/Network/WebSockets/Connection/Options.hs
index 1bf0168..1255c31 100644
--- a/src/Network/WebSockets/Connection/Options.hs
+++ b/src/Network/WebSockets/Connection/Options.hs
@@ -31,6 +31,8 @@ data ConnectionOptions = ConnectionOptions
     { connectionOnPong                :: !(IO ())
       -- ^ Whenever a 'pong' is received, this IO action is executed. It can be
       -- used to tickle connections or fire missiles.
+    , connectionTimeout               :: !Int
+      -- ^ Timeout for connection establishment in seconds. Only used in the client.
     , connectionCompressionOptions    :: !CompressionOptions
       -- ^ Enable 'PermessageDeflate'.
     , connectionStrictUnicode         :: !Bool
@@ -59,9 +61,11 @@ data ConnectionOptions = ConnectionOptions
 -- * Nothing happens when a pong is received.
 -- * Compression is disabled.
 -- * Lenient unicode decoding.
+-- * 30 second timeout for connection establishment.
 defaultConnectionOptions :: ConnectionOptions
 defaultConnectionOptions = ConnectionOptions
     { connectionOnPong                = return ()
+    , connectionTimeout               = 30
     , connectionCompressionOptions    = NoCompression
     , connectionStrictUnicode         = False
     , connectionFramePayloadSizeLimit = mempty
diff --git a/src/Network/WebSockets/Http.hs b/src/Network/WebSockets/Http.hs
index 2194f4e..25d85a7 100644
--- a/src/Network/WebSockets/Http.hs
+++ b/src/Network/WebSockets/Http.hs
@@ -101,6 +101,8 @@ data HandshakeException
     -- | The request was well-formed, but the library user rejected it.
     -- (e.g. "unknown path")
     | RequestRejected RequestHead ResponseHead
+    -- | The connection timed out
+    | ConnectionTimeout
     -- | for example "EOF came too early" (which is actually a parse error)
     -- or for your own errors. (like "unknown path"?)
     | OtherHandshakeException String