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.Concurrent (forkIO)
import Data.String (IsString (fromString))
import qualified Control.Exception as E
sourceSocket :: MonadIO m => Socket -> Source m ByteString
sourceSocket socket =
src
where
src = PipeM pull close
pull = do
bs <- liftIO (recv socket 4096)
return $ if S.null bs then Done Nothing () else HaveOutput src close bs
close = return ()
sinkSocket :: MonadIO m => Socket -> Sink ByteString m ()
sinkSocket socket =
sink
where
sink = NeedInput push (Done Nothing ())
push bs = PipeM (do
liftIO (sendAll socket bs)
return sink) (return ())
type Application m = Source m ByteString
-> Sink ByteString m ()
-> m ()
data ServerSettings = ServerSettings
{ serverPort :: Int
, serverHost :: HostPreference
}
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
}
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 (Show, Eq, Ord)
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'