@@ -13,8 +13,8 @@ module Network.Wai.Handler.CGI
1313import Network.Wai
1414import Network.Wai.Internal
1515import Network.Socket (getAddrInfo, addrAddress)
16+import Data.IORef
1617import Data.Maybe (fromMaybe)
17−import Control.Exception (mask)
1818import qualified Data.ByteString.Char8 as B
1919import qualified Data.ByteString.Lazy as L
2020import Control.Arrow ((***))
@@ -22,15 +22,17 @@ import Data.Char (toLower)
2222import qualified System.IO
2323import qualified Data.String as String
2424import Data.Monoid (mconcat, mempty)
25−import Blaze.ByteString.Builder (fromByteString, toLazyByteString)
25+import Blaze.ByteString.Builder (fromByteString, toLazyByteString, flush)
2626import Blaze.ByteString.Builder.Char8 (fromChar, fromString)
27−import Control.Monad.IO.Class (liftIO)
2827import Data.ByteString.Lazy.Internal (defaultChunkSize)
2928import System.IO (Handle)
3029import Network.HTTP.Types (Status (..))
3130import qualified Network.HTTP.Types as H
3231import qualified Data.CaseInsensitive as CI
3332import Data.Monoid (mappend)
33+import qualified Data.Streaming.Blaze as Blaze
34+import Data.Function (fix)
35+import Control.Monad (unless, void)
3436
3537#if WINDOWS
3638import System.Environment (getEnvironment)
@@ -74,7 +76,7 @@ runSendfile sf app = do
7476-- stick with 'run' or 'runSendfile'.
7577runGeneric
7678 :: [(String, String)] -- ^ all variables
77− -> (Int -> Source IO B.ByteString) -- ^ responseBody of input
79+ -> (Int -> IO (IO B.ByteString)) -- ^ responseBody of input
7880 -> (B.ByteString -> IO ()) -- ^ destination for output
7981 -> Maybe B.ByteString -- ^ does the server support the X-Sendfile header?
8082 -> Application
@@ -96,40 +98,48 @@ runGeneric vars inputH outputH xsendfile app = do
9698 "https" -> True
9799 _ -> False
98100 addrs <- getAddrInfo Nothing (Just remoteHost') Nothing
101+ requestBody' <- inputH contentLength
99102 let addr =
100103 case addrs of
101104 a:_ -> addrAddress a
102105 [] -> error $ "Invalid REMOTE_ADDR or REMOTE_HOST: " ++ remoteHost'
103− mask $ \restore -> do
104− let reqHeaders = map (cleanupVarName *** B.pack) vars
105− env = Request
106− { requestMethod = rmethod
107− , rawPathInfo = B.pack pinfo
108− , pathInfo = H.decodePathSegments $ B.pack pinfo
109− , rawQueryString = B.pack qstring
110− , queryString = H.parseQuery $ B.pack qstring
111− , requestHeaders = reqHeaders
112− , isSecure = isSecure'
113− , remoteHost = addr
114− , httpVersion = H.http11 -- FIXME
115− , requestBody = inputH contentLength
116− , vault = mempty
117− , requestBodyLength = KnownLength $ fromIntegral contentLength
118− , requestHeaderHost = lookup "host" reqHeaders
119− , requestHeaderRange = lookup "range" reqHeaders
120− }
121− -- FIXME worry about exception?
122− res <- restore $ app env
106+ reqHeaders = map (cleanupVarName *** B.pack) vars
107+ env = Request
108+ { requestMethod = rmethod
109+ , rawPathInfo = B.pack pinfo
110+ , pathInfo = H.decodePathSegments $ B.pack pinfo
111+ , rawQueryString = B.pack qstring
112+ , queryString = H.parseQuery $ B.pack qstring
113+ , requestHeaders = reqHeaders
114+ , isSecure = isSecure'
115+ , remoteHost = addr
116+ , httpVersion = H.http11 -- FIXME
117+ , requestBody = requestBody'
118+ , vault = mempty
119+ , requestBodyLength = KnownLength $ fromIntegral contentLength
120+ , requestHeaderHost = lookup "host" reqHeaders
121+ , requestHeaderRange = lookup "range" reqHeaders
122+ }
123+ void $ app env $ \res ->
123124 case (xsendfile, res) of
124− (Just sf, ResponseFile s hs fp Nothing) ->
125− restore $ mapM_ outputH $ L.toChunks $ toLazyByteString $ sfBuilder s hs sf fp
125+ (Just sf, ResponseFile s hs fp Nothing) -> do
126+ mapM_ outputH $ L.toChunks $ toLazyByteString $ sfBuilder s hs sf fp
127+ return ResponseReceived
126128 _ -> do
127− let (s, hs, wb) = responseToSource res
128− wb $ \b ->
129− let src = do
130− yield (Chunk $ headers s hs `mappend` fromChar '\n')
131− b
132− in src $$ builderSink
129+ let (s, hs, wb) = responseToStream res
130+ (blazeRecv, blazeFinish) <- Blaze.newBlazeRecv Blaze.defaultStrategy
131+ wb $ \b -> do
132+ let sendBuilder builder = do
133+ popper <- blazeRecv builder
134+ fix $ \loop -> do
135+ bs <- popper
136+ unless (B.null bs) $ do
137+ outputH bs
138+ loop
139+ sendBuilder $ headers s hs `mappend` fromChar '\n'
140+ b sendBuilder (sendBuilder flush)
141+ blazeFinish >>= maybe (return ()) outputH
142+ return ResponseReceived
133143 where
134144 headers s hs = mconcat (map header $ status s : map header' (fixHeaders hs))
135145 status (Status i m) = (fromByteString "Status", mconcat
@@ -151,13 +161,6 @@ runGeneric vars inputH outputH xsendfile app = do
151161 , fromByteString sf
152162 , fromByteString " not supported"
153163 ]
154− bsSink = await >>= maybe (return ()) push
155− push (Chunk bs) = do
156− liftIO $ outputH bs
157− bsSink
158− -- FIXME actually flush?
159− push Flush = bsSink
160− builderSink = builderToByteStringFlush =$ bsSink
161164 fixHeaders h =
162165 case lookup "content-type" h of
163166 Nothing -> ("Content-Type", "text/html; charset=utf-8") : h
@@ -176,19 +179,19 @@ cleanupVarName s =
176179 helper' (x:rest) = toLower x : helper' rest
177180 helper' [] = []
178181
179−requestBodyHandle :: Handle -> Int -> Source IO B.ByteString
182+requestBodyHandle :: Handle -> Int -> IO (IO B.ByteString)
180183requestBodyHandle h = requestBodyFunc $ \i -> do
181184 bs <- B.hGet h i
182185 return $ if B.null bs then Nothing else Just bs
183186
184−requestBodyFunc :: (Int -> IO (Maybe B.ByteString)) -> Int -> Source IO B.ByteString
185−requestBodyFunc get =
186− loop
187− where
188− loop 0 = return ()
189− loop count = do
190− mbs <- liftIO $ get $ min count defaultChunkSize
191− let count' = count - maybe 0 B.length mbs
192− case mbs of
193− Nothing -> return ()
194− Just bs -> yield bs >> loop count'
187+requestBodyFunc :: (Int -> IO (Maybe B.ByteString)) -> Int -> IO (IO B.ByteString)
188+requestBodyFunc get count0 = do
189+ ref <- newIORef count0
190+ return $ do
191+ count <- readIORef ref
192+ if count <= 0
193+ then return B.empty
194+ else do
195+ mbs <- get $ min count defaultChunkSize
196+ writeIORef ref $ count - maybe 0 B.length mbs
197+ return $ fromMaybe B.empty mbs