sendRawResponse

yesod-core/Yesod/Core/Handler.hs

@@ -89,6 +89,9 @@ module Yesod.Core.Handler
     , sendResponseStatus
     , sendResponseCreated
     , sendWaiResponse
+#if MIN_VERSION_wai(2, 1, 0)
+    , sendRawResponse
+#endif
       -- * Different representations
       -- $representations
     , selectRep
@@ -170,7 +173,7 @@ import qualified Data.ByteString               as S
 import qualified Data.ByteString.Lazy          as L
 import qualified Data.Map                      as Map
 
-import Data.Conduit (Source)
+import Data.Conduit (Source, Sink)
 import           Control.Arrow                 ((***))
 import qualified Data.ByteString.Char8         as S8
 import           Data.Maybe                    (mapMaybe)
@@ -198,6 +201,9 @@ import Data.CaseInsensitive (CI)
 #if MIN_VERSION_wai(2, 0, 0)
 import qualified System.PosixCompat.Files as PC
 #endif
+#if MIN_VERSION_wai(2, 1, 0)
+import Control.Monad.Trans.Control (MonadBaseControl, control)
+#endif
 
 get :: MonadHandler m => m GHState
 get = liftHandlerT $ HandlerT $ I.readIORef . handlerState
@@ -547,6 +553,23 @@ sendResponseCreated url = do
 sendWaiResponse :: MonadHandler m => W.Response -> m b
 sendWaiResponse = handlerError . HCWai
 
+#if MIN_VERSION_wai(2, 1, 0)
+-- | Send a raw response. This is used for cases such as WebSockets. Requires
+-- WAI 2.1 or later, and a web server which supports raw responses (e.g.,
+-- Warp).
+--
+-- Since 1.2.7
+sendRawResponse :: (MonadHandler m, MonadBaseControl IO m)
+                => (Source IO S8.ByteString -> Sink S8.ByteString IO () -> m ())
+                -> m a
+sendRawResponse raw = control $ \runInIO ->
+    runInIO $ sendWaiResponse $ flip W.responseRaw fallback
+    $ \src sink -> runInIO (raw src sink) >> return ()
+  where
+    fallback = W.responseLBS H.status500 [("Content-Type", "text/plain")]
+        "sendRawResponse: backend does not support raw responses"
+#endif
+
 -- | Return a 404 not found page. Also denotes no handler available.
 notFound :: MonadHandler m => m a
 notFound = hcError NotFound

yesod-core/Yesod/Core/Internal/Response.hs

@@ -47,9 +47,20 @@ yarToResponse (YRWai a) _ _ _ is =
     case a of
         ResponseSource s hs w -> return $ ResponseSource s hs $ \f ->
             w f `finally` closeInternalState is
-        _ -> do
+        ResponseBuilder{} -> do
             closeInternalState is
             return a
+        ResponseFile{} -> do
+            closeInternalState is
+            return a
+#if MIN_VERSION_wai(2, 1, 0)
+        -- Ignore the fallback provided, in case it refers to a ResourceT state
+        -- in a ResponseSource.
+        ResponseRaw raw _ -> return $ ResponseRaw
+            (\f -> raw f `finally` closeInternalState is)
+            (responseLBS H.status500 [("Content-Type", "text/plain")]
+                "yarToResponse: backend does not support raw responses")
+#endif
 #else
 yarToResponse (YRWai a) _ _ _ = return a
 #endif

yesod-core/test/YesodCoreTest.hs

@@ -14,6 +14,7 @@ import qualified YesodCoreTest.Redirect as Redirect
 import qualified YesodCoreTest.JsLoader as JsLoader
 import qualified YesodCoreTest.RequestBodySize as RequestBodySize
 import qualified YesodCoreTest.Json as Json
+import qualified YesodCoreTest.RawResponse as RawResponse
 import qualified YesodCoreTest.Streaming as Streaming
 import qualified YesodCoreTest.Reps as Reps
 import qualified YesodCoreTest.Auth as Auth
@@ -37,6 +38,7 @@ specs = do
       JsLoader.specs
       RequestBodySize.specs
       Json.specs
+      RawResponse.specs
       Streaming.specs
       Reps.specs
       Auth.specs

yesod-core/test/YesodCoreTest/RawResponse.hs

@@ -0,0 +1,62 @@
+{-# LANGUAGE OverloadedStrings, TemplateHaskell, QuasiQuotes, TypeFamilies, MultiParamTypeClasses, ScopedTypeVariables #-}
+module YesodCoreTest.RawResponse (specs, Widget) where
+
+import Yesod.Core
+import Test.Hspec
+import qualified Data.Map as Map
+import Network.Wai.Test
+import Data.Text (Text)
+import Data.ByteString.Lazy (ByteString)
+import qualified Data.Conduit.List as CL
+import qualified Data.ByteString.Char8 as S8
+import Data.Conduit
+import qualified Data.Conduit.Binary as CB
+import Data.Char (toUpper)
+import Control.Exception (try, IOException)
+import Data.Conduit.Network
+import Network.Socket (sClose)
+import Control.Concurrent (threadDelay)
+import Control.Concurrent.Async (withAsync)
+import Control.Monad.Trans.Resource (register)
+import Data.IORef
+
+data App = App
+
+mkYesod "App" [parseRoutes|
+/ HomeR GET
+|]
+
+instance Yesod App
+
+getHomeR :: Handler ()
+getHomeR = do
+    ref <- liftIO $ newIORef 0
+    _ <- register $ writeIORef ref 1
+    sendRawResponse $ \src sink -> liftIO $ do
+        val <- readIORef ref
+        yield (S8.pack $ show val) $$ sink
+        src $$ CL.map (S8.map toUpper) =$ sink
+
+getFreePort :: IO Int
+getFreePort = do
+    loop 43124
+  where
+    loop port = do
+        esocket <- try $ bindPort port "*"
+        case esocket of
+            Left (_ :: IOException) -> loop (succ port)
+            Right socket -> do
+                sClose socket
+                return port
+
+specs :: Spec
+specs = describe "RawResponse" $ do
+    it "works" $ do
+        port <- getFreePort
+        withAsync (warp port App) $ \_ -> do
+            threadDelay 100000
+            runTCPClient (clientSettings port "127.0.0.1") $ \ad -> do
+                yield "GET / HTTP/1.1\r\n\r\nhello" $$ appSink ad
+                (appSource ad $$ CB.take 6) >>= (`shouldBe` "0HELLO")
+                yield "WORLd" $$ appSink ad
+                (appSource ad $$ await) >>= (`shouldBe` Just "WORLD")

yesod-core/yesod-core.cabal

@@ -1,5 +1,5 @@
 name:            yesod-core
-version:         1.2.6.7
+version:         1.2.7
 license:         MIT
 license-file:    LICENSE
 author:          Michael Snoyman <michael@snoyman.com>
@@ -122,6 +122,9 @@ test-suite tests
                   , containers
                   , lifted-base
                   , resourcet
+                  , network-conduit
+                  , network
+                  , async
     ghc-options:     -Wall
     extensions: TemplateHaskell