{-# 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
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 ()
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 ()
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
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 ()
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