FileUpload (#358)

yesod-core/Yesod/Handler.hs

@@ -181,6 +181,8 @@ import Control.Monad.Trans.Control
 import Control.Monad.Trans.Resource
 import Control.Monad.Base
 import Yesod.Routes.Class
+import Data.Word (Word64)
+import Data.Conduit (Sink)
 
 class YesodSubRoute s y where
     fromSubRoute :: s -> y -> Route s -> Route y
@@ -193,6 +195,7 @@ data HandlerData sub master = HandlerData
     , handlerRender   :: Route master -> [(Text, Text)] -> Text
     , handlerToMaster :: Route sub -> Route master
     , handlerState    :: I.IORef GHState
+    , handlerUpload   :: Word64 -> FileUpload
     }
 
 handlerSubData :: (Route sub -> Route master)
@@ -322,22 +325,35 @@ hcError = liftIO . throwIO . HCError
 
 runRequestBody :: GHandler s m RequestBodyContents
 runRequestBody = do
+    hd <- ask
+    let getUpload = handlerUpload hd
+        len = reqBodySize $ handlerRequest hd
+        upload = getUpload len
     x <- get
     case ghsRBC x of
         Just rbc -> return rbc
         Nothing -> do
             rr <- waiRequest
-            rbc <- lift $ rbHelper rr
+            rbc <- lift $ rbHelper upload rr
             put x { ghsRBC = Just rbc }
             return rbc
 
-rbHelper :: W.Request -> ResourceT IO RequestBodyContents
+rbHelper :: FileUpload -> W.Request -> ResourceT IO RequestBodyContents
-rbHelper req =
+rbHelper upload =
-    (map fix1 *** map fix2) <$> (NWP.parseRequestBody NWP.lbsSink req) -- FIXME allow control over which backend to use
+    case upload of
+        FileUploadMemory s -> rbHelper' s mkFileInfoLBS
+        FileUploadDisk s -> rbHelper' s mkFileInfoFile
+
+rbHelper' :: Sink S8.ByteString (ResourceT IO) x
+          -> (Text -> Text -> x -> FileInfo)
+          -> W.Request
+          -> ResourceT IO ([(Text, Text)], [(Text, FileInfo)])
+rbHelper' sink mkFI req =
+    (map fix1 *** map fix2) <$> (NWP.parseRequestBody sink req)
   where
     fix1 = go *** go
     fix2 (x, NWP.FileInfo a b c) =
-        (go x, FileInfo (go a) (go b) c)
+        (go x, mkFI (go a) (go b) c)
     go = decodeUtf8With lenientDecode
 
 -- | Get the sub application argument.
@@ -378,8 +394,9 @@ runHandler :: HasReps c
            -> (Route sub -> Route master)
            -> master
            -> sub
+           -> (Word64 -> FileUpload)
            -> YesodApp
-runHandler handler mrender sroute tomr master sub =
+runHandler handler mrender sroute tomr master sub upload =
   YesodApp $ \eh rr cts initSession -> do
     let toErrorHandler e =
             case fromException e of
@@ -400,6 +417,7 @@ runHandler handler mrender sroute tomr master sub =
             , handlerRender = mrender
             , handlerToMaster = tomr
             , handlerState = istate
+            , handlerUpload = upload
             }
     contents' <- catch (fmap Right $ unGHandler handler hd)
         (\e -> return $ Left $ maybe (HCError $ toErrorHandler e) id
@@ -772,6 +790,7 @@ getSession = liftM ghsSession get
 handlerToYAR :: (HasReps a, HasReps b)
              => master -- ^ master site foundation
              -> sub    -- ^ sub site foundation
+             -> (Word64 -> FileUpload)
              -> (Route sub -> Route master)
              -> (Route master -> [(Text, Text)] -> Text) -- route renderer
              -> (ErrorResponse -> GHandler sub master a)
@@ -780,11 +799,11 @@ handlerToYAR :: (HasReps a, HasReps b)
              -> SessionMap
              -> GHandler sub master b
              -> ResourceT IO YesodAppResult
-handlerToYAR y s toMasterRoute render errorHandler rr murl sessionMap h =
+handlerToYAR y s upload toMasterRoute render errorHandler rr murl sessionMap h =
     unYesodApp ya eh' rr types sessionMap
   where
-    ya = runHandler h render murl toMasterRoute y s
+    ya = runHandler h render murl toMasterRoute y s upload
-    eh' er = runHandler (errorHandler' er) render murl toMasterRoute y s
+    eh' er = runHandler (errorHandler' er) render murl toMasterRoute y s upload
     types = httpAccept $ reqWaiRequest rr
     errorHandler' = localNoCurrent . errorHandler
 

yesod-core/Yesod/Internal/Core.hs

@@ -40,6 +40,7 @@ module Yesod.Internal.Core
     , yesodRender
     , resolveApproot
     , Approot (..)
+    , FileUpload (..)
     ) where
 
 import Yesod.Content
@@ -90,6 +91,7 @@ import Data.Aeson (Value (Array, String))
 import Data.Aeson.Encode (encode)
 import qualified Data.Vector as Vector
 import Network.Wai.Middleware.Gzip (GzipSettings, def)
+import Network.Wai.Parse (tempFileSink, lbsSink)
 import qualified Paths_yesod_core
 import Data.Version (showVersion)
 
@@ -325,6 +327,16 @@ $doctype 5
         key <- CS.getKey CS.defaultKeyFile
         return $ Just $ clientSessionBackend key 120
 
+    -- | How to store uploaded files.
+    --
+    -- Default: Whe nthe request body is greater than 50kb, store in a temp
+    -- file. Otherwise, store in memory.
+    fileUpload :: a
+               -> Word64 -- ^ request body size
+               -> FileUpload
+    fileUpload _ size
+        | size > 50000 = FileUploadDisk tempFileSink
+        | otherwise = FileUploadMemory lbsSink
 
 messageLoggerHandler :: Yesod m
                      => Loc -> LogLevel -> Text -> GHandler s m ()
@@ -376,24 +388,18 @@ defaultYesodRunner :: Yesod master
                    -> (Route sub -> Route master)
                    -> Maybe (SessionBackend master)
                    -> W.Application
-defaultYesodRunner _ master _ murl toMaster _ req
+defaultYesodRunner handler master sub murl toMasterRoute msb req
-    | maximumContentLength master (fmap toMaster murl) < len =
+  | maximumContentLength master (fmap toMasterRoute murl) < len =
         return $ W.responseLBS
             (H.Status 413 "Too Large")
             [("Content-Type", "text/plain")]
             "Request body too large to be processed."
-  where
+  | otherwise = do
-    len = fromMaybe 0 $ lookup "content-length" (W.requestHeaders req) >>= readMay
-    readMay s =
-        case reads $ S8.unpack s of
-            [] -> Nothing
-            (x, _):_ -> Just x
-defaultYesodRunner handler master sub murl toMasterRoute msb req = do
     now <- liftIO getCurrentTime
     let dontSaveSession _ _ = return []
     (session, saveSession) <- liftIO $
         maybe (return ([], dontSaveSession)) (\sb -> sbLoadSession sb master req now) msb
-    rr <- liftIO $ parseWaiRequest req session (isJust msb)
+    rr <- liftIO $ parseWaiRequest req session (isJust msb) len
     let h = {-# SCC "h" #-} do
           case murl of
             Nothing -> handler
@@ -413,7 +419,7 @@ defaultYesodRunner handler master sub murl toMasterRoute msb req = do
                 handler
     let sessionMap = Map.fromList . filter ((/=) tokenKey . fst) $ session
     let ra = resolveApproot master req
-    yar <- handlerToYAR master sub toMasterRoute
+    yar <- handlerToYAR master sub (fileUpload master) toMasterRoute
         (yesodRender master ra) errorHandler rr murl sessionMap h
     extraHeaders <- case yar of
         (YARPlain _ _ ct _ newSess) -> do
@@ -425,6 +431,12 @@ defaultYesodRunner handler master sub murl toMasterRoute msb req = do
             return $ ("Content-Type", ct) : map headerToPair sessionHeaders
         _ -> return []
     return $ yarToResponse yar extraHeaders
+  where
+    len = fromMaybe 0 $ lookup "content-length" (W.requestHeaders req) >>= readMay
+    readMay s =
+        case reads $ S8.unpack s of
+            [] -> Nothing
+            (x, _):_ -> Just x
 
 data AuthResult = Authorized | AuthenticationRequired | Unauthorized Text
     deriving (Eq, Show, Read)

yesod-core/Yesod/Internal/Request.hs

@@ -4,7 +4,14 @@ module Yesod.Internal.Request
     ( parseWaiRequest
     , Request (..)
     , RequestBodyContents
-    , FileInfo (..)
+    , FileInfo
+    , fileName
+    , fileContentType
+    , fileSource
+    , fileMove
+    , mkFileInfoLBS
+    , mkFileInfoFile
+    , FileUpload (..)
     -- The below are exported for testing.
     , randomString
     , parseWaiRequest'
@@ -28,6 +35,10 @@ import qualified Data.Set as Set
 import qualified Data.Text as T
 import Data.Text.Encoding (decodeUtf8With)
 import Data.Text.Encoding.Error (lenientDecode)
+import Data.Conduit
+import Data.Conduit.List (sourceList)
+import Data.Conduit.Binary (sourceFile, sinkFile)
+import Data.Word (Word64)
 
 -- | The parsed request information.
 data Request = Request
@@ -38,23 +49,27 @@ data Request = Request
     , reqLangs :: [Text]
       -- | A random, session-specific token used to prevent CSRF attacks.
     , reqToken :: Maybe Text
+      -- | Size of the request body.
+    , reqBodySize :: Word64
     }
 
 parseWaiRequest :: W.Request
                 -> [(Text, ByteString)] -- ^ session
                 -> Bool
+                -> Word64
                 -> IO Request
-parseWaiRequest env session' useToken =
+parseWaiRequest env session' useToken bodySize =
-    parseWaiRequest' env session' useToken <$> newStdGen
+    parseWaiRequest' env session' useToken bodySize <$> newStdGen
 
 parseWaiRequest' :: RandomGen g
                  => W.Request
                  -> [(Text, ByteString)] -- ^ session
                  -> Bool
+                 -> Word64
                  -> g
                  -> Request
-parseWaiRequest' env session' useToken gen = 
+parseWaiRequest' env session' useToken bodySize gen =
-    Request gets'' cookies' env langs'' token
+    Request gets'' cookies' env langs'' token bodySize
   where
     gets' = queryToQueryText $ W.queryString env
     gets'' = map (second $ fromMaybe "") gets'
@@ -116,6 +131,15 @@ type RequestBodyContents =
 data FileInfo = FileInfo
     { fileName :: Text
     , fileContentType :: Text
-    , fileContent :: L.ByteString
+    , fileSource :: Source (ResourceT IO) ByteString
+    , fileMove :: FilePath -> IO ()
     }
-    deriving (Eq, Show)
+
+mkFileInfoLBS :: Text -> Text -> L.ByteString -> FileInfo
+mkFileInfoLBS name ct lbs = FileInfo name ct (sourceList $ L.toChunks lbs) (\fp -> L.writeFile fp lbs)
+
+mkFileInfoFile :: Text -> Text -> FilePath -> FileInfo
+mkFileInfoFile name ct fp = FileInfo name ct (sourceFile fp) (\dst -> runResourceT $ sourceFile fp $$ sinkFile dst)
+
+data FileUpload = FileUploadMemory (Sink ByteString (ResourceT IO) L.ByteString)
+                | FileUploadDisk (Sink ByteString (ResourceT IO) FilePath)

yesod-core/test/YesodCoreTest/ErrorHandling.hs

@@ -49,7 +49,7 @@ postFirstThingR = do
 
 postAfterRunRequestBodyR = do
    x <- runRequestBody
-   _ <- error $ show x
+   _ <- error $ show $ fst x
    getHomeR
 
 errorHandlingTest :: Spec

yesod-core/test/YesodCoreTest/InternalRequest.hs

@@ -40,19 +40,19 @@ tokenSpecs = describe "Yesod.Internal.Request.parseWaiRequest (reqToken)"
 
 noDisabledToken :: Bool
 noDisabledToken = reqToken r == Nothing where
-  r = parseWaiRequest' defaultRequest [] False g
+  r = parseWaiRequest' defaultRequest [] False 0 g
 
 ignoreDisabledToken :: Bool
 ignoreDisabledToken = reqToken r == Nothing where
-  r = parseWaiRequest' defaultRequest [("_TOKEN", "old")] False g
+  r = parseWaiRequest' defaultRequest [("_TOKEN", "old")] False 0 g
 
 useOldToken :: Bool
 useOldToken = reqToken r == Just "old" where
-  r = parseWaiRequest' defaultRequest [("_TOKEN", "old")] True g
+  r = parseWaiRequest' defaultRequest [("_TOKEN", "old")] True 0 g
 
 generateToken :: Bool
 generateToken = reqToken r /= Nothing where
-  r = parseWaiRequest' defaultRequest [("_TOKEN", "old")] True g
+  r = parseWaiRequest' defaultRequest [("_TOKEN", "old")] True 0 g
 
 
 langSpecs :: Spec
@@ -67,21 +67,21 @@ langSpecs = describe "Yesod.Internal.Request.parseWaiRequest (reqLangs)"
 respectAcceptLangs :: Bool
 respectAcceptLangs = reqLangs r == ["en-US", "es", "en"] where
   r = parseWaiRequest' defaultRequest
-        { requestHeaders = [("Accept-Language", "en-US, es")] } [] False g
+        { requestHeaders = [("Accept-Language", "en-US, es")] } [] False 0 g
 
 respectSessionLang :: Bool
 respectSessionLang = reqLangs r == ["en"] where
-  r = parseWaiRequest' defaultRequest [("_LANG", "en")] False g
+  r = parseWaiRequest' defaultRequest [("_LANG", "en")] False 0 g
 
 respectCookieLang :: Bool
 respectCookieLang = reqLangs r == ["en"] where
   r = parseWaiRequest' defaultRequest
         { requestHeaders = [("Cookie", "_LANG=en")]
-        } [] False g
+        } [] False 0 g
 
 respectQueryLang :: Bool
 respectQueryLang = reqLangs r == ["en-US", "en"] where
-  r = parseWaiRequest' defaultRequest { queryString = [("_LANG", Just "en-US")] } [] False g
+  r = parseWaiRequest' defaultRequest { queryString = [("_LANG", Just "en-US")] } [] False 0 g
 
 prioritizeLangs :: Bool
 prioritizeLangs = reqLangs r == ["en-QUERY", "en-COOKIE", "en-SESSION", "en", "es"] where
@@ -90,7 +90,7 @@ prioritizeLangs = reqLangs r == ["en-QUERY", "en-COOKIE", "en-SESSION", "en", "e
                            , ("Cookie", "_LANG=en-COOKIE")
                            ]
         , queryString = [("_LANG", Just "en-QUERY")]
-        } [("_LANG", "en-SESSION")] False g
+        } [("_LANG", "en-SESSION")] False 0 g
 
 
 internalRequestTest :: Spec