{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE OverloadedStrings #-}
module Test.Fetch
( fetch
)
where
import qualified Data.ByteString as BS
import qualified Data.ByteString.Char8 as C8
import qualified Data.ByteString.Lazy as LBS
import Control.Monad.Trans.Resource (runResourceT)
import Data.Conduit (ConduitT, SealedConduitT, Void,
await, sealConduitT, ($$+-))
import qualified Data.Conduit.Binary as CB
import Data.Int (Int64)
import Data.Maybe (fromMaybe)
import Network.Connection (TLSSettings (..))
import Network.HTTP.Client (Request, newManager,
responseBody, responseHeaders,
responseStatus, secure, managerSetProxy, proxyFromRequest)
import Network.HTTP.Client.TLS (mkManagerSettings)
import Network.HTTP.Simple (setRequestManager, withResponse)
import Network.HTTP.Types (hContentLength, statusCode)
import Test.HttpReply
fetch :: Request -> IO HttpReply
fetch req = do
let mSettings = mkManagerSettings (TLSSettingsSimple True False False) Nothing
mSettings' = managerSetProxy proxyFromRequest mSettings
m <- newManager mSettings'
runResourceT $ withResponse (setRequestManager m req) getSrc
where
getSrc resp = do
let mbBodyLength = readInt64 <$> lookup hContentLength (responseHeaders resp)
bodyText <- checkBodySize (sealConduitT $ responseBody resp) mbBodyLength
return $ HttpReply (secure req) (statusCode $ responseStatus resp) (responseHeaders resp) bodyText
bodyCheckBlock :: Int64
bodyCheckBlock = 1000
checkBodySize :: (Monad m) => SealedConduitT () BS.ByteString m () -> Maybe Int64 -> m BS.ByteString
checkBodySize bodySrc Nothing = fmap (BS.concat . LBS.toChunks) $ bodySrc $$+- CB.take $ fromIntegral bodyCheckBlock
checkBodySize bodySrc (Just len)
| len <= bodyCheckBlock = checkBodySize bodySrc Nothing
| otherwise = fromMaybe "Success" <$> (bodySrc $$+- sizeCheckSink len)
-- A pipe that counts the size of each incoming C8.Bytestring, when the last is
-- received, the result is Nothing if the size matches the expected value or an
-- error message if it does not.
sizeCheckSink :: Monad m => Int64 -> ConduitT C8.ByteString Void m (Maybe C8.ByteString)
sizeCheckSink expectedSize = sink 0
where
sink !count = await >>= maybe (closeSink count) (sinkBlock count)
sinkBlock !count bs = sink (count + fromIntegral (BS.length bs))
-- | no more bytes: return Nothing if the count so far matches the expected value
closeSink !count
| count == expectedSize = return Nothing
| otherwise = return $ Just . C8.pack $ "Error : Body length " ++ show count
++ " should have been " ++ show expectedSize ++ "."
readInt64 :: C8.ByteString -> Int64
readInt64 = read . C8.unpack