FileUpload (#358)

2012-07-02 11:15:02 +03:00 · 2012-07-02 11:15:02 +03:00 · 8fac4917b5
commit 8fac4917b5
parent 21a4360f74
5 changed files with 92 additions and 37 deletions
--- a/yesod-core/Yesod/Handler.hs
+++ b/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 req =
-    (map fix1 *** map fix2) <$> (NWP.parseRequestBody NWP.lbsSink req) -- FIXME allow control over which backend to use
+rbHelper :: FileUpload -> W.Request -> ResourceT IO RequestBodyContents
+rbHelper upload =
+    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
-    eh' er = runHandler (errorHandler' er) render murl toMasterRoute y s
+    ya = runHandler h render murl toMasterRoute y s upload
+    eh' er = runHandler (errorHandler' er) render murl toMasterRoute y s upload
    types = httpAccept $ reqWaiRequest rr
    errorHandler' = localNoCurrent . errorHandler

--- a/yesod-core/Yesod/Internal/Core.hs
+++ b/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
-    | maximumContentLength master (fmap toMaster murl) < len =
+defaultYesodRunner handler master sub murl toMasterRoute msb req
+  | 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
-    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
+  | otherwise = 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)
--- a/yesod-core/Yesod/Internal/Request.hs
+++ b/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 <$> newStdGen
+parseWaiRequest env session' useToken bodySize =
+    parseWaiRequest' env session' useToken bodySize <$> newStdGen

 parseWaiRequest' :: RandomGen g
                 => W.Request
                 -> [(Text, ByteString)] -- ^ session
                 -> Bool
+                 -> Word64
                 -> g
                 -> Request
-parseWaiRequest' env session' useToken gen = 
-    Request gets'' cookies' env langs'' token
+parseWaiRequest' env session' useToken bodySize gen =
+    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)
--- a/yesod-core/test/YesodCoreTest/ErrorHandling.hs
+++ b/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
--- a/yesod-core/test/YesodCoreTest/InternalRequest.hs
+++ b/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