{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RecordWildCards #-}
module Main where
import Control.Concurrent (forkIO, threadDelay)
import Control.Concurrent.Async (async, waitAnyCancel)
import Control.Monad (forever, when, void)
import Data.ByteString (ByteString)
import qualified Data.ByteString.Char8 as ByteString
import Data.Default.Class (def)
import Data.Maybe (fromMaybe)
import Data.Time.Clock (diffUTCTime, getCurrentTime)
import qualified Network.HTTP2 as HTTP2
import qualified Network.TLS as TLS
import qualified Network.TLS.Extra.Cipher as TLS
import Options.Applicative
import Data.Monoid ((<>))
import System.IO
import Network.HTTP2.Client
import Network.HTTP2.Client.Helpers
type Path = ByteString
type Verb = ByteString
data PostData = PostBytestring !ByteString | PostFileContent FilePath
deriving Show
postDataForString :: String -> PostData
postDataForString ('@':filepath) = PostFileContent filepath
postDataForString content = PostBytestring (ByteString.pack content)
data ServerPushSwitch = PushEnabled | PushDisabled
deriving Show
data Verbosity = Verbose | NonVerbose
deriving Show
data QueryArgs = QueryArgs {
_host :: !HostName
, _port :: !PortNumber
, _verb :: !Verb
, _path :: !Path
, _extraHeaders :: ![(ByteString, ByteString)]
, _postData :: !(Maybe PostData)
, _interPingDelay :: !Int
, _pingTimeout :: !Int
, _interFlowControlUpdates :: !Int
, _settingsMaxConcurrency :: !Int
, _settingsAllowServerPush :: !ServerPushSwitch
, _settingsMaxFrameSize :: !Int
, _settingsMaxHeaderBlockSize :: !Int
, _settingsInitialWindowSize :: !Int
, _initialWindowKick :: !Int
, _concurrentQueriesCount :: !Int
, _numberQueries :: !Int
, _finalDelay :: !Int
, _finalMessage :: !ByteString
, _encoderBufsize :: !Int
, _decoderBufsize :: !Int
, _downloadPrefix :: !FilePath
, _verboseDebug :: !Verbosity
} deriving Show
clientArgs :: Parser QueryArgs
clientArgs =
QueryArgs
<$> host
<*> port
<*> verb
<*> path
<*> extraHeaders
<*> postData
<*> milliseconds "inter-ping-delay-ms" 0
<*> milliseconds "ping-timeout-ms" 5000
<*> milliseconds "inter-flow-control-updates-ms" 1000
<*> concurrency
<*> allowPush
<*> frameBytes
<*> headersBytes
<*> initialWindowBytes
<*> initialWindowKick
<*> numConcurrentThreads
<*> numQueriesPerThread
<*> milliseconds "delay-before-quitting-ms" 0
<*> kthxByeMessage
<*> encoderBufSize
<*> decoderBufSize
<*> downloadPrefix
<*> verboseDebug
where
bstrOption = fmap ByteString.pack . strOption
milliseconds what base = fmap (*1000) $ option auto (long what <> value base)
keyval kv = let (k,v1) = ByteString.break (== ':') kv in (k, ByteString.drop 1 v1)
host = strOption (long "host" <> value "127.0.0.1")
path = bstrOption (long "path" <> value "/")
port = option auto (long "port" <> value 443)
verb = bstrOption (long "verb" <> value "GET")
extraHeaders = many (fmap keyval $ bstrOption (short 'H'))
postData = optional (fmap postDataForString (strOption (short 'd')))
concurrency = option auto (long "max-concurrency" <> value 100)
allowPush = flag PushEnabled PushDisabled (long "disable-server-push")
frameBytes = option auto (long "max-frame-size" <> value 1048576)
headersBytes = option auto (long "max-headers-list-size" <> value 1048576)
initialWindowBytes = option auto (long "initial-window-size" <> value 10485760)
initialWindowKick = option auto (long "initial-window-kick" <> value 0)
numConcurrentThreads = option auto (long "num-concurrent-threads" <> value 1)
numQueriesPerThread = option auto (long "num-queries-per-thread" <> value 1)
kthxByeMessage = bstrOption (long "exit-greeting" <> value "kthxbye (>;_;<)")
encoderBufSize = option auto (long "hpack-encoder-buffer-size" <> value 4096)
decoderBufSize = option auto (long "hpack-decoder-buffer-size" <> value 4096)
downloadPrefix = strOption (long "push-files-prefix" <> value ":stdout-pp")
verboseDebug = flag NonVerbose Verbose (long "verbose")
main :: IO ()
main = execParser opts >>= client
where
opts = info (helper <*> clientArgs) (mconcat [
fullDesc
, header "http2-client-exe: a CLI HTTP2 client written in Haskell"
])
client :: QueryArgs -> IO ()
client QueryArgs{..} = do
hSetBuffering stdout LineBuffering
-- If we need to post data then we prepare an upload function.
-- Otherwise the function does nothing but we tell the server that we are
-- finished with the stream after HTTP headers.
--
-- Note that this implementation reads the body from files and into memory
-- (better for testing concurrency, worse for large uploads).
(headersFlags, dataPostFunction) <- case _postData of
(Just (PostBytestring dataPayload)) ->
return $ (id, upload dataPayload)
(Just (PostFileContent filepath)) -> do
dataPayload <- ByteString.readFile filepath
return $ (id, upload dataPayload)
Nothing ->
return $ (HTTP2.setEndStream, \_ _ _ _ -> return ())
let headersPairs = [ (":method", _verb)
, (":scheme", "https")
, (":path", _path)
, (":authority", ByteString.pack _host)
] <> _extraHeaders
let ppHandler n idx _ stream ppHdrs streamFlowControl _ = void $ forkIO $ do
let pushpath = fromMaybe "unspecified-path" (lookup ":path" ppHdrs)
timePrint ("push stream started" :: String, pushpath)
ret <- fromStreamResult <$> waitStream stream streamFlowControl
either (\e -> timePrint e) (dump PushPromiseFile pushpath n idx _downloadPrefix) ret
timePrint ("push stream ended" :: String)
let conf = [ (HTTP2.SettingsMaxFrameSize, _settingsMaxFrameSize)
, (HTTP2.SettingsMaxConcurrentStreams, _settingsMaxConcurrency)
, (HTTP2.SettingsMaxHeaderBlockSize, _settingsMaxHeaderBlockSize)
, (HTTP2.SettingsInitialWindowSize, _settingsInitialWindowSize)
, (HTTP2.SettingsEnablePush, case _settingsAllowServerPush of
PushEnabled -> 1 ; PushDisabled -> 0)
]
timePrint conf
frameConn <- newHttp2FrameConnection _host _port tlsParams
let wrappedFrameConn = frameConn {
_makeFrameClientStream = \sid ->
let frameClient = (_makeFrameClientStream frameConn) sid
in frameClient {
_sendFrames = \xs -> do
print $ (">>> "::String, _getStreamId frameClient, map snd xs)
_sendFrames frameClient xs
}
, _serverStream =
let
currentServerStrean = _serverStream frameConn
in
currentServerStrean {
_nextHeaderAndFrame = do
hdrFrame@(hdr,_) <- _nextHeaderAndFrame currentServerStrean
print ("<<< "::String, HTTP2.streamId hdr, hdrFrame)
return hdrFrame
}
}
conn <- case _verboseDebug of
Verbose ->
wrapFrameClient wrappedFrameConn _encoderBufsize _decoderBufsize conf
NonVerbose ->
wrapFrameClient frameConn _encoderBufsize _decoderBufsize conf
linkHttp2Client conn
_addCredit (_incomingFlowControl conn) _initialWindowKick
_ <- forkIO $ forever $ do
updated <- _updateWindow $ _incomingFlowControl conn
when updated $ timePrint ("sending flow-control update" :: String)
threadDelay _interFlowControlUpdates
_ <- forkIO $ when (_interPingDelay > 0) $ forever $ do
threadDelay _interPingDelay
(t0, t1, pingReply) <- ping conn _pingTimeout "pingpong"
timePrint $ ("ping-reply:" :: String, pingReply, diffUTCTime t1 t0)
let go 0 idx = timePrint $ "done worker: " <> show idx
go n idx = do
_ <- (withHttp2Stream conn $ \stream ->
let initStream = do
_ <- async $ onPushPromise stream (ppHandler n idx)
headers stream headersPairs headersFlags
handler streamINFlowControl streamOUTFlowControl = do
timePrint $ "stream started " <> show (idx, n)
_ <- async $ dataPostFunction conn
(_outgoingFlowControl conn)
stream
streamOUTFlowControl
ret <- fromStreamResult <$> waitStream stream streamINFlowControl
either (\e -> timePrint e) (dump MainFile _path n idx _downloadPrefix) ret
timePrint $ "stream ended " <> show (idx, n)
in StreamDefinition initStream handler)
go (n - 1) idx
_ <- waitAnyCancel =<< traverse (async . go _numberQueries) [1 .. _concurrentQueriesCount]
when (_finalDelay > 0) (threadDelay _finalDelay)
_gtfo conn HTTP2.NoError _finalMessage
_close conn
return ()
where
tlsParams = TLS.ClientParams {
TLS.clientWantSessionResume = Nothing
, TLS.clientUseMaxFragmentLength = Nothing
, TLS.clientServerIdentification = ("127.0.0.1", "")
, TLS.clientUseServerNameIndication = True
, TLS.clientShared = def
, TLS.clientHooks = def { TLS.onServerCertificate = \_ _ _ _ -> return []
}
, TLS.clientSupported = def { TLS.supportedCiphers = TLS.ciphersuite_default }
, TLS.clientDebug = def
}
data DumpType = MainFile | PushPromiseFile
dump :: DumpType -> Path -> Int -> Int -> FilePath -> StreamResponse -> IO ()
dump MainFile _ _ _ ":none" (hdrs, _) = do
timePrint hdrs
dump MainFile _ _ _ ":stdout" (hdrs, body) = do
timePrint hdrs
ByteString.putStrLn body
dump PushPromiseFile _ _ _ ":stdout" (hdrs, _) = do
timePrint hdrs
dump MainFile _ _ _ ":stdout-pp" (hdrs, body) = do
timePrint hdrs
ByteString.putStrLn body
dump PushPromiseFile _ _ _ ":stdout-pp" (hdrs, body) = do
timePrint hdrs
ByteString.putStrLn body
dump _ querystring nquery nthread prefix (hdrs, body) = do
timePrint hdrs
ByteString.writeFile filepath body
where
filepath = mconcat [ prefix
, "/"
, show nquery
, "."
, show nthread
, ByteString.unpack $ cleanPath querystring
]
cleanPath = ByteString.takeWhile (/= '?')
. ByteString.map (\c -> if c == '/' then '_' else c)
timePrint :: Show a => a -> IO ()
timePrint x = do
tst <- getCurrentTime
ByteString.hPutStrLn stderr $ ByteString.pack $ show (tst, x)