{-# 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 :: TVar ZREState -> Endpoint -> IO b
beaconRecv TVar ZREState
s Endpoint
e = do
    Socket
sock <- HostName -> PortNumber -> IO Socket
multicastReceiver (ByteString -> HostName
B.unpack (ByteString -> HostName) -> ByteString -> HostName
forall a b. (a -> b) -> a -> b
$ Endpoint -> ByteString
endpointAddr Endpoint
e) (Port -> PortNumber
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Port -> PortNumber) -> Port -> PortNumber
forall a b. (a -> b) -> a -> b
$ Maybe Port -> Port
forall a. HasCallStack => Maybe a -> a
fromJust (Maybe Port -> Port) -> Maybe Port -> Port
forall a b. (a -> b) -> a -> b
$ Endpoint -> Maybe Port
endpointPort Endpoint
e)
    IO () -> IO b
forall (f :: * -> *) a b. Applicative f => f a -> f b
forever (IO () -> IO b) -> IO () -> IO b
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 HostName (ByteString, Integer, UUID, Integer)
parseBeacon ByteString
msg of
          Left HostName
err -> HostName -> IO ()
forall a. Show a => a -> IO ()
print HostName
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 (Integer -> Port
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 (Integer -> Port
forall a b. (Integral a, Num b) => a -> b
fromIntegral Integer
port)
              SockAddr
_ -> () -> IO ()
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 <- STM ZREState -> IO ZREState
forall a. STM a -> IO a
atomically (STM ZREState -> IO ZREState) -> STM ZREState -> IO ZREState
forall a b. (a -> b) -> a -> b
$ TVar ZREState -> STM ZREState
forall a. TVar a -> STM a
readTVar TVar ZREState
s

    if UUID
uuid UUID -> UUID -> Bool
forall a. Eq a => a -> a -> Bool
== ZREState -> UUID
zreUUID ZREState
st
      then () -> IO ()
forall (m :: * -> *) a. Monad m => a -> m a
return () -- our own message
      else do
        case UUID -> Map UUID (TVar Peer) -> Maybe (TVar Peer)
forall k a. Ord k => k -> Map k a -> Maybe a
M.lookup UUID
uuid (Map UUID (TVar Peer) -> Maybe (TVar Peer))
-> Map UUID (TVar Peer) -> Maybe (TVar Peer)
forall a b. (a -> b) -> a -> b
$ ZREState -> Map UUID (TVar Peer)
zrePeers ZREState
st of
          (Just TVar Peer
peer) -> do
            UTCTime
now <- IO UTCTime
getCurrentTime
            STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
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 , ")"]
            IO (TVar Peer) -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO (TVar Peer) -> IO ()) -> IO (TVar Peer) -> IO ()
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 ((UTCTime
  -> UUID
  -> TVar ZREState
  -> STM (TVar Peer, Maybe (IO ()), Maybe (IO ())))
 -> IO (TVar Peer))
-> (UTCTime
    -> UUID
    -> TVar ZREState
    -> STM (TVar Peer, Maybe (IO ()), Maybe (IO ())))
-> IO (TVar Peer)
forall a b. (a -> b) -> a -> b
$ ByteString
-> Port
-> UTCTime
-> UUID
-> TVar ZREState
-> STM (TVar Peer, Maybe (IO ()), Maybe (IO ()))
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
            () -> IO ()
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
    IO () -> IO ()
forall a. IO a -> IO a
withSocketsDo (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
      IO Socket -> (Socket -> IO ()) -> (Socket -> IO ()) -> IO ()
forall a b c. IO a -> (a -> IO b) -> (a -> IO c) -> IO c
bracket (AddrInfo -> IO Socket
getSocket AddrInfo
addrInfo) Socket -> IO ()
close (SockAddr -> ByteString -> Socket -> IO ()
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
      (SocketOption -> IO ()) -> [SocketOption] -> IO ()
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)
      Socket -> IO Socket
forall (m :: * -> *) a. Monad m => a -> m a
return Socket
s
    talk :: SockAddr -> ByteString -> Socket -> IO b
talk SockAddr
addr ByteString
msg Socket
s =
      IO () -> IO b
forall (f :: * -> *) a b. Applicative f => f a -> f b
forever (IO () -> IO b) -> IO () -> IO b
forall a b. (a -> b) -> a -> b
$ do
      IO Port -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void (IO Port -> IO ()) -> IO Port -> IO ()
forall a b. (a -> b) -> a -> b
$ Socket -> ByteString -> SockAddr -> IO Port
sendTo Socket
s ByteString
msg SockAddr
addr
      Port -> IO ()
threadDelay (Port -> IO ()) -> Port -> IO ()
forall a b. (a -> b) -> a -> b
$ Float -> Port
forall a. RealFrac a => a -> Port
sec Float
seconds