Allow mounting websocket apps that have a HTTP fallback

digitallyinduced · Jul 22, 2022 · f799009 · f799009
1 parent d67fb1f
commit f799009
Show file tree

Hide file tree

Showing 2 changed files with 47 additions and 4 deletions.
diff --git a/IHP/ControllerSupport.hs b/IHP/ControllerSupport.hs
@@ -22,6 +22,7 @@ module IHP.ControllerSupport
 , jumpToAction
 , requestBodyJSON
 , startWebSocketApp
+, startWebSocketAppAndFailOnHTTP
 , setHeader
 , addResponseHeaders
 , addResponseHeadersFromContext
@@ -118,8 +119,8 @@ prepareRLSIfNeeded modelContext = do
         Nothing -> pure modelContext
 
 {-# INLINE startWebSocketApp #-}
-startWebSocketApp :: forall webSocketApp application. (?applicationContext :: ApplicationContext, ?context :: RequestContext, InitControllerContext application, ?application :: application, Typeable application, WebSockets.WSApp webSocketApp) => IO ResponseReceived
-startWebSocketApp = do
+startWebSocketApp :: forall webSocketApp application. (?applicationContext :: ApplicationContext, ?context :: RequestContext, InitControllerContext application, ?application :: application, Typeable application, WebSockets.WSApp webSocketApp) => IO ResponseReceived -> IO ResponseReceived
+startWebSocketApp onHTTP = do
     let ?modelContext = ApplicationContext.modelContext ?applicationContext
     let ?requestContext = ?context
     let respond = ?context |> get #respond
@@ -142,7 +143,12 @@ startWebSocketApp = do
         |> WebSockets.websocketsApp WebSockets.defaultConnectionOptions handleConnection
         |> \case
             Just response -> respond response
-            Nothing -> respond $ responseLBS HTTP.status400 [(hContentType, "text/plain")] "This endpoint is only available via a WebSocket"
+            Nothing -> onHTTP
+{-# INLINE startWebSocketAppAndFailOnHTTP #-}
+startWebSocketAppAndFailOnHTTP :: forall webSocketApp application. (?applicationContext :: ApplicationContext, ?context :: RequestContext, InitControllerContext application, ?application :: application, Typeable application, WebSockets.WSApp webSocketApp) => IO ResponseReceived
+startWebSocketAppAndFailOnHTTP = startWebSocketApp @webSocketApp @application (respond $ responseLBS HTTP.status400 [(hContentType, "text/plain")] "This endpoint is only available via a WebSocket")
+    where
+        respond = ?context |> get #respond
 
 
 jumpToAction :: forall action. (Controller action, ?context :: ControllerContext, ?modelContext :: ModelContext) => action -> IO ()

diff --git a/IHP/RouterSupport.hs b/IHP/RouterSupport.hs
@@ -23,6 +23,7 @@ CanRoute (..)
 , parseText
 , webSocketApp
 , webSocketAppWithCustomPath
+, webSocketAppWithHTTPFallback
 , onlyAllowMethods
 , getMethod
 , routeParam
@@ -63,6 +64,7 @@ import qualified Data.Text.Encoding as Text
 import Data.Dynamic
 import IHP.Router.Types
 import IHP.WebSocket (WSApp)
+import qualified IHP.WebSocket as WS
 import GHC.TypeLits as T
 import IHP.Controller.Context
 import IHP.Controller.Param
@@ -683,6 +685,24 @@ webSocketApp = webSocketAppWithCustomPath @webSocketApp typeName
                 |> ByteString.pack
 {-# INLINABLE webSocketApp #-}
 
+webSocketAppWithHTTPFallback :: forall webSocketApp application.
+    ( WSApp webSocketApp
+    , InitControllerContext application
+    , ?application :: application
+    , ?applicationContext :: ApplicationContext
+    , ?context :: RequestContext
+    , Typeable application
+    , Typeable webSocketApp
+    , Controller webSocketApp
+    ) => Parser (IO ResponseReceived)
+webSocketAppWithHTTPFallback = webSocketAppWithCustomPathAndHTTPFallback @webSocketApp @application typeName
+    where
+        typeName :: ByteString
+        typeName = Typeable.typeOf (error "unreachable" :: webSocketApp)
+                |> show
+                |> ByteString.pack
+{-# INLINABLE webSocketAppWithHTTPFallback #-}
+
 -- | Routes to a given WebSocket app if the path matches
 --
 -- __Example:__
@@ -706,9 +726,26 @@ webSocketAppWithCustomPath :: forall webSocketApp application.
 webSocketAppWithCustomPath path = do
         Attoparsec.char '/'
         string path
-        pure (startWebSocketApp @webSocketApp)
+        pure (startWebSocketAppAndFailOnHTTP @webSocketApp)
 {-# INLINABLE webSocketAppWithCustomPath #-}
 
+webSocketAppWithCustomPathAndHTTPFallback :: forall webSocketApp application.
+    ( WSApp webSocketApp
+    , InitControllerContext application
+    , ?application :: application
+    , ?applicationContext :: ApplicationContext
+    , ?context :: RequestContext
+    , Typeable application
+    , Typeable webSocketApp
+    , Controller webSocketApp
+    ) => ByteString -> Parser (IO ResponseReceived)
+webSocketAppWithCustomPathAndHTTPFallback path = do
+        Attoparsec.char '/'
+        string path
+        pure (startWebSocketApp @webSocketApp (runActionWithNewContext (WS.initialState @webSocketApp)))
+{-# INLINABLE webSocketAppWithCustomPathAndHTTPFallback #-}
+
+
 -- | Defines the start page for a router (when @\/@ is requested).
 startPage :: forall action application. (Controller action, InitControllerContext application, ?application::application, ?applicationContext::ApplicationContext, ?context::RequestContext, Typeable application, Typeable action) => action -> Parser (IO ResponseReceived)
 startPage action = get (ByteString.pack (actionPrefix @action)) action