{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ViewPatterns #-}
module Network.HTTP.Client.Connection
    ( connectionReadLine
    , connectionReadLineWith
    , connectionDropTillBlankLine
    , dummyConnection
    , openSocketConnection
    , openSocketConnectionSize
    , makeConnection
    , socketConnection
    , withSocket
    ) where

import Data.ByteString (ByteString, empty)
import Data.IORef
import Control.Monad
import Network.HTTP.Client.Types
import Network.Socket (Socket, HostAddress)
import qualified Network.Socket as NS
import Network.Socket.ByteString (sendAll, recv)
import qualified Control.Exception as E
import qualified Data.ByteString as S
import Data.Word (Word8)
import Data.Function (fix)

connectionReadLine :: Connection -> IO ByteString
connectionReadLine :: Connection -> IO ByteString
connectionReadLine Connection
conn = do
    ByteString
bs <- Connection -> IO ByteString
connectionRead Connection
conn
    Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (ByteString -> Bool
S.null ByteString
bs) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ HttpExceptionContent -> IO ()
forall a. HttpExceptionContent -> IO a
throwHttp HttpExceptionContent
IncompleteHeaders
    Connection -> ByteString -> IO ByteString
connectionReadLineWith Connection
conn ByteString
bs

-- | Keep dropping input until a blank line is found.
connectionDropTillBlankLine :: Connection -> IO ()
connectionDropTillBlankLine :: Connection -> IO ()
connectionDropTillBlankLine Connection
conn = (IO () -> IO ()) -> IO ()
forall a. (a -> a) -> a
fix ((IO () -> IO ()) -> IO ()) -> (IO () -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \IO ()
loop -> do
    ByteString
bs <- Connection -> IO ByteString
connectionReadLine Connection
conn
    Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (ByteString -> Bool
S.null ByteString
bs) IO ()
loop

connectionReadLineWith :: Connection -> ByteString -> IO ByteString
connectionReadLineWith :: Connection -> ByteString -> IO ByteString
connectionReadLineWith Connection
conn ByteString
bs0 =
    ByteString
-> ([ByteString] -> [ByteString]) -> Int -> IO ByteString
go ByteString
bs0 [ByteString] -> [ByteString]
forall a. a -> a
id Int
0
  where
    go :: ByteString
-> ([ByteString] -> [ByteString]) -> Int -> IO ByteString
go ByteString
bs [ByteString] -> [ByteString]
front Int
total =
        case (Word8 -> Bool) -> ByteString -> (ByteString, ByteString)
S.break (Word8 -> Word8 -> Bool
forall a. Eq a => a -> a -> Bool
== Word8
charLF) ByteString
bs of
            (ByteString
_, ByteString
"") -> do
                let total' :: Int
total' = Int
total Int -> Int -> Int
forall a. Num a => a -> a -> a
+ ByteString -> Int
S.length ByteString
bs
                Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
total' Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
4096) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ HttpExceptionContent -> IO ()
forall a. HttpExceptionContent -> IO a
throwHttp HttpExceptionContent
OverlongHeaders
                ByteString
bs' <- Connection -> IO ByteString
connectionRead Connection
conn
                Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (ByteString -> Bool
S.null ByteString
bs') (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ HttpExceptionContent -> IO ()
forall a. HttpExceptionContent -> IO a
throwHttp HttpExceptionContent
IncompleteHeaders
                ByteString
-> ([ByteString] -> [ByteString]) -> Int -> IO ByteString
go ByteString
bs' ([ByteString] -> [ByteString]
front ([ByteString] -> [ByteString])
-> ([ByteString] -> [ByteString]) -> [ByteString] -> [ByteString]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (ByteString
bsByteString -> [ByteString] -> [ByteString]
forall a. a -> [a] -> [a]
:)) Int
total'
            (ByteString
x, Int -> ByteString -> ByteString
S.drop Int
1 -> ByteString
y) -> do
                Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (ByteString -> Bool
S.null ByteString
y) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$! Connection -> ByteString -> IO ()
connectionUnread Connection
conn ByteString
y
                ByteString -> IO ByteString
forall (m :: * -> *) a. Monad m => a -> m a
return (ByteString -> IO ByteString) -> ByteString -> IO ByteString
forall a b. (a -> b) -> a -> b
$! ByteString -> ByteString
killCR (ByteString -> ByteString) -> ByteString -> ByteString
forall a b. (a -> b) -> a -> b
$! [ByteString] -> ByteString
S.concat ([ByteString] -> ByteString) -> [ByteString] -> ByteString
forall a b. (a -> b) -> a -> b
$! [ByteString] -> [ByteString]
front [ByteString
x]

charLF, charCR :: Word8
charLF :: Word8
charLF = Word8
10
charCR :: Word8
charCR = Word8
13

killCR :: ByteString -> ByteString
killCR :: ByteString -> ByteString
killCR ByteString
bs
    | ByteString -> Bool
S.null ByteString
bs = ByteString
bs
    | ByteString -> Word8
S.last ByteString
bs Word8 -> Word8 -> Bool
forall a. Eq a => a -> a -> Bool
== Word8
charCR = ByteString -> ByteString
S.init ByteString
bs
    | Bool
otherwise = ByteString
bs

-- | For testing
dummyConnection :: [ByteString] -- ^ input
                -> IO (Connection, IO [ByteString], IO [ByteString]) -- ^ conn, output, input
dummyConnection :: [ByteString] -> IO (Connection, IO [ByteString], IO [ByteString])
dummyConnection [ByteString]
input0 = do
    IORef [ByteString]
iinput <- [ByteString] -> IO (IORef [ByteString])
forall a. a -> IO (IORef a)
newIORef [ByteString]
input0
    IORef [ByteString]
ioutput <- [ByteString] -> IO (IORef [ByteString])
forall a. a -> IO (IORef a)
newIORef []
    (Connection, IO [ByteString], IO [ByteString])
-> IO (Connection, IO [ByteString], IO [ByteString])
forall (m :: * -> *) a. Monad m => a -> m a
return (Connection :: IO ByteString
-> (ByteString -> IO ())
-> (ByteString -> IO ())
-> IO ()
-> Connection
Connection
        { connectionRead :: IO ByteString
connectionRead = IORef [ByteString]
-> ([ByteString] -> ([ByteString], ByteString)) -> IO ByteString
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef IORef [ByteString]
iinput (([ByteString] -> ([ByteString], ByteString)) -> IO ByteString)
-> ([ByteString] -> ([ByteString], ByteString)) -> IO ByteString
forall a b. (a -> b) -> a -> b
$ \[ByteString]
input ->
            case [ByteString]
input of
                [] -> ([], ByteString
empty)
                ByteString
x:[ByteString]
xs -> ([ByteString]
xs, ByteString
x)
        , connectionUnread :: ByteString -> IO ()
connectionUnread = \ByteString
x -> IORef [ByteString] -> ([ByteString] -> ([ByteString], ())) -> IO ()
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef IORef [ByteString]
iinput (([ByteString] -> ([ByteString], ())) -> IO ())
-> ([ByteString] -> ([ByteString], ())) -> IO ()
forall a b. (a -> b) -> a -> b
$ \[ByteString]
input -> (ByteString
xByteString -> [ByteString] -> [ByteString]
forall a. a -> [a] -> [a]
:[ByteString]
input, ())
        , connectionWrite :: ByteString -> IO ()
connectionWrite = \ByteString
x -> IORef [ByteString] -> ([ByteString] -> ([ByteString], ())) -> IO ()
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef IORef [ByteString]
ioutput (([ByteString] -> ([ByteString], ())) -> IO ())
-> ([ByteString] -> ([ByteString], ())) -> IO ()
forall a b. (a -> b) -> a -> b
$ \[ByteString]
output -> ([ByteString]
output [ByteString] -> [ByteString] -> [ByteString]
forall a. [a] -> [a] -> [a]
++ [ByteString
x], ())
        , connectionClose :: IO ()
connectionClose = () -> IO ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
        }, IORef [ByteString]
-> ([ByteString] -> ([ByteString], [ByteString]))
-> IO [ByteString]
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef IORef [ByteString]
ioutput (([ByteString] -> ([ByteString], [ByteString])) -> IO [ByteString])
-> ([ByteString] -> ([ByteString], [ByteString]))
-> IO [ByteString]
forall a b. (a -> b) -> a -> b
$ \[ByteString]
output -> ([], [ByteString]
output), IORef [ByteString] -> IO [ByteString]
forall a. IORef a -> IO a
readIORef IORef [ByteString]
iinput)

-- | Create a new 'Connection' from a read, write, and close function.
--
-- @since 0.5.3
makeConnection :: IO ByteString -- ^ read
               -> (ByteString -> IO ()) -- ^ write
               -> IO () -- ^ close
               -> IO Connection
makeConnection :: IO ByteString -> (ByteString -> IO ()) -> IO () -> IO Connection
makeConnection IO ByteString
r ByteString -> IO ()
w IO ()
c = do
    IORef [ByteString]
istack <- [ByteString] -> IO (IORef [ByteString])
forall a. a -> IO (IORef a)
newIORef []

    -- it is necessary to make sure we never read from or write to
    -- already closed connection.
    IORef Bool
closedVar <- Bool -> IO (IORef Bool)
forall a. a -> IO (IORef a)
newIORef Bool
False

    let close :: IO ()
close = do
          Bool
closed <- IORef Bool -> (Bool -> (Bool, Bool)) -> IO Bool
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef IORef Bool
closedVar (\Bool
closed -> (Bool
True, Bool
closed))
          Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless Bool
closed (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$
            IO ()
c

    Weak (IORef [ByteString])
_ <- IORef [ByteString] -> IO () -> IO (Weak (IORef [ByteString]))
forall a. IORef a -> IO () -> IO (Weak (IORef a))
mkWeakIORef IORef [ByteString]
istack IO ()
close
    Connection -> IO Connection
forall (m :: * -> *) a. Monad m => a -> m a
return (Connection -> IO Connection) -> Connection -> IO Connection
forall a b. (a -> b) -> a -> b
$! Connection :: IO ByteString
-> (ByteString -> IO ())
-> (ByteString -> IO ())
-> IO ()
-> Connection
Connection
        { connectionRead :: IO ByteString
connectionRead = do
            Bool
closed <- IORef Bool -> IO Bool
forall a. IORef a -> IO a
readIORef IORef Bool
closedVar
            Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
closed (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ HttpExceptionContent -> IO ()
forall a. HttpExceptionContent -> IO a
throwHttp HttpExceptionContent
ConnectionClosed
            IO (IO ByteString) -> IO ByteString
forall (m :: * -> *) a. Monad m => m (m a) -> m a
join (IO (IO ByteString) -> IO ByteString)
-> IO (IO ByteString) -> IO ByteString
forall a b. (a -> b) -> a -> b
$ IORef [ByteString]
-> ([ByteString] -> ([ByteString], IO ByteString))
-> IO (IO ByteString)
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef IORef [ByteString]
istack (([ByteString] -> ([ByteString], IO ByteString))
 -> IO (IO ByteString))
-> ([ByteString] -> ([ByteString], IO ByteString))
-> IO (IO ByteString)
forall a b. (a -> b) -> a -> b
$ \[ByteString]
stack ->
              case [ByteString]
stack of
                  ByteString
x:[ByteString]
xs -> ([ByteString]
xs, ByteString -> IO ByteString
forall (m :: * -> *) a. Monad m => a -> m a
return ByteString
x)
                  [] -> ([], IO ByteString
r)

        , connectionUnread :: ByteString -> IO ()
connectionUnread = \ByteString
x -> do
            Bool
closed <- IORef Bool -> IO Bool
forall a. IORef a -> IO a
readIORef IORef Bool
closedVar
            Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
closed (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ HttpExceptionContent -> IO ()
forall a. HttpExceptionContent -> IO a
throwHttp HttpExceptionContent
ConnectionClosed
            IORef [ByteString] -> ([ByteString] -> ([ByteString], ())) -> IO ()
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef IORef [ByteString]
istack (([ByteString] -> ([ByteString], ())) -> IO ())
-> ([ByteString] -> ([ByteString], ())) -> IO ()
forall a b. (a -> b) -> a -> b
$ \[ByteString]
stack -> (ByteString
xByteString -> [ByteString] -> [ByteString]
forall a. a -> [a] -> [a]
:[ByteString]
stack, ())

        , connectionWrite :: ByteString -> IO ()
connectionWrite = \ByteString
x -> do
            Bool
closed <- IORef Bool -> IO Bool
forall a. IORef a -> IO a
readIORef IORef Bool
closedVar
            Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
closed (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ HttpExceptionContent -> IO ()
forall a. HttpExceptionContent -> IO a
throwHttp HttpExceptionContent
ConnectionClosed
            ByteString -> IO ()
w ByteString
x

        , connectionClose :: IO ()
connectionClose = IO ()
close
        }

-- | Create a new 'Connection' from a 'Socket'.
--
-- @since 0.5.3
socketConnection :: Socket
                 -> Int -- ^ chunk size
                 -> IO Connection
socketConnection :: Socket -> Int -> IO Connection
socketConnection Socket
socket Int
chunksize = IO ByteString -> (ByteString -> IO ()) -> IO () -> IO Connection
makeConnection
    (Socket -> Int -> IO ByteString
recv Socket
socket Int
chunksize)
    (Socket -> ByteString -> IO ()
sendAll Socket
socket)
    (Socket -> IO ()
NS.close Socket
socket)

openSocketConnection :: (Socket -> IO ())
                     -> Maybe HostAddress
                     -> String -- ^ host
                     -> Int -- ^ port
                     -> IO Connection
openSocketConnection :: (Socket -> IO ())
-> Maybe HostAddress -> String -> Int -> IO Connection
openSocketConnection Socket -> IO ()
f = (Socket -> IO ())
-> Int -> Maybe HostAddress -> String -> Int -> IO Connection
openSocketConnectionSize Socket -> IO ()
f Int
8192

openSocketConnectionSize :: (Socket -> IO ())
                         -> Int -- ^ chunk size
                         -> Maybe HostAddress
                         -> String -- ^ host
                         -> Int -- ^ port
                         -> IO Connection
openSocketConnectionSize :: (Socket -> IO ())
-> Int -> Maybe HostAddress -> String -> Int -> IO Connection
openSocketConnectionSize Socket -> IO ()
tweakSocket Int
chunksize Maybe HostAddress
hostAddress' String
host' Int
port' =
    (Socket -> IO ())
-> Maybe HostAddress
-> String
-> Int
-> (Socket -> IO Connection)
-> IO Connection
forall a.
(Socket -> IO ())
-> Maybe HostAddress -> String -> Int -> (Socket -> IO a) -> IO a
withSocket Socket -> IO ()
tweakSocket Maybe HostAddress
hostAddress' String
host' Int
port' ((Socket -> IO Connection) -> IO Connection)
-> (Socket -> IO Connection) -> IO Connection
forall a b. (a -> b) -> a -> b
$ \ Socket
sock ->
        Socket -> Int -> IO Connection
socketConnection Socket
sock Int
chunksize

withSocket :: (Socket -> IO ())
           -> Maybe HostAddress
           -> String -- ^ host
           -> Int -- ^ port
           -> (Socket -> IO a)
           -> IO a
withSocket :: (Socket -> IO ())
-> Maybe HostAddress -> String -> Int -> (Socket -> IO a) -> IO a
withSocket Socket -> IO ()
tweakSocket Maybe HostAddress
hostAddress' String
host' Int
port' Socket -> IO a
f = do
    let hints :: AddrInfo
hints = AddrInfo
NS.defaultHints { addrSocketType :: SocketType
NS.addrSocketType = SocketType
NS.Stream }
    [AddrInfo]
addrs <- case Maybe HostAddress
hostAddress' of
        Maybe HostAddress
Nothing ->
            Maybe AddrInfo -> Maybe String -> Maybe String -> IO [AddrInfo]
NS.getAddrInfo (AddrInfo -> Maybe AddrInfo
forall a. a -> Maybe a
Just AddrInfo
hints) (String -> Maybe String
forall a. a -> Maybe a
Just String
host') (String -> Maybe String
forall a. a -> Maybe a
Just (String -> Maybe String) -> String -> Maybe String
forall a b. (a -> b) -> a -> b
$ Int -> String
forall a. Show a => a -> String
show Int
port')
        Just HostAddress
ha ->
            [AddrInfo] -> IO [AddrInfo]
forall (m :: * -> *) a. Monad m => a -> m a
return
                [AddrInfo :: [AddrInfoFlag]
-> Family
-> SocketType
-> ProtocolNumber
-> SockAddr
-> Maybe String
-> AddrInfo
NS.AddrInfo
                 { addrFlags :: [AddrInfoFlag]
NS.addrFlags = []
                 , addrFamily :: Family
NS.addrFamily = Family
NS.AF_INET
                 , addrSocketType :: SocketType
NS.addrSocketType = SocketType
NS.Stream
                 , addrProtocol :: ProtocolNumber
NS.addrProtocol = ProtocolNumber
6 -- tcp
                 , addrAddress :: SockAddr
NS.addrAddress = PortNumber -> HostAddress -> SockAddr
NS.SockAddrInet (Int -> PortNumber
forall a. Enum a => Int -> a
toEnum Int
port') HostAddress
ha
                 , addrCanonName :: Maybe String
NS.addrCanonName = Maybe String
forall a. Maybe a
Nothing
                 }]

    IO Socket -> (Socket -> IO ()) -> (Socket -> IO a) -> IO a
forall a b c. IO a -> (a -> IO b) -> (a -> IO c) -> IO c
E.bracketOnError ([AddrInfo] -> (AddrInfo -> IO Socket) -> IO Socket
forall a. [AddrInfo] -> (AddrInfo -> IO a) -> IO a
firstSuccessful [AddrInfo]
addrs ((AddrInfo -> IO Socket) -> IO Socket)
-> (AddrInfo -> IO Socket) -> IO Socket
forall a b. (a -> b) -> a -> b
$ (Socket -> IO ()) -> AddrInfo -> IO Socket
forall a. (Socket -> IO a) -> AddrInfo -> IO Socket
openSocket Socket -> IO ()
tweakSocket) Socket -> IO ()
NS.close Socket -> IO a
f

openSocket :: (Socket -> IO a) -> AddrInfo -> IO Socket
openSocket Socket -> IO a
tweakSocket AddrInfo
addr =
    IO Socket
-> (Socket -> IO ()) -> (Socket -> IO Socket) -> IO Socket
forall a b c. IO a -> (a -> IO b) -> (a -> IO c) -> IO c
E.bracketOnError
        (Family -> SocketType -> ProtocolNumber -> IO Socket
NS.socket (AddrInfo -> Family
NS.addrFamily AddrInfo
addr) (AddrInfo -> SocketType
NS.addrSocketType AddrInfo
addr)
                   (AddrInfo -> ProtocolNumber
NS.addrProtocol AddrInfo
addr))
        Socket -> IO ()
NS.close
        (\Socket
sock -> do
            Socket -> SocketOption -> Int -> IO ()
NS.setSocketOption Socket
sock SocketOption
NS.NoDelay Int
1
            Socket -> IO a
tweakSocket Socket
sock
            Socket -> SockAddr -> IO ()
NS.connect Socket
sock (AddrInfo -> SockAddr
NS.addrAddress AddrInfo
addr)
            Socket -> IO Socket
forall (m :: * -> *) a. Monad m => a -> m a
return Socket
sock)

firstSuccessful :: [NS.AddrInfo] -> (NS.AddrInfo -> IO a) -> IO a
firstSuccessful :: [AddrInfo] -> (AddrInfo -> IO a) -> IO a
firstSuccessful []     AddrInfo -> IO a
_  = String -> IO a
forall a. HasCallStack => String -> a
error String
"getAddrInfo returned empty list"
firstSuccessful (AddrInfo
a:[AddrInfo]
as) AddrInfo -> IO a
cb =
    AddrInfo -> IO a
cb AddrInfo
a IO a -> (IOException -> IO a) -> IO a
forall e a. Exception e => IO a -> (e -> IO a) -> IO a
`E.catch` \(IOException
e :: E.IOException) ->
        case [AddrInfo]
as of
            [] -> IOException -> IO a
forall e a. Exception e => e -> IO a
E.throwIO IOException
e
            [AddrInfo]
_  -> [AddrInfo] -> (AddrInfo -> IO a) -> IO a
forall a. [AddrInfo] -> (AddrInfo -> IO a) -> IO a
firstSuccessful [AddrInfo]
as AddrInfo -> IO a
cb