1 {-# LANGUAGE BangPatterns              #-}
    2 {-# LANGUAGE CPP                       #-}
    3 {-# LANGUAGE DeriveDataTypeable        #-}
    4 {-# LANGUAGE ExistentialQuantification #-}
    5 {-# LANGUAGE OverloadedStrings         #-}
    6 {-# LANGUAGE ScopedTypeVariables       #-}
    7 
    8 --------------------------------------------------------------------------------
    9 -- | Helpers for running a 'Snap' web handler with compression.
   10 
   11 module Snap.Util.GZip
   12   ( withCompression
   13   , withCompression'
   14   , noCompression
   15   , BadAcceptEncodingException
   16   ) where
   17 
   18 import           Control.Applicative              (Alternative ((<|>), many), Applicative ((*>), (<*), pure), (<$>))
   19 import           Control.Exception                (Exception, throwIO)
   20 import           Control.Monad                    (Functor (fmap), Monad ((>>), (>>=), return), MonadPlus (mplus), void, when)
   21 import           Control.Monad.IO.Class           (MonadIO (liftIO))
   22 import           Data.Attoparsec.ByteString.Char8 (Parser, char, endOfInput, isAlpha_ascii, isDigit, skipSpace, string, takeWhile, takeWhile1)
   23 import           Data.ByteString.Builder          (Builder)
   24 import           Data.ByteString.Char8            (ByteString)
   25 import qualified Data.ByteString.Char8            as S (takeWhile)
   26 import qualified Data.Char                        as Char (isSpace)
   27 import           Data.Maybe                       (Maybe (Just, Nothing), fromMaybe, isJust, maybe)
   28 import           Data.Set                         (Set)
   29 import qualified Data.Set                         as Set (fromList, member)
   30 import           Data.Typeable                    (Typeable)
   31 import           Prelude                          (Either (..), Eq (..), IO, Show (show), id, not, ($), ($!), (&&), (++), (||))
   32 import           Snap.Core                        (MonadSnap, clearContentLength, finishWith, getHeader, getRequest, getResponse, modifyResponse, modifyResponseBody, setHeader)
   33 import           Snap.Internal.Debug              (debug)
   34 import           Snap.Internal.Parsing            (fullyParse)
   35 import           System.IO.Streams                (OutputStream)
   36 import qualified System.IO.Streams                as Streams (compressBuilder, gzipBuilder)
   37 
   38 ------------------------------------------------------------------------------
   39 -- | Runs a 'Snap' web handler with compression if available.
   40 --
   41 -- If the client has indicated support for @gzip@ or @deflate@ in its
   42 -- @Accept-Encoding@ header, and the @Content-Type@ in the response is one of
   43 -- the following types:
   44 --
   45 --   * @application/x-javascript@
   46 --
   47 --   * @application/json@
   48 --
   49 --   * @text/css@
   50 --
   51 --   * @text/html@
   52 --
   53 --   * @text/javascript@
   54 --
   55 --   * @text/plain@
   56 --
   57 --   * @text/xml@
   58 --
   59 --   * @application/x-font-truetype@
   60 --
   61 -- Then the given handler's output stream will be compressed,
   62 -- @Content-Encoding@ will be set in the output headers, and the
   63 -- @Content-Length@ will be cleared if it was set. (We can't process the
   64 -- stream in O(1) space if the length is known beforehand.)
   65 --
   66 -- The wrapped handler will be run to completion, and then the 'Response'
   67 -- that's contained within the 'Snap' monad state will be passed to
   68 -- 'finishWith' to prevent further processing.
   69 --
   70 --
   71 -- Example:
   72 --
   73 -- @
   74 -- ghci> :set -XOverloadedStrings
   75 -- ghci> import qualified "Data.Map" as M
   76 -- ghci> import qualified "Snap.Test" as T
   77 -- ghci> let r = T.get \"/\" M.empty >> T.addHeader \"Accept-Encoding\" \"gzip,deflate\"
   78 -- ghci> let h = 'Snap.Core.modifyResponse' ('Snap.Core.setContentType' \"text\/plain\") >> 'Snap.Core.writeBS' \"some text\"
   79 -- ghci> T.runHandler r h
   80 -- HTTP\/1.1 200 OK
   81 -- content-type: text\/plain
   82 -- server: Snap\/test
   83 -- date: Fri, 08 Aug 2014 15:40:45 GMT
   84 --
   85 -- some text
   86 -- ghci> T.runHandler r ('withCompression' h)
   87 -- HTTP\/1.1 200 OK
   88 -- content-type: text\/plain
   89 -- vary: Accept-Encoding
   90 -- content-encoding: gzip
   91 -- server: Snap\/test
   92 -- date: Fri, 08 Aug 2014 15:40:10 GMT
   93 --
   94 --
   95 -- @
   96 withCompression :: MonadSnap m
   97                 => m a   -- ^ the web handler to run
   98                 -> m ()
   99 withCompression = withCompression' compressibleMimeTypes
  100 
  101 
  102 ------------------------------------------------------------------------------
  103 -- | The same as 'withCompression', with control over which MIME types to
  104 -- compress.
  105 withCompression' :: MonadSnap m
  106                  => Set ByteString
  107                     -- ^ set of compressible MIME types
  108                  -> m a
  109                     -- ^ the web handler to run
  110                  -> m ()
  111 withCompression' mimeTable action = do
  112     _    <- action
  113     resp <- getResponse
  114 
  115     -- If a content-encoding is already set, do nothing. This prevents
  116     -- "withCompression $ withCompression m" from ruining your day.
  117     when (not $ isJust $ getHeader "Content-Encoding" resp) $ do
  118        let mbCt = fmap chop $ getHeader "Content-Type" resp
  119 
  120        debug $ "withCompression', content-type is " ++ show mbCt
  121 
  122        case mbCt of
  123          (Just ct) -> when (Set.member ct mimeTable) chkAcceptEncoding
  124          _         -> return $! ()
  125 
  126 
  127     getResponse >>= finishWith
  128 
  129   where
  130     chop = S.takeWhile (\c -> c /= ';' && not (Char.isSpace c))
  131 
  132     chkAcceptEncoding = do
  133         req <- getRequest
  134         debug $ "checking accept-encoding"
  135         let mbAcc = getHeader "Accept-Encoding" req
  136         debug $ "accept-encoding is " ++ show mbAcc
  137         let s = fromMaybe "" mbAcc
  138 
  139         types <- liftIO $ parseAcceptEncoding s
  140 
  141         chooseType Nothing types
  142 
  143     chooseType !m []               = maybe (return $! ()) id m
  144     chooseType !_ ("gzip":_)       = gzipCompression "gzip"
  145     chooseType !m ("deflate":xs)   =
  146         chooseType (m `mplus` Just (compressCompression "deflate")) xs
  147 
  148     chooseType !_ ("x-gzip":_)     = gzipCompression "x-gzip"
  149     chooseType !m ("x-deflate":xs) =
  150         chooseType (m `mplus` Just (compressCompression "x-deflate")) xs
  151     chooseType !m (_:xs)           = chooseType m xs
  152 
  153 
  154 ------------------------------------------------------------------------------
  155 -- | Turn off compression by setting \"Content-Encoding: identity\" in the
  156 -- response headers. 'withCompression' is a no-op when a content-encoding is
  157 -- already set.
  158 noCompression :: MonadSnap m => m ()
  159 noCompression = modifyResponse $ setHeader "Content-Encoding" "identity"
  160 
  161 
  162 ------------------------------------------------------------------------------
  163 -- private following
  164 ------------------------------------------------------------------------------
  165 
  166 
  167 ------------------------------------------------------------------------------
  168 compressibleMimeTypes :: Set ByteString
  169 compressibleMimeTypes = Set.fromList [ "application/x-font-truetype"
  170                                      , "application/x-javascript"
  171                                      , "application/json"
  172                                      , "text/css"
  173                                      , "text/html"
  174                                      , "text/javascript"
  175                                      , "text/plain"
  176                                      , "text/xml" ]
  177 
  178 
  179 
  180 
  181 ------------------------------------------------------------------------------
  182 gzipCompression :: MonadSnap m => ByteString -> m ()
  183 gzipCompression ce = modifyResponse f
  184   where
  185     f r = setHeader "Content-Encoding" ce    $
  186           setHeader "Vary" "Accept-Encoding" $
  187           clearContentLength                 $
  188           modifyResponseBody gcompress r
  189 
  190 
  191 ------------------------------------------------------------------------------
  192 compressCompression :: MonadSnap m => ByteString -> m ()
  193 compressCompression ce = modifyResponse f
  194   where
  195     f r = setHeader "Content-Encoding" ce    $
  196           setHeader "Vary" "Accept-Encoding" $
  197           clearContentLength                 $
  198           modifyResponseBody ccompress r
  199 
  200 
  201 ------------------------------------------------------------------------------
  202 gcompress :: (OutputStream Builder -> IO (OutputStream Builder))
  203           -> OutputStream Builder
  204           -> IO (OutputStream Builder)
  205 gcompress body stream = Streams.gzipBuilder 5 stream >>= body
  206 
  207 
  208 ------------------------------------------------------------------------------
  209 ccompress :: (OutputStream Builder -> IO (OutputStream Builder))
  210           -> OutputStream Builder
  211           -> IO (OutputStream Builder)
  212 ccompress body stream = Streams.compressBuilder 5 stream >>= body
  213 
  214 
  215 ------------------------------------------------------------------------------
  216 -- We're not gonna bother with quality values; we'll do gzip or compress in
  217 -- that order.
  218 acceptParser :: Parser [ByteString]
  219 acceptParser = do
  220     xs <- ((:[]) <$> encoding) <|> (return $! [])
  221     ys <- many (char ',' *> encoding)
  222     endOfInput
  223     return $! xs ++ ys
  224   where
  225     encoding = skipSpace *> c <* skipSpace
  226 
  227     c = do
  228         x <- coding
  229         qvalue <|> (return $! ())
  230         return x
  231 
  232     qvalue = do
  233         skipSpace
  234         void $! char ';'
  235         skipSpace
  236         void $! char 'q'
  237         skipSpace
  238         void $! char '='
  239         float
  240         return $! ()
  241 
  242     coding = string "*" <|> takeWhile1 isCodingChar
  243 
  244     isCodingChar ch = isDigit ch || isAlpha_ascii ch || ch == '-' || ch == '_'
  245 
  246     float = takeWhile isDigit >>
  247             (char '.' >> takeWhile isDigit >> (pure $! ())) <|> (pure $! ())
  248 
  249 
  250 ------------------------------------------------------------------------------
  251 -- | Thrown when the 'Accept-Encoding' request header has invalid format.
  252 data BadAcceptEncodingException = BadAcceptEncodingException
  253    deriving (Typeable)
  254 
  255 
  256 ------------------------------------------------------------------------------
  257 instance Show BadAcceptEncodingException where
  258     show BadAcceptEncodingException = "bad 'accept-encoding' header"
  259 
  260 
  261 ------------------------------------------------------------------------------
  262 instance Exception BadAcceptEncodingException
  263 
  264 
  265 ------------------------------------------------------------------------------
  266 parseAcceptEncoding :: ByteString -> IO [ByteString]
  267 parseAcceptEncoding s =
  268     case r of
  269       Left _ -> throwIO BadAcceptEncodingException
  270       Right x -> return x
  271   where
  272     r = fullyParse s acceptParser