diff --git a/yesod-core/Yesod/Core/Handler.hs b/yesod-core/Yesod/Core/Handler.hs index 7c561c52..6b72a4d5 100644 --- a/yesod-core/Yesod/Core/Handler.hs +++ b/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 diff --git a/yesod-core/Yesod/Core/Internal/Response.hs b/yesod-core/Yesod/Core/Internal/Response.hs index 73fe107d..3f06ac23 100644 --- a/yesod-core/Yesod/Core/Internal/Response.hs +++ b/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 diff --git a/yesod-core/test/YesodCoreTest.hs b/yesod-core/test/YesodCoreTest.hs index 60a1cb2e..e0175991 100644 --- a/yesod-core/test/YesodCoreTest.hs +++ b/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 diff --git a/yesod-core/test/YesodCoreTest/RawResponse.hs b/yesod-core/test/YesodCoreTest/RawResponse.hs new file mode 100644 index 00000000..8b768ca2 --- /dev/null +++ b/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") diff --git a/yesod-core/yesod-core.cabal b/yesod-core/yesod-core.cabal index 557362d8..3bc47dd1 100644 --- a/yesod-core/yesod-core.cabal +++ b/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 @@ -122,6 +122,9 @@ test-suite tests , containers , lifted-base , resourcet + , network-conduit + , network + , async ghc-options: -Wall extensions: TemplateHaskell