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