@@ -10,11 +10,14 @@ import Control.Applicative
1010import Control.Concurrent.STM
1111import qualified Control.Exception as E
1212import Control.Monad (void, when)
13+import Data.Bits
1314import qualified Data.ByteString as BS
1415import Data.ByteString.Builder (Builder)
1516import qualified Data.ByteString.Builder.Extra as B
1617import Data.Maybe (isNothing)
17−import Foreign.Ptr
18+import Data.Word (Word8, Word32)
19+import Foreign.Ptr (Ptr, plusPtr)
20+import Foreign.Storable (poke)
1821import Network.HPACK (setLimitForEncoding)
1922import Network.HTTP2
2023import Network.HTTP2.Priority (isEmptySTM, dequeueSTM, Precedence)
@@ -49,14 +52,14 @@ getStreamWindowSize :: Stream -> IO WindowSize
4952getStreamWindowSize Stream{streamWindow} = atomically $ readTVar streamWindow
5053
5154{-# INLINE waitStreamWindowSize #-}
52−waitStreamWindowSize :: Stream -> STM ()
53−waitStreamWindowSize Stream{streamWindow} = do
55+waitStreamWindowSize :: Stream -> IO ()
56+waitStreamWindowSize Stream{streamWindow} = atomically $ do
5457 w <- readTVar streamWindow
5558 check (w > 0)
5659
5760{-# INLINE waitStreaming #-}
58−waitStreaming :: TBQueue a -> STM ()
59−waitStreaming tbq = do
61+waitStreaming :: TBQueue a -> IO ()
62+waitStreaming tbq = atomically $ do
6063 isEmpty <- isEmptyTBQueue tbq
6164 check (isEmpty == False)
6265
@@ -119,22 +122,23 @@ frameSender ctx@Context{outputQ,controlQ,connectionWindow,encodeDynamicTable}
119122 Nothing -> return ()
120123 Just siz -> setLimitForEncoding siz encodeDynamicTable
121124
122− output (ONext strm curr mtbq) off0 lim = do
125+ output (ONext strm curr mtbq tell) off0 lim = do
123126 -- Data frame payload
124127 let !buf = connWriteBuffer `plusPtr` off0
125128 !siz = connBufferSize - off0
126129 Next datPayloadLen mnext <- curr buf siz lim
127− off <- fillDataHeader strm off0 datPayloadLen mnext
128− maybeEnqueueNext strm mtbq mnext
130+ off <- fillDataHeader strm off0 datPayloadLen mnext tell
131+ maybeEnqueueNext strm mtbq mnext tell
129132 return off
130133
131− output (ORspn strm rspn ii) off0 lim = do
134+ output (ORspn strm rspn ii tell) off0 lim = do
132135 -- Header frame and Continuation frame
133− let sid = streamNumber strm
134− endOfStream = case rspn of
136+ let !sid = streamNumber strm
137+ !endOfStream = case rspn of
135138 RspnNobody _ _ -> True
136139 _ -> False
137− kvlen <- headerContinue sid rspn endOfStream off0 ii
140+ ths <- addNecessaryHeaders ctx rspn ii settings
141+ kvlen <- headerContinue sid ths endOfStream off0
138142 off <- sendHeadersIfNecessary $ off0 + frameHeaderLength + kvlen
139143 case rspn of
140144 RspnNobody _ _ -> do
@@ -145,32 +149,48 @@ frameSender ctx@Context{outputQ,controlQ,connectionWindow,encodeDynamicTable}
145149 let payloadOff = off + frameHeaderLength
146150 Next datPayloadLen mnext <-
147151 fillFileBodyGetNext conn ii payloadOff lim path mpart
148− off' <- fillDataHeader strm off datPayloadLen mnext
149− maybeEnqueueNext strm Nothing mnext
152+ off' <- fillDataHeader strm off datPayloadLen mnext tell
153+ maybeEnqueueNext strm Nothing mnext tell
150154 return off'
151155 RspnBuilder _ _ builder -> do
152156 -- Data frame payload
153157 let payloadOff = off + frameHeaderLength
154158 Next datPayloadLen mnext <-
155159 fillBuilderBodyGetNext conn ii payloadOff lim builder
156− off' <- fillDataHeader strm off datPayloadLen mnext
157− maybeEnqueueNext strm Nothing mnext
160+ off' <- fillDataHeader strm off datPayloadLen mnext tell
161+ maybeEnqueueNext strm Nothing mnext tell
158162 return off'
159163 RspnStreaming _ _ tbq -> do
160164 let payloadOff = off + frameHeaderLength
161165 Next datPayloadLen mnext <-
162166 fillStreamBodyGetNext conn payloadOff lim tbq strm
163− off' <- fillDataHeader strm off datPayloadLen mnext
164− maybeEnqueueNext strm (Just tbq) mnext
167+ off' <- fillDataHeader strm off datPayloadLen mnext tell
168+ maybeEnqueueNext strm (Just tbq) mnext tell
165169 return off'
166170
171+ output (OPush strm ths rspn ii tell 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 (ORspn strm rspn ii tell) off lim
178+
179+ output _ _ _ = undefined -- never reach
180+
167181 outputOrEnqueueAgain out off = E.handle resetStream $ do
168182 state <- readIORef $ streamState strm
169183 if isClosed state then
170184 return off
171− else case mtbq of
172− Just tbq -> checkStreaming tbq
173− _ -> checkStreamWindowSize
185+ else case out of
186+ OWait strm' rsp ii wait -> do
187+ -- Checking if all push are done.
188+ let out' = ORspn strm' rsp ii (return ())
189+ forkAndEnqueueWhenReady wait outputQ out' mgr
190+ return off
191+ _ -> case mtbq of
192+ Just tbq -> checkStreaming tbq
193+ _ -> checkStreamWindowSize
174194 where
175195 strm = outputStream out
176196 mtbq = outputMaybeTBQueue out
@@ -202,13 +222,11 @@ frameSender ctx@Context{outputQ,controlQ,connectionWindow,encodeDynamicTable}
202222 flushN :: Int -> IO ()
203223 flushN n = bufferIO connWriteBuffer n connSendAll
204224
205− headerContinue sid rspn endOfStream off ii = do
206− let !s = rspnStatus rspn
207− !h = rspnHeaders rspn
225+ headerContinue sid ths endOfStream off = do
208226 let !offkv = off + frameHeaderLength
209227 let !bufkv = connWriteBuffer `plusPtr` offkv
210228 !limkv = connBufferSize - offkv
211− (hs,kvlen) <- hpackEncodeHeader ctx bufkv limkv ii settings s h
229+ (hs,kvlen) <- hpackEncodeHeader ctx bufkv limkv ths
212230 let flag0 = case hs of
213231 [] -> setEndHeader defaultFlags
214232 _ -> defaultFlags
@@ -221,24 +239,25 @@ frameSender ctx@Context{outputQ,controlQ,connectionWindow,encodeDynamicTable}
221239 !headerPayloadLim = connBufferSize - frameHeaderLength
222240
223241 continue _ kvlen [] = return kvlen
224− continue sid kvlen hs = do
242+ continue sid kvlen ths = do
225243 flushN $ kvlen + frameHeaderLength
226244 -- Now off is 0
227− (hs', kvlen') <- hpackEncodeHeaderLoop ctx bufHeaderPayload headerPayloadLim hs
228− when (hs == hs') $ E.throwIO $ ConnectionError CompressionError "cannot compress the header"
229− let flag = case hs' of
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
230248 [] -> setEndHeader defaultFlags
231249 _ -> defaultFlags
232250 fillFrameHeader FrameContinuation kvlen' sid flag connWriteBuffer
233− continue sid kvlen' hs'
251+ continue sid kvlen' ths'
234252
235253 {-# INLINE maybeEnqueueNext #-}
236254 -- Re-enqueue the stream in the output queue.
237− maybeEnqueueNext :: Stream -> Maybe (TBQueue Sequence) -> Maybe DynaNext -> IO ()
238− maybeEnqueueNext _ _ Nothing = return ()
239− maybeEnqueueNext strm mtbq (Just next) = enqueueOutput outputQ out
255+ maybeEnqueueNext :: Stream -> Maybe (TBQueue Sequence)
256+ -> Maybe DynaNext -> IO () -> IO ()
257+ maybeEnqueueNext _ _ Nothing _ = return ()
258+ maybeEnqueueNext strm mtbq (Just next) tell = enqueueOutput outputQ out
240259 where
241− !out = ONext strm next mtbq
260+ !out = ONext strm next mtbq tell
242261
243262 {-# INLINE sendHeadersIfNecessary #-}
244263 -- Send headers if there is not room for a 1-byte data frame, and return
@@ -250,7 +269,7 @@ frameSender ctx@Context{outputQ,controlQ,connectionWindow,encodeDynamicTable}
250269 flushN off
251270 return 0
252271
253− fillDataHeader strm off datPayloadLen mnext = do
272+ fillDataHeader strm off datPayloadLen mnext tell = do
254273 -- Data frame header
255274 let !sid = streamNumber strm
256275 !buf = connWriteBuffer `plusPtr` off
@@ -259,11 +278,27 @@ frameSender ctx@Context{outputQ,controlQ,connectionWindow,encodeDynamicTable}
259278 flag | done = setEndStream defaultFlags
260279 | otherwise = defaultFlags
261280 fillFrameHeader FrameData datPayloadLen sid flag buf
262− when done $ closed ctx strm Finished
281+ when done $ do
282+ void $ tell
283+ closed ctx strm Finished
263284 atomically $ modifyTVar' connectionWindow (subtract datPayloadLen)
264285 atomically $ modifyTVar' (streamWindow strm) (subtract datPayloadLen)
265286 return off'
266287
288+ pushPromise pid sid ths off = do
289+ let !offsid = off + frameHeaderLength
290+ !bufsid = connWriteBuffer `plusPtr` offsid
291+ poke32 bufsid $ fromIntegral sid
292+ let !offkv = offsid + 4
293+ !bufkv = connWriteBuffer `plusPtr` offkv
294+ !limkv = connBufferSize - offkv
295+ (_,kvlen) <- hpackEncodeHeader ctx bufkv limkv ths
296+ let !flag = setEndHeader defaultFlags -- No EndStream flag
297+ !buf = connWriteBuffer `plusPtr` off
298+ !len = kvlen + 4
299+ fillFrameHeader FramePushPromise len pid flag buf
300+ return len
301+
267302 {-# INLINE fillFrameHeader #-}
268303 fillFrameHeader ftyp len sid flag buf = encodeFrameHeaderBuf ftyp hinfo buf
269304 where
@@ -470,3 +505,20 @@ mini :: Int -> Integer -> Int
470505mini i n
471506 | fromIntegral i < n = i
472507 | otherwise = fromIntegral n
508+
509+
510+----------------------------------------------------------------
511+
512+poke32 :: Ptr Word8 -> Word32 -> IO ()
513+poke32 ptr i = do
514+ poke ptr w0
515+ poke8 ptr 1 w1
516+ poke8 ptr 2 w2
517+ poke8 ptr 3 w3
518+ where
519+ w0 = fromIntegral ((i `shiftR` 24) .&. 0xff)
520+ w1 = fromIntegral ((i `shiftR` 16) .&. 0xff)
521+ w2 = fromIntegral ((i `shiftR` 8) .&. 0xff)
522+ w3 = fromIntegral (i .&. 0xff)
523+ poke8 :: Ptr Word8 -> Int -> Word8 -> IO ()
524+ poke8 ptr0 n w = poke (ptr0 `plusPtr` n) w