@@ -37,6 +37,14 @@ import Blaze.ByteString.Builder (fromByteString)
3737import Control.Exception (try, SomeException)
3838import qualified Data.Set as Set
3939import 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)
4048
4149data GzipSettings = GzipSettings
4250 { gzipFiles :: GzipFiles
@@ -69,20 +77,19 @@ defaultCheckMime bs =
6977--
7078-- * Only compress if the response is above a certain size.
7179gzip :: GzipSettings -> Middleware
72−gzip set app env = do
73− res <- app env
80+gzip set app env sendResponse = app env $ \res ->
7481 case res of
75− ResponseFile{} | gzipFiles set == GzipIgnore -> return res
82+ ResponseFile{} | gzipFiles set == GzipIgnore -> sendResponse res
7683 _ -> if "gzip" `elem` enc && not isMSIE6 && not (isEncoded res)
7784 then
7885 case (res, gzipFiles set) of
7986 (ResponseFile s hs file Nothing, GzipCacheFolder cache) ->
8087 case lookup "content-type" hs of
8188 Just m
82− | gzipCheckMime set m -> liftIO $ compressFile s hs file cache
83− _ -> return res
84− _ -> return $ compressE set res
85− 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
8693 where
8794 enc = fromMaybe [] $ (splitCommas . S8.unpack)
8895 `fmap` lookup "Accept-Encoding" (requestHeaders env)
@@ -90,22 +97,38 @@ gzip set app env = do
9097 isMSIE6 = "MSIE 6" `S.isInfixOf` ua
9198 isEncoded res = isJust $ lookup "Content-Encoding" $ responseHeaders res
9299
93−compressFile :: Status -> [Header] -> FilePath -> FilePath -> IO Response
94−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
95102 e <- doesFileExist tmpfile
96103 if e
97104 then onSucc
98105 else do
99106 createDirectoryIfMissing True cache
100− x <-
101− try $ runResourceT $ CB.sourceFile file
102− C.$$ CZ.gzip C.=$ CB.sinkFile tmpfile
103− 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!
104128 where
105− onSucc = return $ ResponseFile s (fixHeaders hs) tmpfile Nothing
129+ onSucc = sendResponse $ responseFile s (fixHeaders hs) tmpfile Nothing
106130
107− onErr :: SomeException -> IO Response
108− 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
109132
110133 tmpfile = cache ++ '/' : map safe file
111134 safe c
@@ -118,18 +141,42 @@ compressFile s hs file cache = do
118141
119142compressE :: GzipSettings
120143 -> Response
121− -> Response
122−compressE set res =
144+ -> (Response -> IO a)
145+ -> IO a
146+compressE set res sendResponse =
123147 case lookup "content-type" hs of
124148 Just m | gzipCheckMime set m ->
125149 let hs' = fixHeaders hs
126− in ResponseSource s hs' $ \f -> wb $ \b -> f $
127− b C.$= builderToByteStringFlush
128− C.$= CZ.compressFlush 1 (CZ.WindowBits 31)
129− C.$= CL.map (fmap fromByteString)
130− _ -> 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
131178 where
132− (s, hs, wb) = responseToSource res
179+ (s, hs, wb) = responseToStream res
133180
134181-- Remove Content-Length header, since we will certainly have a
135182-- different length after gzip compression.