-----------------------------------------------------------------------------
-- |
-- Module      :  Network.Multicast
-- Copyright   :  (c) Audrey Tang 2008
-- License     :  MIT License
-- 
-- Maintainer  :  audreyt@audreyt.org
-- Stability   :  experimental
-- Portability :  portable
--
-- The "Network.Multicast" module is for sending UDP datagrams over multicast
-- (class D) addresses.
--
-----------------------------------------------------------------------------
#include <sys/types.h>
#include <netinet/in.h>

module Network.Multicast (
    -- * Simple sending and receiving
      multicastSender, multicastReceiver
    -- * Additional Socket operations
    , addMembership, dropMembership, setLoopbackMode
    -- * Loopback flags
    , LoopbackMode, enableLoopback, noLoopback
) where
import Network.BSD
import Network.Socket
import Foreign.C.Types
import Foreign.C.Error
import Foreign.Storable
import Foreign.Marshal
import Foreign.Ptr

type LoopbackMode = Bool

enableLoopback, noLoopback :: LoopbackMode
enableLoopback = True
noLoopback     = False

-- | Calling 'multicastSender' creates a client side UDP socket for sending
-- multicast datagrams to the specified host and port.
--
-- Minimal example:
--
-- > import Network.Socket
-- > import Network.Multicast
-- > main = withSocketsDo $ do
-- >     (sock, addr) <- multicastSender "224.0.0.99" 9999 noLoopback
-- >     let loop = do
-- >         sendTo sock "Hello, world" addr
-- >         loop in loop
--
multicastSender :: HostName -> PortNumber -> LoopbackMode -> IO (Socket, SockAddr)
multicastSender host port loop = do
    proto <- getProtocolNumber "udp"
    let hints = defaultHints
            { addrFlags      = [AI_ADDRCONFIG]
            , addrProtocol   = proto
            , addrSocketType = Datagram
            }
    (addr:_) <- getAddrInfo (Just hints) (Just host) (Just $ show port)
    sock     <- socket (addrFamily addr) (addrSocketType addr) (addrProtocol addr)
    if loop then return () else setLoopbackMode sock loop
    return (sock, addrAddress addr)

-- | Calling 'multicastReceiver' creates and binds a UDP socket for listening
-- multicast datagrams on the specified host and port.
--
-- Minimal example:
--
-- > import Network.Socket
-- > import Network.Multicast
-- > main = withSocketsDo $ do
-- >     sock <- multicastReceiver "224.0.0.99" 9999
-- >     let loop = do
-- >         (msg, _, addr) <- recvFrom sock 1024
-- >         print (msg, addr) in loop
--
multicastReceiver :: HostName -> PortNumber -> IO Socket
multicastReceiver host port = do
    proto <- getProtocolNumber "udp"
    let hints = defaultHints
            { addrFlags      = [AI_PASSIVE]
            , addrProtocol   = proto
            , addrSocketType = Datagram
            }
    (addr:_) <- getAddrInfo (Just hints) (Just host) (Just $ show port)
    sock     <- socket (addrFamily addr) (addrSocketType addr) (addrProtocol addr)
    addMembership sock host
    setSocketOption sock ReusePort 1
    bindSocket sock (addrAddress addr)
    return sock

-- | Enable or disable the loopback mode on a socket created by 'multicastSender'.
-- Loopback is enabled by default; disabling it may improve performance a little bit.
setLoopbackMode :: Socket -> LoopbackMode -> IO ()
setLoopbackMode (MkSocket s _ _ _ _) mode = maybeIOError "setLoopbackMode" $
    alloca $ \loopPtr -> do
        let loop = if mode then 1 else 0 :: CUChar
        poke loopPtr loop
        c_setsockopt s _IPPROTO_IP _IP_MULTICAST_LOOP (castPtr loopPtr) (toEnum $ sizeOf loop)

-- | Make the socket listen on multicast datagrams sent by the specified 'HostName'.
addMembership :: Socket -> HostName -> IO ()
addMembership s = maybeIOError "addMembership" . doMulticastGroup _IP_ADD_MEMBERSHIP s

-- | Stop the socket from listening on multicast datagrams sent by the specified 'HostName'.
dropMembership :: Socket -> HostName -> IO ()
dropMembership s = maybeIOError "dropMembership" . doMulticastGroup _IP_DROP_MEMBERSHIP s

maybeIOError :: String -> IO CInt -> IO ()
maybeIOError name f = f >>= \err -> case err of
    0 -> return ()
    _ -> ioError (errnoToIOError name (Errno (fromIntegral err)) Nothing Nothing)

doMulticastGroup :: CInt -> Socket -> HostName -> IO CInt
doMulticastGroup flag (MkSocket s _ _ _ _) host = allocaBytes #{size struct ip_mreq} $ \mReqPtr -> do
    addr <- inet_addr host
    #{poke struct ip_mreq, imr_multiaddr} mReqPtr addr
    #{poke struct ip_mreq, imr_interface} mReqPtr (#{const INADDR_ANY} `asTypeOf` addr)
    c_setsockopt s _IPPROTO_IP flag (castPtr mReqPtr) (#{size struct ip_mreq})

foreign import ccall unsafe "setsockopt"
    c_setsockopt :: CInt -> CInt -> CInt -> Ptr CInt -> CInt -> IO CInt

_IPPROTO_IP :: CInt
_IPPROTO_IP = #const IPPROTO_IP

_IP_ADD_MEMBERSHIP, _IP_DROP_MEMBERSHIP, _IP_MULTICAST_LOOP :: CInt
_IP_ADD_MEMBERSHIP  = #const IP_ADD_MEMBERSHIP
_IP_DROP_MEMBERSHIP = #const IP_DROP_MEMBERSHIP
_IP_MULTICAST_LOOP  = #const IP_MULTICAST_LOOP