From 8fac4917b58afe84beba58e8f00cac5f14413160 Mon Sep 17 00:00:00 2001 From: Michael Snoyman Date: Mon, 2 Jul 2012 11:15:02 +0300 Subject: [PATCH] FileUpload (#358) --- yesod-core/Yesod/Handler.hs | 37 +++++++++++++----- yesod-core/Yesod/Internal/Core.hs | 34 +++++++++++------ yesod-core/Yesod/Internal/Request.hs | 38 +++++++++++++++---- .../test/YesodCoreTest/ErrorHandling.hs | 2 +- .../test/YesodCoreTest/InternalRequest.hs | 18 ++++----- 5 files changed, 92 insertions(+), 37 deletions(-) diff --git a/yesod-core/Yesod/Handler.hs b/yesod-core/Yesod/Handler.hs index 68ec0a7e..d371d16a 100644 --- 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 diff --git a/yesod-core/Yesod/Internal/Core.hs b/yesod-core/Yesod/Internal/Core.hs index 4a1e7764..7ead0d1c 100644 --- 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) diff --git a/yesod-core/Yesod/Internal/Request.hs b/yesod-core/Yesod/Internal/Request.hs index a05d69cb..58238a0a 100644 --- 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) diff --git a/yesod-core/test/YesodCoreTest/ErrorHandling.hs b/yesod-core/test/YesodCoreTest/ErrorHandling.hs index 3ec1df30..8267a133 100644 --- 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 diff --git a/yesod-core/test/YesodCoreTest/InternalRequest.hs b/yesod-core/test/YesodCoreTest/InternalRequest.hs index 2158f262..b2a1e179 100644 --- 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