{-# 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 hiding (accept, send, sendTo, recv, recvFrom)
import Network.Socket.ByteString
import Network.SockAddr
import Network.Multicast

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

beaconRecv :: TVar ZREState -> Endpoint -> IO b
beaconRecv :: TVar ZREState -> Endpoint -> IO b
beaconRecv s :: TVar ZREState
s e :: 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
        (msg :: ByteString
msg, addr :: SockAddr
addr) <- Socket -> Port -> IO (ByteString, SockAddr)
recvFrom Socket
sock 22
        case ByteString -> Either HostName (ByteString, Integer, UUID, Integer)
parseBeacon ByteString
msg of
          Left err :: HostName
err -> HostName -> IO ()
forall a. Show a => a -> IO ()
print HostName
err
          Right (_lead :: ByteString
_lead, _ver :: Integer
_ver, uuid :: UUID
uuid, port :: Integer
port) -> do
            case SockAddr
addr of
              x :: SockAddr
x@(SockAddrInet _hisport :: PortNumber
_hisport _host :: 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 _hisport :: PortNumber
_hisport _ _host :: HostAddress6
_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)
              _ -> () -> 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 -> B.ByteString -> UUID -> Int -> IO ()
beaconHandle :: TVar ZREState -> ByteString -> UUID -> Port -> IO ()
beaconHandle s :: TVar ZREState
s addr :: ByteString
addr uuid :: UUID
uuid port :: 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 peer :: 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
          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 ()


-- sends udp multicast beacons
beacon :: AddrInfo -> B.ByteString -> Port -> IO a
beacon :: AddrInfo -> ByteString -> Port -> IO a
beacon addrInfo :: AddrInfo
addrInfo uuid :: ByteString
uuid port :: Port
port = do
    IO a -> IO a
forall a. IO a -> IO a
withSocketsDo (IO a -> IO a) -> IO a -> IO a
forall a b. (a -> b) -> a -> b
$ do
      IO Socket -> (Socket -> IO ()) -> (Socket -> IO a) -> IO a
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 a
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 addr :: 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_ (\x :: SocketOption
x -> Socket -> SocketOption -> Port -> IO ()
setSocketOption Socket
s SocketOption
x 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 addr :: SockAddr
addr msg :: ByteString
msg s :: 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
zreBeaconMs