Move common WebSocket setup to the moduleHEADmaster

3 files changed, 16 insertions, 13 deletions

diff --git a/src/Attach.hs b/src/Attach.hs
index c15b0a5..cf6bd1c 100644
--- a/src/Attach.hs
+++ b/src/Attach.hs
@@ -30,7 +30,7 @@ import Text.Blaze.Html5 qualified as H
 import Text.Blaze.Html5.Attributes qualified as A
 
 import JavaScript
-import WebSocket (startClient, receiveMessage)
+import WebSocket (startDefaultWebSocketConnections)
 
 
 data AttachState = AttachState
@@ -120,11 +120,7 @@ initiateAttachConnection baseHead dgst = do
             Nothing -> return ()
         locationAssign "./"
 
-    startClient asServer "a.discovery.erebosprotocol.net" 443 "" $ \conn -> do
-        void $ forkIO $ forever $ do
-            msg <- receiveMessage conn
-            receivedFromCustomAddress asServer conn msg
-        void $ serverPeerCustom asServer conn
+    startDefaultWebSocketConnections asServer
 
     runExceptT (discoverySearch asServer dgst) >>= \case
         Right _ -> return ()
diff --git a/src/Main.hs b/src/Main.hs
index c8a2e82..547aaca 100644
--- a/src/Main.hs
+++ b/src/Main.hs
@@ -45,7 +45,7 @@ import JavaScript as JS
 import Storage.Cache
 import Storage.IndexedDB
 import Version
-import WebSocket (startClient, receiveMessage)
+import WebSocket (startDefaultWebSocketConnections)
 
 main :: IO ()
 main = error "unused"
@@ -288,11 +288,7 @@ setup = do
     maybe (return ()) (watchPeers gs server) =<<
         JS.getElementById "peer_list"
 
-    startClient server "a.discovery.erebosprotocol.net" 443 "" $ \conn -> do
-        void $ forkIO $ forever $ do
-            msg <- receiveMessage conn
-            receivedFromCustomAddress server conn msg
-        void $ serverPeerCustom server conn
+    startDefaultWebSocketConnections server
 
     Just inviteGenerateInput <- JS.getElementById "invite_name"
     Just inviteGenerateForm <- JS.getElementById "invite_generate"
diff --git a/src/WebSocket.hs b/src/WebSocket.hs
index 0492c81..4179cfa 100644
--- a/src/WebSocket.hs
+++ b/src/WebSocket.hs
@@ -1,11 +1,13 @@
 module WebSocket (
     Connection,
     startClient,
+    startDefaultWebSocketConnections,
     sendMessage,
     receiveMessage,
 ) where
 
-import Control.Concurrent.Chan
+import Control.Concurrent
+import Control.Monad
 
 import Data.ByteString (ByteString)
 import Data.ByteString.Unsafe
@@ -70,6 +72,15 @@ startClient server addr port path fun = do
         dropPeerAddress server $ CustomPeerAddress conn
 
 
+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
+
+
 sendMessage :: Connection -> ByteString -> IO ()
 sendMessage Connection {..} bs = do
     unsafeUseAsCStringLen bs $ \( ptr, len ) -> do