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