Initial WAI 3.0 support, one test still failing

yesod-core/Yesod/Core/Dispatch.hs

@@ -92,8 +92,13 @@ toWaiAppYre yre req =
   where
     site = yreSite yre
     sendRedirect :: Yesod master => master -> [Text] -> W.Application
+#if MIN_VERSION_wai(3, 0, 0)
+    sendRedirect y segments' env sendResponse =
+         sendResponse $ W.responseLBS status301
+#else
     sendRedirect y segments' env =
          return $ W.responseLBS status301
+#endif
                 [ ("Content-Type", "text/plain")
                 , ("Location", Blaze.ByteString.Builder.toByteString dest')
                 ] "Redirecting"

yesod-core/Yesod/Core/Handler.hs

@@ -198,6 +198,8 @@ import Control.Exception (throwIO)
 import Blaze.ByteString.Builder (Builder)
 import Safe (headMay)
 import Data.CaseInsensitive (CI)
+import qualified Data.Conduit.List as CL
+import Control.Monad (unless)
 import           Control.Monad.Trans.Resource  (MonadResource, InternalState, runResourceT, withInternalState, getInternalState, liftResourceT, resourceForkIO
 #if MIN_VERSION_wai(2, 0, 0)
 #else
@@ -368,7 +370,11 @@ handlerToIO =
           where
             oldReq    = handlerRequest oldHandlerData
             oldWaiReq = reqWaiRequest oldReq
+#if MIN_VERSION_wai(3, 0, 0)
+            newWaiReq = oldWaiReq { W.requestBody = return mempty
+#else
             newWaiReq = oldWaiReq { W.requestBody = mempty
+#endif
                                   , W.requestBodyLength = W.KnownLength 0
                                   }
         oldEnv = handlerEnv oldHandlerData
@@ -585,6 +591,19 @@ sendWaiResponse = handlerError . HCWai
 sendRawResponse :: (MonadHandler m, MonadBaseControl IO m)
                 => (Source IO S8.ByteString -> Sink S8.ByteString IO () -> m ())
                 -> m a
+#if MIN_VERSION_wai(3, 0, 0)
+sendRawResponse raw = control $ \runInIO ->
+    runInIO $ sendWaiResponse $ flip W.responseRaw fallback
+    $ \src sink -> runInIO (raw (src' src) (CL.mapM_ sink)) >> return ()
+  where
+    fallback = W.responseLBS H.status500 [("Content-Type", "text/plain")]
+        "sendRawResponse: backend does not support raw responses"
+    src' src = do
+        bs <- liftIO src
+        unless (S.null bs) $ do
+            yield bs
+            src' src
+#else
 sendRawResponse raw = control $ \runInIO ->
     runInIO $ sendWaiResponse $ flip W.responseRaw fallback
     $ \src sink -> runInIO (raw src sink) >> return ()
@@ -592,6 +611,7 @@ sendRawResponse raw = control $ \runInIO ->
     fallback = W.responseLBS H.status500 [("Content-Type", "text/plain")]
         "sendRawResponse: backend does not support raw responses"
 #endif
+#endif
 
 -- | Return a 404 not found page. Also denotes no handler available.
 notFound :: MonadHandler m => m a
@@ -1068,6 +1088,14 @@ provideRepType ct handler =
 rawRequestBody :: MonadHandler m => Source m S.ByteString
 rawRequestBody = do
     req <- lift waiRequest
+#if MIN_VERSION_wai(3, 0, 0)
+    let loop = do
+            bs <- liftIO $ W.requestBody req
+            unless (S.null bs) $ do
+                yield bs
+                loop
+    loop
+#else
     transPipe
 #if MIN_VERSION_wai(2, 0, 0)
         liftIO
@@ -1075,6 +1103,7 @@ rawRequestBody = do
         liftResourceT
 #endif
         (W.requestBody req)
+#endif
 
 -- | Stream the data from the file. Since Yesod 1.2, this has been generalized
 -- to work in any @MonadResource@.

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

@@ -44,11 +44,28 @@ import Control.Monad.Trans.Resource (runResourceT, ResourceT)
 import Control.Exception (throwIO)
 import Yesod.Core.Types
 import qualified Data.Map as Map
+import Data.IORef
 
 -- | Impose a limit on the size of the request body.
-limitRequestBody :: Word64 -> W.Request -> W.Request
+limitRequestBody :: Word64 -> W.Request -> IO W.Request
+#if MIN_VERSION_wai(3, 0, 0)
+limitRequestBody maxLen req = do
+    ref <- newIORef maxLen
+    return req
+        { W.requestBody = do
+            bs <- W.requestBody req
+            remaining <- readIORef ref
+            let len = fromIntegral $ S8.length bs
+                remaining' = remaining - len
+            if remaining < len
+                then throwIO $ HCWai tooLargeResponse
+                else do
+                    writeIORef ref remaining'
+                    return bs
+        }
+#else
 limitRequestBody maxLen req =
-    req { W.requestBody = W.requestBody req $= limit maxLen }
+    return req { W.requestBody = W.requestBody req $= limit maxLen }
   where
     tooLarge = liftIO $ throwIO $ HCWai tooLargeResponse
 
@@ -63,6 +80,7 @@ limitRequestBody maxLen req =
                 else do
                     yield bs
                     limit $ remaining - len
+#endif
 
 tooLargeResponse :: W.Response
 tooLargeResponse = W.responseLBS
@@ -75,7 +93,7 @@ parseWaiRequest :: RandomGen g
                 -> SessionMap
                 -> Bool
                 -> Maybe Word64 -- ^ max body size
-                -> (Either YesodRequest (g -> YesodRequest))
+                -> (Either (IO YesodRequest) (g -> IO YesodRequest))
 parseWaiRequest env session useToken mmaxBodySize =
     -- In most cases, we won't need to generate any random values. Therefore,
     -- we split our results: if we need a random generator, return a Right
@@ -85,17 +103,19 @@ parseWaiRequest env session useToken mmaxBodySize =
         Left token -> Left $ mkRequest token
         Right mkToken -> Right $ mkRequest . mkToken
   where
-    mkRequest token' = YesodRequest
+    mkRequest token' = do
-        { reqGetParams  = gets
+        envLimited <- maybe return limitRequestBody mmaxBodySize env
-        , reqCookies    = cookies
+        return YesodRequest
-        , reqWaiRequest = maybe id limitRequestBody mmaxBodySize env
+            { reqGetParams  = gets
-        , reqLangs      = langs''
+            , reqCookies    = cookies
-        , reqToken      = token'
+            , reqWaiRequest = envLimited
-        , reqSession    = if useToken
+            , reqLangs      = langs''
-                            then Map.delete tokenKey session
+            , reqToken      = token'
-                            else session
+            , reqSession    = if useToken
-        , reqAccept     = httpAccept env
+                                then Map.delete tokenKey session
-        }
+                                else session
+            , reqAccept     = httpAccept env
+            }
     gets = textQueryString env
     reqCookie = lookup "Cookie" $ W.requestHeaders env
     cookies = maybe [] parseCookiesText reqCookie

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

@@ -33,7 +33,49 @@ import qualified Data.ByteString.Lazy         as L
 import qualified Data.Map                     as Map
 import           Yesod.Core.Internal.Request  (tokenKey)
 import           Data.Text.Encoding           (encodeUtf8)
+import           Data.Conduit                 (Flush (..), ($$))
+import qualified Data.Conduit.List            as CL
 
+#if MIN_VERSION_wai(3, 0, 0)
+yarToResponse :: YesodResponse
+              -> (SessionMap -> IO [Header]) -- ^ save session
+              -> YesodRequest
+              -> Request
+              -> InternalState
+              -> IO Response
+yarToResponse (YRWai a) _ _ _ _ = return a
+yarToResponse (YRPlain s' hs ct c newSess) saveSession yreq req is = do
+    extraHeaders <- do
+        let nsToken = maybe
+                newSess
+                (\n -> Map.insert tokenKey (encodeUtf8 n) newSess)
+                (reqToken yreq)
+        sessionHeaders <- saveSession nsToken
+        return $ ("Content-Type", ct) : map headerToPair sessionHeaders
+    let finalHeaders = extraHeaders ++ map headerToPair hs
+        finalHeaders' len = ("Content-Length", S8.pack $ show len)
+                          : finalHeaders
+
+    let go (ContentBuilder b mlen) = do
+            let hs' = maybe finalHeaders finalHeaders' mlen
+            return $ ResponseBuilder s hs' b
+        go (ContentFile fp p) = do
+            return $ ResponseFile s finalHeaders fp p
+        go (ContentSource body) = return $ responseStream s finalHeaders
+            $ \sendChunk flush -> do
+                transPipe (flip runInternalState is) body
+                $$ CL.mapM_ (\mchunk ->
+                    case mchunk of
+                        Flush -> flush
+                        Chunk builder -> sendChunk builder)
+        go (ContentDontEvaluate c') = go c'
+    go c
+  where
+    s
+        | s' == defaultStatus = H.status200
+        | otherwise = s'
+
+#else
 yarToResponse :: YesodResponse
               -> (SessionMap -> IO [Header]) -- ^ save session
               -> YesodRequest
@@ -106,6 +148,7 @@ yarToResponse (YRPlain s' hs ct c newSess) saveSession yreq req
     s
         | s' == defaultStatus = H.status200
         | otherwise = s'
+#endif
 
 -- | Indicates that the user provided no specific status code to be used, and
 -- therefore the default status code should be used. For normal responses, this

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

@@ -230,7 +230,11 @@ runFakeHandler fakeSessionMap logger site handler = liftIO $ do
           , remoteHost     = error "runFakeHandler-remoteHost"
           , pathInfo       = ["runFakeHandler", "pathInfo"]
           , queryString    = []
+#if MIN_VERSION_wai(3, 0, 0)
+          , requestBody    = return mempty
+#else
           , requestBody    = mempty
+#endif
           , vault          = mempty
           , requestBodyLength = KnownLength 0
           }
@@ -253,8 +257,13 @@ yesodRunner :: (ToTypedContent res, Yesod site)
             -> YesodRunnerEnv site
             -> Maybe (Route site)
             -> Application
+#if MIN_VERSION_wai(3, 0, 0)
+yesodRunner handler' YesodRunnerEnv {..} route req sendResponse
+  | Just maxLen <- mmaxLen, KnownLength len <- requestBodyLength req, maxLen < len = sendResponse tooLargeResponse
+#else
 yesodRunner handler' YesodRunnerEnv {..} route req
   | Just maxLen <- mmaxLen, KnownLength len <- requestBodyLength req, maxLen < len = return tooLargeResponse
+#endif
   | otherwise = do
     let dontSaveSession _ = return []
     (session, saveSession) <- liftIO $ do
@@ -281,6 +290,16 @@ yesodRunner handler' YesodRunnerEnv {..} route req
         rhe = rheSafe
             { rheOnError = runHandler rheSafe . errorHandler
             }
+#if MIN_VERSION_wai(3, 0, 0)
+
+    E.bracket createInternalState closeInternalState $ \is -> do
+        yreq' <- yreq
+        yar <- runInternalState (runHandler rhe handler yreq') is
+        res <- yarToResponse yar saveSession yreq' req is
+        sendResponse res
+
+#else
+
 #if MIN_VERSION_wai(2, 0, 0)
     bracketOnError createInternalState closeInternalState $ \is -> do
         yar <- runInternalState (runHandler rhe handler yreq) is
@@ -288,6 +307,7 @@ yesodRunner handler' YesodRunnerEnv {..} route req
 #else
     yar <- runHandler rhe handler yreq
     liftIO $ yarToResponse yar saveSession yreq req
+#endif
 #endif
   where
     mmaxLen = maximumContentLength yreSite route

yesod-core/test/YesodCoreTest/CleanPath.hs

@@ -1,6 +1,7 @@
 {-# LANGUAGE QuasiQuotes, TypeFamilies, TemplateHaskell, MultiParamTypeClasses #-}
 {-# LANGUAGE OverloadedStrings #-}
 {-# LANGUAGE FlexibleInstances #-}
+{-# LANGUAGE CPP #-}
 module YesodCoreTest.CleanPath (cleanPathTest, Widget) where
 
 import Test.Hspec
@@ -32,7 +33,11 @@ instance ParseRoute Subsite where
     parseRoute (x, _) = Just $ SubsiteRoute x
 
 instance YesodSubDispatch Subsite master where
+#if MIN_VERSION_wai(3, 0, 0)
+    yesodSubDispatch _ req f = f $ responseLBS
+#else
     yesodSubDispatch _ req = return $ responseLBS
+#endif
         status200
         [ ("Content-Type", "SUBSITE")
         ] $ L8.pack $ show (pathInfo req)

yesod-core/test/YesodCoreTest/ErrorHandling.hs

@@ -128,7 +128,7 @@ errorHandlingTest = describe "Test.ErrorHandling" $ do
       it "builder includes content-length" caseGoodBuilder
       forM_ [1..10] $ \i -> it ("error case " ++ show i) (caseError i)
 
-runner :: Session () -> IO ()
+runner :: Session a -> IO a
 runner f = toWaiApp App >>= runSession f
 
 caseNotFound :: IO ()
@@ -175,11 +175,10 @@ caseErrorInBody = runner $ do
 caseErrorInBodyNoEval :: IO ()
 caseErrorInBodyNoEval = do
     eres <- try $ runner $ do
-        _ <- request defaultRequest { pathInfo = ["error-in-body-noeval"] }
+        request defaultRequest { pathInfo = ["error-in-body-noeval"] }
-        return ()
     case eres of
         Left (_ :: SomeException) -> return ()
-        Right _ -> error "Expected an exception"
+        Right x -> error $ "Expected an exception, got: " ++ show x
 
 caseOverrideStatus :: IO ()
 caseOverrideStatus = runner $ do

yesod-core/test/YesodCoreTest/InternalRequest.hs

@@ -12,6 +12,7 @@ import Data.Monoid (mempty)
 import Data.Map (singleton)
 import Yesod.Core
 import Data.Word (Word64)
+import System.IO.Unsafe (unsafePerformIO)
 
 randomStringSpecs :: Spec
 randomStringSpecs = describe "Yesod.Internal.Request.randomString" $ do
@@ -36,7 +37,7 @@ parseWaiRequest' :: Request
                  -> Bool
                  -> Word64
                  -> YesodRequest
-parseWaiRequest' a b c d =
+parseWaiRequest' a b c d = unsafePerformIO $ -- ugly hack, just to ease migration, should be removed
     case parseWaiRequest a b c (Just d) of
         Left yreq -> yreq
         Right needGen -> needGen g

yesod-core/test/YesodCoreTest/RequestBodySize.hs

@@ -83,8 +83,8 @@ specs :: Spec
 specs = describe "Test.RequestBodySize" $ do
     caseHelper "lookupPostParam- large" "post" "foobarbaz=bin" 413 413
     caseHelper "lookupPostParam- small" "post" "foo=bin" 200 200
-    caseHelper "consume- large" "consume" "this is longer than 10" 413 413
+    caseHelper "total consume- large" "consume" "this is longer than 10" 413 413
-    caseHelper "consume- small" "consume" "smaller" 200 200
+    caseHelper "total consume- small" "consume" "smaller" 200 200
     caseHelper "partial consume- large" "partial-consume" "this is longer than 10" 200 413
     caseHelper "partial consume- small" "partial-consume" "smaller" 200 200
     caseHelper "unused- large" "unused" "this is longer than 10" 200 413

yesod-core/test/YesodCoreTest/WaiSubsite.hs

@@ -1,4 +1,4 @@
-{-# LANGUAGE QuasiQuotes, TemplateHaskell, TypeFamilies, MultiParamTypeClasses, OverloadedStrings #-}
+{-# LANGUAGE CPP, QuasiQuotes, TemplateHaskell, TypeFamilies, MultiParamTypeClasses, OverloadedStrings #-}
 module YesodCoreTest.WaiSubsite (specs, Widget) where
 
 import YesodCoreTest.YesodTest
@@ -6,7 +6,11 @@ import Yesod.Core
 import qualified Network.HTTP.Types as H
 
 myApp :: Application
+#if MIN_VERSION_wai(3, 0, 0)
+myApp _ f = f $ responseLBS H.status200 [("Content-type", "text/plain")] "WAI"
+#else
 myApp _ = return $ responseLBS H.status200 [("Content-type", "text/plain")] "WAI"
+#endif
 
 getApp :: a -> WaiSubsite
 getApp _ = WaiSubsite myApp

yesod-core/yesod-core.cabal

@@ -131,6 +131,7 @@ test-suite tests
                   , conduit-extra
                   , shakespeare
                   , streaming-commons
+                  , wai-extra
     ghc-options:     -Wall
     extensions: TemplateHaskell