Server.hs

{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RankNTypes #-}

module HTTP3.Server (
    setup,
    setupWith,
    setupShortTimeout,
    server,
    countingServer,
    recordedBody,
    teardown,
    trailersMaker,
    firstTrailerValue,
    CH.hashInit,
) where

import Control.Concurrent

-- cryptonite

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 Data.IP ()
import Network.HPACK
import Network.HTTP.Types
import Network.HTTP3.Server
import qualified Network.QUIC.Server as QUIC
import Network.Socket ()
import System.IO.Unsafe (unsafePerformIO)
import qualified System.TimeManager as T
import Test.Hspec

import HTTP3.Config

setup :: Server -> Int -> IO ThreadId
setup = setupWith id

-- | 'setup', with the HTTP/3 configuration changed as given.
setupWith :: (Config -> Config) -> Server -> Int -> IO ThreadId
setupWith modify svr siz = do
    sc <- makeTestServerConfig
    tid <- forkIO $ QUIC.run sc loop
    threadDelay 500000 -- give enough time to the server
    return tid
  where
    loop conn = E.bracket allocSimpleConfig freeSimpleConfig $ \conf0 -> do
        let conf =
                conf0
                    { confQEncoderConfig = defaultQEncoderConfig{ecMaxTableCapacity = siz}
                    , confQDecoderConfig = defaultQDecoderConfig{dcMaxTableCapacity = siz}
                    }

        run conn (modify conf) svr

-- | 'setup', with handlers timing out after a second of not sending.
setupShortTimeout :: Server -> IO ThreadId
setupShortTimeout svr = do
    sc <- makeTestServerConfig
    tid <- forkIO $ QUIC.run sc loop
    threadDelay 500000 -- give enough time to the server
    return tid
  where
    loop conn = T.withManager 1000000 $ \mgr ->
        E.bracket allocSimpleConfig freeSimpleConfig $ \conf ->
            run conn conf{confTimeoutManager = mgr} svr

teardown :: ThreadId -> IO ()
teardown tid = killThread tid

-- | 'server', but counting how many times it is handed a request.
--
-- A request the server has already rejected must not reach the application at
-- all, and the only place that shows is here.
countingServer :: IORef Int -> Server
countingServer ref req aux sendResponse = do
    atomicModifyIORef' ref $ \n -> (n + 1, ())
    server req aux sendResponse

server :: Server
server req aux sendResponse = case requestMethod req of
    Just "GET" -> case requestPath req of
        Just "/" -> sendResponse responseHello []
        Just "/sockaddr" -> sendResponse (responseSockAddr aux) []
        -- A response that is reset after one DATA frame, between two
        -- frames: the body throws, and the stream is reset.
        Just "/reset" ->
            sendResponse (responseStreaming ok200 [] resetAfterOne) []
        -- No content, and a content-length describing content it does not
        -- have: allowed for a 304 (RFC 9114, section 4.1.2).
        Just "/not-modified" ->
            sendResponse (responseNoBody notModified304 bigLength) []
        -- The same for a 200, which is meant to have content: malformed.
        Just "/no-content" -> sendResponse (responseNoBody ok200 bigLength) []
        -- Some of the body, then nothing for longer than the timeout of
        -- 'setupShortTimeout': the handler is killed part-way.
        Just "/stall" -> sendResponse (responseStreaming ok200 [] stall) []
        -- A response whose header section is some 2K.
        Just "/bigheader" ->
            sendResponse (responseNoBody ok200 [("x-big", B.replicate 2000 0x61)]) []
        _ -> sendResponse response404 []
    Just "POST" -> case requestPath req of
        Just "/echo" -> sendResponse (responseEcho req) []
        -- Reads the body and says nothing about it, so that a test can be
        -- about what reading the body does rather than about trailers.
        Just "/drain" -> do
            let loop = do
                    bs <- getRequestBodyChunk req
                    unless (B.null bs) loop
            loop
            sendResponse responseHello []
        -- What the handler made of the body: how long it was, or what it
        -- was told instead of its end ('recordedBody').
        Just "/record" -> do
            let loop n = do
                    bs <- getRequestBodyChunk req
                    if B.null bs then return n else loop (n + B.length bs)
            r <- E.try $ loop (0 :: Int)
            putMVar recordedBody $ either (\e -> Left (show (e :: E.SomeException))) Right r
            sendResponse responseHello []
        -- Answers with the length of the body it read.
        Just "/length" -> do
            let loop n = do
                    bs <- getRequestBodyChunk req
                    if B.null bs then return n else loop (n + B.length bs)
            n <- loop (0 :: Int)
            sendResponse (responseLength n) []
        _ -> sendResponse responseHello []
    -- No content, as a response to HEAD has, and the content-length of the
    -- GET.
    Just "HEAD" -> sendResponse (responseNoBody ok200 bigLength) []
    _ -> sendResponse response405 []

-- | What the handler of /record made of the body it was sent.
recordedBody :: MVar (Either String Int)
recordedBody = unsafePerformIO newEmptyMVar
{-# NOINLINE recordedBody #-}

bigLength :: ResponseHeaders
bigLength = [("content-length", "1000")]

responseLength :: Int -> Response
responseLength n = responseBuilder ok200 header body
  where
    header = [("Content-Type", "text/plain")]
    body = byteString $ C8.pack $ show n

stall :: (Builder -> IO ()) -> IO () -> IO ()
stall write flush = do
    write $ byteString "partial"
    flush
    threadDelay 10000000

resetAfterOne :: (Builder -> IO ()) -> IO () -> IO ()
resetAfterOne write flush = do
    write $ byteString "partial"
    flush
    -- A RESET_STREAM can go out ahead of stream data queued before it, and
    -- then the client sees neither HEADERS nor DATA.  Let them arrive first.
    threadDelay 200000
    E.throwIO $ userError "reset after one DATA frame"

responseHello :: Response
responseHello = responseBuilder ok200 header body
  where
    header = [("Content-Type", "text/plain")]
    body = byteString "Hello, world!\n"

-- | Hands back the two addresses the server was given, so that a test can
-- check they are not the same one twice over.
responseSockAddr :: Aux -> Response
responseSockAddr aux = responseBuilder ok200 header body
  where
    header = [("Content-Type", "text/plain")]
    body =
        byteString $
            C8.pack $
                show (auxMySockAddr aux) ++ " " ++ show (auxPeerSockAddr aux)

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")]
    streamingBody write _flush = do
        loop
        mt <- getRequestTrailers req
        firstTrailerValue <$> mt
            `shouldBe` Just "b0870457df2b8cae06a88657a198d9b52f8e2b0a"
      where
        loop = do
            bs <- getRequestBodyChunk req
            unless (B.null bs) $ do
                void $ write $ byteString bs
                loop
    maker = trailersMaker CH.hashInit

-- 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

firstTrailerValue :: TokenHeaderTable -> FieldValue
firstTrailerValue tbl = case fst tbl of
    [] -> error "firstTrailerValue"
    x : _ -> snd x