network-enumerator-0.1: Network/Socket/Enumerator.hs

-----------------------------------------------------------------------------
-- |
-- 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