module Network.HTTP.Client.Manager
( ManagerSettings (..)
, newManager
, closeManager
, withManager
, getConn
, failedConnectionException
, defaultManagerSettings
, rawConnectionModifySocket
, proxyFromRequest
, noProxy
, useProxy
, proxyEnvironment
, proxyEnvironmentNamed
, defaultProxy
) where
#if !MIN_VERSION_base(4,6,0)
import Prelude hiding (catch)
#endif
import Data.Monoid (mappend)
import System.IO (hClose, hFlush, IOMode(..))
import qualified Data.IORef as I
import qualified Data.Map as Map
import qualified Data.ByteString.Char8 as S8
import qualified Data.ByteString.Lazy as L
import qualified Blaze.ByteString.Builder as Blaze
import Data.Text (Text)
import qualified Data.Text as T
import Data.Text.Read (decimal)
import Control.Monad.IO.Class (MonadIO, liftIO)
import Control.Monad (unless, join, when, void, mplus)
import Control.Exception (mask_, SomeException, bracket, catch, throwIO, fromException, mask, IOException, Exception (..), handle)
import Control.Concurrent (forkIO, threadDelay)
import Data.Time (UTCTime (..), Day (..), DiffTime, getCurrentTime, addUTCTime)
import Control.DeepSeq (deepseq)
import qualified Network.Socket as NS
import Data.Maybe (mapMaybe)
import System.IO (Handle)
import System.Mem.Weak (Weak, deRefWeak)
import Network.HTTP.Types (status200)
import Network.HTTP.Client.Types
import Network.HTTP.Client.Connection
import Network.HTTP.Client.Headers (parseStatusHeaders)
import Control.Concurrent.MVar (MVar, takeMVar, tryPutMVar, newEmptyMVar)
import System.Environment (getEnvironment)
import qualified Network.URI as U
import Control.Monad (guard)
rawConnectionModifySocket :: (NS.Socket -> IO ())
-> IO (Maybe NS.HostAddress -> String -> Int -> IO Connection)
rawConnectionModifySocket = return . openSocketConnection
rawConnectionModifySocketSize :: (NS.Socket -> IO ())
-> IO (Int -> Maybe NS.HostAddress -> String -> Int -> IO Connection)
rawConnectionModifySocketSize = return . openSocketConnectionSize
defaultManagerSettings :: ManagerSettings
defaultManagerSettings = ManagerSettings
{ managerConnCount = 10
, managerRawConnection = return $ openSocketConnection (const $ return ())
, managerTlsConnection = return $ \_ _ _ -> throwIO TlsNotSupported
, managerTlsProxyConnection = return $ \_ _ _ _ _ _ -> throwIO TlsNotSupported
, managerResponseTimeout = Just 30000000
, managerRetryableException = \e ->
case fromException e of
Just (_ :: IOException) -> True
_ ->
case fromException e of
Just NoResponseDataReceived -> True
Just IncompleteHeaders -> True
_ -> False
, managerWrapIOException =
let wrapper se =
case fromException se of
Just e -> toException $ InternalIOException e
Nothing -> se
in handle $ throwIO . wrapper
, managerIdleConnectionCount = 512
, managerModifyRequest = return
, managerProxyInsecure = defaultProxy
, managerProxySecure = defaultProxy
}
takeSocket :: Manager -> ConnKey -> IO (Maybe Connection)
takeSocket man key =
I.atomicModifyIORef (mConns man) go
where
go ManagerClosed = (ManagerClosed, Nothing)
go mcOrig@(ManagerOpen idleCount m) =
case Map.lookup key m of
Nothing -> (mcOrig, Nothing)
Just (One a _) ->
let mc = ManagerOpen (idleCount 1) (Map.delete key m)
in mc `seq` (mc, Just a)
Just (Cons a _ _ rest) ->
let mc = ManagerOpen (idleCount 1) (Map.insert key rest m)
in mc `seq` (mc, Just a)
putSocket :: Manager -> ConnKey -> Connection -> IO ()
putSocket man key ci = do
now <- getCurrentTime
join $ I.atomicModifyIORef (mConns man) (go now)
void $ tryPutMVar (mConnsBaton man) ()
where
go _ ManagerClosed = (ManagerClosed , connectionClose ci)
go now mc@(ManagerOpen idleCount m)
| idleCount >= mIdleConnectionCount man = (mc, connectionClose ci)
| otherwise = case Map.lookup key m of
Nothing ->
let cnt' = idleCount + 1
m' = ManagerOpen cnt' (Map.insert key (One ci now) m)
in m' `seq` (m', return ())
Just l ->
let (l', mx) = addToList now (mMaxConns man) ci l
cnt' = idleCount + maybe 0 (const 1) mx
m' = ManagerOpen cnt' (Map.insert key l' m)
in m' `seq` (m', maybe (return ()) connectionClose mx)
addToList :: UTCTime -> Int -> a -> NonEmptyList a -> (NonEmptyList a, Maybe a)
addToList _ i x l | i <= 1 = (l, Just x)
addToList now _ x l@One{} = (Cons x 2 now l, Nothing)
addToList now maxCount x l@(Cons _ currCount _ _)
| maxCount > currCount = (Cons x (currCount + 1) now l, Nothing)
| otherwise = (l, Just x)
newManager :: ManagerSettings -> IO Manager
newManager ms = do
NS.withSocketsDo $ return ()
rawConnection <- managerRawConnection ms
tlsConnection <- managerTlsConnection ms
tlsProxyConnection <- managerTlsProxyConnection ms
mapRef <- I.newIORef $! ManagerOpen 0 Map.empty
baton <- newEmptyMVar
wmapRef <- I.mkWeakIORef mapRef $ closeManager' mapRef
httpProxy <- runProxyOverride (managerProxyInsecure ms) False
httpsProxy <- runProxyOverride (managerProxySecure ms) True
_ <- forkIO $ reap baton wmapRef
let manager = Manager
{ mConns = mapRef
, mConnsBaton = baton
, mMaxConns = managerConnCount ms
, mResponseTimeout = managerResponseTimeout ms
, mRawConnection = rawConnection
, mTlsConnection = tlsConnection
, mTlsProxyConnection = tlsProxyConnection
, mRetryableException = managerRetryableException ms
, mWrapIOException = managerWrapIOException ms
, mIdleConnectionCount = managerIdleConnectionCount ms
, mModifyRequest = managerModifyRequest ms
, mSetProxy = \req ->
if secure req
then httpsProxy req
else httpProxy req
}
return manager
reap :: MVar () -> Weak (I.IORef ConnsMap) -> IO ()
reap baton wmapRef =
mask_ loop
where
loop = do
threadDelay (5 * 1000 * 1000)
mmapRef <- deRefWeak wmapRef
case mmapRef of
Nothing -> return ()
Just mapRef -> goMapRef mapRef
goMapRef mapRef = do
now <- getCurrentTime
let isNotStale time = 30 `addUTCTime` time >= now
(newMap, toDestroy) <- I.atomicModifyIORef mapRef $ \m ->
let (newMap, toDestroy) = findStaleWrap isNotStale m
in (newMap, (newMap, toDestroy))
mapM_ safeConnClose toDestroy
case newMap of
ManagerOpen _ m | not $ Map.null m -> return ()
_ -> takeMVar baton
loop
findStaleWrap _ ManagerClosed = (ManagerClosed, [])
findStaleWrap isNotStale (ManagerOpen idleCount m) =
let (x, y) = findStale isNotStale m
in (ManagerOpen (idleCount length y) x, y)
findStale isNotStale =
findStale' id id . Map.toList
where
findStale' destroy keep [] = (Map.fromList $ keep [], destroy [])
findStale' destroy keep ((connkey, nelist):rest) =
findStale' destroy' keep' rest
where
(notStale, stale) = span (isNotStale . fst) $ neToList nelist
destroy' = destroy . (map snd stale++)
keep' =
case neFromList notStale of
Nothing -> keep
Just x -> keep . ((connkey, x):)
flushStaleCerts now =
Map.fromList . mapMaybe flushStaleCerts' . Map.toList
where
flushStaleCerts' (host', inner) =
case mapMaybe flushStaleCerts'' $ Map.toList inner of
[] -> Nothing
pairs ->
let x = take 10 pairs
in x `seqPairs` Just (host', Map.fromList x)
flushStaleCerts'' (certs, expires)
| expires > now = Just (certs, expires)
| otherwise = Nothing
seqPairs :: [(L.ByteString, UTCTime)] -> b -> b
seqPairs [] b = b
seqPairs (p:ps) b = p `seqPair` ps `seqPairs` b
seqPair :: (L.ByteString, UTCTime) -> b -> b
seqPair (lbs, utc) b = lbs `seqLBS` utc `seqUTC` b
seqLBS :: L.ByteString -> b -> b
seqLBS lbs b = L.length lbs `seq` b
seqUTC :: UTCTime -> b -> b
seqUTC (UTCTime day dt) b = day `seqDay` dt `seqDT` b
seqDay :: Day -> b -> b
seqDay (ModifiedJulianDay i) b = i `deepseq` b
seqDT :: DiffTime -> b -> b
seqDT = seq
neToList :: NonEmptyList a -> [(UTCTime, a)]
neToList (One a t) = [(t, a)]
neToList (Cons a _ t nelist) = (t, a) : neToList nelist
neFromList :: [(UTCTime, a)] -> Maybe (NonEmptyList a)
neFromList [] = Nothing
neFromList [(t, a)] = Just (One a t)
neFromList xs =
Just . snd . go $ xs
where
go [] = error "neFromList.go []"
go [(t, a)] = (2, One a t)
go ((t, a):rest) =
let (i, rest') = go rest
i' = i + 1
in i' `seq` (i', Cons a i t rest')
closeManager :: Manager -> IO ()
closeManager = closeManager' . mConns
closeManager' :: I.IORef ConnsMap
-> IO ()
closeManager' connsRef = mask_ $ do
!m <- I.atomicModifyIORef connsRef $ \x -> (ManagerClosed, x)
case m of
ManagerClosed -> return ()
ManagerOpen _ m -> mapM_ (nonEmptyMapM_ safeConnClose) $ Map.elems m
withManager :: ManagerSettings -> (Manager -> IO a) -> IO a
withManager settings = bracket (newManager settings) closeManager
safeConnClose :: Connection -> IO ()
safeConnClose ci = connectionClose ci `catch` \(_ :: IOException) -> return ()
nonEmptyMapM_ :: Monad m => (a -> m ()) -> NonEmptyList a -> m ()
nonEmptyMapM_ f (One x _) = f x
nonEmptyMapM_ f (Cons x _ _ l) = f x >> nonEmptyMapM_ f l
getManagedConn
:: Manager
-> ConnKey
-> IO Connection
-> IO (ConnRelease, Connection, ManagedConn)
getManagedConn man key open = mask $ \restore -> do
mci <- takeSocket man key
(ci, isManaged) <-
case mci of
Nothing -> do
ci <- restore open
return (ci, Fresh)
Just ci -> return (ci, Reused)
toReuseRef <- I.newIORef DontReuse
wasReleasedRef <- I.newIORef False
let connRelease r = do
I.writeIORef toReuseRef r
releaseHelper
releaseHelper = mask $ \restore -> do
wasReleased <- I.atomicModifyIORef wasReleasedRef $ \x -> (True, x)
unless wasReleased $ do
toReuse <- I.readIORef toReuseRef
restore $ case toReuse of
Reuse -> putSocket man key ci
DontReuse -> connectionClose ci
return (connRelease, ci, isManaged)
failedConnectionException :: Request -> HttpException
failedConnectionException req =
FailedConnectionException host' port'
where
(_, host', port') = getConnDest req
getConnDest :: Request -> (Bool, String, Int)
getConnDest req =
case proxy req of
Just p -> (True, S8.unpack (proxyHost p), proxyPort p)
Nothing -> (False, S8.unpack $ host req, port req)
getConn :: Request
-> Manager
-> IO (ConnRelease, Connection, ManagedConn)
getConn req m
| S8.null h = throwIO $ InvalidDestinationHost h
| otherwise =
getManagedConn m (ConnKey connKeyHost connport (host req) (port req) (secure req)) $
wrapConnectExc $ go connaddr connhost connport
where
h = host req
(useProxy, connhost, connport) = getConnDest req
(connaddr, connKeyHost) =
case (hostAddress req, useProxy) of
(Just ha, False) -> (Just ha, HostAddress ha)
_ -> (Nothing, HostName $ T.pack connhost)
wrapConnectExc = handle $ \e ->
throwIO $ FailedConnectionException2 connhost connport (secure req)
(toException (e :: IOException))
go =
case (secure req, useProxy) of
(False, _) -> mRawConnection m
(True, False) -> mTlsConnection m
(True, True) ->
let ultHost = host req
ultPort = port req
connstr = S8.concat
[ "CONNECT "
, ultHost
, ":"
, S8.pack $ show ultPort
, " HTTP/1.1\r\n\r\n"
]
parse conn = do
sh@(StatusHeaders status _ _) <- parseStatusHeaders conn
unless (status == status200) $
throwIO $ ProxyConnectException ultHost ultPort $ Left $ S8.pack $ show sh
in mTlsProxyConnection m connstr parse (S8.unpack ultHost)
proxyFromRequest :: ProxyOverride
proxyFromRequest = ProxyOverride $ const $ return id
noProxy :: ProxyOverride
noProxy = ProxyOverride $ const $ return $ \req -> req { proxy = Nothing }
useProxy :: Proxy -> ProxyOverride
useProxy p = ProxyOverride $ const $ return $ \req -> req { proxy = Just p }
proxyEnvironment :: Maybe Proxy
-> ProxyOverride
proxyEnvironment mp = ProxyOverride $ \secure ->
envHelper (envName secure) $ maybe EHNoProxy EHUseProxy mp
envName :: Bool
-> Text
envName False = "http_proxy"
envName True = "https_proxy"
proxyEnvironmentNamed
:: Text
-> Maybe Proxy
-> ProxyOverride
proxyEnvironmentNamed name =
ProxyOverride . const . envHelper name
. maybe EHNoProxy EHUseProxy
defaultProxy :: ProxyOverride
defaultProxy = ProxyOverride $ \secure ->
envHelper (envName secure) EHFromRequest
data EnvHelper = EHFromRequest
| EHNoProxy
| EHUseProxy Proxy
envHelper :: Text -> EnvHelper -> IO (Request -> Request)
envHelper name eh = do
env <- getEnvironment
case lookup (T.unpack name) env of
Nothing -> return $
case eh of
EHFromRequest -> id
EHNoProxy -> \req -> req { proxy = Nothing }
EHUseProxy p -> \req -> req { proxy = Just p }
Just str -> do
let invalid = throwIO $ InvalidProxyEnvironmentVariable name (T.pack str)
p <- maybe invalid return $ do
uri <- case U.parseURI str of
Just u | U.uriScheme u == "http:" -> return u
_ -> U.parseURI $ "http://" ++ str
guard $ U.uriScheme uri == "http:"
guard $ null (U.uriPath uri) || U.uriPath uri == "/"
guard $ null $ U.uriQuery uri
guard $ null $ U.uriFragment uri
auth <- U.uriAuthority uri
guard $ null $ U.uriUserInfo auth
port <-
case U.uriPort auth of
"" -> Just 80
':':rest ->
case decimal $ T.pack rest of
Right (p, "") -> Just p
_ -> Nothing
_ -> Nothing
Just $ Proxy (S8.pack $ U.uriRegName auth) port
return $ \req -> req { proxy = Just p }