{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE CPP #-}
module System.IO.Streams.TCP
( TCPConnection
, connect
, connectSocket
, socketToConnection
, defaultChunkSize
, bindAndListen
, bindAndListenWith
, accept
, acceptWith
) where
import qualified Control.Exception as E
import Control.Monad
import Data.Connection
import qualified Data.ByteString as B
import qualified Data.ByteString.Lazy.Internal as L
import qualified Network.Socket as N
import qualified Network.Socket.ByteString as NB
import qualified Network.Socket.ByteString.Lazy as NL
import qualified System.IO.Streams as S
import Foreign.Storable (sizeOf)
addrAny :: N.HostAddress
#if MIN_VERSION_network(2,7,0)
addrAny = N.tupleToHostAddress (0,0,0,0)
#else
addrAny = N.iNADDR_ANY
#endif
type TCPConnection = Connection (N.Socket, N.SockAddr)
defaultChunkSize :: Int
defaultChunkSize = 32 * k - chunkOverhead
where
k = 1024
chunkOverhead = 2 * sizeOf (undefined :: Int)
connectSocket :: N.HostName
-> N.PortNumber
-> IO (N.Socket, N.SockAddr)
connectSocket host port = do
(family, socketType, protocol, addr) <- resolveAddrInfo host port
E.bracketOnError (N.socket family socketType protocol)
N.close
(\sock -> do N.connect sock addr
N.setSocketOption sock N.NoDelay 1
return (sock, addr)
)
where
resolveAddrInfo host port = do
(addrInfo:_) <- N.getAddrInfo (Just hints) (Just host) (Just $ show port)
let family = N.addrFamily addrInfo
let socketType = N.addrSocketType addrInfo
let protocol = N.addrProtocol addrInfo
let addr = N.addrAddress addrInfo
return (family, socketType, protocol, addr)
where
hints = N.defaultHints {
N.addrFlags = [N.AI_ADDRCONFIG, N.AI_NUMERICSERV]
, N.addrSocketType = N.Stream
}
{-# INLINABLE resolveAddrInfo #-}
socketToConnection
:: Int
-> (N.Socket, N.SockAddr)
-> IO TCPConnection
socketToConnection bufsiz (sock, addr) = do
is <- S.makeInputStream $ do
s <- NB.recv sock bufsiz
return $! if B.null s then Nothing else Just s
return (Connection is (send sock) (N.close sock) (sock, addr))
where
send _ (L.Empty) = return ()
send sock (L.Chunk bs L.Empty) = unless (B.null bs) (NB.sendAll sock bs)
send sock lbs = NL.sendAll sock lbs
connect :: N.HostName
-> N.PortNumber
-> IO TCPConnection
connect host port = connectSocket host port >>= socketToConnection defaultChunkSize
bindAndListen :: Int
-> N.PortNumber
-> IO N.Socket
bindAndListen = bindAndListenWith $ \ sock -> do
N.setSocketOption sock N.ReuseAddr 1
N.setSocketOption sock N.NoDelay 1
bindAndListenWith :: (N.Socket -> IO ())
-> Int
-> N.PortNumber
-> IO N.Socket
bindAndListenWith f maxc port =
E.bracketOnError (N.socket N.AF_INET N.Stream 0)
N.close
(\sock -> do f sock
N.bind sock (N.SockAddrInet port addrAny)
N.listen sock maxc
return sock
)
accept :: N.Socket -> IO TCPConnection
accept = acceptWith (socketToConnection defaultChunkSize)
acceptWith :: ((N.Socket, N.SockAddr) -> IO TCPConnection)
-> N.Socket
-> IO TCPConnection
acceptWith f = f <=< N.accept