content: {-# LANGUAGE LambdaCase #-}
  {-# LANGUAGE NamedFieldPuns #-}
  {-# LANGUAGE OverloadedStrings #-}
  {-# LANGUAGE RecordWildCards #-}
  {-# LANGUAGE ScopedTypeVariables #-}
  
  module Network.HTTP2.H2.Sender (
      frameSender,
  ) where
  
  import Control.Concurrent.STM
  import qualified Control.Exception as E
  import Data.IORef (modifyIORef', readIORef, writeIORef)
  import Data.IntMap.Strict (IntMap)
  import Foreign.Ptr (minusPtr, plusPtr)
  import Network.ByteOrder
  import Network.HTTP.Semantics.Client
  import Network.HTTP.Semantics.IO
  
  import Imports
  import Network.HPACK (setLimitForEncoding, toTokenHeaderTable)
  import Network.HTTP2.Frame
  import Network.HTTP2.H2.Context
  import Network.HTTP2.H2.EncodeFrame
  import Network.HTTP2.H2.HPACK
  import Network.HTTP2.H2.Queue
  import Network.HTTP2.H2.Settings
  import Network.HTTP2.H2.Stream
  import Network.HTTP2.H2.StreamTable
  import Network.HTTP2.H2.Types
  import Network.HTTP2.H2.Window
  
  ----------------------------------------------------------------
  
  data Switch
      = C Control
      | O Output
      | Flush
  
  -- Peer SETTINGS_INITIAL_WINDOW_SIZE
  -- Adjusting initial window size for streams
  updatePeerSettings :: Context -> SettingsList -> IO ()
  updatePeerSettings Context{peerSettings, oddStreamTable, evenStreamTable} peerAlist = do
      oldws <- initialWindowSize <$> readIORef peerSettings
      modifyIORef' peerSettings $ \old -> fromSettingsList old peerAlist
      newws <- initialWindowSize <$> readIORef peerSettings
      -- FIXME: race condition
      -- 1) newOddStream reads old peerSettings and
      --    insert it to its stream table after adjusting.
      -- 2) newOddStream reads new peerSettings and
      --    insert it to its stream table before adjusting.
      let dif = newws - oldws
      when (dif /= 0) $ do
          getOddStreams oddStreamTable >>= updateAllStreamTxFlow dif
          getEvenStreams evenStreamTable >>= updateAllStreamTxFlow dif
    where
      updateAllStreamTxFlow :: WindowSize -> IntMap Stream -> IO ()
      updateAllStreamTxFlow siz strms =
          forM_ strms $ \strm -> increaseStreamWindowSize strm siz
  
  frameSender :: Context -> Config -> IO ()
  frameSender
      ctx@Context{outputQ, controlQ, encodeDynamicTable, outputBufferLimit}
      Config{..} = do
          labelMe "H2 sender"
          loop 0
        where
          ----------------------------------------------------------------
          loop :: Offset -> IO ()
          loop off = do
              x <- atomically $ dequeue off
              case x of
                  C ctl -> flushN off >> control ctl >> loop 0
                  O out -> outputAndSync out off >>= flushIfNecessary >>= loop
                  Flush -> flushN off >> loop 0
  
          -- Flush the connection buffer to the socket, where the first 'n' bytes of
          -- the buffer are filled.
          flushN :: Offset -> IO ()
          flushN 0 = return ()
          flushN n = bufferIO confWriteBuffer n confSendAll
  
          flushIfNecessary :: Offset -> IO Offset
          flushIfNecessary off = do
              buflim <- readIORef outputBufferLimit
              if off <= buflim - 512
                  then return off
                  else do
                      flushN off
                      return 0
  
          dequeue :: Offset -> STM Switch
          dequeue off = do
              isEmptyC <- isEmptyTQueue controlQ
              if isEmptyC
                  then do
                      -- FLOW CONTROL: WINDOW_UPDATE 0: send: respecting peer's limit
                      waitConnectionWindowSize ctx
                      isEmptyO <- isEmptyTQueue outputQ
                      if isEmptyO
                          then if off /= 0 then return Flush else retry
                          else O <$> readTQueue outputQ
                  else C <$> readTQueue controlQ
  
          ----------------------------------------------------------------
          copyAll [] buf = return buf
          copyAll (x : xs) buf = copy buf x >>= copyAll xs
  
          -- called with off == 0
          control :: Control -> IO ()
          control (CFrames ms xs) = do
              buf <- copyAll xs confWriteBuffer
              let off = buf `minusPtr` confWriteBuffer
              flushN off
              case ms of
                  Nothing -> return ()
                  Just peerAlist -> do
                      -- Peer SETTINGS_INITIAL_WINDOW_SIZE
                      updatePeerSettings ctx peerAlist
                      -- Peer SETTINGS_MAX_FRAME_SIZE
                      case lookup SettingsMaxFrameSize peerAlist of
                          Nothing -> return ()
                          Just payloadLen -> do
                              let dlim = payloadLen + frameHeaderLength
                                  buflim
                                      | confBufferSize >= dlim = dlim
                                      | otherwise = confBufferSize
                              writeIORef outputBufferLimit buflim
                      -- Peer SETTINGS_HEADER_TABLE_SIZE
                      case lookup SettingsTokenHeaderTableSize peerAlist of
                          Nothing -> return ()
                          Just siz -> setLimitForEncoding siz encodeDynamicTable
  
          ----------------------------------------------------------------
          -- INVARIANT
          --
          -- Both the stream window and the connection window are open.
          ----------------------------------------------------------------
          outputAndSync :: Output -> Offset -> IO Offset
          outputAndSync out@(Output strm otyp sync) off = E.handle (\e -> resetStream strm InternalError e >> return off) $ do
              state <- readStreamState strm
              if isHalfClosedLocal state
                  then return off
                  else case otyp of
                      OHeader hdr mnext tlrmkr -> do
                          (off', mout') <- outputHeader strm hdr mnext tlrmkr sync off
                          sync mout'
                          return off'
                      _ -> do
                          sws <- getStreamWindowSize strm
                          cws <- getConnectionWindowSize ctx -- not 0
                          let lim = min cws sws
                          (off', mout') <- output out off lim
                          sync mout'
                          return off'
  
          resetStream :: Stream -> ErrorCode -> E.SomeException -> IO ()
          resetStream strm err e
              | isAsyncException e = E.throwIO e
              | otherwise = do
                  closed ctx strm (ResetByMe e)
                  let rst = resetFrame err $ streamNumber strm
                  enqueueControl controlQ $ CFrames Nothing [rst]
  
          ----------------------------------------------------------------
          outputHeader
              :: Stream
              -> [Header]
              -> Maybe DynaNext
              -> TrailersMaker
              -> (Maybe Output -> IO ())
              -> Offset
              -> IO (Offset, Maybe Output)
          outputHeader strm hdr mnext tlrmkr sync off0 = do
              -- Header frame and Continuation frame
              let sid = streamNumber strm
                  endOfStream = isNothing mnext
              (ths, _) <- toTokenHeaderTable $ fixHeaders hdr
              off' <- headerContinue sid ths endOfStream off0
              -- halfClosedLocal calls closed which removes
              -- the stream from stream table.
              off <- flushIfNecessary off'
              case mnext of
                  Nothing -> do
                      -- endOfStream
                      halfClosedLocal ctx strm Finished
                      return (off, Nothing)
                  Just next -> do
                      let out' = Output strm (ONext next tlrmkr) sync
                      return (off, Just out')
  
          ----------------------------------------------------------------
          output :: Output -> Offset -> WindowSize -> IO (Offset, Maybe Output)
          output out@(Output strm (ONext curr tlrmkr) _) off0 lim = do
              -- Data frame payload
              buflim <- readIORef outputBufferLimit
              let payloadOff = off0 + frameHeaderLength
                  datBuf = confWriteBuffer `plusPtr` payloadOff
                  datBufSiz = buflim - payloadOff
              curr datBuf (min datBufSiz lim) >>= \case
                  Next datPayloadLen reqflush mnext -> do
                      NextTrailersMaker tlrmkr' <- runTrailersMaker tlrmkr datBuf datPayloadLen
                      fillDataHeader
                          strm
                          off0
                          datPayloadLen
                          mnext
                          tlrmkr'
                          out
                          reqflush
                  CancelNext mErr -> do
                      -- Stream cancelled
                      --
                      -- At this point, the headers have already been sent.
                      -- Therefore, the stream cannot be in the 'Idle' state, so we
                      -- are justified in sending @RST_STREAM@.
                      --
                      -- By the invariant on the 'outputQ', there are no other
                      -- outputs for this stream already enqueued. Therefore, we can
                      -- safely cancel it knowing that we won't try and send any
                      -- more data frames on this stream.
                      case mErr of
                          Just err ->
                              resetStream strm InternalError err
                          Nothing ->
                              resetStream strm Cancel (E.toException CancelledStream)
                      return (off0, Nothing)
          output (Output strm (OPush ths pid) _) off0 _lim = do
              -- Creating a push promise header
              -- Frame id should be associated stream id from the client.
              let sid = streamNumber strm
              len <- pushPromise pid sid ths off0
              off <- flushIfNecessary $ off0 + frameHeaderLength + len
              return (off, Nothing)
          output _ _ _ = undefined -- never reached
  
          ----------------------------------------------------------------
          headerContinue :: StreamId -> TokenHeaderList -> Bool -> Offset -> IO Offset
          headerContinue sid ths0 endOfStream off0 = do
              buflim <- readIORef outputBufferLimit
              let offkv = off0 + frameHeaderLength
                  bufkv = confWriteBuffer `plusPtr` offkv
                  limkv = buflim - offkv
              (ths, kvlen) <- hpackEncodeHeader ctx bufkv limkv ths0
              if kvlen == 0
                  then continue off0 ths FrameHeaders
                  else do
                      let flag = getFlag ths
                          buf = confWriteBuffer `plusPtr` off0
                          off = offkv + kvlen
                      fillFrameHeader FrameHeaders kvlen sid flag buf
                      continue off ths FrameContinuation
            where
              eos = if endOfStream then setEndStream else id
              getFlag [] = eos $ setEndHeader defaultFlags
              getFlag _ = eos defaultFlags
  
              continue :: Offset -> TokenHeaderList -> FrameType -> IO Offset
              continue off [] _ = return off
              continue off ths ft = do
                  flushN off
                  -- Now off is 0
                  buflim <- readIORef outputBufferLimit
                  let bufHeaderPayload = confWriteBuffer `plusPtr` frameHeaderLength
  
                      headerPayloadLim = buflim - frameHeaderLength
                  (ths', kvlen') <-
                      hpackEncodeHeaderLoop ctx bufHeaderPayload headerPayloadLim ths
                  when (ths == ths') $
                      E.throwIO $
                          ConnectionErrorIsSent CompressionError sid "cannot compress the header"
                  let flag = getFlag ths'
                      off' = frameHeaderLength + kvlen'
                  fillFrameHeader ft kvlen' sid flag confWriteBuffer
                  continue off' ths' FrameContinuation
  
          ----------------------------------------------------------------
          fillDataHeader
              :: Stream
              -> Offset
              -> Int
              -> Maybe DynaNext
              -> (Maybe ByteString -> IO NextTrailersMaker)
              -> Output
              -> Bool
              -> IO (Offset, Maybe Output)
          fillDataHeader
              strm@Stream{streamNumber}
              off
              datPayloadLen
              Nothing
              tlrmkr
              _
              reqflush = do
                  let buf = confWriteBuffer `plusPtr` off
                  (mtrailers, flag) <- do
                      Trailers trailers <- tlrmkr Nothing
                      if null trailers
                          then return (Nothing, setEndStream defaultFlags)
                          else return (Just trailers, defaultFlags)
                  -- Avoid sending an empty data frame before trailers at the end
                  -- of a stream
                  off' <-
                      if datPayloadLen /= 0 || isNothing mtrailers
                          then do
                              decreaseWindowSize ctx strm datPayloadLen
                              fillFrameHeader FrameData datPayloadLen streamNumber flag buf
                              return $ off + frameHeaderLength + datPayloadLen
                          else
                              return off
                  off'' <- handleTrailers mtrailers off'
                  halfClosedLocal ctx strm Finished
                  if reqflush
                      then do
                          flushN off''
                          return (0, Nothing)
                      else return (off'', Nothing)
                where
                  handleTrailers Nothing off0 = return off0
                  handleTrailers (Just trailers) off0 = do
                      (ths, _) <- toTokenHeaderTable trailers
                      headerContinue streamNumber ths True {- endOfStream -} off0
          fillDataHeader
              _
              off
              0
              (Just next)
              tlrmkr
              out
              reqflush = do
                  let out' = out{outputType = ONext next tlrmkr}
                  if reqflush
                      then do
                          flushN off
                          return (0, Just out')
                      else return (off, Just out')
          fillDataHeader
              strm@Stream{streamNumber}
              off
              datPayloadLen
              (Just next)
              tlrmkr
              out
              reqflush = do
                  let buf = confWriteBuffer `plusPtr` off
                      off' = off + frameHeaderLength + datPayloadLen
                      flag = defaultFlags
                  fillFrameHeader FrameData datPayloadLen streamNumber flag buf
                  decreaseWindowSize ctx strm datPayloadLen
                  let out' = out{outputType = ONext next tlrmkr}
                  if reqflush
                      then do
                          flushN off'
                          return (0, Just out')
                      else return (off', Just out')
  
          ----------------------------------------------------------------
          pushPromise :: StreamId -> StreamId -> TokenHeaderList -> Offset -> IO Int
          pushPromise pid sid ths off = do
              let offsid = off + frameHeaderLength -- checkme
                  bufsid = confWriteBuffer `plusPtr` offsid
              poke32 (fromIntegral sid) bufsid 0
              let offkv = offsid + 4
                  bufkv = confWriteBuffer `plusPtr` offkv
                  limkv = confBufferSize - offkv
              (_, kvlen) <- hpackEncodeHeader ctx bufkv limkv ths
              let flag = setEndHeader defaultFlags -- No EndStream flag
                  buf = confWriteBuffer `plusPtr` off
                  len = kvlen + 4
              fillFrameHeader FramePushPromise len pid flag buf
              return len
  
          ----------------------------------------------------------------
          {-# INLINE fillFrameHeader #-}
          fillFrameHeader :: FrameType -> Int -> StreamId -> FrameFlags -> Buffer -> IO ()
          fillFrameHeader ftyp len sid flag buf = encodeFrameHeaderBuf ftyp hinfo buf
            where
              hinfo =
                  FrameHeader
                      { payloadLength = len
                      , flags = flag
                      , streamId = sid
                      }
  