Drop websocket peer when connection is closed

1 files changed, 8 insertions, 2 deletions

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