1 files changed, 81 insertions, 0 deletions
diff --git a/main/WebSocket.hs b/main/WebSocket.hs
new file mode 100644
index 0000000..79cb141
--- /dev/null
+++ b/main/WebSocket.hs
@@ -0,0 +1,81 @@
+{-# LANGUAGE OverloadedStrings #-}
+
+module WebSocket (
+    WebSocketAddress(..),
+    WebSocketOptions(..), defaultWebSocketOptions,
+    startWebsocketServer,
+) where
+
+import Control.Concurrent
+import Control.Exception
+import Control.Monad
+
+import Data.ByteString.Char8 qualified as BC
+import Data.ByteString.Lazy qualified as BL
+import Data.Unique
+
+import Erebos.Network
+
+import Network.WebSockets qualified as WS
+
+
+data WebSocketAddress = WebSocketAddress Unique (Maybe String) WS.Connection
+
+instance Eq WebSocketAddress where
+    WebSocketAddress u _ _ == WebSocketAddress u' _ _ = u == u'
+
+instance Ord WebSocketAddress where
+    compare (WebSocketAddress u _ _) (WebSocketAddress u' _ _) = compare u u'
+
+instance Show WebSocketAddress where
+    show (WebSocketAddress _ Nothing _) = "websocket"
+    show (WebSocketAddress _ (Just addr) _) = "websocket " <> addr
+
+instance PeerAddressType WebSocketAddress where
+    sendBytesToAddress (WebSocketAddress _ _ conn) msg = do
+        WS.sendDataMessage conn $ WS.Binary $ BL.fromStrict msg
+    connectionToAddressClosed (WebSocketAddress _ _ conn) = do
+        WS.sendClose conn BL.empty `catch` \e -> if
+            | Just WS.ConnectionClosed <- fromException e -> return ()
+            | otherwise -> throwIO e
+
+
+data WebSocketOptions = WebSocketOptions
+    { wsAddress :: String
+    , wsPort :: Int
+    , wsDebugLog :: Bool
+    }
+
+defaultWebSocketOptions :: WebSocketOptions
+defaultWebSocketOptions = WebSocketOptions
+    { wsAddress = "::"
+    , wsPort = 80
+    , wsDebugLog = False
+    }
+
+
+startWebsocketServer :: Server -> (String -> IO ()) -> WebSocketOptions -> IO ()
+startWebsocketServer server logd WebSocketOptions {..} = do
+    void $ forkIO $ do
+        WS.runServer wsAddress wsPort $ \pending -> do
+            when wsDebugLog $ do
+                logd $ "WebSocket request: " <> show (WS.pendingRequest pending)
+            conn <- WS.acceptRequest pending
+            u <- newUnique
+            let mbaddr = fmap BC.unpack $ lookup "X-Real-IP" $ WS.requestHeaders $ WS.pendingRequest pending
+            let paddr = WebSocketAddress u mbaddr conn
+            void $ serverPeerCustom server paddr
+
+            let handler e
+                    | Just WS.CloseRequest {} <- fromException e = do
+                        dropPeerAddress server $ CustomPeerAddress paddr
+                    | Just WS.ConnectionClosed <- fromException e = do
+                        dropPeerAddress server $ CustomPeerAddress paddr
+                    | otherwise = do
+                        logd $ "WebSocket thread exception: " ++ show e
+            handle handler $ do
+                WS.withPingThread conn 30 (return ()) $ do
+                    forever $ do
+                        WS.receiveDataMessage conn >>= \case
+                            WS.Binary msg -> receivedFromCustomAddress server paddr $ BL.toStrict msg
+                            WS.Text {} -> logd $ "unexpected websocket text message"