HPACK.hs

{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RecordWildCards #-}

module Network.HTTP2.H2.HPACK (
    hpackEncodeHeader,
    hpackEncodeHeaderRest,
    hpackDecodeHeader,
    hpackDecodeTrailer,
    hpackDiscardHeader,
    just,
    fixHeaders,
) where

import qualified Control.Exception as E
import qualified Data.ByteString as BS
import Data.ByteString.Internal (create)
import qualified Data.ByteString.Lazy as BS.Lazy
import Foreign.Marshal.Alloc (free, mallocBytes)
import Foreign.Marshal.Utils (copyBytes)
import Network.ByteOrder
import Network.HTTP.Semantics
import Network.HTTP.Types

import Imports
import Network.HPACK
import Network.HTTP2.Frame
import Network.HTTP2.H2.Context
import Network.HTTP2.H2.Types

-- $setup
-- >>> :set -XOverloadedStrings

----------------------------------------------------------------

fixHeaders :: ResponseHeaders -> ResponseHeaders
fixHeaders hdr = deleteUnnecessaryHeaders hdr

deleteUnnecessaryHeaders :: ResponseHeaders -> ResponseHeaders
deleteUnnecessaryHeaders hdr = filter del hdr
  where
    del (k, _) = k `notElem` headersToBeRemoved

headersToBeRemoved :: [HeaderName]
headersToBeRemoved =
    [ hConnection
    , "Transfer-Encoding"
    -- Keep-Alive
    -- Proxy-Connection
    -- Upgrade
    ]

----------------------------------------------------------------

strategy :: EncodeStrategy
strategy = EncodeStrategy{compressionAlgo = Linear, useHuffman = False}

-- Set-Cookie: contains only one cookie value.
-- So, we don't need to split it.
hpackEncodeHeader
    :: Context
    -> Buffer
    -> BufferSize
    -> TokenHeaderList
    -> IO (TokenHeaderList, Int)
hpackEncodeHeader Context{..} buf siz ths =
    encodeTokenHeader buf siz strategy True encodeDynamicTable ths

hpackEncodeHeaderLoop
    :: Context
    -> Buffer
    -> BufferSize
    -> TokenHeaderList
    -> IO (TokenHeaderList, Int)
hpackEncodeHeaderLoop Context{..} buf siz hs =
    encodeTokenHeader buf siz strategy False encodeDynamicTable hs

-- | Encode the rest of a header block whose start 'hpackEncodeHeader' wrote
--
-- For a block that did not fit where it was being written.  Grows the buffer
-- as needed: a header that does not fit is retried with a larger buffer (the
-- encoder does not modify the dynamic table for a header it could not write).
hpackEncodeHeaderRest
    :: Context
    -> BufferSize
    -- ^ Initial buffer size
    -> TokenHeaderList
    -> IO BS.Lazy.ByteString
hpackEncodeHeaderRest ctx = go []
  where
    go acc _ [] = return $ BS.Lazy.fromChunks (reverse acc)
    go acc siz ths = do
        (chunk, ths') <- E.bracket (mallocBytes siz) free $ \buf -> do
            (ths', len) <- hpackEncodeHeaderLoop ctx buf siz ths
            chunk <- create len $ \p -> copyBytes p buf len
            return (chunk, ths')
        if BS.null chunk
            then go acc (siz * 2) ths -- no progress: grow
            else go (chunk : acc) siz ths'

----------------------------------------------------------------

hpackDecodeHeader
    :: HeaderBlockFragment -> StreamId -> Context -> IO TokenHeaderTable
hpackDecodeHeader hdrblk sid ctx = do
    tbl@(_, vt) <- hpackDecode "illegal header" hdrblk sid ctx
    if isClient ctx || checkRequestHeader vt
        then return tbl
        else E.throwIO $ StreamErrorIsSent ProtocolError sid "illegal header"

hpackDecodeTrailer
    :: HeaderBlockFragment -> StreamId -> Context -> IO TokenHeaderTable
hpackDecodeTrailer = hpackDecode "illegal trailer"

-- | Decode a field block for a stream we no longer have, and discard the result
--
-- The block must still be decoded: it may modify the dynamic table.
hpackDiscardHeader :: HeaderBlockFragment -> StreamId -> Context -> IO ()
hpackDiscardHeader hdrblk sid ctx =
    void (hpackDecode "illegal header" hdrblk sid ctx) `E.catch` ignore
  where
    -- Decoded to the end, so there is nothing to refuse: the message goes
    -- nowhere anyway.
    ignore StreamErrorIsSent{} = return ()
    ignore e = E.throwIO e

-- | Decode a field block, reporting a block we could not get through as a
-- connection error, and a malformed message in a block decoded to the end
-- as a stream error.
--
-- The first argument says which kind of block it was, since the peer reads
-- this in the GOAWAY and "illegal trailer" about a request's headers is a
-- confusing thing to be told.
hpackDecode
    :: ReasonPhrase
    -> HeaderBlockFragment
    -> StreamId
    -> Context
    -> IO TokenHeaderTable
hpackDecode illegal hdrblk sid Context{..} =
    decodeTokenHeader decodeDynamicTable hdrblk `E.catch` handl
  where
    -- A malformed message: 'decodeTokenHeader' says so only once it has
    -- decoded the whole block, so the dynamic table is up to date and the
    -- connection can go on.  A stream error, as RFC 9113 section 8.1.1 has
    -- it.  This used to be a connection error, as the block was abandoned
    -- at the malformed field; one request with an upper-case field name
    -- closed the connection, every other stream on it with it.
    handl IllegalHeaderName =
        E.throwIO $ StreamErrorIsSent ProtocolError sid illegal
    handl TooLargeHeader =
        E.throwIO $ StreamErrorIsSent ProtocolError sid "too many fields"
    -- A block we could not get through: our dynamic table now holds the
    -- entries decoded before the throw and nothing after them -- no longer
    -- what the peer's encoder believes we have.  Section 10.5.1: "The field
    -- block MUST be processed to ensure a consistent connection state,
    -- unless the connection is closed."  We did not, so it must be.
    handl e = do
        let msg = fromString $ show e
        E.throwIO $ ConnectionErrorIsSent CompressionError sid msg

{-# INLINE checkRequestHeader #-}
checkRequestHeader :: ValueTable -> Bool
checkRequestHeader reqvt
    | just mMethod (== "CONNECT") = isNothing mPath && isNothing mScheme
    | isJust mStatus = False
    | isNothing mMethod = False
    | isNothing mScheme = False
    | isNothing mPath = False
    | mPath == Just "" = False
    | isJust mConnection = False
    | just mTE (/= "trailers") = False
    | otherwise = checkAuth mAuthority mHost
  where
    mStatus = getFieldValue tokenStatus reqvt
    mScheme = getFieldValue tokenScheme reqvt
    mPath = getFieldValue tokenPath reqvt
    mMethod = getFieldValue tokenMethod reqvt
    mConnection = getFieldValue tokenConnection reqvt
    mTE = getFieldValue tokenTE reqvt
    mAuthority = getFieldValue tokenAuthority reqvt
    mHost = getFieldValue tokenHost reqvt

checkAuth :: Maybe ByteString -> Maybe ByteString -> Bool
checkAuth Nothing Nothing = False
checkAuth (Just a) (Just h) | a /= h = False
checkAuth _ _ = True

{-# INLINE just #-}
just :: Maybe a -> (a -> Bool) -> Bool
just Nothing _ = False
just (Just x) p
    | p x = True
    | otherwise = False