diff --git a/Data/Conduit/Network.hs b/Data/Conduit/Network.hs --- a/Data/Conduit/Network.hs +++ b/Data/Conduit/Network.hs @@ -1,6 +1,5 @@ -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE TypeFamilies #-} {-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE RankNTypes #-} module Data.Conduit.Network ( -- * Basic utilities sourceSocket @@ -26,14 +25,15 @@ 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.Exception (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 +import Data.Conduit.Network.Utils (HostPreference) +import qualified Data.Conduit.Network.Utils as Utils + -- | Stream data from the socket. -- -- This function does /not/ automatically close the socket. @@ -120,13 +120,7 @@ -- Since 0.2.1 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) + (sock, addr) <- Utils.getSocket host' port' NS.Stream ee <- try' $ NS.connect sock (NS.addrAddress addr) case ee of Left e -> NS.sClose sock >> throwIO e @@ -135,81 +129,12 @@ try' :: IO a -> IO (Either SomeException a) try' = try --- | Which host to bind. --- --- Note: The @IsString@ instance recognizes the following special values: --- --- * @*@ means @HostAny@ --- --- * @*4@ means @HostIPv4@ --- --- * @*6@ means @HostIPv6@ -data HostPreference = - HostAny - | HostIPv4 - | HostIPv6 - | Host String - deriving (Eq, Ord, Show, Read) - -instance IsString HostPreference where - -- The funny code coming up is to get around some irritating warnings from - -- GHC. I should be able to just write: - {- - fromString "*" = HostAny - fromString "*4" = HostIPv4 - fromString "*6" = HostIPv6 - -} - fromString s'@('*':s) = - case s of - [] -> HostAny - ['4'] -> HostIPv4 - ['6'] -> HostIPv6 - _ -> Host s' - fromString s = Host s - -- | Attempt to bind a listening @Socket@ on the given host/port. If no host is -- given, will use the first address available. -- -- Since 0.3.0 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 - -- Choose an IPv6 socket if exists. This ensures the socket can - -- handle both IPv4 and IPv6 if v6only is false. - 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' + sock <- Utils.bindPort p s NS.Stream + NS.listen sock NS.maxListenQueue + return sock diff --git a/Data/Conduit/Network/UDP.hs b/Data/Conduit/Network/UDP.hs new file mode 100644 --- /dev/null +++ b/Data/Conduit/Network/UDP.hs @@ -0,0 +1,101 @@ +{-# LANGUAGE RankNTypes #-} +module Data.Conduit.Network.UDP + ( -- * UDP message representation + Message (..) + -- * Basic utilities + , sourceSocket + , sinkSocket + , sinkAllSocket + , sinkToSocket + , sinkAllToSocket + -- * Helper Utilities + , HostPreference (..) + , bindPort + , getSocket + ) where + +import Data.Conduit +import Network.Socket (AddrInfo, SockAddr, Socket) +import qualified Network.Socket as NS +import Network.Socket.ByteString (recvFrom, send, sendAll, sendTo, sendAllTo) +import Data.ByteString (ByteString) +import Control.Monad.IO.Class (MonadIO (liftIO)) +import Control.Monad (void) +import Control.Monad.Trans.Class (lift) + +import Data.Conduit.Network.Utils (HostPreference) +import qualified Data.Conduit.Network.Utils as Utils + +-- | Representation of a single message +data Message = Message { msgData :: {-# UNPACK #-} !ByteString + , msgSender :: !SockAddr + } + +-- | Stream messages from the socket. +-- +-- The given @len@ defines the maximum packet size. Every produced item +-- contains the message payload and the origin address. +-- +-- This function does /not/ automatically close the socket. +sourceSocket :: MonadIO m => Socket -> Int -> GSource m Message +sourceSocket socket len = loop + where + loop = do + (bs, addr) <- lift $ liftIO $ recvFrom socket len + yield (Message bs addr) >> loop + +-- | Stream messages to the connected socket. +-- +-- The payload is sent using @send@, so some of it might be lost. +-- +-- This function does /not/ automatically close the socket. +sinkSocket :: MonadIO m => Socket -> GInfSink ByteString m +sinkSocket = sinkSocketHelper (\sock bs -> void $ send sock bs) + +-- | Stream messages to the connected socket. +-- +-- The payload is sent using @sendAll@, so it might end up in multiple packets. +-- +-- This function does /not/ automatically close the socket. +sinkAllSocket :: MonadIO m => Socket -> GInfSink ByteString m +sinkAllSocket = sinkSocketHelper sendAll + +-- | Stream messages to the socket. +-- +-- Every handled item contains the message payload and the destination +-- address. The payload is sent using @sendTo@, so some of it might be +-- lost. +-- +-- This function does /not/ automatically close the socket. +sinkToSocket :: MonadIO m => Socket -> GInfSink Message m +sinkToSocket = sinkSocketHelper (\sock (Message bs addr) -> void $ sendTo sock bs addr) + +-- | Stream messages to the socket. +-- +-- Every handled item contains the message payload and the destination +-- address. The payload is sent using @sendAllTo@, so it might end up in +-- multiple packets. +-- +-- This function does /not/ automatically close the socket. +sinkAllToSocket :: MonadIO m => Socket -> GInfSink Message m +sinkAllToSocket = sinkSocketHelper (\sock (Message bs addr) -> sendAllTo sock bs addr) + +-- | Attempt to connect to the given host/port. +getSocket :: String -> Int -> IO (Socket, AddrInfo) +getSocket host' port' = Utils.getSocket host' port' NS.Datagram + +-- | Attempt to bind a listening @Socket@ on the given host/port. If no host is +-- given, will use the first address available. +bindPort :: Int -> HostPreference -> IO Socket +bindPort p s = Utils.bindPort p s NS.Datagram + +-- Internal +sinkSocketHelper :: MonadIO m => (Socket -> a -> IO ()) + -> Socket + -> GInfSink a m +sinkSocketHelper act socket = loop + where + loop = awaitE >>= either + return + (\a -> lift (liftIO $ act socket a) >> loop) +{-# INLINE sinkSocketHelper #-} diff --git a/Data/Conduit/Network/Utils.hs b/Data/Conduit/Network/Utils.hs new file mode 100644 --- /dev/null +++ b/Data/Conduit/Network/Utils.hs @@ -0,0 +1,101 @@ +{-# LANGUAGE ScopedTypeVariables #-} +module Data.Conduit.Network.Utils + ( -- * Helper utilities + HostPreference (..) + , bindPort + , getSocket + ) where + +import Network.Socket (AddrInfo, Socket, SocketType) +import qualified Network.Socket as NS +import Data.String (IsString (fromString)) +import Control.Exception (bracketOnError, IOException) +import qualified Control.Exception as E + +-- | Attempt to connect to the given host/port using given @SocketType@. +getSocket :: String -> Int -> SocketType -> IO (Socket, AddrInfo) +getSocket host' port' sockettype = do + let hints = NS.defaultHints { + NS.addrFlags = [NS.AI_ADDRCONFIG] + , NS.addrSocketType = sockettype + } + (addr:_) <- NS.getAddrInfo (Just hints) (Just host') (Just $ show port') + sock <- NS.socket (NS.addrFamily addr) (NS.addrSocketType addr) + (NS.addrProtocol addr) + return (sock, addr) + +-- | Which host to bind. +-- +-- Note: The @IsString@ instance recognizes the following special values: +-- +-- * @*@ means @HostAny@ +-- +-- * @*4@ means @HostIPv4@ +-- +-- * @*6@ means @HostIPv6@ +data HostPreference = + HostAny + | HostIPv4 + | HostIPv6 + | Host String + deriving (Eq, Ord, Show, Read) + +instance IsString HostPreference where + -- The funny code coming up is to get around some irritating warnings from + -- GHC. I should be able to just write: + {- + fromString "*" = HostAny + fromString "*4" = HostIPv4 + fromString "*6" = HostIPv6 + -} + fromString s'@('*':s) = + case s of + [] -> HostAny + ['4'] -> HostIPv4 + ['6'] -> HostIPv6 + _ -> Host s' + fromString s = Host s + +-- | Attempt to bind a listening @Socket@ on the given host/port using given +-- @SocketType@. If no host is given, will use the first address available. +bindPort :: Int -> HostPreference -> SocketType -> IO Socket +bindPort p s sockettype = do + let hints = NS.defaultHints + { NS.addrFlags = [ NS.AI_PASSIVE + , NS.AI_NUMERICSERV + , NS.AI_NUMERICHOST + ] + , NS.addrSocketType = sockettype + } + host = + case s of + Host s' -> Just s' + _ -> Nothing + port = Just . show $ p + addrs <- NS.getAddrInfo (Just hints) host port + -- Choose an IPv6 socket if exists. This ensures the socket can + -- handle both IPv4 and IPv6 if v6only is false. + 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) + return sock + ) + tryAddrs addrs' diff --git a/network-conduit.cabal b/network-conduit.cabal --- a/network-conduit.cabal +++ b/network-conduit.cabal @@ -1,5 +1,5 @@ Name: network-conduit -Version: 0.5.0 +Version: 0.5.0.1 Synopsis: Stream socket data using conduits. Description: Stream socket data using conduits. License: BSD3 @@ -17,6 +17,8 @@ Library Exposed-modules: Data.Conduit.Network + , Data.Conduit.Network.UDP + other-modules: Data.Conduit.Network.Utils Build-depends: base >= 4 && < 5 , transformers >= 0.2.2 && < 0.4 , bytestring >= 0.9