@@ -22,12 +22,15 @@ import qualified Data.CaseInsensitive as CI
2222import Data.Conduit
2323import Data.Conduit.Blaze (builderToByteString)
2424import qualified Data.Conduit.List as CL
25+import Data.Function (on)
26+import Data.List (deleteBy)
2527import Data.Maybe (isJust, listToMaybe)
2628import Data.Monoid (mappend)
2729import Data.Version (showVersion)
2830import Network.HTTP.Attoparsec (parseByteRanges)
2931import qualified Network.HTTP.Types as H
3032import Network.Wai
33+import qualified Network.Wai.Handler.Warp.Date as D
3134import Network.Wai.Handler.Warp.Header
3235import Network.Wai.Handler.Warp.ResponseHeader
3336import qualified Network.Wai.Handler.Warp.Timeout as T
@@ -37,6 +40,9 @@ import Numeric (showInt)
3740import qualified Paths_warp
3841import qualified System.PosixCompat.Files as P
3942
43+-- $setup
44+-- >>> :set -XOverloadedStrings
45+
4046----------------------------------------------------------------
4147
4248fileRange :: H.Status -> H.ResponseHeaders -> FilePath
@@ -135,69 +141,81 @@ sendResponse :: Connection
135141 -> IndexedHeader -- ^ Indexed header of HTTP request.
136142 -> Response -- ^ HTTP response including status code and response header.
137143 -> IO Bool -- ^ Returing True if the connection is persistent.
138−
139−----------------------------------------------------------------
140−
141−sendResponse conn ii restore req reqidxhdr (ResponseFile s0 hs0 path mPart)
142− | hasBody s0 req = restore $ do
143− ex <- fileRange s0 hs path mPart mRange
144− case ex of
145− Left _ -> sendResponse conn ii restore req reqidxhdr notFound
146− Right (s, hs1, beg, len) -> do
147− lheader <- composeHeader version s hs1
148− connSendFile conn path beg len (T.tickle th) [lheader]
149− T.tickle th
150− return isPersist
151− | otherwise = restore $ sendResponseNoBody conn th version s0 hs isPersist
144+sendResponse conn ii restore req reqidxhdr response = restore $ do
145+ hs <- addServerAndDate hs0
146+ if hasBody s req then do
147+ sendRsp conn ver s hs rsp
148+ T.tickle th
149+ return ret
150+ else do
151+ sendResponseNoBody conn ver s hs response
152+ T.tickle th
153+ return isPersist
152154 where
153− hs = addServer rspidxhdr $ addAcceptRanges hs0
155+ ver = httpVersion req
156+ s = responseStatus response
157+ hs0 = responseHeaders response
154158 rspidxhdr = indexResponseHeader hs0
155159 th = threadHandle ii
156− version = httpVersion req
157− (isPersist,_) = infoFromRequest req reqidxhdr
160+ dc = dateCacher ii
161+ addServerAndDate = addDate dc . addServer rspidxhdr
158162 mRange = reqidxhdr ! idxRange
159− notFound = responseLBS H.status404 [(H.hContentType, "text/plain")] "File not found"
163+ reqinfo@(isPersist,_) = infoFromRequest req reqidxhdr
164+ (isKeepAlive, needsChunked) = infoFromResponse rspidxhdr reqinfo
165+ rsp = case response of
166+ ResponseFile _ _ path mPart -> RspFile path mPart mRange (T.tickle th)
167+ ResponseBuilder _ _ b -> RspBuilder b needsChunked
168+ ResponseSource _ _ fb -> RspSource fb needsChunked th
169+ ret = case response of
170+ ResponseFile _ _ _ _ -> isPersist
171+ ResponseBuilder _ _ _ -> isKeepAlive
172+ ResponseSource _ _ _ -> isKeepAlive
160173
161174----------------------------------------------------------------
162175
163−sendResponse conn ii restore req reqidxhdr (ResponseBuilder s hs0 b)
164− | hasBody s req = restore $ do
165− header <- composeHeaderBuilder version s hs needsChunked
166− let body
167− | needsChunked = header `mappend` chunkedTransferEncoding b
168− `mappend` chunkedTransferTerminator
169− | otherwise = header `mappend` b
170− flip toByteStringIO body $ \bs -> do
171− connSendAll conn bs
172− T.tickle th
173− return isKeepAlive
174− | otherwise = restore $ sendResponseNoBody conn th version s hs isPersist
176+data Rsp = RspFile FilePath (Maybe FilePart) (Maybe HeaderValue) (IO ())
177+ | RspBuilder Builder Bool
178+ | RspSource (forall b. WithSource IO (Flush Builder) b) Bool T.Handle
179+
180+----------------------------------------------------------------
181+
182+sendRsp :: Connection
183+ -> H.HttpVersion
184+ -> H.Status
185+ -> H.ResponseHeaders
186+ -> Rsp
187+ -> IO ()
188+sendRsp conn ver s0 hs0 (RspFile path mPart mRange hook) = do
189+ ex <- fileRange s0 hs path mPart mRange
190+ case ex of
191+ Left _ -> sendRsp conn ver s2 hs2 (RspBuilder body True)
192+ Right (s, hs1, beg, len) -> do
193+ lheader <- composeHeader ver s hs1
194+ connSendFile conn path beg len hook [lheader]
175195 where
176− hs = addServer rspidxhdr hs0
177− rspidxhdr = indexResponseHeader hs0
178− th = threadHandle ii
179− version = httpVersion req
180− reqinfo@(isPersist,_) = infoFromRequest req reqidxhdr
181− (isKeepAlive, needsChunked) = infoFromResponse rspidxhdr reqinfo
196+ hs = addAcceptRanges hs0
197+ s2 = H.status404
198+ hs2 = replaceHeader H.hContentType "text/plain" hs0
199+ body = fromByteString "File not found"
200+
201+----------------------------------------------------------------
202+
203+sendRsp conn ver s hs (RspBuilder b needsChunked) = do
204+ header <- composeHeaderBuilder ver s hs needsChunked
205+ let body
206+ | needsChunked = header `mappend` chunkedTransferEncoding b
207+ `mappend` chunkedTransferTerminator
208+ | otherwise = header `mappend` b
209+ flip toByteStringIO body $ \bs -> do
210+ connSendAll conn bs
182211
183212----------------------------------------------------------------
184213
185−sendResponse conn ii restore req reqidxhdr (ResponseSource s hs0 withBodyFlush)
186− | hasBody s req = withBodyFlush $ \bodyFlush -> restore $ do
187− header <- composeHeaderBuilder version s hs needsChunked
188− let src = CL.sourceList [header] `mappend` cbody bodyFlush
189− src $$ builderToByteString =$ connSink conn th
190− return isKeepAlive
191− | otherwise = withBodyFlush $ \_bodyFlush ->
192− -- make sure any cleanup is called
193− restore $ sendResponseNoBody conn th version s hs isPersist
214+sendRsp conn ver s hs (RspSource withBodyFlush needsChunked th) = withBodyFlush $ \bodyFlush -> do
215+ header <- composeHeaderBuilder ver s hs needsChunked
216+ let src = CL.sourceList [header] `mappend` cbody bodyFlush
217+ src $$ builderToByteString =$ connSink conn th
194218 where
195− hs = addServer rspidxhdr hs0
196− rspidxhdr = indexResponseHeader hs0
197− th = threadHandle ii
198− version = httpVersion req
199− reqinfo@(isPersist,_) = infoFromRequest req reqidxhdr
200− (isKeepAlive, needsChunked) = infoFromResponse rspidxhdr reqinfo
201219 cbody bodyFlush = if needsChunked then body $= chunk else body
202220 where
203221 body = mapOutput (\x -> case x of
@@ -211,21 +229,24 @@ sendResponse conn ii restore req reqidxhdr (ResponseSource s hs0 withBodyFlush)
211229
212230----------------------------------------------------------------
213231
214−sendResponseNoBody :: Connection -> T.Handle
215− -> H.HttpVersion -> H.Status -> H.ResponseHeaders
216− -> Bool -> IO Bool
217−sendResponseNoBody conn th version s hs isPersist = do
218− composeHeader version s hs >>= connSendAll conn
219− T.tickle th
220− return isPersist
232+sendResponseNoBody :: Connection
233+ -> H.HttpVersion
234+ -> H.Status
235+ -> H.ResponseHeaders
236+ -> Response
237+ -> IO ()
238+sendResponseNoBody conn ver s hs (ResponseSource _ _ withBodyFlush) =
239+ withBodyFlush $ \_bodyFlush ->
240+ composeHeader ver s hs >>= connSendAll conn
241+sendResponseNoBody conn ver s hs _ =
242+ composeHeader ver s hs >>= connSendAll conn
221243
222244----------------------------------------------------------------
223245----------------------------------------------------------------
224246
225247-- | Use 'connSendAll' to send this data while respecting timeout rules.
226248connSink :: Connection -> T.Handle -> Sink ByteString IO ()
227−connSink Connection { connSendAll = send } th =
228− sink
249+connSink Connection { connSendAll = send } th = sink
229250 where
230251 sink = await >>= maybe close push
231252 close = liftIO (T.resume th)
@@ -315,6 +336,11 @@ addContentRange beg end total hdrs = (hContentRange, range) : hdrs
315336 ( '/'
316337 : showInt total ""))
317338
339+addDate :: D.DateCache -> H.ResponseHeaders -> IO H.ResponseHeaders
340+addDate dc hdrs = do
341+ gmtdate <- D.getDate dc
342+ return $ (H.hDate, gmtdate) : hdrs
343+
318344----------------------------------------------------------------
319345
320346-- | The version of Warp.
@@ -331,6 +357,15 @@ addServer rspidxhdr hdrs = case rspidxhdr ! idxServer of
331357
332358----------------------------------------------------------------
333359
360+-- |
361+--
362+-- >>> replaceHeader "Content-Type" "new" [("content-type","old")]
363+-- [("Content-Type","new")]
364+replaceHeader :: H.HeaderName -> HeaderValue -> H.ResponseHeaders -> H.ResponseHeaders
365+replaceHeader k v hdrs = (k,v) : deleteBy ((==) `on` fst) (k,v) hdrs
366+
367+----------------------------------------------------------------
368+
334369composeHeaderBuilder :: H.HttpVersion -> H.Status -> H.ResponseHeaders -> Bool -> IO Builder
335370composeHeaderBuilder ver s hs True =
336371 fromByteString <$> composeHeader ver s (addTransferEncoding hs)