Drop websocket peer when connection is closed

2 files changed, 9 insertions, 3 deletions

diff --git a/src/Main.hs b/src/Main.hs
index 0d6fc45..31b1436 100644
--- a/src/Main.hs
+++ b/src/Main.hs
@@ -205,7 +205,7 @@ setup = do
     peerList <- JS.getElementById "peer_list"
     watchPeers gs server peerList
 
-    startClient "a.discovery.erebosprotocol.net" 443 "" $ \conn -> do
+    startClient server "a.discovery.erebosprotocol.net" 443 "" $ \conn -> do
         void $ forkIO $ forever $ do
             msg <- receiveMessage conn
             receivedFromCustomAddress server conn msg
diff --git a/src/WebSocket.hs b/src/WebSocket.hs
index 6adbfd9..a1196ca 100644
--- a/src/WebSocket.hs
+++ b/src/WebSocket.hs
@@ -43,8 +43,8 @@ instance PeerAddressType Connection where
     sendBytesToAddress = sendMessage
     connectionToAddressClosed = closeConnection
 
-startClient :: String -> Int -> String -> (Connection -> IO ()) -> IO ()
-startClient addr port path fun = do
+startClient :: Server -> String -> Int -> String -> (Connection -> IO ()) -> IO ()
+startClient server addr port path fun = do
     connUnique <- newUnique
     let connAddress = "wss://" <> addr <> ":" <> show port <> "/" <> path
     connJS <- js_initWebSocket (toJSString connAddress)
@@ -62,6 +62,12 @@ startClient addr port path fun = do
         bs <- unsafePackCStringFinalizer ptr len (free ptr)
         writeChan connInQueue bs
 
+    JS.addEventListener connJS "close" $ \_ -> do
+        dropPeerAddress server $ CustomPeerAddress conn
+
+    JS.addEventListener connJS "error" $ \_ -> do
+        dropPeerAddress server $ CustomPeerAddress conn
+
 
 sendMessage :: Connection -> ByteString -> IO ()
 sendMessage Connection {..} bs = do