{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RecordWildCards #-}
module Main where
import Control.Concurrent (forkIO, myThreadId, throwTo, 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.Monoid ((<>))
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 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 UseTLS = UseTLS | PlainText
deriving Show
data QueryArgs = QueryArgs {
_host :: !HostName
, _port :: !PortNumber
, _verb :: !Verb
, _path :: !Path
, _extraHeaders :: ![(ByteString, ByteString)]
, _extraTrailers :: ![(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
, _useTLS :: !UseTLS
} deriving Show
clientArgs :: Parser QueryArgs
clientArgs =
QueryArgs
<$> host
<*> port
<*> verb
<*> path
<*> extraHeaders
<*> extraTrailers
<*> 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
<*> useTLS
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'))
extraTrailers = many (fmap keyval $ bstrOption (short 'T'))
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")
useTLS = flag UseTLS PlainText (long "plain-text")
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 (_extraTrailers, _postData) of
([], (Just (PostBytestring dataPayload))) ->
return $ (id, upload dataPayload HTTP2.setEndStream)
([], (Just (PostFileContent filepath))) -> do
dataPayload <- ByteString.readFile filepath
return $ (id, upload dataPayload HTTP2.setEndStream)
([], Nothing) ->
return $ (HTTP2.setEndStream, \_ _ _ _ -> return ())
(_, (Just (PostBytestring dataPayload))) ->
return $ (id, (\c ofc s sofc -> do
upload dataPayload id c ofc s sofc
trailers s _extraTrailers HTTP2.setEndStream
))
(_, (Just (PostFileContent filepath))) -> do
dataPayload <- ByteString.readFile filepath
return $ (id, (\c ofc s sofc -> do
upload dataPayload id c ofc s sofc
trailers s _extraTrailers HTTP2.setEndStream
))
(_, Nothing) ->
return $ (id, (\_ _ s _ -> do
trailers s _extraTrailers HTTP2.setEndStream
))
let headersPairs = [ (":method", _verb)
, (":scheme", "https")
, (":path", _path)
, (":authority", ByteString.pack _host)
] <> _extraHeaders <> [("Trailer", ByteString.unwords $ fmap fst _extraTrailers)]
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 (ppHandler n idx)
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
let tlsSetting = case _useTLS of
UseTLS ->
Just tlsParams
PlainText ->
Nothing
frameConn <- newHttp2FrameConnection _host _port tlsSetting
let wrappedFrameConn = frameConn {
_makeFrameClientStream = \sid ->
let frameClient = (_makeFrameClientStream frameConn) sid
in frameClient {
_sendFrames = \mkFrames -> do
xs <- mkFrames
print $ (">>> "::String, _getStreamId frameClient, map snd xs)
_sendFrames frameClient (pure xs)
}
, _serverStream =
let
currentServerStrean = _serverStream frameConn
in
currentServerStrean {
_nextHeaderAndFrame = do
hdrFrame@(hdr,_) <- _nextHeaderAndFrame currentServerStrean
print ("<<< "::String, HTTP2.streamId hdr, hdrFrame)
return hdrFrame
}
}
parentThread <- myThreadId
let withConn = case _verboseDebug of
Verbose ->
runHttp2Client wrappedFrameConn _encoderBufsize _decoderBufsize conf defaultGoAwayHandler printUnhandledFrame
NonVerbose ->
runHttp2Client frameConn _encoderBufsize _decoderBufsize conf (throwTo parentThread) ignoreFallbackHandler
withConn $ \conn -> do
_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 =
headers stream headersPairs headersFlags
handler streamINFlowControl streamOUTFlowControl = do
timePrint $ "stream started " <> show (idx, n)
_ <- async $ do
dataPostFunction conn
(_outgoingFlowControl conn)
stream
streamOUTFlowControl
ret <- fromStreamResult <$> waitStream stream streamINFlowControl (ppHandler n idx)
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
where
tlsParams = TLS.ClientParams {
TLS.clientWantSessionResume = Nothing
, TLS.clientUseMaxFragmentLength = Nothing
, TLS.clientServerIdentification = (_host, ByteString.pack $ show _port)
, 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, _, trls) = do
timePrint hdrs
timePrint trls
dump MainFile _ _ _ ":stdout" (hdrs, body, trls) = do
timePrint hdrs
ByteString.putStrLn body
timePrint trls
dump PushPromiseFile _ _ _ ":stdout" (hdrs, _, trls) = do
timePrint hdrs
timePrint trls
dump MainFile _ _ _ ":stdout-pp" (hdrs, body, trls) = do
timePrint hdrs
ByteString.putStrLn body
timePrint trls
dump PushPromiseFile _ _ _ ":stdout-pp" (hdrs, body, trls) = do
timePrint hdrs
ByteString.putStrLn body
timePrint trls
dump _ querystring nquery nthread prefix (hdrs, body, trls) = do
timePrint hdrs
ByteString.writeFile filepath body
timePrint trls
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)
printUnhandledFrame :: FallBackFrameHandler
printUnhandledFrame (fh,fp) = timePrint ("UNHANDLED:"::String, fh, fp)