← Back to branch summary

~ubuntu-branches/ubuntu/wily/haskell-wai-extra/wily-proposed

« back to all changes in this revision

Viewing changes to Network/Wai/Middleware/Gzip.hs

  • Committer: Package Import Robot
  • Author(s): Clint Adams
  • Date: 2014-06-06 11:40:45 UTC
  • mfrom: (1.1.13)
  • Revision ID: package-import@ubuntu.com-20140606114045-er7ij5963lk64k70
Tags: 3.0.0-1
New upstream version.

Show diffs side-by-side

added added

removed removed

Lines of Context:
Network/Wai/Middleware/Gzip.hs
28
28
import Data.Maybe (fromMaybe, isJust)
29
29
import qualified Data.ByteString.Char8 as S8
30
30
import qualified Data.ByteString as S
31
 
import Data.Default
 
31
import Data.Default.Class
32
32
import Network.HTTP.Types (Status, Header)
33
33
import Control.Monad.IO.Class (liftIO)
 
34
import Control.Monad.Trans.Resource (runResourceT)
34
35
import System.Directory (doesFileExist, createDirectoryIfMissing)
35
 
import qualified Data.Conduit as C
36
 
import qualified Data.Conduit.Zlib as CZ
37
 
import qualified Data.Conduit.Binary as CB
38
 
import qualified Data.Conduit.List as CL
39
 
import Data.Conduit.Blaze (builderToByteStringFlush)
40
36
import Blaze.ByteString.Builder (fromByteString)
41
37
import Control.Exception (try, SomeException)
42
38
import qualified Data.Set as Set
43
39
import Network.Wai.Internal
 
40
import qualified Data.Streaming.Blaze as B
 
41
import qualified Data.Streaming.Zlib as Z
 
42
import qualified Blaze.ByteString.Builder as Blaze
 
43
import Control.Monad (unless)
 
44
import Data.Function (fix)
 
45
import Control.Exception (throwIO)
 
46
import qualified System.IO as IO
 
47
import Data.ByteString.Lazy.Internal (defaultChunkSize)
44
48
 
45
49
data GzipSettings = GzipSettings
46
50
    { gzipFiles :: GzipFiles
73
77
--
74
78
-- * Only compress if the response is above a certain size.
75
79
gzip :: GzipSettings -> Middleware
76
 
gzip set app env = do
77
 
    res <- app env
 
80
gzip set app env sendResponse = app env $ \res ->
78
81
    case res of
79
 
        ResponseFile{} | gzipFiles set == GzipIgnore -> return res
 
82
        ResponseFile{} | gzipFiles set == GzipIgnore -> sendResponse res
80
83
        _ -> if "gzip" `elem` enc && not isMSIE6 && not (isEncoded res)
81
84
                then
82
85
                    case (res, gzipFiles set) of
83
86
                        (ResponseFile s hs file Nothing, GzipCacheFolder cache) ->
84
87
                            case lookup "content-type" hs of
85
88
                                Just m
86
 
                                    | gzipCheckMime set m -> liftIO $ compressFile s hs file cache
87
 
                                _ -> return res
88
 
                        _ -> return $ compressE set res
89
 
                else return res
 
89
                                    | gzipCheckMime set m -> compressFile s hs file cache sendResponse
 
90
                                _ -> sendResponse res
 
91
                        _ -> compressE set res sendResponse
 
92
                else sendResponse res
90
93
  where
91
94
    enc = fromMaybe [] $ (splitCommas . S8.unpack)
92
95
                    `fmap` lookup "Accept-Encoding" (requestHeaders env)
94
97
    isMSIE6 = "MSIE 6" `S.isInfixOf` ua
95
98
    isEncoded res = isJust $ lookup "Content-Encoding" $ responseHeaders res
96
99
 
97
 
compressFile :: Status -> [Header] -> FilePath -> FilePath -> IO Response
98
 
compressFile s hs file cache = do
 
100
compressFile :: Status -> [Header] -> FilePath -> FilePath -> (Response -> IO a) -> IO a
 
101
compressFile s hs file cache sendResponse = do
99
102
    e <- doesFileExist tmpfile
100
103
    if e
101
104
        then onSucc
102
105
        else do
103
106
            createDirectoryIfMissing True cache
104
 
            x <-
105
 
               try $ C.runResourceT $ CB.sourceFile file
106
 
                C.$$ CZ.gzip C.=$ CB.sinkFile tmpfile
107
 
            either onErr (const onSucc) x
 
107
            x <- try $
 
108
                 IO.withBinaryFile file IO.ReadMode $ \inH ->
 
109
                 IO.withBinaryFile tmpfile IO.WriteMode $ \outH -> do
 
110
                    deflate <- Z.initDeflate 7 $ Z.WindowBits 31
 
111
                    -- FIXME this code should write to a temporary file, then
 
112
                    -- rename to the final file
 
113
                    let goPopper popper = fix $ \loop -> do
 
114
                            res <- popper
 
115
                            case res of
 
116
                                Z.PRDone -> return ()
 
117
                                Z.PRNext bs -> do
 
118
                                    S.hPut outH bs
 
119
                                    loop
 
120
                                Z.PRError e -> throwIO e
 
121
                    fix $ \loop -> do
 
122
                        bs <- S.hGetSome inH defaultChunkSize
 
123
                        unless (S.null bs) $ do
 
124
                            Z.feedDeflate deflate bs >>= goPopper
 
125
                            loop
 
126
                    goPopper $ Z.finishDeflate deflate
 
127
            either onErr (const onSucc) (x :: Either SomeException ()) -- FIXME bad! don't catch all exceptions like that!
108
128
  where
109
 
    onSucc = return $ ResponseFile s (fixHeaders hs) tmpfile Nothing
 
129
    onSucc = sendResponse $ responseFile s (fixHeaders hs) tmpfile Nothing
110
130
 
111
 
    onErr :: SomeException -> IO Response
112
 
    onErr = const $ return $ ResponseFile s hs file Nothing -- FIXME log the error message
 
131
    onErr _ = sendResponse $ responseFile s hs file Nothing -- FIXME log the error message
113
132
 
114
133
    tmpfile = cache ++ '/' : map safe file
115
134
    safe c
122
141
 
123
142
compressE :: GzipSettings
124
143
          -> Response
125
 
          -> Response
126
 
compressE set res =
 
144
          -> (Response -> IO a)
 
145
          -> IO a
 
146
compressE set res sendResponse =
127
147
    case lookup "content-type" hs of
128
148
        Just m | gzipCheckMime set m ->
129
149
            let hs' = fixHeaders hs
130
 
             in ResponseSource s hs' $ \f -> wb $ \b -> f $
131
 
                                       b C.$= builderToByteStringFlush
132
 
                                         C.$= CZ.compressFlush 1 (CZ.WindowBits 31)
133
 
                                         C.$= CL.map (fmap fromByteString)
134
 
        _ -> res
 
150
             in wb $ \body -> sendResponse $ responseStream s hs' $ \sendChunk flush -> do
 
151
                    (blazeRecv, blazeFinish) <- B.newBlazeRecv B.defaultStrategy
 
152
                    deflate <- Z.initDeflate 1 (Z.WindowBits 31)
 
153
                    let sendBuilder builder = do
 
154
                            popper <- blazeRecv builder
 
155
                            fix $ \loop -> do
 
156
                                bs <- popper
 
157
                                unless (S.null bs) $ do
 
158
                                    sendBS bs
 
159
                                    loop
 
160
                        sendBS bs = Z.feedDeflate deflate bs >>= deflatePopper
 
161
                        flushBuilder = do
 
162
                            sendBuilder Blaze.flush
 
163
                            deflatePopper $ Z.flushDeflate deflate
 
164
                            flush
 
165
                        deflatePopper popper = fix $ \loop -> do
 
166
                            res <- popper
 
167
                            case res of
 
168
                                Z.PRDone -> return ()
 
169
                                Z.PRNext bs' -> do
 
170
                                    sendChunk $ fromByteString bs'
 
171
                                    loop
 
172
                                Z.PRError e -> throwIO e
 
173
 
 
174
                    body sendBuilder flushBuilder
 
175
                    sendBuilder Blaze.flush
 
176
                    deflatePopper $ Z.finishDeflate deflate
 
177
        _ -> sendResponse res
135
178
  where
136
 
    (s, hs, wb) = responseToSource res
 
179
    (s, hs, wb) = responseToStream res
137
180
 
138
181
-- Remove Content-Length header, since we will certainly have a
139
182
-- different length after gzip compression.