{-# LANGUAGE BangPatterns              #-}
{-# LANGUAGE CPP                       #-}
{-# LANGUAGE DeriveDataTypeable        #-}
{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE OverloadedStrings         #-}
{-# LANGUAGE ScopedTypeVariables       #-}

--------------------------------------------------------------------------------
-- | Helpers for running a 'Snap' web handler with compression.

module Snap.Util.GZip
  ( withCompression
  , withCompression'
  , noCompression
  , BadAcceptEncodingException
  , compressibleMimeTypes
  ) where

import           Control.Applicative              (Alternative ((<|>), many), Applicative ((*>), (<*), pure), (<$>))
import           Control.Exception                (Exception, throwIO)
import           Control.Monad                    (Functor (fmap), Monad ((>>), (>>=), return), MonadPlus (mplus), void, when)
import           Control.Monad.IO.Class           (MonadIO (liftIO))
import           Data.Attoparsec.ByteString.Char8 (Parser, char, endOfInput, isAlpha_ascii, isDigit, skipSpace, string, takeWhile, takeWhile1)
import           Data.ByteString.Builder          (Builder)
import           Data.ByteString.Char8            (ByteString)
import qualified Data.ByteString.Char8            as S (takeWhile)
import qualified Data.Char                        as Char (isSpace)
import           Data.Maybe                       (Maybe (Just, Nothing), fromMaybe, isJust, maybe)
import           Data.Set                         (Set)
import qualified Data.Set                         as Set (fromList, member)
import           Data.Typeable                    (Typeable)
import           Prelude                          (Either (..), Eq (..), IO, Show (show), id, not, ($), ($!), (&&), (++), (||))
import           Snap.Core                        (MonadSnap, clearContentLength, finishWith, getHeader, getRequest, getResponse, modifyResponse, modifyResponseBody, setHeader)
import           Snap.Internal.Debug              (debug)
import           Snap.Internal.Parsing            (fullyParse)
import           System.IO.Streams                (OutputStream)
import qualified System.IO.Streams                as Streams (compressBuilder, gzipBuilder)

------------------------------------------------------------------------------
-- | Runs a 'Snap' web handler with compression if available.
--
-- If the client has indicated support for @gzip@ or @deflate@ in its
-- @Accept-Encoding@ header, and the @Content-Type@ in the response is one of
-- the following types:
--
--   * @application/x-javascript@
--
--   * @application/json@
--
--   * @text/css@
--
--   * @text/html@
--
--   * @text/javascript@
--
--   * @text/plain@
--
--   * @text/xml@
--
--   * @application/x-font-truetype@
--
-- Then the given handler's output stream will be compressed,
-- @Content-Encoding@ will be set in the output headers, and the
-- @Content-Length@ will be cleared if it was set. (We can't process the
-- stream in O(1) space if the length is known beforehand.)
--
-- The wrapped handler will be run to completion, and then the 'Response'
-- that's contained within the 'Snap' monad state will be passed to
-- 'finishWith' to prevent further processing.
--
--
-- Example:
--
-- @
-- ghci> :set -XOverloadedStrings
-- ghci> import qualified "Data.Map" as M
-- ghci> import qualified "Snap.Test" as T
-- ghci> let r = T.get \"/\" M.empty >> T.addHeader \"Accept-Encoding\" \"gzip,deflate\"
-- ghci> let h = 'Snap.Core.modifyResponse' ('Snap.Core.setContentType' \"text\/plain\") >> 'Snap.Core.writeBS' \"some text\"
-- ghci> T.runHandler r h
-- HTTP\/1.1 200 OK
-- content-type: text\/plain
-- server: Snap\/test
-- date: Fri, 08 Aug 2014 15:40:45 GMT
--
-- some text
-- ghci> T.runHandler r ('withCompression' h)
-- HTTP\/1.1 200 OK
-- content-type: text\/plain
-- vary: Accept-Encoding
-- content-encoding: gzip
-- server: Snap\/test
-- date: Fri, 08 Aug 2014 15:40:10 GMT
--
--
-- @
withCompression :: MonadSnap m
                => m a   -- ^ the web handler to run
                -> m ()
withCompression :: forall (m :: * -> *) a. MonadSnap m => m a -> m ()
withCompression = forall (m :: * -> *) a.
MonadSnap m =>
Set ByteString -> m a -> m ()
withCompression' Set ByteString
compressibleMimeTypes


------------------------------------------------------------------------------
-- | The same as 'withCompression', with control over which MIME types to
-- compress.
withCompression' :: MonadSnap m
                 => Set ByteString
                    -- ^ set of compressible MIME types
                 -> m a
                    -- ^ the web handler to run
                 -> m ()
withCompression' :: forall (m :: * -> *) a.
MonadSnap m =>
Set ByteString -> m a -> m ()
withCompression' Set ByteString
mimeTable m a
action = do
    a
_    <- m a
action
    Response
resp <- forall (m :: * -> *). MonadSnap m => m Response
getResponse

    -- If a content-encoding is already set, do nothing. This prevents
    -- "withCompression $ withCompression m" from ruining your day.
    forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Bool -> Bool
not forall a b. (a -> b) -> a -> b
$ forall a. Maybe a -> Bool
isJust forall a b. (a -> b) -> a -> b
$ forall a. HasHeaders a => CI ByteString -> a -> Maybe ByteString
getHeader CI ByteString
"Content-Encoding" Response
resp) forall a b. (a -> b) -> a -> b
$ do
       let mbCt :: Maybe ByteString
mbCt = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ByteString -> ByteString
chop forall a b. (a -> b) -> a -> b
$ forall a. HasHeaders a => CI ByteString -> a -> Maybe ByteString
getHeader CI ByteString
"Content-Type" Response
resp

       forall (m :: * -> *). MonadIO m => [Char] -> m ()
debug forall a b. (a -> b) -> a -> b
$ [Char]
"withCompression', content-type is " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> [Char]
show Maybe ByteString
mbCt

       case Maybe ByteString
mbCt of
         (Just ByteString
ct) -> forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (forall a. Ord a => a -> Set a -> Bool
Set.member ByteString
ct Set ByteString
mimeTable) m ()
chkAcceptEncoding
         Maybe ByteString
_         -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$! ()


    forall (m :: * -> *). MonadSnap m => m Response
getResponse forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= forall (m :: * -> *) a. MonadSnap m => Response -> m a
finishWith

  where
    chop :: ByteString -> ByteString
chop = (Char -> Bool) -> ByteString -> ByteString
S.takeWhile (\Char
c -> Char
c forall a. Eq a => a -> a -> Bool
/= Char
';' Bool -> Bool -> Bool
&& Bool -> Bool
not (Char -> Bool
Char.isSpace Char
c))

    chkAcceptEncoding :: m ()
chkAcceptEncoding = do
        Request
req <- forall (m :: * -> *). MonadSnap m => m Request
getRequest
        forall (m :: * -> *). MonadIO m => [Char] -> m ()
debug forall a b. (a -> b) -> a -> b
$ [Char]
"checking accept-encoding"
        let mbAcc :: Maybe ByteString
mbAcc = forall a. HasHeaders a => CI ByteString -> a -> Maybe ByteString
getHeader CI ByteString
"Accept-Encoding" Request
req
        forall (m :: * -> *). MonadIO m => [Char] -> m ()
debug forall a b. (a -> b) -> a -> b
$ [Char]
"accept-encoding is " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> [Char]
show Maybe ByteString
mbAcc
        let s :: ByteString
s = forall a. a -> Maybe a -> a
fromMaybe ByteString
"" Maybe ByteString
mbAcc

        [ByteString]
types <- forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO forall a b. (a -> b) -> a -> b
$ ByteString -> IO [ByteString]
parseAcceptEncoding ByteString
s

        forall {m :: * -> *} {a}.
(Eq a, IsString a, MonadSnap m) =>
Maybe (m ()) -> [a] -> m ()
chooseType forall a. Maybe a
Nothing [ByteString]
types

    chooseType :: Maybe (m ()) -> [a] -> m ()
chooseType !Maybe (m ())
m []               = forall b a. b -> (a -> b) -> Maybe a -> b
maybe (forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$! ()) forall a. a -> a
id Maybe (m ())
m
    chooseType !Maybe (m ())
_ (a
"gzip":[a]
_)       = forall (m :: * -> *). MonadSnap m => ByteString -> m ()
gzipCompression ByteString
"gzip"
    chooseType !Maybe (m ())
m (a
"deflate":[a]
xs)   =
        Maybe (m ()) -> [a] -> m ()
chooseType (Maybe (m ())
m forall (m :: * -> *) a. MonadPlus m => m a -> m a -> m a
`mplus` forall a. a -> Maybe a
Just (forall (m :: * -> *). MonadSnap m => ByteString -> m ()
compressCompression ByteString
"deflate")) [a]
xs

    chooseType !Maybe (m ())
_ (a
"x-gzip":[a]
_)     = forall (m :: * -> *). MonadSnap m => ByteString -> m ()
gzipCompression ByteString
"x-gzip"
    chooseType !Maybe (m ())
m (a
"x-deflate":[a]
xs) =
        Maybe (m ()) -> [a] -> m ()
chooseType (Maybe (m ())
m forall (m :: * -> *) a. MonadPlus m => m a -> m a -> m a
`mplus` forall a. a -> Maybe a
Just (forall (m :: * -> *). MonadSnap m => ByteString -> m ()
compressCompression ByteString
"x-deflate")) [a]
xs
    chooseType !Maybe (m ())
m (a
_:[a]
xs)           = Maybe (m ()) -> [a] -> m ()
chooseType Maybe (m ())
m [a]
xs


------------------------------------------------------------------------------
-- | Turn off compression by setting \"Content-Encoding: identity\" in the
-- response headers. 'withCompression' is a no-op when a content-encoding is
-- already set.
noCompression :: MonadSnap m => m ()
noCompression :: forall (m :: * -> *). MonadSnap m => m ()
noCompression = forall (m :: * -> *). MonadSnap m => (Response -> Response) -> m ()
modifyResponse forall a b. (a -> b) -> a -> b
$ forall a. HasHeaders a => CI ByteString -> ByteString -> a -> a
setHeader CI ByteString
"Content-Encoding" ByteString
"identity"


------------------------------------------------------------------------------
-- private following
------------------------------------------------------------------------------


------------------------------------------------------------------------------
compressibleMimeTypes :: Set ByteString
compressibleMimeTypes :: Set ByteString
compressibleMimeTypes = forall a. Ord a => [a] -> Set a
Set.fromList [ ByteString
"application/x-font-truetype"
                                     , ByteString
"application/x-javascript"
                                     , ByteString
"application/json"
                                     , ByteString
"text/css"
                                     , ByteString
"text/html"
                                     , ByteString
"text/javascript"
                                     , ByteString
"text/plain"
                                     , ByteString
"text/xml" ]




------------------------------------------------------------------------------
gzipCompression :: MonadSnap m => ByteString -> m ()
gzipCompression :: forall (m :: * -> *). MonadSnap m => ByteString -> m ()
gzipCompression ByteString
ce = forall (m :: * -> *). MonadSnap m => (Response -> Response) -> m ()
modifyResponse Response -> Response
f
  where
    f :: Response -> Response
f Response
r = forall a. HasHeaders a => CI ByteString -> ByteString -> a -> a
setHeader CI ByteString
"Content-Encoding" ByteString
ce    forall a b. (a -> b) -> a -> b
$
          forall a. HasHeaders a => CI ByteString -> ByteString -> a -> a
setHeader CI ByteString
"Vary" ByteString
"Accept-Encoding" forall a b. (a -> b) -> a -> b
$
          Response -> Response
clearContentLength                 forall a b. (a -> b) -> a -> b
$
          ((OutputStream Builder -> IO (OutputStream Builder))
 -> OutputStream Builder -> IO (OutputStream Builder))
-> Response -> Response
modifyResponseBody (OutputStream Builder -> IO (OutputStream Builder))
-> OutputStream Builder -> IO (OutputStream Builder)
gcompress Response
r


------------------------------------------------------------------------------
compressCompression :: MonadSnap m => ByteString -> m ()
compressCompression :: forall (m :: * -> *). MonadSnap m => ByteString -> m ()
compressCompression ByteString
ce = forall (m :: * -> *). MonadSnap m => (Response -> Response) -> m ()
modifyResponse Response -> Response
f
  where
    f :: Response -> Response
f Response
r = forall a. HasHeaders a => CI ByteString -> ByteString -> a -> a
setHeader CI ByteString
"Content-Encoding" ByteString
ce    forall a b. (a -> b) -> a -> b
$
          forall a. HasHeaders a => CI ByteString -> ByteString -> a -> a
setHeader CI ByteString
"Vary" ByteString
"Accept-Encoding" forall a b. (a -> b) -> a -> b
$
          Response -> Response
clearContentLength                 forall a b. (a -> b) -> a -> b
$
          ((OutputStream Builder -> IO (OutputStream Builder))
 -> OutputStream Builder -> IO (OutputStream Builder))
-> Response -> Response
modifyResponseBody (OutputStream Builder -> IO (OutputStream Builder))
-> OutputStream Builder -> IO (OutputStream Builder)
ccompress Response
r


------------------------------------------------------------------------------
gcompress :: (OutputStream Builder -> IO (OutputStream Builder))
          -> OutputStream Builder
          -> IO (OutputStream Builder)
gcompress :: (OutputStream Builder -> IO (OutputStream Builder))
-> OutputStream Builder -> IO (OutputStream Builder)
gcompress OutputStream Builder -> IO (OutputStream Builder)
body OutputStream Builder
stream = CompressionLevel
-> OutputStream Builder -> IO (OutputStream Builder)
Streams.gzipBuilder CompressionLevel
5 OutputStream Builder
stream forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= OutputStream Builder -> IO (OutputStream Builder)
body


------------------------------------------------------------------------------
ccompress :: (OutputStream Builder -> IO (OutputStream Builder))
          -> OutputStream Builder
          -> IO (OutputStream Builder)
ccompress :: (OutputStream Builder -> IO (OutputStream Builder))
-> OutputStream Builder -> IO (OutputStream Builder)
ccompress OutputStream Builder -> IO (OutputStream Builder)
body OutputStream Builder
stream = CompressionLevel
-> OutputStream Builder -> IO (OutputStream Builder)
Streams.compressBuilder CompressionLevel
5 OutputStream Builder
stream forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= OutputStream Builder -> IO (OutputStream Builder)
body


------------------------------------------------------------------------------
-- We're not gonna bother with quality values; we'll do gzip or compress in
-- that order.
acceptParser :: Parser [ByteString]
acceptParser :: Parser [ByteString]
acceptParser = do
    [ByteString]
xs <- ((forall a. a -> [a] -> [a]
:[]) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Parser ByteString ByteString
encoding) forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> (forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$! [])
    [ByteString]
ys <- forall (f :: * -> *) a. Alternative f => f a -> f [a]
many (Char -> Parser Char
char Char
',' forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> Parser ByteString ByteString
encoding)
    forall t. Chunk t => Parser t ()
endOfInput
    forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$! [ByteString]
xs forall a. [a] -> [a] -> [a]
++ [ByteString]
ys
  where
    encoding :: Parser ByteString ByteString
encoding = Parser ()
skipSpace forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> Parser ByteString ByteString
c forall (f :: * -> *) a b. Applicative f => f a -> f b -> f a
<* Parser ()
skipSpace

    c :: Parser ByteString ByteString
c = do
        ByteString
x <- Parser ByteString ByteString
coding
        Parser ()
qvalue forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> (forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$! ())
        forall (m :: * -> *) a. Monad m => a -> m a
return ByteString
x

    qvalue :: Parser ()
qvalue = do
        Parser ()
skipSpace
        forall (f :: * -> *) a. Functor f => f a -> f ()
void forall a b. (a -> b) -> a -> b
$! Char -> Parser Char
char Char
';'
        Parser ()
skipSpace
        forall (f :: * -> *) a. Functor f => f a -> f ()
void forall a b. (a -> b) -> a -> b
$! Char -> Parser Char
char Char
'q'
        Parser ()
skipSpace
        forall (f :: * -> *) a. Functor f => f a -> f ()
void forall a b. (a -> b) -> a -> b
$! Char -> Parser Char
char Char
'='
        Parser ()
float
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$! ()

    coding :: Parser ByteString ByteString
coding = ByteString -> Parser ByteString ByteString
string ByteString
"*" forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> (Char -> Bool) -> Parser ByteString ByteString
takeWhile1 Char -> Bool
isCodingChar

    isCodingChar :: Char -> Bool
isCodingChar Char
ch = Char -> Bool
isDigit Char
ch Bool -> Bool -> Bool
|| Char -> Bool
isAlpha_ascii Char
ch Bool -> Bool -> Bool
|| Char
ch forall a. Eq a => a -> a -> Bool
== Char
'-' Bool -> Bool -> Bool
|| Char
ch forall a. Eq a => a -> a -> Bool
== Char
'_'

    float :: Parser ()
float = (Char -> Bool) -> Parser ByteString ByteString
takeWhile Char -> Bool
isDigit forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>>
            (Char -> Parser Char
char Char
'.' forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> (Char -> Bool) -> Parser ByteString ByteString
takeWhile Char -> Bool
isDigit forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> (forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$! ())) forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> (forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$! ())


------------------------------------------------------------------------------
-- | Thrown when the 'Accept-Encoding' request header has invalid format.
data BadAcceptEncodingException = BadAcceptEncodingException
   deriving (Typeable)


------------------------------------------------------------------------------
instance Show BadAcceptEncodingException where
    show :: BadAcceptEncodingException -> [Char]
show BadAcceptEncodingException
BadAcceptEncodingException = [Char]
"bad 'accept-encoding' header"


------------------------------------------------------------------------------
instance Exception BadAcceptEncodingException


------------------------------------------------------------------------------
parseAcceptEncoding :: ByteString -> IO [ByteString]
parseAcceptEncoding :: ByteString -> IO [ByteString]
parseAcceptEncoding ByteString
s =
    case Either [Char] [ByteString]
r of
      Left [Char]
_ -> forall e a. Exception e => e -> IO a
throwIO BadAcceptEncodingException
BadAcceptEncodingException
      Right [ByteString]
x -> forall (m :: * -> *) a. Monad m => a -> m a
return [ByteString]
x
  where
    r :: Either [Char] [ByteString]
r = forall a. ByteString -> Parser a -> Either [Char] a
fullyParse ByteString
s Parser [ByteString]
acceptParser