{-|
Description:
  Handler for HTTP range requests

HTTP range requests include a header specifying the range of bytes expected in
the response. For example:

> Range: bytes=0-1023

See this <https://developer.mozilla.org/en-US/docs/Web/HTTP/Range_requests
documentation> for more details on range requests.
-}
{-# LANGUAGE BangPatterns        #-}
{-# LANGUAGE CPP                 #-}
{-# LANGUAGE OverloadedStrings   #-}
{-# LANGUAGE ScopedTypeVariables #-}
module Snap.Util.FileServe.Stream where

------------------------------------------------------------------------------
import Control.Applicative
import Control.Monad
import Control.Monad.IO.Class
import Data.Attoparsec.ByteString.Char8 hiding (char8)
import Data.ByteString.Builder
import Data.ByteString.Char8 (ByteString)
import qualified Data.ByteString.Char8 as S
import qualified Data.ByteString.Lazy.Char8 as L
import Data.Maybe (fromMaybe, isNothing)
import Data.Word (Word64)
import Prelude
import Snap.Core
import Snap.Internal.Parsing (fullyParse, parseNum)
import System.IO.Streams (OutputStream)


------------------------------------------------------------------------------
-- | Serves a file, with support for range requests and explicit mime type
-- specification.
--
-- This function can be used in the presence or absence of @Range@ headers: it
-- can be used to serve partial files and whole files.
--
-- Similar to 'serveFile'.
serveStreamAs :: MonadSnap m
              => ByteString
              -- ^ MIME type
              -> Word64
              -- ^ The size of the file being streamed
              -> (Word64 -> Word64 -> OutputStream Builder -> IO ())
              -- ^ If a partial range is requested, this function is used to send that range.
              -> (OutputStream Builder -> IO ())
              -- ^ If a partial range is not requested, this function is used to send the whole file.
              -> m ()
serveStreamAs :: ByteString
-> Word64
-> (Word64 -> Word64 -> OutputStream Builder -> IO ())
-> (OutputStream Builder -> IO ())
-> m ()
serveStreamAs ByteString
mime Word64
sz Word64 -> Word64 -> OutputStream Builder -> IO ()
stream OutputStream Builder -> IO ()
streamAll = do
    Request
reqOrig <- m Request
forall (m :: * -> *). MonadSnap m => m Request
getRequest

    -- If-Range header must be ignored if there is no Range: header in the
    -- request (RFC 2616 section 14.27)
    let req :: Request
req = if Maybe ByteString -> Bool
forall a. Maybe a -> Bool
isNothing (Maybe ByteString -> Bool) -> Maybe ByteString -> Bool
forall a b. (a -> b) -> a -> b
$ CI ByteString -> Request -> Maybe ByteString
forall a. HasHeaders a => CI ByteString -> a -> Maybe ByteString
getHeader CI ByteString
"range" Request
reqOrig
                then CI ByteString -> Request -> Request
forall a. HasHeaders a => CI ByteString -> a -> a
deleteHeader CI ByteString
"if-range" Request
reqOrig
                else Request
reqOrig

    -- ok, at this point we know the last-modified time and the
    -- content-type. set those.
    (Response -> Response) -> m ()
forall (m :: * -> *). MonadSnap m => (Response -> Response) -> m ()
modifyResponse ((Response -> Response) -> m ()) -> (Response -> Response) -> m ()
forall a b. (a -> b) -> a -> b
$ CI ByteString -> ByteString -> Response -> Response
forall a. HasHeaders a => CI ByteString -> ByteString -> a -> a
setHeader CI ByteString
"Accept-Ranges" ByteString
"bytes"
                   (Response -> Response)
-> (Response -> Response) -> Response -> Response
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> Response -> Response
setContentType ByteString
mime

    -- checkRangeReq checks for a Range: header in the request and sends a
    -- partial response if it matches.
    Bool
wasRange <- Snap Bool -> m Bool
forall (m :: * -> *) a. MonadSnap m => Snap a -> m a
liftSnap (Snap Bool -> m Bool) -> Snap Bool -> m Bool
forall a b. (a -> b) -> a -> b
$ Request
-> (Word64 -> Word64 -> OutputStream Builder -> IO ())
-> Word64
-> Snap Bool
forall (m :: * -> *).
MonadSnap m =>
Request
-> (Word64 -> Word64 -> OutputStream Builder -> IO ())
-> Word64
-> m Bool
checkRangeReq Request
req Word64 -> Word64 -> OutputStream Builder -> IO ()
stream Word64
sz

    -- if we didn't have a range request, we just do normal sendfile
    Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless Bool
wasRange (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ do
      (Response -> Response) -> m ()
forall (m :: * -> *). MonadSnap m => (Response -> Response) -> m ()
modifyResponse ((Response -> Response) -> m ()) -> (Response -> Response) -> m ()
forall a b. (a -> b) -> a -> b
$ Int -> Response -> Response
setResponseCode Int
200
                     (Response -> Response)
-> (Response -> Response) -> Response -> Response
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Word64 -> Response -> Response
setContentLength Word64
sz
      (OutputStream Builder -> IO (OutputStream Builder)) -> m ()
forall (m :: * -> *).
MonadSnap m =>
(OutputStream Builder -> IO (OutputStream Builder)) -> m ()
addToOutput ((OutputStream Builder -> IO (OutputStream Builder)) -> m ())
-> (OutputStream Builder -> IO (OutputStream Builder)) -> m ()
forall a b. (a -> b) -> a -> b
$ \OutputStream Builder
str -> IO () -> IO ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (OutputStream Builder -> IO ()
streamAll OutputStream Builder
str) IO () -> IO (OutputStream Builder) -> IO (OutputStream Builder)
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> OutputStream Builder -> IO (OutputStream Builder)
forall (m :: * -> *) a. Monad m => a -> m a
return OutputStream Builder
str

------------------------------------------------------------------------------
data RangeReq = RangeReq !Word64 !(Maybe Word64)
              | SuffixRangeReq !Word64


------------------------------------------------------------------------------
rangeParser :: Parser RangeReq
rangeParser :: Parser RangeReq
rangeParser = ByteString -> Parser ByteString
string ByteString
"bytes=" Parser ByteString -> Parser RangeReq -> Parser RangeReq
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*>
              (Parser RangeReq
byteRangeSpec Parser RangeReq -> Parser RangeReq -> Parser RangeReq
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> Parser RangeReq
suffixByteRangeSpec) Parser RangeReq -> Parser ByteString () -> Parser RangeReq
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f a
<*
              Parser ByteString ()
forall t. Chunk t => Parser t ()
endOfInput
  where
    byteRangeSpec :: Parser RangeReq
byteRangeSpec = do
        Word64
start <- Int64 -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int64 -> Word64)
-> Parser ByteString Int64 -> Parser ByteString Word64
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Parser ByteString Int64
parseNum
        Parser ByteString Char -> Parser ByteString ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (Parser ByteString Char -> Parser ByteString ())
-> Parser ByteString Char -> Parser ByteString ()
forall a b. (a -> b) -> a -> b
$! Char -> Parser ByteString Char
char Char
'-'
        Maybe Int64
end   <- Maybe Int64
-> Parser ByteString (Maybe Int64)
-> Parser ByteString (Maybe Int64)
forall (f :: * -> *) a. Alternative f => a -> f a -> f a
option Maybe Int64
forall a. Maybe a
Nothing (Parser ByteString (Maybe Int64)
 -> Parser ByteString (Maybe Int64))
-> Parser ByteString (Maybe Int64)
-> Parser ByteString (Maybe Int64)
forall a b. (a -> b) -> a -> b
$ (Int64 -> Maybe Int64)
-> Parser ByteString Int64 -> Parser ByteString (Maybe Int64)
forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
liftM Int64 -> Maybe Int64
forall a. a -> Maybe a
Just Parser ByteString Int64
parseNum

        RangeReq -> Parser RangeReq
forall (m :: * -> *) a. Monad m => a -> m a
return (RangeReq -> Parser RangeReq) -> RangeReq -> Parser RangeReq
forall a b. (a -> b) -> a -> b
$! Word64 -> Maybe Word64 -> RangeReq
RangeReq Word64
start (Int64 -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int64 -> Word64) -> Maybe Int64 -> Maybe Word64
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe Int64
end)

    suffixByteRangeSpec :: Parser RangeReq
suffixByteRangeSpec =
        (Int64 -> RangeReq) -> Parser ByteString Int64 -> Parser RangeReq
forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
liftM (Word64 -> RangeReq
SuffixRangeReq (Word64 -> RangeReq) -> (Int64 -> Word64) -> Int64 -> RangeReq
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int64 -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral) (Parser ByteString Int64 -> Parser RangeReq)
-> Parser ByteString Int64 -> Parser RangeReq
forall a b. (a -> b) -> a -> b
$ Char -> Parser ByteString Char
char Char
'-' Parser ByteString Char
-> Parser ByteString Int64 -> Parser ByteString Int64
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> Parser ByteString Int64
parseNum


------------------------------------------------------------------------------
checkRangeReq :: (MonadSnap m)
              => Request
              -> (Word64 -> Word64 -> OutputStream Builder -> IO ())
              -> Word64
              -> m Bool
checkRangeReq :: Request
-> (Word64 -> Word64 -> OutputStream Builder -> IO ())
-> Word64
-> m Bool
checkRangeReq Request
req Word64 -> Word64 -> OutputStream Builder -> IO ()
stream Word64
sz = do
    -- TODO/FIXME: multiple ranges
    m Bool -> (ByteString -> m Bool) -> Maybe ByteString -> m Bool
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (Bool -> m Bool
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
False)
          (\ByteString
s -> (String -> m Bool)
-> (RangeReq -> m Bool) -> Either String RangeReq -> m Bool
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either (m Bool -> String -> m Bool
forall a b. a -> b -> a
const (m Bool -> String -> m Bool) -> m Bool -> String -> m Bool
forall a b. (a -> b) -> a -> b
$ Bool -> m Bool
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
False)
                        RangeReq -> m Bool
withRange
                        (ByteString -> Parser RangeReq -> Either String RangeReq
forall a. ByteString -> Parser a -> Either String a
fullyParse ByteString
s Parser RangeReq
rangeParser))
          (CI ByteString -> Request -> Maybe ByteString
forall a. HasHeaders a => CI ByteString -> a -> Maybe ByteString
getHeader CI ByteString
"range" Request
req)

  where
    withRange :: RangeReq -> m Bool
withRange (RangeReq Word64
start Maybe Word64
mend) = do
        let end :: Word64
end = Word64 -> Maybe Word64 -> Word64
forall a. a -> Maybe a -> a
fromMaybe (Word64
szWord64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
-Word64
1) Maybe Word64
mend

        if Word64
start Word64 -> Word64 -> Bool
forall a. Ord a => a -> a -> Bool
< Word64
0 Bool -> Bool -> Bool
|| Word64
end Word64 -> Word64 -> Bool
forall a. Ord a => a -> a -> Bool
< Word64
start Bool -> Bool -> Bool
|| Word64
start Word64 -> Word64 -> Bool
forall a. Ord a => a -> a -> Bool
>= Word64
sz Bool -> Bool -> Bool
|| Word64
end Word64 -> Word64 -> Bool
forall a. Ord a => a -> a -> Bool
>= Word64
sz
           then m Bool
send416
           else Word64 -> Word64 -> m Bool
forall (m :: * -> *). MonadSnap m => Word64 -> Word64 -> m Bool
send206 Word64
start Word64
end

    withRange (SuffixRangeReq Word64
nbytes) = do
        let end :: Word64
end   = Word64
szWord64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
-Word64
1
        let start :: Word64
start = Word64
sz Word64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
- Word64
nbytes

        if Word64
start Word64 -> Word64 -> Bool
forall a. Ord a => a -> a -> Bool
< Word64
0 Bool -> Bool -> Bool
|| Word64
end Word64 -> Word64 -> Bool
forall a. Ord a => a -> a -> Bool
< Word64
start Bool -> Bool -> Bool
|| Word64
start Word64 -> Word64 -> Bool
forall a. Ord a => a -> a -> Bool
>= Word64
sz Bool -> Bool -> Bool
|| Word64
end Word64 -> Word64 -> Bool
forall a. Ord a => a -> a -> Bool
>= Word64
sz
           then m Bool
send416
           else Word64 -> Word64 -> m Bool
forall (m :: * -> *). MonadSnap m => Word64 -> Word64 -> m Bool
send206 Word64
start Word64
end

    -- note: start and end INCLUSIVE here
    send206 :: Word64 -> Word64 -> m Bool
send206 Word64
start Word64
end = do
        let !len :: Word64
len = Word64
endWord64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
-Word64
startWord64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
+Word64
1
        let crng :: ByteString
crng = [ByteString] -> ByteString
S.concat ([ByteString] -> ByteString)
-> (ByteString -> [ByteString]) -> ByteString -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> [ByteString]
L.toChunks (ByteString -> ByteString) -> ByteString -> ByteString
forall a b. (a -> b) -> a -> b
$
                   Builder -> ByteString
toLazyByteString (Builder -> ByteString) -> Builder -> ByteString
forall a b. (a -> b) -> a -> b
$
                   [Builder] -> Builder
forall a. Monoid a => [a] -> a
mconcat [ ByteString -> Builder
byteString ByteString
"bytes "
                           , Word64 -> Builder
forall a. Show a => a -> Builder
fromShow Word64
start
                           , Char -> Builder
char8 Char
'-'
                           , Word64 -> Builder
forall a. Show a => a -> Builder
fromShow Word64
end
                           , Char -> Builder
char8 Char
'/'
                           , Word64 -> Builder
forall a. Show a => a -> Builder
fromShow Word64
sz ]

        (Response -> Response) -> m ()
forall (m :: * -> *). MonadSnap m => (Response -> Response) -> m ()
modifyResponse ((Response -> Response) -> m ()) -> (Response -> Response) -> m ()
forall a b. (a -> b) -> a -> b
$ Int -> Response -> Response
setResponseCode Int
206
                       (Response -> Response)
-> (Response -> Response) -> Response -> Response
forall b c a. (b -> c) -> (a -> b) -> a -> c
. CI ByteString -> ByteString -> Response -> Response
forall a. HasHeaders a => CI ByteString -> ByteString -> a -> a
setHeader CI ByteString
"Content-Range" ByteString
crng
                       (Response -> Response)
-> (Response -> Response) -> Response -> Response
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Word64 -> Response -> Response
setContentLength Word64
len

        -- end here was inclusive, sendFilePartial is exclusive
        (OutputStream Builder -> IO (OutputStream Builder)) -> m ()
forall (m :: * -> *).
MonadSnap m =>
(OutputStream Builder -> IO (OutputStream Builder)) -> m ()
addToOutput ((OutputStream Builder -> IO (OutputStream Builder)) -> m ())
-> (OutputStream Builder -> IO (OutputStream Builder)) -> m ()
forall a b. (a -> b) -> a -> b
$ \OutputStream Builder
str -> IO () -> IO ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (Word64 -> Word64 -> OutputStream Builder -> IO ()
stream Word64
start (Word64
endWord64 -> Word64 -> Word64
forall a. Num a => a -> a -> a
+Word64
1) OutputStream Builder
str) IO () -> IO (OutputStream Builder) -> IO (OutputStream Builder)
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> OutputStream Builder -> IO (OutputStream Builder)
forall (m :: * -> *) a. Monad m => a -> m a
return OutputStream Builder
str
        Bool -> m Bool
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
True


    send416 :: m Bool
send416 = do
        -- if there's an "If-Range" header in the request, then we just send
        -- back 200
        if CI ByteString -> Request -> Maybe ByteString
forall a. HasHeaders a => CI ByteString -> a -> Maybe ByteString
getHeader CI ByteString
"If-Range" Request
req Maybe ByteString -> Maybe ByteString -> Bool
forall a. Eq a => a -> a -> Bool
/= Maybe ByteString
forall a. Maybe a
Nothing
           then Bool -> m Bool
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
False
           else do
               let crng :: ByteString
crng = [ByteString] -> ByteString
S.concat ([ByteString] -> ByteString)
-> (ByteString -> [ByteString]) -> ByteString -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> [ByteString]
L.toChunks (ByteString -> ByteString) -> ByteString -> ByteString
forall a b. (a -> b) -> a -> b
$
                          Builder -> ByteString
toLazyByteString (Builder -> ByteString) -> Builder -> ByteString
forall a b. (a -> b) -> a -> b
$
                          [Builder] -> Builder
forall a. Monoid a => [a] -> a
mconcat [ ByteString -> Builder
byteString ByteString
"bytes */"
                                  , Word64 -> Builder
forall a. Show a => a -> Builder
fromShow Word64
sz ]

               (Response -> Response) -> m ()
forall (m :: * -> *). MonadSnap m => (Response -> Response) -> m ()
modifyResponse ((Response -> Response) -> m ()) -> (Response -> Response) -> m ()
forall a b. (a -> b) -> a -> b
$ Int -> Response -> Response
setResponseCode Int
416
                              (Response -> Response)
-> (Response -> Response) -> Response -> Response
forall b c a. (b -> c) -> (a -> b) -> a -> c
. CI ByteString -> ByteString -> Response -> Response
forall a. HasHeaders a => CI ByteString -> ByteString -> a -> a
setHeader CI ByteString
"Content-Range" ByteString
crng
                              (Response -> Response)
-> (Response -> Response) -> Response -> Response
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Word64 -> Response -> Response
setContentLength Word64
0
                              (Response -> Response)
-> (Response -> Response) -> Response -> Response
forall b c a. (b -> c) -> (a -> b) -> a -> c
. CI ByteString -> Response -> Response
forall a. HasHeaders a => CI ByteString -> a -> a
deleteHeader CI ByteString
"Content-Type"
                              (Response -> Response)
-> (Response -> Response) -> Response -> Response
forall b c a. (b -> c) -> (a -> b) -> a -> c
. CI ByteString -> Response -> Response
forall a. HasHeaders a => CI ByteString -> a -> a
deleteHeader CI ByteString
"Content-Encoding"
                              (Response -> Response)
-> (Response -> Response) -> Response -> Response
forall b c a. (b -> c) -> (a -> b) -> a -> c
. CI ByteString -> Response -> Response
forall a. HasHeaders a => CI ByteString -> a -> a
deleteHeader CI ByteString
"Transfer-Encoding"
                              (Response -> Response)
-> (Response -> Response) -> Response -> Response
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (OutputStream Builder -> IO (OutputStream Builder))
-> Response -> Response
setResponseBody (OutputStream Builder -> IO (OutputStream Builder)
forall (m :: * -> *) a. Monad m => a -> m a
return (OutputStream Builder -> IO (OutputStream Builder))
-> (OutputStream Builder -> OutputStream Builder)
-> OutputStream Builder
-> IO (OutputStream Builder)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. OutputStream Builder -> OutputStream Builder
forall a. a -> a
id)

               Bool -> m Bool
forall (m :: * -> *) a. Monad m => a -> m a
return Bool
True

------------------------------------------------------------------------------
fromShow :: Show a => a -> Builder
fromShow :: a -> Builder
fromShow = String -> Builder
stringUtf8 (String -> Builder) -> (a -> String) -> a -> Builder
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> String
forall a. Show a => a -> String
show