From f92cdd5c41010ca0954ae8c0a3e55e9393bd10c4 Mon Sep 17 00:00:00 2001
From: Michael Snoyman <michael@snoyman.com>
Date: Sun, 15 Jun 2014 15:48:57 +0300
Subject: [PATCH] sendWaiApplication

---
 yesod-core/Yesod/Core/Handler.hs             |  7 ++
 yesod-core/Yesod/Core/Internal/Response.hs   | 15 +++--
 yesod-core/Yesod/Core/Internal/Run.hs        |  4 +-
 yesod-core/Yesod/Core/Types.hs               |  3 +
 yesod-core/test/YesodCoreTest/RawResponse.hs | 67 +++++++++++++++++---
 yesod-core/yesod-core.cabal                  |  4 +-
 6 files changed, 80 insertions(+), 20 deletions(-)

diff --git a/yesod-core/Yesod/Core/Handler.hs b/yesod-core/Yesod/Core/Handler.hs
index 2e5d7cb4..aa2f6195 100644
--- a/yesod-core/Yesod/Core/Handler.hs
+++ b/yesod-core/Yesod/Core/Handler.hs
@@ -91,6 +91,7 @@ module Yesod.Core.Handler
     , sendResponseStatus
     , sendResponseCreated
     , sendWaiResponse
+    , sendWaiApplication
 #if MIN_VERSION_wai(2, 1, 0)
     , sendRawResponse
 #endif
@@ -585,6 +586,12 @@ sendResponseCreated url = do
 sendWaiResponse :: MonadHandler m => W.Response -> m b
 sendWaiResponse = handlerError . HCWai
 
+-- | Switch over to handling the current request with a WAI @Application@.
+--
+-- Since 1.2.17
+sendWaiApplication :: MonadHandler m => W.Application -> m b
+sendWaiApplication = handlerError . HCWaiApp
+
 #if MIN_VERSION_wai(3, 0, 0)
 -- | Send a raw response without conduit. This is used for cases such as
 -- WebSockets. Requires WAI 3.0 or later, and a web server which supports raw
diff --git a/yesod-core/Yesod/Core/Internal/Response.hs b/yesod-core/Yesod/Core/Internal/Response.hs
index 7336ef1d..affa1a7b 100644
--- a/yesod-core/Yesod/Core/Internal/Response.hs
+++ b/yesod-core/Yesod/Core/Internal/Response.hs
@@ -42,9 +42,11 @@ yarToResponse :: YesodResponse
               -> YesodRequest
               -> Request
               -> InternalState
-              -> IO Response
-yarToResponse (YRWai a) _ _ _ _ = return a
-yarToResponse (YRPlain s' hs ct c newSess) saveSession yreq req is = do
+              -> (Response -> IO ResponseReceived)
+              -> IO ResponseReceived
+yarToResponse (YRWai a) _ _ _ _ sendResponse = sendResponse a
+yarToResponse (YRWaiApp app) _ _ req _ sendResponse = app req sendResponse
+yarToResponse (YRPlain s' hs ct c newSess) saveSession yreq req is sendResponse = do
     extraHeaders <- do
         let nsToken = maybe
                 newSess
@@ -58,10 +60,10 @@ yarToResponse (YRPlain s' hs ct c newSess) saveSession yreq req is = do
 
     let go (ContentBuilder b mlen) = do
             let hs' = maybe finalHeaders finalHeaders' mlen
-            return $ ResponseBuilder s hs' b
+            sendResponse $ ResponseBuilder s hs' b
         go (ContentFile fp p) = do
-            return $ ResponseFile s finalHeaders fp p
-        go (ContentSource body) = return $ responseStream s finalHeaders
+            sendResponse $ ResponseFile s finalHeaders fp p
+        go (ContentSource body) = sendResponse $ responseStream s finalHeaders
             $ \sendChunk flush -> do
                 transPipe (flip runInternalState is) body
                 $$ CL.mapM_ (\mchunk ->
@@ -85,6 +87,7 @@ yarToResponse :: YesodResponse
 #endif
               -> IO Response
 #if MIN_VERSION_wai(2, 0, 0)
+yarToResponse (YRWaiApp app) _ _ req _ = app req
 yarToResponse (YRWai a) _ _ _ is =
     case a of
         ResponseSource s hs w -> return $ ResponseSource s hs $ \f ->
diff --git a/yesod-core/Yesod/Core/Internal/Run.hs b/yesod-core/Yesod/Core/Internal/Run.hs
index 09b4609e..311f2088 100644
--- a/yesod-core/Yesod/Core/Internal/Run.hs
+++ b/yesod-core/Yesod/Core/Internal/Run.hs
@@ -152,6 +152,7 @@ runHandler rhe@RunHandlerEnv {..} handler yreq = withInternalState $ \resState -
                 emptyContent
                 finalSession
         HCWai r -> return $ YRWai r
+        HCWaiApp a -> return $ YRWaiApp a
 
 safeEh :: (Loc -> LogSource -> LogLevel -> LogStr -> IO ())
        -> ErrorResponse
@@ -295,8 +296,7 @@ yesodRunner handler' YesodRunnerEnv {..} route req
     E.bracket createInternalState closeInternalState $ \is -> do
         yreq' <- yreq
         yar <- runInternalState (runHandler rhe handler yreq') is
-        res <- yarToResponse yar saveSession yreq' req is
-        sendResponse res
+        yarToResponse yar saveSession yreq' req is sendResponse
 
 #else
 
diff --git a/yesod-core/Yesod/Core/Types.hs b/yesod-core/Yesod/Core/Types.hs
index 7e3fd0dd..09d274f8 100644
--- a/yesod-core/Yesod/Core/Types.hs
+++ b/yesod-core/Yesod/Core/Types.hs
@@ -126,6 +126,7 @@ data YesodRequest = YesodRequest
 -- or a higher-level data structure which Yesod will turn into a @Response@.
 data YesodResponse
     = YRWai !W.Response
+    | YRWaiApp !W.Application
     | YRPlain !H.Status ![Header] !ContentType !Content !SessionMap
 
 -- | A tuple containing both the POST parameters and submitted files.
@@ -372,6 +373,7 @@ data HandlerContents =
     | HCRedirect H.Status Text
     | HCCreated Text
     | HCWai W.Response
+    | HCWaiApp W.Application
     deriving Typeable
 
 instance Show HandlerContents where
@@ -381,6 +383,7 @@ instance Show HandlerContents where
     show (HCRedirect s t) = "HCRedirect " ++ show (s, t)
     show (HCCreated t) = "HCCreated " ++ show t
     show (HCWai _) = "HCWai"
+    show (HCWaiApp _) = "HCWaiApp"
 instance Exception HandlerContents
 
 -- Instances for WidgetT
diff --git a/yesod-core/test/YesodCoreTest/RawResponse.hs b/yesod-core/test/YesodCoreTest/RawResponse.hs
index 83bd686a..e4977afb 100644
--- a/yesod-core/test/YesodCoreTest/RawResponse.hs
+++ b/yesod-core/test/YesodCoreTest/RawResponse.hs
@@ -5,6 +5,7 @@ import Yesod.Core
 import Test.Hspec
 import qualified Data.Map as Map
 import Network.Wai.Test
+import Network.Wai (responseStream)
 import Data.Text (Text)
 import Data.ByteString.Lazy (ByteString)
 import qualified Data.Conduit.List as CL
@@ -20,11 +21,15 @@ import Control.Concurrent.Async (withAsync)
 import Control.Monad.Trans.Resource (register)
 import Data.IORef
 import Data.Streaming.Network (bindPortTCP)
+import Network.HTTP.Types (status200)
+import Blaze.ByteString.Builder (fromByteString)
 
 data App = App
 
 mkYesod "App" [parseRoutes|
 / HomeR GET
+/wai-stream WaiStreamR GET
+/wai-app-stream WaiAppStreamR GET
 |]
 
 instance Yesod App
@@ -38,6 +43,20 @@ getHomeR = do
         yield (S8.pack $ show val) $$ sink
         src $$ CL.map (S8.map toUpper) =$ sink
 
+getWaiStreamR :: Handler ()
+getWaiStreamR = sendWaiResponse $ responseStream status200 [] $ \send flush -> do
+    flush
+    send $ fromByteString "hello"
+    flush
+    send $ fromByteString " world"
+
+getWaiAppStreamR :: Handler ()
+getWaiAppStreamR = sendWaiApplication $ \_ f -> f $ responseStream status200 [] $ \send flush -> do
+    flush
+    send $ fromByteString "hello"
+    flush
+    send $ fromByteString " world"
+
 getFreePort :: IO Int
 getFreePort = do
     loop 43124
@@ -51,13 +70,41 @@ getFreePort = do
                 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")
+specs = do
+    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")
+
+    let body req = do
+            port <- getFreePort
+            withAsync (warp port App) $ \_ -> do
+                threadDelay 100000
+                runTCPClient (clientSettings port "127.0.0.1") $ \ad -> do
+                    yield req $$ appSink ad
+                    appSource ad $$ CB.lines =$ do
+                        let loop = do
+                                x <- await
+                                case x of
+                                    Nothing -> return ()
+                                    Just "\r" -> return ()
+                                    _ -> loop
+                        loop
+
+                        Just "0005\r" <- await
+                        Just "hello\r" <- await
+
+                        Just "0006\r" <- await
+                        Just " world\r" <- await
+
+                        return ()
+    it "sendWaiResponse + responseStream" $ do
+        body "GET /wai-stream HTTP/1.1\r\n\r\n"
+    it "sendWaiApplication + responseStream" $ do
+        body "GET /wai-app-stream HTTP/1.1\r\n\r\n"
diff --git a/yesod-core/yesod-core.cabal b/yesod-core/yesod-core.cabal
index 124883b9..eac143da 100644
--- a/yesod-core/yesod-core.cabal
+++ b/yesod-core/yesod-core.cabal
@@ -1,5 +1,5 @@
 name:            yesod-core
-version:         1.2.16.1
+version:         1.2.17
 license:         MIT
 license-file:    LICENSE
 author:          Michael Snoyman <michael@snoyman.com>
@@ -108,7 +108,7 @@ test-suite tests
     build-depends: base
                   ,hspec >= 1.3
                   ,wai-test >= 1.3.0.5
-                  ,wai
+                  ,wai >= 3.0
                   ,yesod-core
                   ,bytestring
                   ,hamlet