{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RecordWildCards #-}
module Data.ZRE (
    zreVer
  , newZRE
  , parseZRE
  , encodeZRE
  , zreBeacon
  , parseBeacon
  , Name
  , Headers
  , Content
  , Group
  , Groups
  , Seq
  , GroupSeq
  , ZREMsg(..)
  , ZRECmd(..)) where
import Prelude hiding (putStrLn, take)
import qualified Data.ByteString.Char8 as B
import qualified Data.ByteString.Lazy as BL

import GHC.Word

import qualified Data.Map as M
import qualified Data.Set as S
import Data.UUID
import Data.Time.Clock

import System.ZMQ4.Endpoint
import Data.ZMQParse

zreVer :: Int
zreVer :: Int
zreVer = 2
zreSig :: Word16
zreSig :: Word16
zreSig = 0xAAA1

type Seq = Int
type GroupSeq = Int
type Group = B.ByteString
type Groups = S.Set Group
type Name = B.ByteString
type Headers = M.Map B.ByteString B.ByteString
type Content = [B.ByteString]

data ZREMsg = ZREMsg {
    ZREMsg -> Maybe UUID
msgFrom :: Maybe UUID
  , ZREMsg -> Int
msgSeq :: Seq
  , ZREMsg -> Maybe UTCTime
msgTime :: Maybe UTCTime
  , ZREMsg -> ZRECmd
msgCmd :: ZRECmd
  } deriving (Int -> ZREMsg -> ShowS
[ZREMsg] -> ShowS
ZREMsg -> String
(Int -> ZREMsg -> ShowS)
-> (ZREMsg -> String) -> ([ZREMsg] -> ShowS) -> Show ZREMsg
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [ZREMsg] -> ShowS
$cshowList :: [ZREMsg] -> ShowS
show :: ZREMsg -> String
$cshow :: ZREMsg -> String
showsPrec :: Int -> ZREMsg -> ShowS
$cshowsPrec :: Int -> ZREMsg -> ShowS
Show, ZREMsg -> ZREMsg -> Bool
(ZREMsg -> ZREMsg -> Bool)
-> (ZREMsg -> ZREMsg -> Bool) -> Eq ZREMsg
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: ZREMsg -> ZREMsg -> Bool
$c/= :: ZREMsg -> ZREMsg -> Bool
== :: ZREMsg -> ZREMsg -> Bool
$c== :: ZREMsg -> ZREMsg -> Bool
Eq, Eq ZREMsg
Eq ZREMsg =>
(ZREMsg -> ZREMsg -> Ordering)
-> (ZREMsg -> ZREMsg -> Bool)
-> (ZREMsg -> ZREMsg -> Bool)
-> (ZREMsg -> ZREMsg -> Bool)
-> (ZREMsg -> ZREMsg -> Bool)
-> (ZREMsg -> ZREMsg -> ZREMsg)
-> (ZREMsg -> ZREMsg -> ZREMsg)
-> Ord ZREMsg
ZREMsg -> ZREMsg -> Bool
ZREMsg -> ZREMsg -> Ordering
ZREMsg -> ZREMsg -> ZREMsg
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 :: ZREMsg -> ZREMsg -> ZREMsg
$cmin :: ZREMsg -> ZREMsg -> ZREMsg
max :: ZREMsg -> ZREMsg -> ZREMsg
$cmax :: ZREMsg -> ZREMsg -> ZREMsg
>= :: ZREMsg -> ZREMsg -> Bool
$c>= :: ZREMsg -> ZREMsg -> Bool
> :: ZREMsg -> ZREMsg -> Bool
$c> :: ZREMsg -> ZREMsg -> Bool
<= :: ZREMsg -> ZREMsg -> Bool
$c<= :: ZREMsg -> ZREMsg -> Bool
< :: ZREMsg -> ZREMsg -> Bool
$c< :: ZREMsg -> ZREMsg -> Bool
compare :: ZREMsg -> ZREMsg -> Ordering
$ccompare :: ZREMsg -> ZREMsg -> Ordering
$cp1Ord :: Eq ZREMsg
Ord)

data ZRECmd =
    Hello Endpoint Groups GroupSeq Name Headers
  | Whisper Content
  | Shout Group Content
  | Join Group GroupSeq
  | Leave Group GroupSeq
  | Ping
  | PingOk
  deriving (Int -> ZRECmd -> ShowS
[ZRECmd] -> ShowS
ZRECmd -> String
(Int -> ZRECmd -> ShowS)
-> (ZRECmd -> String) -> ([ZRECmd] -> ShowS) -> Show ZRECmd
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [ZRECmd] -> ShowS
$cshowList :: [ZRECmd] -> ShowS
show :: ZRECmd -> String
$cshow :: ZRECmd -> String
showsPrec :: Int -> ZRECmd -> ShowS
$cshowsPrec :: Int -> ZRECmd -> ShowS
Show, ZRECmd -> ZRECmd -> Bool
(ZRECmd -> ZRECmd -> Bool)
-> (ZRECmd -> ZRECmd -> Bool) -> Eq ZRECmd
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: ZRECmd -> ZRECmd -> Bool
$c/= :: ZRECmd -> ZRECmd -> Bool
== :: ZRECmd -> ZRECmd -> Bool
$c== :: ZRECmd -> ZRECmd -> Bool
Eq, Eq ZRECmd
Eq ZRECmd =>
(ZRECmd -> ZRECmd -> Ordering)
-> (ZRECmd -> ZRECmd -> Bool)
-> (ZRECmd -> ZRECmd -> Bool)
-> (ZRECmd -> ZRECmd -> Bool)
-> (ZRECmd -> ZRECmd -> Bool)
-> (ZRECmd -> ZRECmd -> ZRECmd)
-> (ZRECmd -> ZRECmd -> ZRECmd)
-> Ord ZRECmd
ZRECmd -> ZRECmd -> Bool
ZRECmd -> ZRECmd -> Ordering
ZRECmd -> ZRECmd -> ZRECmd
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 :: ZRECmd -> ZRECmd -> ZRECmd
$cmin :: ZRECmd -> ZRECmd -> ZRECmd
max :: ZRECmd -> ZRECmd -> ZRECmd
$cmax :: ZRECmd -> ZRECmd -> ZRECmd
>= :: ZRECmd -> ZRECmd -> Bool
$c>= :: ZRECmd -> ZRECmd -> Bool
> :: ZRECmd -> ZRECmd -> Bool
$c> :: ZRECmd -> ZRECmd -> Bool
<= :: ZRECmd -> ZRECmd -> Bool
$c<= :: ZRECmd -> ZRECmd -> Bool
< :: ZRECmd -> ZRECmd -> Bool
$c< :: ZRECmd -> ZRECmd -> Bool
compare :: ZRECmd -> ZRECmd -> Ordering
$ccompare :: ZRECmd -> ZRECmd -> Ordering
$cp1Ord :: Eq ZRECmd
Ord)

zreBeacon :: B.ByteString -> Port -> B.ByteString
zreBeacon :: ByteString -> Int -> ByteString
zreBeacon uuid :: ByteString
uuid port :: Int
port = Put -> ByteString
runPut (Put -> ByteString) -> Put -> ByteString
forall a b. (a -> b) -> a -> b
$ do
  ByteString -> Put
putByteString "ZRE"
  -- XXX: for compatibility with zyre implementation
  -- this should use 0x01 instead, but why when
  -- we can stick zre version there and use it for filtering?
  -- for now leave in compat mode as we don't
  -- assert this but zyre does
  Int8 -> Put
putInt8 (Int8 -> Put) -> Int8 -> Put
forall a b. (a -> b) -> a -> b
$ Int -> Int8
forall a b. (Integral a, Num b) => a -> b
fromIntegral (0x01 :: Int) -- compat
  --putInt8 $ fromIntegral zreVer -- non-compat
  ByteString -> Put
putByteString ByteString
uuid
  Int16 -> Put
putInt16be (Int16 -> Put) -> Int16 -> Put
forall a b. (a -> b) -> a -> b
$ Int -> Int16
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
port

parseUUID :: Get UUID
parseUUID :: Get UUID
parseUUID =  do
  Maybe UUID
muuid <- ByteString -> Maybe UUID
fromByteString (ByteString -> Maybe UUID)
-> (ByteString -> ByteString) -> ByteString -> Maybe UUID
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> ByteString
BL.fromStrict (ByteString -> Maybe UUID) -> Get ByteString -> Get (Maybe UUID)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> Get ByteString
getByteString 16
  case Maybe UUID
muuid of
    Just uuid :: UUID
uuid -> UUID -> Get UUID
forall (m :: * -> *) a. Monad m => a -> m a
return UUID
uuid
    Nothing -> String -> Get UUID
forall (m :: * -> *) a. MonadFail m => String -> m a
fail "Unable to parse UUID"

parseBeacon :: B.ByteString
            -> (Either String (B.ByteString, Integer, UUID, Integer))
parseBeacon :: ByteString -> Either String (ByteString, Integer, UUID, Integer)
parseBeacon = Get (ByteString, Integer, UUID, Integer)
-> ByteString -> Either String (ByteString, Integer, UUID, Integer)
forall a. Get a -> ByteString -> Either String a
runGet (Get (ByteString, Integer, UUID, Integer)
 -> ByteString
 -> Either String (ByteString, Integer, UUID, Integer))
-> Get (ByteString, Integer, UUID, Integer)
-> ByteString
-> Either String (ByteString, Integer, UUID, Integer)
forall a b. (a -> b) -> a -> b
$ do
  ByteString
lead <- Int -> Get ByteString
getByteString 3
  Integer
ver <- Get Integer
forall a. Integral a => Get a
getInt8
  UUID
uuid <- Get UUID
parseUUID
  Integer
port <- Get Integer
forall a. Integral a => Get a
getInt16
  (ByteString, Integer, UUID, Integer)
-> Get (ByteString, Integer, UUID, Integer)
forall (m :: * -> *) a. Monad m => a -> m a
return (ByteString
lead, Integer
ver, UUID
uuid, Integer
port)

cmdCode :: ZRECmd -> Word8
cmdCode :: ZRECmd -> Word8
cmdCode (Hello _ _ _ _ _) = 0x01
cmdCode (Whisper _)       = 0x02
cmdCode (Shout _ _)       = 0x03
cmdCode (Join _ _)        = 0x04
cmdCode (Leave _ _)       = 0x05
cmdCode Ping              = 0x06
cmdCode PingOk            = 0x07

getContent :: ZRECmd -> Content
getContent :: ZRECmd -> Content
getContent (Whisper c :: Content
c) = Content
c
getContent (Shout _ c :: Content
c) = Content
c
getContent _ = []

newZRE :: Seq -> ZRECmd -> ZREMsg
newZRE :: Int -> ZRECmd -> ZREMsg
newZRE seqNum :: Int
seqNum cmd :: ZRECmd
cmd = Maybe UUID -> Int -> Maybe UTCTime -> ZRECmd -> ZREMsg
ZREMsg Maybe UUID
forall a. Maybe a
Nothing Int
seqNum Maybe UTCTime
forall a. Maybe a
Nothing ZRECmd
cmd

encodeZRE :: ZREMsg -> [B.ByteString]
encodeZRE :: ZREMsg -> Content
encodeZRE ZREMsg{..} = ByteString
msgByteString -> Content -> Content
forall a. a -> [a] -> [a]
:(ZRECmd -> Content
getContent ZRECmd
msgCmd)
  where
    msg :: ByteString
msg = Put -> ByteString
runPut (Put -> ByteString) -> Put -> ByteString
forall a b. (a -> b) -> a -> b
$ do
      Word16 -> Put
putWord16be Word16
zreSig
      Word8 -> Put
putWord8 (Word8 -> Put) -> Word8 -> Put
forall a b. (a -> b) -> a -> b
$ ZRECmd -> Word8
cmdCode ZRECmd
msgCmd
      Int8 -> Put
putInt8 (Int8 -> Put) -> Int8 -> Put
forall a b. (a -> b) -> a -> b
$ Int -> Int8
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
zreVer
      Int16 -> Put
putInt16be (Int16 -> Put) -> Int16 -> Put
forall a b. (a -> b) -> a -> b
$ Int -> Int16
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
msgSeq
      ZRECmd -> Put
encodeCmd ZRECmd
msgCmd

encodeCmd :: ZRECmd -> PutM ()
encodeCmd :: ZRECmd -> Put
encodeCmd (Hello endpoint :: Endpoint
endpoint groups :: Groups
groups statusSeq :: Int
statusSeq name :: ByteString
name headers :: Headers
headers) = do
  ByteString -> Put
putByteStringLen (Endpoint -> ByteString
pEndpoint Endpoint
endpoint)
  Groups -> Put
forall (t :: * -> *). Foldable t => t ByteString -> Put
putByteStrings Groups
groups
  Int8 -> Put
putInt8 (Int8 -> Put) -> Int8 -> Put
forall a b. (a -> b) -> a -> b
$ Int -> Int8
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
statusSeq
  ByteString -> Put
putByteStringLen ByteString
name
  Headers -> Put
putMap Headers
headers
encodeCmd (Shout group :: ByteString
group _content :: Content
_content) = ByteString -> Put
putByteStringLen ByteString
group
encodeCmd (Join group :: ByteString
group statusSeq :: Int
statusSeq) = do
  ByteString -> Put
putByteStringLen ByteString
group
  Int8 -> Put
putInt8 (Int8 -> Put) -> Int8 -> Put
forall a b. (a -> b) -> a -> b
$ Int -> Int8
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
statusSeq
encodeCmd (Leave group :: ByteString
group statusSeq :: Int
statusSeq) = do
  ByteString -> Put
putByteStringLen ByteString
group
  Int8 -> Put
putInt8 (Int8 -> Put) -> Int8 -> Put
forall a b. (a -> b) -> a -> b
$ Int -> Int8
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
statusSeq
encodeCmd _ = () -> Put
forall (m :: * -> *) a. Monad m => a -> m a
return ()

parseHello :: Get ZRECmd
parseHello :: Get ZRECmd
parseHello = Endpoint -> Groups -> Int -> ByteString -> Headers -> ZRECmd
Hello
  (Endpoint -> Groups -> Int -> ByteString -> Headers -> ZRECmd)
-> Get Endpoint
-> Get (Groups -> Int -> ByteString -> Headers -> ZRECmd)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Get Endpoint
parseEndpoint'
  Get (Groups -> Int -> ByteString -> Headers -> ZRECmd)
-> Get Groups -> Get (Int -> ByteString -> Headers -> ZRECmd)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (Content -> Groups) -> Get Content -> Get Groups
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Content -> Groups
forall a. Ord a => [a] -> Set a
S.fromList Get Content
parseStrings
  Get (Int -> ByteString -> Headers -> ZRECmd)
-> Get Int -> Get (ByteString -> Headers -> ZRECmd)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Get Int
forall a. Integral a => Get a
getInt8
  Get (ByteString -> Headers -> ZRECmd)
-> Get ByteString -> Get (Headers -> ZRECmd)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Get ByteString
parseString
  Get (Headers -> ZRECmd) -> Get Headers -> Get ZRECmd
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Get Headers
parseMap
  where
    parseEndpoint' :: Get Endpoint
parseEndpoint' = do
      ByteString
s <- Get ByteString
parseString
      case ByteString -> Either String Endpoint
parseAttoEndpoint ByteString
s of
        (Left err :: String
err) -> String -> Get Endpoint
forall (m :: * -> *) a. MonadFail m => String -> m a
fail (String -> Get Endpoint) -> String -> Get Endpoint
forall a b. (a -> b) -> a -> b
$ "Unable to parse endpoint: " String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
err
        (Right endpoint :: Endpoint
endpoint) -> Endpoint -> Get Endpoint
forall (m :: * -> *) a. Monad m => a -> m a
return Endpoint
endpoint

parseShout :: Content -> Get ZRECmd
parseShout :: Content -> Get ZRECmd
parseShout frames :: Content
frames = ByteString -> Content -> ZRECmd
Shout (ByteString -> Content -> ZRECmd)
-> Get ByteString -> Get (Content -> ZRECmd)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Get ByteString
parseString Get (Content -> ZRECmd) -> Get Content -> Get ZRECmd
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Content -> Get Content
forall (f :: * -> *) a. Applicative f => a -> f a
pure Content
frames

parseJoin :: Get ZRECmd
parseJoin :: Get ZRECmd
parseJoin = ByteString -> Int -> ZRECmd
Join (ByteString -> Int -> ZRECmd)
-> Get ByteString -> Get (Int -> ZRECmd)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Get ByteString
parseString Get (Int -> ZRECmd) -> Get Int -> Get ZRECmd
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Get Int
forall a. Integral a => Get a
getInt8

parseLeave :: Get ZRECmd
parseLeave :: Get ZRECmd
parseLeave = ByteString -> Int -> ZRECmd
Leave (ByteString -> Int -> ZRECmd)
-> Get ByteString -> Get (Int -> ZRECmd)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Get ByteString
parseString Get (Int -> ZRECmd) -> Get Int -> Get ZRECmd
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Get Int
forall a. Integral a => Get a
getInt8

parseCmd :: B.ByteString -> Content -> Get ZREMsg
parseCmd :: ByteString -> Content -> Get ZREMsg
parseCmd from :: ByteString
from frames :: Content
frames = do
    Int
cmd <- (Get Int
forall a. Integral a => Get a
getInt8 :: Get Int)
    Int
ver <- Get Int
forall a. Integral a => Get a
getInt8
    Int
sqn <- Get Int
forall a. Integral a => Get a
getInt16

    case Get UUID -> ByteString -> Either String UUID
forall a. Get a -> ByteString -> Either String a
runGet Get UUID
parseUUID ByteString
from of
      (Left err :: String
err) -> String -> Get ZREMsg
forall (m :: * -> *) a. MonadFail m => String -> m a
fail (String -> Get ZREMsg) -> String -> Get ZREMsg
forall a b. (a -> b) -> a -> b
$ "No UUID: " String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
err
      (Right uuid :: UUID
uuid)-> do
        if Int
ver Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
zreVer
          then String -> Get ZREMsg
forall (m :: * -> *) a. MonadFail m => String -> m a
fail "Protocol version mismatch"
          else do

            ZRECmd
zcmd <- case Int
cmd of
              0x01 -> Get ZRECmd
parseHello
              0x02 -> ZRECmd -> Get ZRECmd
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ZRECmd -> Get ZRECmd) -> ZRECmd -> Get ZRECmd
forall a b. (a -> b) -> a -> b
$ Content -> ZRECmd
Whisper Content
frames -- parseWhisper
              0x03 -> Content -> Get ZRECmd
parseShout Content
frames
              0x04 -> Get ZRECmd
parseJoin
              0x05 -> Get ZRECmd
parseLeave
              0x06 -> ZRECmd -> Get ZRECmd
forall (f :: * -> *) a. Applicative f => a -> f a
pure ZRECmd
Ping
              0x07 -> ZRECmd -> Get ZRECmd
forall (f :: * -> *) a. Applicative f => a -> f a
pure ZRECmd
PingOk
              _    -> String -> Get ZRECmd
forall (m :: * -> *) a. MonadFail m => String -> m a
fail "Unknown command"

            ZREMsg -> Get ZREMsg
forall (m :: * -> *) a. Monad m => a -> m a
return (ZREMsg -> Get ZREMsg) -> ZREMsg -> Get ZREMsg
forall a b. (a -> b) -> a -> b
$ Maybe UUID -> Int -> Maybe UTCTime -> ZRECmd -> ZREMsg
ZREMsg (UUID -> Maybe UUID
forall a. a -> Maybe a
Just UUID
uuid) Int
sqn Maybe UTCTime
forall a. Maybe a
Nothing ZRECmd
zcmd

parseZRE :: [B.ByteString] -> Either String ZREMsg
parseZRE :: Content -> Either String ZREMsg
parseZRE (from :: ByteString
from:msg :: ByteString
msg:rest :: Content
rest) = ByteString -> ByteString -> Content -> Either String ZREMsg
parseZre ByteString
from ByteString
msg Content
rest
parseZRE _ = String -> Either String ZREMsg
forall a b. a -> Either a b
Left "empty message"

parseZre :: B.ByteString -> B.ByteString -> Content -> Either String ZREMsg
parseZre :: ByteString -> ByteString -> Content -> Either String ZREMsg
parseZre from :: ByteString
from msg :: ByteString
msg frames :: Content
frames = (Get ZREMsg -> ByteString -> Either String ZREMsg)
-> ByteString -> Get ZREMsg -> Either String ZREMsg
forall a b c. (a -> b -> c) -> b -> a -> c
flip Get ZREMsg -> ByteString -> Either String ZREMsg
forall a. Get a -> ByteString -> Either String a
runGet ByteString
msg (Get ZREMsg -> Either String ZREMsg)
-> Get ZREMsg -> Either String ZREMsg
forall a b. (a -> b) -> a -> b
$ do
  Word16
sig <- Get Word16
forall a. Integral a => Get a
getInt16
  if Word16
sig Word16 -> Word16 -> Bool
forall a. Eq a => a -> a -> Bool
/= Word16
zreSig
    then String -> Get ZREMsg
forall (m :: * -> *) a. MonadFail m => String -> m a
fail "Signature mismatch"
    else do
      -- we need to drop 1st byte of from string which is '1':UUID (17 bytes)
      ZREMsg
res <- ByteString -> Content -> Get ZREMsg
parseCmd (ByteString -> ByteString
B.tail ByteString
from) Content
frames
      ZREMsg -> Get ZREMsg
forall (m :: * -> *) a. Monad m => a -> m a
return ZREMsg
res