MultipartFormData.hs

{-# LANGUAGE CPP, OverloadedStrings #-}
-- | This module handles building multipart/form-data. Example usage:
--
-- > {-# LANGUAGE OverloadedStrings #-}
-- > import Network
-- > import Network.HTTP.Conduit
-- > import Network.HTTP.Conduit.MultipartFormData
-- >
-- > import Data.Text.Encoding as TE
-- >
-- > import Control.Monad
-- >
-- > main = withSocketsDo $ withManager $ \m -> do
-- >     Response{responseBody=cat} <- flip httpLbs m $ fromJust $ parseUrl "http://random-cat-photo.net/cat.jpg"
-- >     flip httpLbs m =<<
-- >         (formDataBody [partBS "title" "Bleaurgh"
-- >                       ,partBS "text" $ TE.encodeUtf8 "矢田矢田矢田矢田矢田"
-- >                       ,partFileSource "file1" "/home/friedrich/Photos/MyLittlePony.jpg"
-- >                       ,partFileRequestBody "file2" "cat.jpg" $ RequestBodyLBS cat]
-- >             $ fromJust $ parseUrl "http://example.org/~friedrich/blog/addPost.hs")
module Network.HTTP.Conduit.MultipartFormData
    (Part(..)
    ,partBS
    ,partLBS
    ,partFile
    ,partFileSource
    ,partFileSourceChunked
    ,partFileRequestBody
    ,partFileRequestBodyM
    ,formDataBody
    ,formDataBodyPure
    ,formDataBodyWithBoundary
    ,webkitBoundary
    ) where

import Network.HTTP.Conduit
import Network.Mime
import Network.HTTP.Types (hContentType, methodPost)

import Blaze.ByteString.Builder
import qualified Data.Conduit.List as CL
import qualified Data.Conduit.Binary as CB
import Data.Conduit

import Data.Text
import qualified Data.Text.Encoding as TE
#if MIN_VERSION_bytestring(0,10,0)
import Data.ByteString.Lazy (fromStrict)
#endif
import qualified Data.ByteString.Lazy as BL
import qualified Data.ByteString as BS

import Control.Monad.Trans.Resource
import Control.Monad.IO.Class
import System.FilePath
import System.Random
import Data.Array.Base
import System.IO
import Data.Bits
import Data.Int
import Data.Word
import Data.Functor.Identity
import Data.Monoid
import Control.Monad
import Control.Applicative

#if !MIN_VERSION_bytestring(0,10,0)
{-# INLINE fromStrict #-}
fromStrict :: BS.ByteString -> BL.ByteString
fromStrict x = BL.fromChunks [x]
#endif

#if !MIN_VERSION_base(4,5,0)
{-# INLINE (<>) #-}
infixr 5 <>
(<>) :: Monoid m => m -> m -> m
(<>) = mappend
#endif

{-# INLINE sourceSingle #-}
sourceSingle :: Monad m => a -> Pipe l i a u m ()
sourceSingle = CL.sourceList . return

instance Monad m => Monoid (RequestBody m) where
    mempty = RequestBodyLBS mempty
    mappend (RequestBodySourceChunked a) b =
        RequestBodySourceChunked (a <> toChunked b)
    mappend a (RequestBodySourceChunked b) =
        RequestBodySourceChunked (toChunked a <> b)
    
    mappend (RequestBodySource l1 a) b =
        let (l2, b') = toSource b in RequestBodySource (l1 + l2) (a <> b')
    mappend a (RequestBodySource l2 b) =
        let (l1, a') = toSource a in RequestBodySource (l1 + l2) (a' <> b)
    
    mappend (RequestBodyBuilder l1 a) b =
        let (l2, b') = toBuilder b in RequestBodyBuilder (l1 + l2) (a <> b')
    mappend a (RequestBodyBuilder l2 b) =
        let (l1, a') = toBuilder a in RequestBodyBuilder (l1 + l2) (a' <> b)
    
    mappend (RequestBodyLBS a) b = RequestBodyLBS (a <> toLBS b)
    mappend a (RequestBodyLBS b) = RequestBodyLBS (toLBS a <> b)
    
    mappend (RequestBodyBS a) (RequestBodyBS b) = RequestBodyLBS (BL.fromChunks [a,b])

toChunked :: Monad m => RequestBody m -> Source m Builder
toChunked (RequestBodyBS a) = sourceSingle $ fromByteString a
toChunked (RequestBodyLBS a) = sourceSingle $ fromLazyByteString a
toChunked (RequestBodyBuilder _ a) = sourceSingle a
toChunked (RequestBodySource _ a) = a
toChunked (RequestBodySourceChunked a) = a

toSource :: Monad m => RequestBody m -> (Int64, Source m Builder)
toSource (RequestBodyBS a) = (fromIntegral $ BS.length a, sourceSingle $ fromByteString a)
toSource (RequestBodyLBS a) = (BL.length a, sourceSingle $ fromLazyByteString a)
toSource (RequestBodyBuilder l a) = (l, sourceSingle a)
toSource (RequestBodySource l a) = (l, a)

toBuilder :: RequestBody m -> (Int64, Builder)
toBuilder (RequestBodyBS a) = (fromIntegral $ BS.length a, fromByteString a)
toBuilder (RequestBodyLBS a) = (BL.length a, fromLazyByteString a)
toBuilder (RequestBodyBuilder l a) = (l, a)

toLBS :: RequestBody m -> BL.ByteString
toLBS (RequestBodyBS a) = fromStrict a
toLBS (RequestBodyLBS a) = a

-- | A single part of a multipart message.
data Part m m' = Part
    { partName :: Text -- ^ Name of the corresponding \<input\>
    , partFilename :: Maybe String -- ^ A file name, if this is an attached file
    , partContentType :: Maybe MimeType -- ^ Content type
    , partGetBody :: m (RequestBody m') -- ^ Action in m which returns the body
                                        -- of a message.
    }

partBS :: (Monad m, Monad m') => Text -> BS.ByteString -> Part m m'
partBS n b = Part n mempty mempty $ return $ RequestBodyBS b

partLBS :: (Monad m, Monad m') => Text -> BL.ByteString -> Part m m'
partLBS n b = Part n mempty mempty $ return $ RequestBodyLBS b

-- | Make a 'Part' from a file, the entire file will reside in memory at once.
-- If you want constant memory usage use 'partFileSource'
partFile :: (Functor m, MonadIO m, Monad m') => Text -> FilePath -> Part m m'
partFile n f =
    partFileRequestBodyM n f $ do
        fmap RequestBodyBS $ liftIO $ BS.readFile f

-- | Stream 'Part' from a file.
partFileSource :: (Functor m, MonadIO m, MonadResource m') => Text -> FilePath -> Part m m'
partFileSource n f =
    partFileRequestBodyM n f $ do
        size <- fmap fromInteger $ liftIO $ withBinaryFile f ReadMode hFileSize
        return $ RequestBodySource size $ CB.sourceFile f $= CL.map fromByteString

-- | 'partFileSourceChunked' will read a file and send it in chunks.
--
-- Note that not all servers support this. Only use 'partFileSourceChunked'
-- if you know the server you're sending to supports chunked request bodies. 
partFileSourceChunked :: (Monad m, MonadResource m') => Text -> FilePath -> Part m m'
partFileSourceChunked n f =
    partFileRequestBody n f $ do
        RequestBodySourceChunked $ CB.sourceFile f $= CL.map fromByteString

-- | Construct a 'Part' from form name, filepath and a 'RequestBody'
--
-- > partFileRequestBody "who_calls" "caller.json" $ RequestBodyBS "{\"caller\":\"Jason J Jason\"}"
partFileRequestBody :: (Monad m, Monad m') => Text -> FilePath -> RequestBody m' -> Part m m'
partFileRequestBody n f rqb =
    partFileRequestBodyM n f $ return rqb

-- | Construct a 'Part' from action returning the 'RequestBody'
--
-- > partFileRequestBodyM "cat_photo" "haskell-the-cat.jpg" $ do
-- >     size <- fromInteger <$> withBinaryFile "haskell-the-cat.jpg" ReadMode hFileSize
-- >     return $ RequestBodySource size $ CB.sourceFile "haskell-the-cat.jpg" $= CL.map fromByteString
partFileRequestBodyM :: Monad m' => Text -> FilePath -> m (RequestBody m') -> Part m m'
partFileRequestBodyM n f rqb =
    Part n (Just f) (Just $ defaultMimeLookup $ pack f) rqb

{-# INLINABLE cp #-}
cp :: BS.ByteString -> RequestBody m
cp bs = RequestBodyBuilder (fromIntegral $ BS.length bs) $ copyByteString bs

renderPart :: (Functor m, Monad m') => BS.ByteString -> Part m m' -> m (RequestBody m')
renderPart boundary (Part name mfilename mcontenttype get) = fmap render get
  where render renderBody =
            cp "--" <> cp boundary <> cp "\r\n"
         <> cp "Content-Disposition: form-data; name=\""
         <> RequestBodyBS (TE.encodeUtf8 name)
         <> (case mfilename of
                 Just f -> cp "\"; filename=\""
                        <> RequestBodyBS (TE.encodeUtf8 $ pack $ takeFileName f)
                 _ -> mempty)
         <> cp "\""
         <> (case mcontenttype of
                Just ct -> cp "\r\n"
                        <> cp "Content-Type: "
                        <> cp ct
                _ -> mempty)
         <> cp "\r\n\r\n"
         <> renderBody <> cp "\r\n"

renderParts :: (Functor m, Monad m, Monad m') => BS.ByteString -> [Part m m'] -> m (RequestBody m')
renderParts boundary parts = fin . mconcat <$> mapM (renderPart boundary) parts
  where fin = (<> cp "--" <> cp boundary <> cp "--\r\n")

-- | Generate a boundary simillar to those generated by WebKit-based browsers.
webkitBoundary :: IO BS.ByteString
webkitBoundary = do
    fmap (BS.append prefix . BS.pack . Prelude.concat) $ replicateM 4 $ do
        randomness <- randomIO :: IO Int
        return [unsafeAt alphaNumericEncodingMap $ randomness `shiftR` 24 .&. 0x3F
               ,unsafeAt alphaNumericEncodingMap $ randomness `shiftR` 16 .&. 0x3F
               ,unsafeAt alphaNumericEncodingMap $ randomness `shiftR` 8 .&. 0x3F
               ,unsafeAt alphaNumericEncodingMap $ randomness .&. 0x3F]
  where
    prefix = "----WebKitFormBoundary"
    alphaNumericEncodingMap :: UArray Int Word8
    alphaNumericEncodingMap = listArray (0, 63)
        [0x41, 0x42, 0x43, 0x44, 0x45, 0x46, 0x47, 0x48,
         0x49, 0x4A, 0x4B, 0x4C, 0x4D, 0x4E, 0x4F, 0x50,
         0x51, 0x52, 0x53, 0x54, 0x55, 0x56, 0x57, 0x58,
         0x59, 0x5A, 0x61, 0x62, 0x63, 0x64, 0x65, 0x66,
         0x67, 0x68, 0x69, 0x6A, 0x6B, 0x6C, 0x6D, 0x6E,
         0x6F, 0x70, 0x71, 0x72, 0x73, 0x74, 0x75, 0x76,
         0x77, 0x78, 0x79, 0x7A, 0x30, 0x31, 0x32, 0x33,
         0x34, 0x35, 0x36, 0x37, 0x38, 0x39, 0x41, 0x42]

-- | Add form data to the 'Request'.
--
-- This sets a new 'requestBody', adds a content-type request header and changes the method to POST.
formDataBody :: (Functor m, MonadIO m, Monad m') => [Part m m'] -> Request m' -> m (Request m')
formDataBody a b = do
    boundary <- liftIO webkitBoundary
    formDataBodyWithBoundary boundary a b

{-# INLINE formDataBodyPure #-}
-- | Add form data to request without doing any IO. Your form data should only
-- contain pure parts ('partBS', 'partLBS', 'partFileRequestBody'). You'll have
-- to supply your own boundary (for example one generated by 'webkitBoundary')
formDataBodyPure :: Monad m => BS.ByteString -> [Part Identity m] -> Request m -> Request m
formDataBodyPure = \boundary parts req ->
    runIdentity $ formDataBodyWithBoundary boundary parts req

-- | Add form data with supplied boundary
formDataBodyWithBoundary :: (Functor m, Monad m, Monad m') => BS.ByteString -> [Part m m'] -> Request m' -> m (Request m')
formDataBodyWithBoundary boundary parts req = do
    body <- renderParts boundary parts
    return $ req
        { method = methodPost
        , requestHeaders =
            (hContentType, "multipart/form-data; boundary=" <> boundary)
          : Prelude.filter (\(x, _) -> x /= hContentType) (requestHeaders req)
        , requestBody = body
        }

instance Show (RequestBody m) where
    showsPrec d (RequestBodyBS a) =
        showParen (d>=11) $ showString "RequestBodyBS " . showsPrec 11 a
    showsPrec d (RequestBodyLBS a) =
        showParen (d>=11) $ showString "RequestBodyLBS " . showsPrec 11 a
    showsPrec d (RequestBodyBuilder l _) =
        showParen (d>=11) $ showString "RequestBodyBuilder " . showsPrec 11 l .
            showString " " . showString "<Builder>"
    showsPrec d (RequestBodySource l _) =
        showParen (d>=11) $ showString "RequestBodySource " . showsPrec 11 l .
            showString " <Source m Builder>"
    showsPrec d (RequestBodySourceChunked _) =
        showParen (d>=11) $ showString "RequestBodySource <Source m Builder>"

instance Show (Part m m') where
    showsPrec d (Part n f c _) =
        showParen (d>=11) $ showString "Part "
            . showsPrec 11 n
            . showString " "
            . showsPrec 11 f
            . showString " "
            . showsPrec 11 c
            . showString " "
            . showString "<m (RequestBody m)>"