content: {-# LANGUAGE BangPatterns #-}
  {-# LANGUAGE OverloadedStrings #-}
  {-# LANGUAGE RankNTypes #-}
  {-# LANGUAGE RecordWildCards #-}
  
  module HTTP2.ServerSpec (spec) where
  
  import Control.Concurrent
  import Control.Concurrent.Async
  import qualified Control.Exception as E
  import Control.Monad
  import Crypto.Hash (Context, SHA1)
  import qualified Crypto.Hash as CH
  import Data.ByteString (ByteString)
  import qualified Data.ByteString as B
  import Data.ByteString.Builder (Builder, byteString)
  import qualified Data.ByteString.Char8 as C8
  import Data.IORef
  import Network.HTTP.Semantics
  import Network.HTTP.Types
  import Network.Run.TCP
  import Network.Socket
  import Network.Socket.ByteString
  import System.IO
  import System.IO.Unsafe
  import System.Random
  import Test.Hspec
  
  import Network.HPACK
  import Network.HPACK.Internal
  import qualified Network.HTTP2.Client as C
  import qualified Network.HTTP2.Client.Internal as C
  import Network.HTTP2.Frame
  import Network.HTTP2.Server
  
  port :: String
  port = show $ unsafePerformIO (randomPort <$> getStdGen)
    where
      randomPort = fst . randomR (43124 :: Int, 44320)
  
  host :: String
  host = "127.0.0.1"
  
  spec :: Spec
  spec = do
      describe "server" $ do
          it "handles normal cases" $
              E.bracket (forkIO runServer) killThread $ \_ -> do
                  threadDelay 10000
                  runClient allocSimpleConfig
  
          it "should always send the connection preface first" $ do
              prefaceVar <- newEmptyMVar
              E.bracket (forkIO (runFakeServer prefaceVar)) killThread $ \_ -> do
                  threadDelay 10000
                  E.catch (runClient allocSlowPrefaceConfig) ignoreHTTP2Error
  
              preface <- takeMVar prefaceVar
              preface `shouldBe` connectionPreface
  
          it "prevents attacks" $
              E.bracket (forkIO runServer) killThread $ \_ -> do
                  threadDelay 10000
                  runAttack rapidSettings `shouldThrow` connectionError "too many settings"
                  runAttack rapidPing `shouldThrow` connectionError "too many ping"
                  runAttack rapidEmptyHeader
                      `shouldThrow` connectionError "too many empty headers"
                  runAttack rapidEmptyData `shouldThrow` connectionError "too many empty data"
                  runAttack rapidRst `shouldThrow` connectionError "too many rst_stream"
  
  ignoreHTTP2Error :: C.HTTP2Error -> IO ()
  ignoreHTTP2Error _ = pure ()
  
  runServer :: IO ()
  runServer = runTCPServer (Just host) port runHTTP2Server
    where
      runHTTP2Server s =
          E.bracket
              (allocSimpleConfig s 32768)
              freeSimpleConfig
              (\conf -> run defaultServerConfig conf server)
  
  runFakeServer :: MVar ByteString -> IO ()
  runFakeServer prefaceVar = do
      runTCPServer (Just host) port $ \s -> do
          ref <- newIORef Nothing
  
          -- send settings
          sendAll s $
              "\x00\x00\x12\x04\x00\x00\x00\x00\x00"
                  `mappend` "\x00\x03\x00\x00\x00\x80\x00\x04\x00"
                  `mappend` "\x01\x00\x00\x00\x05\x00\xff\xff\xff"
  
          -- receive preface
          value <- defaultReadN s ref (B.length connectionPreface)
          putMVar prefaceVar value
  
          -- send goaway frame
          sendAll s "\x00\x00\x08\x07\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x01"
  
          -- wait for a few ms to make sure the client has a chance to close the
          -- socket on its end
          threadDelay 10000
  
  server :: Server
  server req _aux sendResponse = case requestMethod req of
      Just "GET" -> case requestPath req of
          Just "/" -> sendResponse responseHello []
          Just "/stream" -> sendResponse responseInfinite []
          Just "/push" -> do
              let pp = pushPromise "/push-pp" responsePP 0
              sendResponse responseHello [pp]
          _ -> sendResponse response404 []
      Just "POST" -> case requestPath req of
          Just "/echo" -> sendResponse (responseEcho req) []
          _ -> sendResponse responseHello []
      _ -> sendResponse response405 []
  
  responseHello :: Response
  responseHello = responseBuilder ok200 header body
    where
      header = [("Content-Type", "text/plain")]
      body = byteString "Hello, world!\n"
  
  responsePP :: Response
  responsePP = responseBuilder ok200 header body
    where
      header =
          [ ("Content-Type", "text/plain")
          , ("x-push", "True")
          ]
      body = byteString "Push\n"
  
  responseInfinite :: Response
  responseInfinite = responseStreaming ok200 header body
    where
      header = [("Content-Type", "text/plain")]
      body :: (Builder -> IO ()) -> IO () -> IO ()
      body write flush = do
          let go n = write (byteString (C8.pack (show n)) `mappend` "\n") *> flush *> go (succ n)
          go (0 :: Int)
  
  response404 :: Response
  response404 = responseNoBody notFound404 []
  
  response405 :: Response
  response405 = responseNoBody methodNotAllowed405 []
  
  responseEcho :: Request -> Response
  responseEcho req = setResponseTrailersMaker h2rsp maker
    where
      h2rsp = responseStreaming ok200 header streamingBody
      header = [("Content-Type", "text/plain")]
      mhx = getFieldValue (toToken "X-Tag") (snd (requestHeaders req))
      streamingBody write _flush = do
          loop
          mt <- getRequestTrailers req
          firstTrailerValue <$> mt `shouldBe` mhx
        where
          loop = do
              bs <- getRequestBodyChunk req
              when (bs /= "") $ do
                  void $ write $ byteString bs
                  loop
      maker = trailersMaker (CH.hashInit :: Context SHA1)
  
  -- Strictness is important for Context.
  trailersMaker :: Context SHA1 -> Maybe ByteString -> IO NextTrailersMaker
  trailersMaker ctx Nothing = return $ Trailers [("X-SHA1", sha1)]
    where
      !sha1 = C8.pack $ show $ CH.hashFinalize ctx
  trailersMaker ctx (Just bs) = return $ NextTrailersMaker $ trailersMaker ctx'
    where
      !ctx' = CH.hashUpdate ctx bs
  
  runClient :: (Socket -> BufferSize -> IO Config) -> IO ()
  runClient allocConfig =
      runTCPClient host port runHTTP2Client
    where
      auth = host
      cliconf = C.defaultClientConfig{C.authority = auth}
      runHTTP2Client s =
          E.bracket
              (allocConfig s 4096)
              freeSimpleConfig
              (\conf -> C.run cliconf conf client)
  
      client :: C.Client ()
      client sendRequest aux =
          foldr1
              concurrently_
              [ client0 sendRequest aux
              , client1 sendRequest aux
              , client2 sendRequest aux
              , client3 sendRequest aux
              , client3' sendRequest aux
              , client3'' sendRequest aux
              , client4 sendRequest aux
              , client5 sendRequest aux
              ]
  
  -- delay sending preface to be able to test if it is always sent first
  allocSlowPrefaceConfig :: Socket -> BufferSize -> IO Config
  allocSlowPrefaceConfig s size = do
      config <- allocSimpleConfig s size
      pure config{confSendAll = slowPrefaceSend (confSendAll config)}
    where
      slowPrefaceSend :: (ByteString -> IO ()) -> ByteString -> IO ()
      slowPrefaceSend orig chunk = do
          when (C8.pack "PRI" `C8.isPrefixOf` chunk) $ do
              threadDelay 10000
          orig chunk
  
  client0 :: C.Client ()
  client0 sendRequest _aux = do
      let req = C.requestNoBody methodGet "/" []
      sendRequest req $ \rsp -> do
          C.responseStatus rsp `shouldBe` Just ok200
          fmap statusMessage (C.responseStatus rsp) `shouldBe` Just "OK"
  
  client1 :: C.Client ()
  client1 sendRequest _aux = do
      let req = C.requestNoBody methodGet "/push-pp" []
      sendRequest req $ \rsp -> do
          C.responseStatus rsp `shouldBe` Just notFound404
  
  client2 :: C.Client ()
  client2 sendRequest _aux = do
      let req = C.requestNoBody methodPut "/" []
      sendRequest req $ \rsp -> do
          C.responseStatus rsp `shouldBe` Just methodNotAllowed405
  
  client3 :: C.Client ()
  client3 sendRequest _aux = do
      let hx = "b0870457df2b8cae06a88657a198d9b52f8e2b0a"
          req0 =
              C.requestFile methodPost "/echo" [("X-Tag", hx)] $
                  FileSpec "test/inputFile" 0 1012731
          req = C.setRequestTrailersMaker req0 maker
      sendRequest req $ \rsp -> do
          let consumeBody = do
                  bs <- C.getResponseBodyChunk rsp
                  when (bs /= "") consumeBody
          consumeBody
          mt <- C.getResponseTrailers rsp
          firstTrailerValue <$> mt `shouldBe` Just hx
    where
      !maker = trailersMaker (CH.hashInit :: Context SHA1)
  
  client3' :: C.Client ()
  client3' sendRequest _aux = do
      let hx = "b0870457df2b8cae06a88657a198d9b52f8e2b0a"
          req0 = C.requestStreaming methodPost "/echo" [("X-Tag", hx)] $ \write _flush -> do
              let sendFile h = do
                      bs <- B.hGet h 1024
                      when (bs /= "") $ do
                          write $ byteString bs
                          sendFile h
              withFile "test/inputFile" ReadMode sendFile
          req = C.setRequestTrailersMaker req0 maker
      sendRequest req $ \rsp -> do
          let consumeBody = do
                  bs <- C.getResponseBodyChunk rsp
                  when (bs /= "") consumeBody
          consumeBody
          mt <- C.getResponseTrailers rsp
          firstTrailerValue <$> mt `shouldBe` Just hx
    where
      !maker = trailersMaker (CH.hashInit :: Context SHA1)
  
  client3'' :: C.Client ()
  client3'' sendRequest _axu = do
      let hx = "59f82dfddc0adf5bdf7494b8704f203a67e25d4a"
          req0 = C.requestStreaming methodPost "/echo" [("X-Tag", hx)] $ \write _flush -> do
              let chunk = C8.replicate (16384 * 2) 'c'
                  tag = C8.replicate 16 't'
              -- I don't think 9 is important here, this is just what I have, the client hangs on receiving the last one
              replicateM_ 9 $ write $ byteString chunk
              write $ byteString tag
          req = C.setRequestTrailersMaker req0 maker
      sendRequest req $ \rsp -> do
          let consumeBody = do
                  bs <- C.getResponseBodyChunk rsp
                  when (bs /= "") consumeBody
          consumeBody
          mt <- C.getResponseTrailers rsp
          firstTrailerValue <$> mt `shouldBe` Just hx
    where
      !maker = trailersMaker (CH.hashInit :: Context SHA1)
  
  client4 :: C.Client ()
  client4 sendRequest _aux = do
      let req0 = C.requestNoBody methodGet "/push" []
      sendRequest req0 $ \rsp -> do
          C.responseStatus rsp `shouldBe` Just ok200
      let req1 = C.requestNoBody methodGet "/push-pp" []
      sendRequest req1 $ \rsp -> do
          C.responseStatus rsp `shouldBe` Just ok200
  
  client5 :: C.Client ()
  client5 sendRequest _aux = do
      let req0 = C.requestNoBody methodGet "/stream" []
      sendRequest req0 $ \rsp -> do
          C.responseStatus rsp `shouldBe` Just ok200
          let go n
                  | n > 0 = do
                      _ <- C.getResponseBodyChunk rsp
                      go (pred n)
                  | otherwise = pure ()
          go (100 :: Int)
  
  firstTrailerValue :: TokenHeaderTable -> FieldValue
  firstTrailerValue tbl = case fst tbl of
      [] -> error "firstTrailerValue"
      x : _ -> snd x
  
  runAttack :: (C.ClientIO -> IO ()) -> IO ()
  runAttack attack =
      runTCPClient host port runHTTP2Client
    where
      auth = host
      cliconf = C.defaultClientConfig{C.authority = auth}
      runHTTP2Client s =
          E.bracket
              (allocSimpleConfig s 4096)
              freeSimpleConfig
              (\conf -> C.runIO cliconf conf client)
      client cconf = return $ do
          attack cconf
          threadDelay 1000000
  
  rapidSettings :: C.ClientIO -> IO ()
  rapidSettings C.ClientIO{..} = do
      let einfo = EncodeInfo defaultFlags 0 Nothing
          bs = encodeFrame einfo $ SettingsFrame [(SettingsEnablePush, 0)]
      cioWriteBytes bs
      cioWriteBytes bs
      cioWriteBytes bs
      cioWriteBytes bs
      cioWriteBytes bs
      cioWriteBytes bs
      cioWriteBytes bs
      cioWriteBytes bs
      cioWriteBytes bs
  
  rapidPing :: C.ClientIO -> IO ()
  rapidPing C.ClientIO{..} = do
      let einfo = EncodeInfo defaultFlags 0 Nothing
          opaque64 = "01234567"
          bs = encodeFrame einfo $ PingFrame opaque64
      replicateM_ 20 $ cioWriteBytes bs
  
  rapidEmptyHeader :: C.ClientIO -> IO ()
  rapidEmptyHeader C.ClientIO{..} = do
      (sid, _) <- cioCreateStream
      let einfo = EncodeInfo defaultFlags sid Nothing
          bs = encodeFrame einfo $ HeadersFrame Nothing ""
      cioWriteBytes bs
      cioWriteBytes bs
      cioWriteBytes bs
      cioWriteBytes bs
      cioWriteBytes bs
      cioWriteBytes bs
      cioWriteBytes bs
      cioWriteBytes bs
      cioWriteBytes bs
  
  rapidEmptyData :: C.ClientIO -> IO ()
  rapidEmptyData C.ClientIO{..} = do
      (sid, _) <- cioCreateStream
      let einfoH = EncodeInfo (setEndHeader defaultFlags) sid Nothing
          hdr =
              hpackEncode
                  [ (":scheme", "http")
                  , (":authority", "127.0.0.1")
                  , (":path", "/")
                  , (":method", "GET")
                  ]
          bsH = encodeFrame einfoH $ HeadersFrame Nothing hdr
      cioWriteBytes bsH
      let einfoD = EncodeInfo defaultFlags sid Nothing
          bsD = encodeFrame einfoD $ DataFrame ""
      cioWriteBytes bsD
      cioWriteBytes bsD
      cioWriteBytes bsD
      cioWriteBytes bsD
      cioWriteBytes bsD
      cioWriteBytes bsD
      cioWriteBytes bsD
      cioWriteBytes bsD
  
  rapidRst :: C.ClientIO -> IO ()
  rapidRst C.ClientIO{..} = do
      reset
      reset
      reset
      reset
      reset
      reset
      reset
      reset
    where
      reset = do
          (sid, _) <- cioCreateStream
          -- setEndStream for HalfClosedRemote
          let einfoH = EncodeInfo (setEndStream $ setEndHeader defaultFlags) sid Nothing
              hdr =
                  hpackEncode
                      [ (":scheme", "http")
                      , (":authority", "127.0.0.1")
                      , (":path", "/")
                      , (":method", "GET")
                      ]
              bsH = encodeFrame einfoH $ HeadersFrame Nothing hdr
          cioWriteBytes bsH
          let einfoR = EncodeInfo defaultFlags sid Nothing
              -- Only (HalfClosedRemote, NoError) is accepted.
              -- Otherwise, a stream error terminates the connection.
              bsR = encodeFrame einfoR $ RSTStreamFrame NoError
          cioWriteBytes bsR
  
  connectionError :: C.ReasonPhrase -> C.HTTP2Error -> Bool
  connectionError phrase (C.ConnectionErrorIsReceived _ _ p)
      | phrase == p = True
  connectionError _ _ = False
  
  hpackEncode :: [(ByteString, ByteString)] -> ByteString
  hpackEncode kvs = foldr cat "" kvs
    where
      (k, v) `cat` b =
          B.singleton 0x10
              <> unsafePerformIO (encodeInteger 7 (B.length k))
              <> k
              <> unsafePerformIO (encodeInteger 7 (B.length v))
              <> v
              <> b
  