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)
enumSocket :: MonadIO m
=> Integer
-> 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
enumSocketFrom :: MonadIO m
=> Integer
-> 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
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)
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