module Data.Conduit.Network
(
sourceSocket
, sinkSocket
, Application
, ServerSettings (..)
, runTCPServer
, ClientSettings (..)
, runTCPClient
, HostPreference (..)
, bindPort
, getSocket
) where
import Data.Conduit
import qualified Network.Socket as NS
import Network.Socket (Socket)
import Network.Socket.ByteString (sendAll, recv)
import Data.ByteString (ByteString)
import qualified Data.ByteString as S
import Control.Monad.IO.Class (MonadIO (liftIO))
import Control.Exception (bracketOnError, IOException, throwIO, SomeException, try, finally, bracket)
import Control.Monad (forever)
import Control.Monad.Trans.Control (MonadBaseControl, control)
import Control.Monad.Trans.Class (lift)
import Control.Concurrent (forkIO)
import Data.String (IsString (fromString))
import qualified Control.Exception as E
sourceSocket :: MonadIO m => Socket -> GSource m ByteString
sourceSocket socket =
loop
where
loop = do
bs <- lift $ liftIO $ recv socket 4096
if S.null bs
then return ()
else yield bs >> loop
sinkSocket :: MonadIO m => Socket -> GInfSink ByteString m
sinkSocket socket =
loop
where
loop = awaitE >>= either return (\bs -> lift (liftIO $ sendAll socket bs) >> loop)
type Application m = Source m ByteString
-> Sink ByteString m ()
-> m ()
data ServerSettings = ServerSettings
{ serverPort :: Int
, serverHost :: HostPreference
}
deriving (Eq, Show, Read)
runTCPServer :: (MonadIO m, MonadBaseControl IO m) => ServerSettings -> Application m -> m ()
runTCPServer (ServerSettings port host) app = control $ \run -> bracket
(liftIO $ bindPort port host)
(liftIO . NS.sClose)
(run . forever . serve)
where
serve lsocket = do
(socket, _addr) <- liftIO $ NS.accept lsocket
let src = sourceSocket socket
sink = sinkSocket socket
app' run = run (app src sink) >> return ()
appClose run = app' run `finally` NS.sClose socket
control $ \run -> forkIO (appClose run) >> run (return ())
data ClientSettings = ClientSettings
{ clientPort :: Int
, clientHost :: String
}
deriving (Eq, Show, Read)
runTCPClient :: (MonadIO m, MonadBaseControl IO m) => ClientSettings -> Application m -> m ()
runTCPClient (ClientSettings port host) app = control $ \run -> bracket
(getSocket host port)
NS.sClose
(\s -> run $ app (sourceSocket s) (sinkSocket s))
getSocket :: String -> Int -> IO NS.Socket
getSocket host' port' = do
let hints = NS.defaultHints {
NS.addrFlags = [NS.AI_ADDRCONFIG]
, NS.addrSocketType = NS.Stream
}
(addr:_) <- NS.getAddrInfo (Just hints) (Just host') (Just $ show port')
sock <- NS.socket (NS.addrFamily addr) (NS.addrSocketType addr)
(NS.addrProtocol addr)
ee <- try' $ NS.connect sock (NS.addrAddress addr)
case ee of
Left e -> NS.sClose sock >> throwIO e
Right () -> return sock
where
try' :: IO a -> IO (Either SomeException a)
try' = try
data HostPreference =
HostAny
| HostIPv4
| HostIPv6
| Host String
deriving (Eq, Ord, Show, Read)
instance IsString HostPreference where
fromString s'@('*':s) =
case s of
[] -> HostAny
['4'] -> HostIPv4
['6'] -> HostIPv6
_ -> Host s'
fromString s = Host s
bindPort :: Int -> HostPreference -> IO Socket
bindPort p s = do
let hints = NS.defaultHints
{ NS.addrFlags = [ NS.AI_PASSIVE
, NS.AI_NUMERICSERV
, NS.AI_NUMERICHOST
]
, NS.addrSocketType = NS.Stream
}
host =
case s of
Host s' -> Just s'
_ -> Nothing
port = Just . show $ p
addrs <- NS.getAddrInfo (Just hints) host port
let addrs4 = filter (\x -> NS.addrFamily x /= NS.AF_INET6) addrs
addrs6 = filter (\x -> NS.addrFamily x == NS.AF_INET6) addrs
addrs' =
case s of
HostIPv4 -> addrs4 ++ addrs6
HostIPv6 -> addrs6 ++ addrs4
_ -> addrs
tryAddrs (addr1:rest@(_:_)) =
E.catch
(theBody addr1)
(\(_ :: IOException) -> tryAddrs rest)
tryAddrs (addr1:[]) = theBody addr1
tryAddrs _ = error "bindPort: addrs is empty"
theBody addr =
bracketOnError
(NS.socket (NS.addrFamily addr) (NS.addrSocketType addr) (NS.addrProtocol addr))
NS.sClose
(\sock -> do
NS.setSocketOption sock NS.ReuseAddr 1
NS.bindSocket sock (NS.addrAddress addr)
NS.listen sock NS.maxListenQueue
return sock
)
tryAddrs addrs'