{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RecordWildCards #-}

-- zgossip protocol https://github.com/zeromq/czmq/blob/master/src/zgossip_msg.bnf
-- client sends HELLO, recieves all stored tuples, forwards to other clients

module Data.ZGossip (
    newZGS
  , parseZGS
  , encodeZGS
  , Key
  , Value
  , TTL
  , Peer
  , ZGSCmd(..)
  , ZGSMsg(..)
  ) where

import Prelude hiding (putStrLn, take)
import Data.ByteString (ByteString)

import GHC.Word

import Data.ZMQParse

-- | Version of the ZGossip protocol
zgsVer :: Int
zgsVer :: TTL
zgsVer = TTL
1

-- | Signature of the ZGossip protocol
zgsSig :: Word16
zgsSig :: Word16
zgsSig = Word16
0xAAA0

type Peer  = ByteString
type Key   = ByteString
type Value = ByteString
type TTL   = Int

data ZGSMsg = ZGSMsg {
    ZGSMsg -> Maybe ByteString
zgsFrom :: Maybe ByteString
  , ZGSMsg -> ZGSCmd
zgsCmd :: ZGSCmd
  } deriving (TTL -> ZGSMsg -> ShowS
[ZGSMsg] -> ShowS
ZGSMsg -> String
forall a.
(TTL -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [ZGSMsg] -> ShowS
$cshowList :: [ZGSMsg] -> ShowS
show :: ZGSMsg -> String
$cshow :: ZGSMsg -> String
showsPrec :: TTL -> ZGSMsg -> ShowS
$cshowsPrec :: TTL -> ZGSMsg -> ShowS
Show, ZGSMsg -> ZGSMsg -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: ZGSMsg -> ZGSMsg -> Bool
$c/= :: ZGSMsg -> ZGSMsg -> Bool
== :: ZGSMsg -> ZGSMsg -> Bool
$c== :: ZGSMsg -> ZGSMsg -> Bool
Eq, Eq ZGSMsg
ZGSMsg -> ZGSMsg -> Bool
ZGSMsg -> ZGSMsg -> Ordering
ZGSMsg -> ZGSMsg -> ZGSMsg
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
min :: ZGSMsg -> ZGSMsg -> ZGSMsg
$cmin :: ZGSMsg -> ZGSMsg -> ZGSMsg
max :: ZGSMsg -> ZGSMsg -> ZGSMsg
$cmax :: ZGSMsg -> ZGSMsg -> ZGSMsg
>= :: ZGSMsg -> ZGSMsg -> Bool
$c>= :: ZGSMsg -> ZGSMsg -> Bool
> :: ZGSMsg -> ZGSMsg -> Bool
$c> :: ZGSMsg -> ZGSMsg -> Bool
<= :: ZGSMsg -> ZGSMsg -> Bool
$c<= :: ZGSMsg -> ZGSMsg -> Bool
< :: ZGSMsg -> ZGSMsg -> Bool
$c< :: ZGSMsg -> ZGSMsg -> Bool
compare :: ZGSMsg -> ZGSMsg -> Ordering
$ccompare :: ZGSMsg -> ZGSMsg -> Ordering
Ord)

data ZGSCmd =
    Hello
  | Publish Key Value TTL
  | Ping
  | PingOk
  | Invalid
  deriving (TTL -> ZGSCmd -> ShowS
[ZGSCmd] -> ShowS
ZGSCmd -> String
forall a.
(TTL -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [ZGSCmd] -> ShowS
$cshowList :: [ZGSCmd] -> ShowS
show :: ZGSCmd -> String
$cshow :: ZGSCmd -> String
showsPrec :: TTL -> ZGSCmd -> ShowS
$cshowsPrec :: TTL -> ZGSCmd -> ShowS
Show, ZGSCmd -> ZGSCmd -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: ZGSCmd -> ZGSCmd -> Bool
$c/= :: ZGSCmd -> ZGSCmd -> Bool
== :: ZGSCmd -> ZGSCmd -> Bool
$c== :: ZGSCmd -> ZGSCmd -> Bool
Eq, Eq ZGSCmd
ZGSCmd -> ZGSCmd -> Bool
ZGSCmd -> ZGSCmd -> Ordering
ZGSCmd -> ZGSCmd -> ZGSCmd
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
min :: ZGSCmd -> ZGSCmd -> ZGSCmd
$cmin :: ZGSCmd -> ZGSCmd -> ZGSCmd
max :: ZGSCmd -> ZGSCmd -> ZGSCmd
$cmax :: ZGSCmd -> ZGSCmd -> ZGSCmd
>= :: ZGSCmd -> ZGSCmd -> Bool
$c>= :: ZGSCmd -> ZGSCmd -> Bool
> :: ZGSCmd -> ZGSCmd -> Bool
$c> :: ZGSCmd -> ZGSCmd -> Bool
<= :: ZGSCmd -> ZGSCmd -> Bool
$c<= :: ZGSCmd -> ZGSCmd -> Bool
< :: ZGSCmd -> ZGSCmd -> Bool
$c< :: ZGSCmd -> ZGSCmd -> Bool
compare :: ZGSCmd -> ZGSCmd -> Ordering
$ccompare :: ZGSCmd -> ZGSCmd -> Ordering
Ord)

cmdCode :: ZGSCmd -> Word8
cmdCode :: ZGSCmd -> Word8
cmdCode ZGSCmd
Hello           = Word8
0x01
cmdCode (Publish ByteString
_ ByteString
_ TTL
_) = Word8
0x02
cmdCode ZGSCmd
Ping            = Word8
0x03
cmdCode ZGSCmd
PingOk          = Word8
0x04
cmdCode ZGSCmd
Invalid         = Word8
0x05

newZGS :: ZGSCmd -> ZGSMsg
newZGS :: ZGSCmd -> ZGSMsg
newZGS ZGSCmd
cmd = Maybe ByteString -> ZGSCmd -> ZGSMsg
ZGSMsg forall a. Maybe a
Nothing ZGSCmd
cmd

encodeZGS :: ZGSMsg -> ByteString
encodeZGS :: ZGSMsg -> ByteString
encodeZGS ZGSMsg{Maybe ByteString
ZGSCmd
zgsCmd :: ZGSCmd
zgsFrom :: Maybe ByteString
zgsCmd :: ZGSMsg -> ZGSCmd
zgsFrom :: ZGSMsg -> Maybe ByteString
..} = ByteString
msg
  where
    msg :: ByteString
msg = Put -> ByteString
runPut forall a b. (a -> b) -> a -> b
$ do
      Word16 -> Put
putWord16be Word16
zgsSig
      Word8 -> Put
putWord8 forall a b. (a -> b) -> a -> b
$ ZGSCmd -> Word8
cmdCode ZGSCmd
zgsCmd
      Int8 -> Put
putInt8 forall a b. (a -> b) -> a -> b
$ forall a b. (Integral a, Num b) => a -> b
fromIntegral TTL
zgsVer
      ZGSCmd -> Put
encodeCmd ZGSCmd
zgsCmd

encodeCmd :: ZGSCmd -> PutM ()
encodeCmd :: ZGSCmd -> Put
encodeCmd (Publish ByteString
k ByteString
v TTL
ttl) = do
  ByteString -> Put
putByteStringLen ByteString
k
  ByteString -> Put
putLongByteStringLen ByteString
v
  Int32 -> Put
putInt32be forall a b. (a -> b) -> a -> b
$ forall a b. (Integral a, Num b) => a -> b
fromIntegral TTL
ttl
encodeCmd ZGSCmd
_ = forall (m :: * -> *) a. Monad m => a -> m a
return ()

parsePublish :: Get ZGSCmd
parsePublish :: Get ZGSCmd
parsePublish = ByteString -> ByteString -> TTL -> ZGSCmd
Publish
  forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Get ByteString
parseString
  forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Get ByteString
parseLongString
  forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall a. Integral a => Get a
getInt32

parseCmd :: ByteString -> Get ZGSMsg
parseCmd :: ByteString -> Get ZGSMsg
parseCmd ByteString
from = do
    TTL
cmd <- (forall a. Integral a => Get a
getInt8 :: Get Int)
    TTL
ver <- forall a. Integral a => Get a
getInt8

    if TTL
ver forall a. Eq a => a -> a -> Bool
/= TTL
zgsVer
      then forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"Protocol version mismatch"
      else do

        ZGSCmd
zcmd <- case TTL
cmd of
          TTL
0x01 -> forall (f :: * -> *) a. Applicative f => a -> f a
pure ZGSCmd
Hello
          TTL
0x02 -> Get ZGSCmd
parsePublish
          TTL
0x03 -> forall (f :: * -> *) a. Applicative f => a -> f a
pure ZGSCmd
Ping
          TTL
0x04 -> forall (f :: * -> *) a. Applicative f => a -> f a
pure ZGSCmd
PingOk
          TTL
0x05 -> forall (f :: * -> *) a. Applicative f => a -> f a
pure ZGSCmd
Invalid
          TTL
_    -> forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"Unknown command"

        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ Maybe ByteString -> ZGSCmd -> ZGSMsg
ZGSMsg (forall a. a -> Maybe a
Just ByteString
from) ZGSCmd
zcmd

parseZGS :: [ByteString] -> Either String ZGSMsg
parseZGS :: [ByteString] -> Either String ZGSMsg
parseZGS [ByteString
from, ByteString
msg] = ByteString -> ByteString -> Either String ZGSMsg
parseZgs ByteString
from ByteString
msg
parseZGS [ByteString]
_ = forall a b. a -> Either a b
Left String
"empty message"

parseZgs :: ByteString -> ByteString -> Either String ZGSMsg
parseZgs :: ByteString -> ByteString -> Either String ZGSMsg
parseZgs ByteString
from ByteString
msg = forall a b c. (a -> b -> c) -> b -> a -> c
flip forall a. Get a -> ByteString -> Either String a
runGet ByteString
msg forall a b. (a -> b) -> a -> b
$ do
  Word16
sig <- forall a. Integral a => Get a
getInt16
  if Word16
sig forall a. Eq a => a -> a -> Bool
/= Word16
zgsSig
    then forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"Signature mismatch"
    else do
      ZGSMsg
res <- ByteString -> Get ZGSMsg
parseCmd ByteString
from
      forall (m :: * -> *) a. Monad m => a -> m a
return ZGSMsg
res