----------------------------------------------------------------------------- -- | -- Module: Network.Socket.Enumerator -- Copyright: 2010 John Millikin -- License: MIT -- -- Maintainer: jmillikin@gmail.com -- Portability: portable -- ----------------------------------------------------------------------------- module Network.Socket.Enumerator ( enumSocket , enumSocketFrom , iterSocket , iterSocketTo ) where import qualified Control.Exception as Exc import Control.Monad.IO.Class (MonadIO, liftIO) import qualified Data.ByteString as B import Data.Enumerator ((>>==)) import qualified Data.Enumerator as E import qualified Network.Socket as S import Network.Socket.ByteString (sendMany, sendManyTo, recv, recvFrom) -- | Enumerate binary data from a 'S.Socket', using 'recv'. The socket -- must be connected. enumSocket :: MonadIO m => Integer -- ^ Buffer size -> S.Socket -> E.Enumerator B.ByteString m b enumSocket bufferSize sock = loop where intSize = fromInteger bufferSize loop (E.Continue k) = do bytes <- try (recv sock intSize) if B.null bytes then E.continue k else k (E.Chunks [bytes]) >>== loop loop step = E.returnI step -- | Enumerate binary data from a 'S.Socket', using 'recvFrom'. The socket -- does not have to be connected. Each chunk of data received will be paired -- with its address. enumSocketFrom :: MonadIO m => Integer -- ^ Buffer size -> S.Socket -> E.Enumerator (B.ByteString, S.SockAddr) m b enumSocketFrom bufferSize sock = loop where intSize = fromInteger bufferSize loop (E.Continue k) = do (bytes, addr) <- try (recvFrom sock intSize) if B.null bytes then E.continue k else k (E.Chunks [(bytes, addr)]) >>== loop loop step = E.returnI step -- | Write data to a 'S.Socket', using 'sendMany'. The socket must be connected. iterSocket :: MonadIO m => S.Socket -> E.Iteratee B.ByteString m () iterSocket sock = E.continue step where step E.EOF = E.yield () E.EOF step (E.Chunks []) = E.continue step step (E.Chunks xs) = try (sendMany sock xs) -- | Write data to a 'S.Socket', using 'sendManyTo'. The socket does not -- have to be connected. iterSocketTo :: MonadIO m => S.Socket -> S.SockAddr -> E.Iteratee B.ByteString m () iterSocketTo sock addr = E.continue step where step E.EOF = E.yield () E.EOF step (E.Chunks []) = E.continue step step (E.Chunks xs) = try (sendManyTo sock xs addr) try :: MonadIO m => IO b -> E.Iteratee a m b try io = do tried <- liftIO (Exc.try io) case tried of Left err -> E.throwError (err :: Exc.SomeException) Right b -> return b