@@ -0,0 +1,538 @@
1+{-# LANGUAGE OverloadedStrings #-}
2+{-# LANGUAGE BangPatterns, CPP #-}
3+{-# LANGUAGE NamedFieldPuns #-}
4+
5+module Network.Wai.Handler.Warp.HTTP2.Sender (frameSender) where
6+
7+import Control.Concurrent.STM
8+import qualified Control.Exception as E
9+import Control.Monad (void, when)
10+import Data.Bits
11+import qualified Data.ByteString as BS
12+import Data.ByteString.Builder (Builder)
13+import qualified Data.ByteString.Builder.Extra as B
14+import Data.IORef
15+import Data.Maybe (isNothing)
16+import Data.Word (Word8, Word32)
17+import Foreign.Ptr (Ptr, plusPtr)
18+import Foreign.Storable (poke)
19+import Network.HPACK (setLimitForEncoding, toHeaderTable)
20+import Network.HTTP2
21+import Network.HTTP2.Priority (isEmptySTM, dequeueSTM, Precedence)
22+import Network.Wai
23+import Network.Wai.Handler.Warp.Buffer
24+import Network.Wai.Handler.Warp.HTTP2.EncodeFrame
25+import Network.Wai.Handler.Warp.HTTP2.HPACK
26+import Network.Wai.Handler.Warp.HTTP2.Manager (Manager)
27+import Network.Wai.Handler.Warp.HTTP2.Types
28+import qualified Network.Wai.Handler.Warp.Settings as S
29+import Network.Wai.Handler.Warp.Types
30+
31+#ifdef WINDOWS
32+import qualified System.IO as IO
33+#else
34+import Network.Wai.Handler.Warp.FdCache
35+import Network.Wai.Handler.Warp.SendFile (positionRead)
36+import qualified Network.Wai.Handler.Warp.Timeout as T
37+#endif
38+
39+----------------------------------------------------------------
40+
41+data Leftover = LZero
42+ | LOne B.BufferWriter
43+ | LTwo BS.ByteString B.BufferWriter
44+
45+----------------------------------------------------------------
46+
47+{-# INLINE getStreamWindowSize #-}
48+getStreamWindowSize :: Stream -> IO WindowSize
49+getStreamWindowSize Stream{streamWindow} = atomically $ readTVar streamWindow
50+
51+{-# INLINE waitStreamWindowSize #-}
52+waitStreamWindowSize :: Stream -> IO ()
53+waitStreamWindowSize Stream{streamWindow} = atomically $ do
54+ w <- readTVar streamWindow
55+ check (w > 0)
56+
57+{-# INLINE waitStreaming #-}
58+waitStreaming :: TBQueue a -> IO ()
59+waitStreaming tbq = atomically $ do
60+ isEmpty <- isEmptyTBQueue tbq
61+ check (isEmpty == False)
62+
63+data Switch = C Control
64+ | O (StreamId,Precedence,Output)
65+ | Flush
66+
67+frameSender :: Context -> Connection -> S.Settings -> Manager -> IO ()
68+frameSender ctx@Context{outputQ,controlQ,connectionWindow,encodeDynamicTable}
69+ conn@Connection{connWriteBuffer,connBufferSize,connSendAll}
70+ settings mgr = loop 0 `E.catch` ignore
71+ where
72+ dequeue off = do
73+ isEmpty <- isEmptyTQueue controlQ
74+ if isEmpty then do
75+ w <- readTVar connectionWindow
76+ check (w > 0)
77+ emp <- isEmptySTM outputQ
78+ if emp then
79+ if off /= 0 then return Flush else retry
80+ else
81+ O <$> dequeueSTM outputQ
82+ else
83+ C <$> readTQueue controlQ
84+
85+ loop off = do
86+ x <- atomically $ dequeue off
87+ case x of
88+ C ctl -> do
89+ when (off /= 0) $ flushN off
90+ off' <- control ctl off
91+ when (off' >= 0) $ loop off'
92+ O (_,pre,out) -> do
93+ let strm = outputStream out
94+ writeIORef (streamPrecedence strm) pre
95+ off' <- outputOrEnqueueAgain out off
96+ case off' of
97+ 0 -> loop 0
98+ _ | off' > 15872 -> flushN off' >> loop 0 -- fixme: hard-coding
99+ | otherwise -> loop off'
100+ Flush -> flushN off >> loop 0
101+
102+ control CFinish _ = return (-1)
103+ control (CGoaway frame) _ = connSendAll frame >> return (-1)
104+ control (CFrame frame) _ = connSendAll frame >> return 0
105+ control (CSettings frame alist) _ = do
106+ connSendAll frame
107+ setLimit alist
108+ return 0
109+ control (CSettings0 frame1 frame2 alist) off = do -- off == 0, just in case
110+ let !buf = connWriteBuffer `plusPtr` off
111+ !off' = off + BS.length frame1 + BS.length frame2
112+ buf' <- copy buf frame1
113+ void $ copy buf' frame2
114+ setLimit alist
115+ return off'
116+
117+ {-# INLINE setLimit #-}
118+ setLimit alist = case lookup SettingsHeaderTableSize alist of
119+ Nothing -> return ()
120+ Just siz -> setLimitForEncoding siz encodeDynamicTable
121+
122+ output out@(Output strm _ _ tell getH2D (ONext curr)) off0 lim = do
123+ -- Data frame payload
124+ let !buf = connWriteBuffer `plusPtr` off0
125+ !siz = connBufferSize - off0
126+ Next datPayloadLen mnext <- curr buf siz lim
127+ off <- fillDataHeader strm off0 datPayloadLen mnext tell getH2D
128+ maybeEnqueueNext out mnext
129+ return off
130+
131+ output out@(Output strm rspn ii tell getH2D ORspn) off0 lim = do
132+ -- Header frame and Continuation frame
133+ let !sid = streamNumber strm
134+ !endOfStream = case rspn of
135+ RspnNobody _ _ -> True
136+ _ -> False
137+ ths <- addNecessaryHeaders ctx rspn ii settings
138+ kvlen <- headerContinue sid ths endOfStream off0
139+ off <- sendHeadersIfNecessary $ off0 + frameHeaderLength + kvlen
140+ case rspn of
141+ RspnNobody _ _ -> do
142+ closed ctx strm Finished
143+ return off
144+ RspnFile _ _ path mpart -> do
145+ -- Data frame payload
146+ let payloadOff = off + frameHeaderLength
147+ Next datPayloadLen mnext <-
148+ fillFileBodyGetNext conn ii payloadOff lim path mpart
149+ off' <- fillDataHeader strm off datPayloadLen mnext tell getH2D
150+ maybeEnqueueNext out mnext
151+ return off'
152+ RspnBuilder _ _ builder -> do
153+ -- Data frame payload
154+ let payloadOff = off + frameHeaderLength
155+ Next datPayloadLen mnext <-
156+ fillBuilderBodyGetNext conn ii payloadOff lim builder
157+ off' <- fillDataHeader strm off datPayloadLen mnext tell getH2D
158+ maybeEnqueueNext out mnext
159+ return off'
160+ RspnStreaming _ _ tbq -> do
161+ let payloadOff = off + frameHeaderLength
162+ Next datPayloadLen mnext <-
163+ fillStreamBodyGetNext conn payloadOff lim tbq strm
164+ off' <- fillDataHeader strm off datPayloadLen mnext tell getH2D
165+ maybeEnqueueNext out mnext
166+ return off'
167+
168+ output out@(Output strm _ _ _ _ (OPush ths pid)) off0 lim = do
169+ -- Creating a push promise header
170+ -- Frame id should be associated stream id from the client.
171+ let !sid = streamNumber strm
172+ len <- pushPromise pid sid ths off0
173+ off <- sendHeadersIfNecessary $ off0 + frameHeaderLength + len
174+ output out{ outputType = ORspn } off lim
175+
176+ output _ _ _ = undefined -- never reach
177+
178+ outputOrEnqueueAgain out off = E.handle resetStream $ do
179+ state <- readIORef $ streamState strm
180+ if isClosed state then
181+ return off
182+ else case out of
183+ Output _ _ _ wait _ OWait -> do
184+ -- Checking if all push are done.
185+ let out' = out {
186+ outputHook = return ()
187+ , outputType = ORspn
188+ }
189+ forkAndEnqueueWhenReady wait outputQ out' mgr
190+ return off
191+ _ -> case mtbq of
192+ Just tbq -> checkStreaming tbq
193+ _ -> checkStreamWindowSize
194+ where
195+ strm = outputStream out
196+ mtbq = outputMaybeTBQueue out
197+ checkStreaming tbq = do
198+ isEmpty <- atomically $ isEmptyTBQueue tbq
199+ if isEmpty then do
200+ forkAndEnqueueWhenReady (waitStreaming tbq) outputQ out mgr
201+ return off
202+ else
203+ checkStreamWindowSize
204+ checkStreamWindowSize = do
205+ sws <- getStreamWindowSize strm
206+ if sws == 0 then do
207+ forkAndEnqueueWhenReady (waitStreamWindowSize strm) outputQ out mgr
208+ return off
209+ else do
210+ cws <- atomically $ readTVar connectionWindow -- not 0
211+ let !lim = min cws sws
212+ output out off lim
213+ resetStream e = do
214+ closed ctx strm (ResetByMe e)
215+ let !rst = resetFrame InternalError $ streamNumber strm
216+ enqueueControl controlQ $ CFrame rst
217+ return off
218+
219+ {-# INLINE flushN #-}
220+ -- Flush the connection buffer to the socket, where the first 'n' bytes of
221+ -- the buffer are filled.
222+ flushN :: Int -> IO ()
223+ flushN n = bufferIO connWriteBuffer n connSendAll
224+
225+ headerContinue sid ths endOfStream off = do
226+ let !offkv = off + frameHeaderLength
227+ let !bufkv = connWriteBuffer `plusPtr` offkv
228+ !limkv = connBufferSize - offkv
229+ (hs,kvlen) <- hpackEncodeHeader ctx bufkv limkv ths
230+ let flag0 = case hs of
231+ [] -> setEndHeader defaultFlags
232+ _ -> defaultFlags
233+ flag = if endOfStream then setEndStream flag0 else flag0
234+ let buf = connWriteBuffer `plusPtr` off
235+ fillFrameHeader FrameHeaders kvlen sid flag buf
236+ continue sid kvlen hs
237+
238+ !bufHeaderPayload = connWriteBuffer `plusPtr` frameHeaderLength
239+ !headerPayloadLim = connBufferSize - frameHeaderLength
240+
241+ continue _ kvlen [] = return kvlen
242+ continue sid kvlen ths = do
243+ flushN $ kvlen + frameHeaderLength
244+ -- Now off is 0
245+ (ths', kvlen') <- hpackEncodeHeaderLoop ctx bufHeaderPayload headerPayloadLim ths
246+ when (ths == ths') $ E.throwIO $ ConnectionError CompressionError "cannot compress the header"
247+ let flag = case ths' of
248+ [] -> setEndHeader defaultFlags
249+ _ -> defaultFlags
250+ fillFrameHeader FrameContinuation kvlen' sid flag connWriteBuffer
251+ continue sid kvlen' ths'
252+
253+ {-# INLINE maybeEnqueueNext #-}
254+ -- Re-enqueue the stream in the output queue.
255+ maybeEnqueueNext :: Output -> Maybe DynaNext -> IO ()
256+ maybeEnqueueNext _ Nothing = return ()
257+ maybeEnqueueNext out (Just next) = enqueueOutput outputQ out'
258+ where
259+ !out' = out { outputType = ONext next }
260+
261+ {-# INLINE sendHeadersIfNecessary #-}
262+ -- Send headers if there is not room for a 1-byte data frame, and return
263+ -- the offset of the next frame's first header byte.
264+ sendHeadersIfNecessary off
265+ -- True if the connection buffer has room for a 1-byte data frame.
266+ | off + frameHeaderLength < connBufferSize = return off
267+ | otherwise = do
268+ flushN off
269+ return 0
270+
271+ fillDataHeader strm@Stream{streamWindow,streamNumber}
272+ off datPayloadLen mnext tell getH2D = do
273+ -- Data frame header
274+ mh2d <- getH2D
275+ let (!trailers,!noTrailers) = case http2dataTrailers <$> mh2d of
276+ Nothing -> ([], True)
277+ Just ts -> (ts, null ts)
278+ !buf = connWriteBuffer `plusPtr` off
279+ !off' = off + frameHeaderLength + datPayloadLen
280+ !noMoreBody = isNothing mnext
281+ flag | noMoreBody && noTrailers = setEndStream defaultFlags
282+ | otherwise = defaultFlags
283+ fillFrameHeader FrameData datPayloadLen streamNumber flag buf
284+ off'' <- handleEndOfBody noMoreBody off' noTrailers trailers
285+ atomically $ modifyTVar' connectionWindow (subtract datPayloadLen)
286+ atomically $ modifyTVar' streamWindow (subtract datPayloadLen)
287+ return off''
288+ where
289+ handleTrailers True off0 _ = return off0
290+ handleTrailers _ off0 trailers = do
291+ (ths,_) <- toHeaderTable trailers
292+ kvlen <- headerContinue streamNumber ths True off0
293+ sendHeadersIfNecessary $ off0 + frameHeaderLength + kvlen
294+ handleEndOfBody True off0 noTrailers trailers = do
295+ off1 <- handleTrailers noTrailers off0 trailers
296+ void $ tell
297+ closed ctx strm Finished
298+ return off1
299+ handleEndOfBody False off0 _ _ = return off0
300+
301+
302+ pushPromise pid sid ths off = do
303+ let !offsid = off + frameHeaderLength
304+ !bufsid = connWriteBuffer `plusPtr` offsid
305+ poke32 bufsid $ fromIntegral sid
306+ let !offkv = offsid + 4
307+ !bufkv = connWriteBuffer `plusPtr` offkv
308+ !limkv = connBufferSize - offkv
309+ (_,kvlen) <- hpackEncodeHeader ctx bufkv limkv ths
310+ let !flag = setEndHeader defaultFlags -- No EndStream flag
311+ !buf = connWriteBuffer `plusPtr` off
312+ !len = kvlen + 4
313+ fillFrameHeader FramePushPromise len pid flag buf
314+ return len
315+
316+ {-# INLINE fillFrameHeader #-}
317+ fillFrameHeader ftyp len sid flag buf = encodeFrameHeaderBuf ftyp hinfo buf
318+ where
319+ hinfo = FrameHeader len flag sid
320+
321+ {-# INLINE ignore #-}
322+ ignore :: E.SomeException -> IO ()
323+ ignore _ = return ()
324+
325+----------------------------------------------------------------
326+
327+{-
328+ResponseFile Status ResponseHeaders FilePath (Maybe FilePart)
329+ResponseBuilder Status ResponseHeaders Builder
330+ResponseStream Status ResponseHeaders StreamingBody
331+ResponseRaw (IO ByteString -> (ByteString -> IO ()) -> IO ()) Response
332+-}
333+
334+fillBuilderBodyGetNext :: Connection -> InternalInfo -> Int -> WindowSize -> Builder -> IO Next
335+fillBuilderBodyGetNext Connection{connWriteBuffer,connBufferSize}
336+ _ off lim bb = do
337+ let datBuf = connWriteBuffer `plusPtr` off
338+ room = min (connBufferSize - off) lim
339+ (len, signal) <- B.runBuilder bb datBuf room
340+ return $ nextForBuilder len signal
341+
342+fillFileBodyGetNext :: Connection -> InternalInfo -> Int -> WindowSize -> FilePath -> Maybe FilePart -> IO Next
343+#ifdef WINDOWS
344+fillFileBodyGetNext Connection{connWriteBuffer,connBufferSize}
345+ _ off lim path mpart = do
346+ let datBuf = connWriteBuffer `plusPtr` off
347+ room = min (connBufferSize - off) lim
348+ (start, bytes) <- fileStartEnd path mpart
349+ -- fixme: how to close Handle? GC does it at this moment.
350+ hdl <- IO.openBinaryFile path IO.ReadMode
351+ IO.hSeek hdl IO.AbsoluteSeek start
352+ len <- IO.hGetBufSome hdl datBuf (mini room bytes)
353+ let bytes' = bytes - fromIntegral len
354+ -- fixme: connWriteBuffer connBufferSize
355+ return $ nextForFile len hdl bytes' (return ())
356+#else
357+fillFileBodyGetNext Connection{connWriteBuffer,connBufferSize}
358+ ii off lim path mpart = do
359+ (mfd, refresh') <- getFd ii path
360+ (fd, refresh) <- case mfd of
361+ Nothing -> do
362+ fd' <- openFile path
363+ th <- T.register (timeoutManager ii) (closeFile fd')
364+ return (fd', T.tickle th)
365+ Just fd -> return (fd, refresh')
366+ let datBuf = connWriteBuffer `plusPtr` off
367+ room = min (connBufferSize - off) lim
368+ (start, bytes) <- fileStartEnd path mpart
369+ len <- positionRead fd datBuf (mini room bytes) start
370+ refresh
371+ let len' = fromIntegral len
372+ return $ nextForFile len fd (start + len') (bytes - len') refresh
373+#endif
374+
375+fileStartEnd :: FilePath -> Maybe FilePart -> IO (Integer, Integer)
376+fileStartEnd _ (Just part) =
377+ return (filePartOffset part, filePartByteCount part)
378+fileStartEnd _ _ = error "fileStartEnd"
379+
380+----------------------------------------------------------------
381+
382+fillStreamBodyGetNext :: Connection -> Int -> WindowSize -> TBQueue Sequence -> Stream -> IO Next
383+fillStreamBodyGetNext Connection{connWriteBuffer,connBufferSize}
384+ off lim sq strm = do
385+ let datBuf = connWriteBuffer `plusPtr` off
386+ room = min (connBufferSize - off) lim
387+ (leftover, cont, len) <- runStreamBuilder datBuf room sq
388+ return $ nextForStream sq strm leftover cont len
389+
390+----------------------------------------------------------------
391+
392+fillBufBuilder :: Leftover -> DynaNext
393+fillBufBuilder leftover buf0 siz0 lim = do
394+ let payloadBuf = buf0 `plusPtr` frameHeaderLength
395+ room = min (siz0 - frameHeaderLength) lim
396+ case leftover of
397+ LZero -> error "fillBufBuilder: LZero"
398+ LOne writer -> do
399+ (len, signal) <- writer payloadBuf room
400+ getNext len signal
401+ LTwo bs writer
402+ | BS.length bs <= room -> do
403+ buf1 <- copy payloadBuf bs
404+ let len1 = BS.length bs
405+ (len2, signal) <- writer buf1 (room - len1)
406+ getNext (len1 + len2) signal
407+ | otherwise -> do
408+ let (bs1,bs2) = BS.splitAt room bs
409+ void $ copy payloadBuf bs1
410+ getNext room (B.Chunk bs2 writer)
411+ where
412+ getNext l s = return $ nextForBuilder l s
413+
414+nextForBuilder :: BytesFilled -> B.Next -> Next
415+nextForBuilder len B.Done
416+ = Next len Nothing
417+nextForBuilder len (B.More _ writer)
418+ = Next len $ Just (fillBufBuilder (LOne writer))
419+nextForBuilder len (B.Chunk bs writer)
420+ = Next len $ Just (fillBufBuilder (LTwo bs writer))
421+
422+----------------------------------------------------------------
423+
424+runStreamBuilder :: Buffer -> BufSize -> TBQueue Sequence
425+ -> IO (Leftover, Bool, BytesFilled)
426+runStreamBuilder buf0 room0 sq = loop buf0 room0 0
427+ where
428+ loop !buf !room !total = do
429+ mbuilder <- atomically $ tryReadTBQueue sq
430+ case mbuilder of
431+ Nothing -> return (LZero, True, total)
432+ Just (SBuilder builder) -> do
433+ (len, signal) <- B.runBuilder builder buf room
434+ let !total' = total + len
435+ case signal of
436+ B.Done -> loop (buf `plusPtr` len) (room - len) total'
437+ B.More _ writer -> return (LOne writer, True, total')
438+ B.Chunk bs writer -> return (LTwo bs writer, True, total')
439+ Just SFlush -> return (LZero, True, total)
440+ Just SFinish -> return (LZero, False, total)
441+
442+fillBufStream :: Leftover -> TBQueue Sequence -> Stream -> DynaNext
443+fillBufStream leftover0 sq strm buf0 siz0 lim0 = do
444+ let payloadBuf = buf0 `plusPtr` frameHeaderLength
445+ room0 = min (siz0 - frameHeaderLength) lim0
446+ case leftover0 of
447+ LZero -> do
448+ (leftover, cont, len) <- runStreamBuilder payloadBuf room0 sq
449+ getNext leftover cont len
450+ LOne writer -> write writer payloadBuf room0 0
451+ LTwo bs writer
452+ | BS.length bs <= room0 -> do
453+ buf1 <- copy payloadBuf bs
454+ let len = BS.length bs
455+ write writer buf1 (room0 - len) len
456+ | otherwise -> do
457+ let (bs1,bs2) = BS.splitAt room0 bs
458+ void $ copy payloadBuf bs1
459+ getNext (LTwo bs2 writer) True room0
460+ where
461+ getNext l b r = return $ nextForStream sq strm l b r
462+ write writer1 buf room sofar = do
463+ (len, signal) <- writer1 buf room
464+ case signal of
465+ B.Done -> do
466+ (leftover, cont, extra) <- runStreamBuilder (buf `plusPtr` len) (room - len) sq
467+ let !total = sofar + len + extra
468+ getNext leftover cont total
469+ B.More _ writer -> do
470+ let !total = sofar + len
471+ getNext (LOne writer) True total
472+ B.Chunk bs writer -> do
473+ let !total = sofar + len
474+ getNext (LTwo bs writer) True total
475+
476+nextForStream :: TBQueue Sequence -> Stream
477+ -> Leftover -> Bool -> BytesFilled
478+ -> Next
479+nextForStream _ _ _ False len = Next len Nothing
480+nextForStream sq strm leftOrZero True len =
481+ Next len $ Just (fillBufStream leftOrZero sq strm)
482+
483+----------------------------------------------------------------
484+
485+#ifdef WINDOWS
486+fillBufFile :: IO.Handle -> Integer -> IO () -> DynaNext
487+fillBufFile h bytes refresh buf siz lim = do
488+ let payloadBuf = buf `plusPtr` frameHeaderLength
489+ room = min (siz - frameHeaderLength) lim
490+ len <- IO.hGetBufSome h payloadBuf room
491+ refresh
492+ let bytes' = bytes - fromIntegral len
493+ return $ nextForFile len h bytes' refresh
494+
495+nextForFile :: BytesFilled -> IO.Handle -> Integer -> IO () -> Next
496+nextForFile 0 _ _ _ = Next 0 Nothing
497+nextForFile len _ 0 _ = Next len Nothing
498+nextForFile len h bytes refresh =
499+ Next len $ Just (fillBufFile h bytes refresh)
500+#else
501+fillBufFile :: Fd -> Integer -> Integer -> IO () -> DynaNext
502+fillBufFile fd start bytes refresh buf siz lim = do
503+ let payloadBuf = buf `plusPtr` frameHeaderLength
504+ room = min (siz - frameHeaderLength) lim
505+ len <- positionRead fd payloadBuf (mini room bytes) start
506+ let len' = fromIntegral len
507+ refresh
508+ return $ nextForFile len fd (start + len') (bytes - len') refresh
509+
510+nextForFile :: BytesFilled -> Fd -> Integer -> Integer -> IO () -> Next
511+nextForFile 0 _ _ _ _ = Next 0 Nothing
512+nextForFile len _ _ 0 _ = Next len Nothing
513+nextForFile len fd start bytes refresh =
514+ Next len $ Just (fillBufFile fd start bytes refresh)
515+#endif
516+
517+{-# INLINE mini #-}
518+mini :: Int -> Integer -> Int
519+mini i n
520+ | fromIntegral i < n = i
521+ | otherwise = fromIntegral n
522+
523+
524+----------------------------------------------------------------
525+
526+poke32 :: Ptr Word8 -> Word32 -> IO ()
527+poke32 ptr i = do
528+ poke ptr w0
529+ poke8 ptr 1 w1
530+ poke8 ptr 2 w2
531+ poke8 ptr 3 w3
532+ where
533+ w0 = fromIntegral ((i `shiftR` 24) .&. 0xff)
534+ w1 = fromIntegral ((i `shiftR` 16) .&. 0xff)
535+ w2 = fromIntegral ((i `shiftR` 8) .&. 0xff)
536+ w3 = fromIntegral (i .&. 0xff)
537+ poke8 :: Ptr Word8 -> Int -> Word8 -> IO ()
538+ poke8 ptr0 n w = poke (ptr0 `plusPtr` n) w