Try reconnecting to the websocket connection when it's lost

1 files changed, 19 insertions, 5 deletions

diff --git a/src/WebSocket.hs b/src/WebSocket.hs
index 71ff6d6..685ed90 100644
--- a/src/WebSocket.hs
+++ b/src/WebSocket.hs
@@ -75,11 +75,25 @@ startClient server addr port path fun = do
 
 startDefaultWebSocketConnections :: Server -> IO ()
 startDefaultWebSocketConnections server = do
-    startClient server "a.discovery.erebosprotocol.net" 443 "" $ \conn -> do
-        void $ forkIO $ forever $ do
-            msg <- receiveMessage conn
-            receivedFromCustomAddress server conn msg
-        void $ serverPeerCustom server conn
+    peerVar <- newEmptyMVar
+    let doStart = do
+            startClient server "a.discovery.erebosprotocol.net" 443 "" $ \conn -> do
+                putMVar peerVar =<< serverPeerCustom server conn
+                void $ forkIO $ forever $ do
+                    msg <- receiveMessage conn
+                    receivedFromCustomAddress server conn msg
+
+    void $ forkIO $ void $ forever $ do
+        peer <- getNextPeerChange server
+        isPeerDropped peer >>= \case
+            True -> do
+                wsPeer <- takeMVar peerVar
+                if peer == wsPeer
+                  then doStart
+                  else putMVar peerVar wsPeer
+            False -> return ()
+
+    doStart
 
 
 sendMessage :: Connection -> ByteString -> IO ()