{-# LANGUAGE OverloadedStrings #-}
module Network.ZRE.Beacon (
    beacon
  , beaconRecv
  ) where

import Control.Monad
import Control.Exception
import Control.Concurrent
import Control.Concurrent.STM
import Network.Socket
import Network.Socket.ByteString
import Network.SockAddr
import Network.Multicast

import Data.ByteString (ByteString)
import Data.Maybe
import Data.UUID
import Data.Time.Clock
import qualified Data.Map as M
import qualified Data.ByteString.Char8 as B

import Data.ZRE
import Network.ZRE.Peer
import Network.ZRE.Types
import System.ZMQ4.Endpoint

-- | Receive beacons from UDP multicast
beaconRecv :: TVar ZREState -> Endpoint -> IO b
beaconRecv :: forall b. TVar ZREState -> Endpoint -> IO b
beaconRecv TVar ZREState
s Endpoint
e = do
    Socket
sock <- String -> PortNumber -> IO Socket
multicastReceiver (ByteString -> String
B.unpack forall a b. (a -> b) -> a -> b
$ Endpoint -> ByteString
endpointAddr Endpoint
e) (forall a b. (Integral a, Num b) => a -> b
fromIntegral forall a b. (a -> b) -> a -> b
$ forall a. HasCallStack => Maybe a -> a
fromJust forall a b. (a -> b) -> a -> b
$ Endpoint -> Maybe Port
endpointPort Endpoint
e)
    forall (f :: * -> *) a b. Applicative f => f a -> f b
forever forall a b. (a -> b) -> a -> b
$ do
        (ByteString
msg, SockAddr
addr) <- Socket -> Port -> IO (ByteString, SockAddr)
recvFrom Socket
sock Port
22
        case ByteString -> Either String (ByteString, Integer, UUID, Integer)
parseBeacon ByteString
msg of
          Left String
err -> forall a. Show a => a -> IO ()
print String
err
          Right (ByteString
_lead, Integer
_ver, UUID
uuid, Integer
port) -> do
            case SockAddr
addr of
              x :: SockAddr
x@(SockAddrInet PortNumber
_hisport HostAddress
_host) -> do
                TVar ZREState -> ByteString -> UUID -> Port -> IO ()
beaconHandle TVar ZREState
s (SockAddr -> ByteString
showSockAddrBS SockAddr
x) UUID
uuid (forall a b. (Integral a, Num b) => a -> b
fromIntegral Integer
port)
              x :: SockAddr
x@(SockAddrInet6 PortNumber
_hisport HostAddress
_ HostAddress6
_host HostAddress
_) -> do
                TVar ZREState -> ByteString -> UUID -> Port -> IO ()
beaconHandle TVar ZREState
s (SockAddr -> ByteString
showSockAddrBS SockAddr
x) UUID
uuid (forall a b. (Integral a, Num b) => a -> b
fromIntegral Integer
port)
              SockAddr
_ -> forall (m :: * -> *) a. Monad m => a -> m a
return ()

-- | Handle messages received on beacon
-- * creates new peers
-- * updates peers last heard
beaconHandle :: TVar ZREState -> ByteString -> UUID -> Int -> IO ()
beaconHandle :: TVar ZREState -> ByteString -> UUID -> Port -> IO ()
beaconHandle TVar ZREState
s ByteString
addr UUID
uuid Port
port = do
    ZREState
st <- forall a. STM a -> IO a
atomically forall a b. (a -> b) -> a -> b
$ forall a. TVar a -> STM a
readTVar TVar ZREState
s

    if UUID
uuid forall a. Eq a => a -> a -> Bool
== ZREState -> UUID
zreUUID ZREState
st
      then forall (m :: * -> *) a. Monad m => a -> m a
return () -- our own message
      else do
        case forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup UUID
uuid forall a b. (a -> b) -> a -> b
$ ZREState -> Peers
zrePeers ZREState
st of
          (Just TVar Peer
peer) -> do
            UTCTime
now <- IO UTCTime
getCurrentTime
            forall a. STM a -> IO a
atomically forall a b. (a -> b) -> a -> b
$ TVar Peer -> UTCTime -> STM ()
updateLastHeard TVar Peer
peer UTCTime
now
          Maybe (TVar Peer)
Nothing -> do
            -- B.putStrLn $ B.concat ["New peer from beacon ", B.pack $ show uuid, " (", addr, ":", B.pack $ show port , ")"]
            forall (f :: * -> *) a. Functor f => f a -> f ()
void forall a b. (a -> b) -> a -> b
$ TVar ZREState
-> UUID
-> (UTCTime
    -> UUID
    -> TVar ZREState
    -> STM (TVar Peer, Maybe (IO ()), Maybe (IO ())))
-> IO (TVar Peer)
makePeer TVar ZREState
s UUID
uuid forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a b.
MonadIO m =>
ByteString
-> Port
-> UTCTime
-> UUID
-> TVar ZREState
-> STM (TVar Peer, Maybe (m a), Maybe (IO b))
newPeerFromBeacon ByteString
addr Port
port
            forall (m :: * -> *) a. Monad m => a -> m a
return ()


-- | Send UDP multicast beacons periodically
beacon :: Float -> AddrInfo -> ByteString -> Port -> IO ()
beacon :: Float -> AddrInfo -> ByteString -> Port -> IO ()
beacon Float
seconds AddrInfo
addrInfo ByteString
uuid Port
port = do
    forall a. IO a -> IO a
withSocketsDo forall a b. (a -> b) -> a -> b
$ do
      forall a b c. IO a -> (a -> IO b) -> (a -> IO c) -> IO c
bracket (AddrInfo -> IO Socket
getSocket AddrInfo
addrInfo) Socket -> IO ()
close (forall {b}. SockAddr -> ByteString -> Socket -> IO b
talk (AddrInfo -> SockAddr
addrAddress AddrInfo
addrInfo) (ByteString -> Port -> ByteString
zreBeacon ByteString
uuid Port
port))
  where
    getSocket :: AddrInfo -> IO Socket
getSocket AddrInfo
addr = do
      Socket
s <- Family -> SocketType -> ProtocolNumber -> IO Socket
socket (AddrInfo -> Family
addrFamily AddrInfo
addr) SocketType
Datagram ProtocolNumber
defaultProtocol
      forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (\SocketOption
x -> Socket -> SocketOption -> Port -> IO ()
setSocketOption Socket
s SocketOption
x Port
1) [SocketOption
Broadcast, SocketOption
ReuseAddr, SocketOption
ReusePort]
      Socket -> SockAddr -> IO ()
bind Socket
s (AddrInfo -> SockAddr
addrAddress AddrInfo
addr)
      forall (m :: * -> *) a. Monad m => a -> m a
return Socket
s
    talk :: SockAddr -> ByteString -> Socket -> IO b
talk SockAddr
addr ByteString
msg Socket
s =
      forall (f :: * -> *) a b. Applicative f => f a -> f b
forever forall a b. (a -> b) -> a -> b
$ do
      forall (f :: * -> *) a. Functor f => f a -> f ()
void forall a b. (a -> b) -> a -> b
$ Socket -> ByteString -> SockAddr -> IO Port
sendTo Socket
s ByteString
msg SockAddr
addr
      Port -> IO ()
threadDelay forall a b. (a -> b) -> a -> b
$ forall a. RealFrac a => a -> Port
sec Float
seconds