module Data.Conduit.Network
(
sourceSocket
, sinkSocket
, Application
, AppData
, appSource
, appSink
, appSockAddr
, appLocalAddr
, ServerSettings
, serverSettings
, serverPort
, serverHost
, serverAfterBind
, serverNeedLocalAddr
, runTCPServer
, ClientSettings
, clientSettings
, clientPort
, clientHost
, runTCPClient
, HostPreference (..)
, bindPort
, getSocket
, acceptSafe
) where
import Prelude hiding (catch)
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 qualified Data.ByteString.Char8 as S8
import Control.Monad.IO.Class (MonadIO (liftIO))
import Control.Exception (throwIO, SomeException, try, finally, bracket, IOException, catch)
import Control.Monad (forever)
import Control.Monad.Trans.Control (MonadBaseControl, control)
import Control.Monad.Trans.Class (lift)
import Control.Concurrent (forkIO, threadDelay)
import Data.Conduit.Network.Internal
import Data.Conduit.Network.Utils (HostPreference)
import qualified Data.Conduit.Network.Utils as Utils
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 = AppData m -> m ()
serverSettings :: Monad m
=> Int
-> HostPreference
-> ServerSettings m
serverSettings port host = ServerSettings
{ serverPort = port
, serverHost = host
, serverAfterBind = const $ return ()
, serverNeedLocalAddr = False
}
runTCPServer :: (MonadIO m, MonadBaseControl IO m) => ServerSettings m -> Application m -> m ()
runTCPServer (ServerSettings port host afterBind needLocalAddr) app = control $ \run -> bracket
(liftIO $ bindPort port host)
(liftIO . NS.sClose)
(\socket -> run $ do
afterBind socket
forever $ serve socket)
where
serve lsocket = do
(socket, addr) <- liftIO $ acceptSafe lsocket
mlocal <- if needLocalAddr
then fmap Just $ liftIO (NS.getSocketName socket)
else return Nothing
let ad = AppData
{ appSource = sourceSocket socket
, appSink = sinkSocket socket
, appSockAddr = addr
, appLocalAddr = mlocal
}
app' run = run (app ad) >> return ()
appClose run = app' run `finally` NS.sClose socket
control $ \run -> forkIO (appClose run) >> run (return ())
clientSettings :: Monad m
=> Int
-> ByteString
-> ClientSettings m
clientSettings port host = ClientSettings
{ clientPort = port
, clientHost = host
}
runTCPClient :: (MonadIO m, MonadBaseControl IO m) => ClientSettings m -> Application m -> m ()
runTCPClient (ClientSettings port host) app = control $ \run -> bracket
(getSocket host port)
(NS.sClose . fst)
(\(s, address) -> run $ app AppData
{ appSource = sourceSocket s
, appSink = sinkSocket s
, appSockAddr = address
, appLocalAddr = Nothing
})
getSocket :: ByteString -> Int -> IO (NS.Socket, NS.SockAddr)
getSocket host' port' = do
(sock, addr) <- Utils.getSocket (S8.unpack host') port' NS.Stream
ee <- try' $ NS.connect sock (NS.addrAddress addr)
case ee of
Left e -> NS.sClose sock >> throwIO e
Right () -> return (sock, NS.addrAddress addr)
where
try' :: IO a -> IO (Either SomeException a)
try' = try
bindPort :: Int -> HostPreference -> IO Socket
bindPort p s = do
sock <- Utils.bindPort p s NS.Stream
NS.listen sock (max 2048 NS.maxListenQueue)
return sock
acceptSafe :: Socket -> IO (Socket, NS.SockAddr)
acceptSafe socket =
loop
where
loop =
NS.accept socket `catch` \(_ :: IOException) -> do
threadDelay 1000000
loop