nuthatch/test/Nuthatch/Probe/Http.hs

module Nuthatch.Probe.Http
  ( CookieSession
  , HttpException
  , ProbeClient
  , ProbeResponse (..)
  , RequestBody (..)
  , RequestOptions (..)
  , cookieHeader
  , cookieValue
  , defaultRequestOptions
  , downloadTo
  , get
  , getWith
  , headWith
  , headerValue
  , headerValues
  , isConnectionFailure
  , newCookieSession
  , newProbeClient
  , postWith
  , putWith
  ) where

import           Control.Concurrent.MVar
import qualified Data.ByteString as BS
import qualified Data.ByteString.Char8 as B8
import qualified Data.ByteString.Lazy as BL
import           Data.List (find)
import qualified Data.CaseInsensitive as CI
import           Network.Connection (TLSSettings (..))
import           Network.HTTP.Client
  ( Cookie (..)
  , CookieJar
  , HttpException (..)
  , HttpExceptionContent (..)
  , Manager
  , Request
  , Response
  , brRead
  , applyBasicAuth
  , checkResponse
  , cookieJar
  , createCookieJar
  , destroyCookieJar
  , httpLbs
  , method
  , parseRequest
  , redirectCount
  , requestBody
  , requestHeaders
  , responseBody
  , responseCookieJar
  , responseHeaders
  , responseStatus
  , responseTimeout
  , responseTimeoutMicro
  , withResponse
  )
import qualified Network.HTTP.Client as HTTP
import           Network.HTTP.Client.TLS (mkManagerSettings)
import           Network.HTTP.Types
  ( Header
  , Method
  , hContentType
  , methodGet
  , methodHead
  , methodPost
  , methodPut
  , statusCode
  )
import           Network.HTTP.Types.URI (renderSimpleQuery)
import           Network.TLS (defaultSupported)
import           System.IO (IOMode (WriteMode), withBinaryFile)

newtype ProbeClient = ProbeClient Manager

newtype CookieSession = CookieSession (MVar CookieJar)

data ProbeResponse = ProbeResponse
  { probeStatus  :: Int
  , probeHeaders :: [Header]
  , probeBody    :: BS.ByteString
  }
  deriving stock (Eq, Show)

data RequestBody
  = NoBody
  | ByteBody BS.ByteString
  | FileBody FilePath
  | FormBody [(BS.ByteString, BS.ByteString)]
  deriving stock (Eq, Show)

data RequestOptions = RequestOptions
  { optionHeaders       :: [Header]
  , optionBody          :: RequestBody
  , optionCookies       :: Maybe CookieSession
  , optionBasicAuth     :: Maybe (BS.ByteString, BS.ByteString)
  , optionQuery         :: [(BS.ByteString, Maybe BS.ByteString)]
  , optionTimeoutMicros :: Maybe Int
  }

defaultRequestOptions :: RequestOptions
defaultRequestOptions = RequestOptions
  { optionHeaders = []
  , optionBody = NoBody
  , optionCookies = Nothing
  , optionBasicAuth = Nothing
  , optionQuery = []
  , optionTimeoutMicros = Nothing
  }

newProbeClient :: IO ProbeClient
newProbeClient =
  ProbeClient <$> HTTP.newManager
    (mkManagerSettings
      TLSSettingsSimple
        { settingDisableCertificateValidation = True
        , settingDisableSession = False
        , settingUseServerName = True
        , settingClientSupported = defaultSupported
        }
      Nothing)

newCookieSession :: IO CookieSession
newCookieSession = CookieSession <$> newMVar (createCookieJar [])

get :: ProbeClient -> String -> IO ProbeResponse
get client = getWith client defaultRequestOptions

getWith :: ProbeClient -> RequestOptions -> String -> IO ProbeResponse
getWith client = perform client methodGet

postWith :: ProbeClient -> RequestOptions -> String -> IO ProbeResponse
postWith client = perform client methodPost

putWith :: ProbeClient -> RequestOptions -> String -> IO ProbeResponse
putWith client = perform client methodPut

headWith :: ProbeClient -> RequestOptions -> String -> IO ProbeResponse
headWith client = perform client methodHead

downloadTo
  :: ProbeClient
  -> RequestOptions
  -> FilePath
  -> String
  -> IO ProbeResponse
downloadTo (ProbeClient manager) options outputPath url = do
  request <- buildRequest methodGet options url
  withResponse request manager $ \response -> do
    updateCookies options response
    withBinaryFile outputPath WriteMode $ \handle ->
      drain handle (responseBody response)
    pure (responseWithoutBody response)
  where
    drain handle reader = do
      chunk <- brRead reader
      if BS.null chunk
        then pure ()
        else BS.hPut handle chunk >> drain handle reader

perform
  :: ProbeClient
  -> Method
  -> RequestOptions
  -> String
  -> IO ProbeResponse
perform (ProbeClient manager) requestMethod options url = do
  request <- buildRequest requestMethod options url
  response <- httpLbs request manager
  updateCookies options response
  pure ProbeResponse
    { probeStatus = statusCode (responseStatus response)
    , probeHeaders = responseHeaders response
    , probeBody = BL.toStrict (responseBody response)
    }

buildRequest :: Method -> RequestOptions -> String -> IO Request
buildRequest requestMethod options url = do
  initial <- parseRequest url
  body <- toRequestBody (optionBody options)
  storedCookies <- traverse readCookies (optionCookies options)
  let formHeaders =
        case optionBody options of
          FormBody _ -> [(hContentType, "application/x-www-form-urlencoded")]
          _ -> []
      queried =
        if null (optionQuery options)
          then initial
          else HTTP.setQueryString (optionQuery options) initial
  let configured =
        queried
          { method = requestMethod
          , requestHeaders = formHeaders <> optionHeaders options
          , requestBody = body
          , cookieJar = storedCookies
          , redirectCount = 0
          , checkResponse = \_ _ -> pure ()
          , responseTimeout =
              maybe
                (responseTimeout initial)
                responseTimeoutMicro
                (optionTimeoutMicros options)
          }
  pure $ maybe configured
    (\(username, password) -> applyBasicAuth username password configured)
    (optionBasicAuth options)

toRequestBody :: RequestBody -> IO HTTP.RequestBody
toRequestBody NoBody = pure (HTTP.RequestBodyBS BS.empty)
toRequestBody (ByteBody body) = pure (HTTP.RequestBodyBS body)
toRequestBody (FileBody path) = HTTP.streamFile path
toRequestBody (FormBody fields) =
  pure . HTTP.RequestBodyBS $ renderSimpleQuery False fields

updateCookies :: RequestOptions -> Response body -> IO ()
updateCookies options response =
  case optionCookies options of
    Nothing -> pure ()
    Just (CookieSession storage) -> do
      _ <- swapMVar storage (responseCookieJar response)
      pure ()

readCookies :: CookieSession -> IO CookieJar
readCookies (CookieSession storage) = readMVar storage

cookieValue :: CookieSession -> BS.ByteString -> IO BS.ByteString
cookieValue session expectedName = do
  cookies <- destroyCookieJar <$> readCookies session
  maybe
    (fail ("cookie not found: " <> B8.unpack expectedName))
    (pure . cookie_value)
    (find ((== expectedName) . cookie_name) cookies)

cookieHeader :: CookieSession -> IO BS.ByteString
cookieHeader session = do
  cookies <- destroyCookieJar <$> readCookies session
  pure . B8.intercalate "; " $
    [cookie_name cookie <> "=" <> cookie_value cookie | cookie <- cookies]

headerValue :: BS.ByteString -> ProbeResponse -> Maybe BS.ByteString
headerValue name response =
  snd <$> find ((== CI.mk name) . fst) (probeHeaders response)

headerValues :: BS.ByteString -> ProbeResponse -> [BS.ByteString]
headerValues name response =
  map snd (filter ((== CI.mk name) . fst) (probeHeaders response))

responseWithoutBody :: Response body -> ProbeResponse
responseWithoutBody response = ProbeResponse
  { probeStatus = statusCode (responseStatus response)
  , probeHeaders = responseHeaders response
  , probeBody = BS.empty
  }

isConnectionFailure :: HttpException -> Bool
isConnectionFailure (HttpExceptionRequest _ (ConnectionFailure _)) = True
isConnectionFailure _ = False