agent entity types (#148)

* agent entity types (WIP - fails)

* agent entities (tests pass, TODO - reduce boilerplate)

* simplify test patters

* simplify test patterns 2

* refactor testEquality for ACommand

* stricter entity parsing and correct serialization, updated tests

* fix check of entity ID

* remove unused instance Eq (Entity t)

package.yaml

@@ -30,6 +30,7 @@ dependencies:
   - base >= 4.7 && < 5
   - base64-bytestring >= 1.0 && < 1.3
   - bytestring == 0.10.*
+  - constraints == 0.12.*
   - containers == 0.6.*
   - cryptonite == 0.27.*
   - directory == 1.3.*

src/Simplex/Messaging/Agent.hs

@@ -7,6 +7,7 @@
 {-# LANGUAGE OverloadedStrings #-}
 {-# LANGUAGE RankNTypes #-}
 {-# LANGUAGE ScopedTypeVariables #-}
+{-# LANGUAGE TypeApplications #-}
 
 -- |
 -- Module      : Simplex.Messaging.Agent
@@ -49,7 +50,7 @@ import Simplex.Messaging.Agent.Store
 import Simplex.Messaging.Agent.Store.SQLite (SQLiteStore, connectSQLiteStore)
 import Simplex.Messaging.Client (SMPServerTransmission)
 import qualified Simplex.Messaging.Crypto as C
-import Simplex.Messaging.Protocol (CorrId (..), MsgBody, SenderPublicKey)
+import Simplex.Messaging.Protocol (MsgBody, SenderPublicKey)
 import qualified Simplex.Messaging.Protocol as SMP
 import Simplex.Messaging.Transport (ATransport (..), TProxy, Transport (..), runTransportServer)
 import Simplex.Messaging.Util (bshow)
@@ -105,12 +106,14 @@ runSMPAgentClient c = do
   race_ (subscriber c s1) (client c s2)
 
 receive :: forall c m. (Transport c, MonadUnliftIO m) => c -> AgentClient -> m ()
-receive h c@AgentClient {rcvQ, sndQ} = forever $ do
-  (corrId, cAlias, cmdOrErr) <- tGet SClient h
-  case cmdOrErr of
-    Right cmd -> write rcvQ (corrId, cAlias, cmd)
-    Left e -> write sndQ (corrId, cAlias, ERR e)
+receive h c@AgentClient {rcvQ, sndQ} = forever loop
   where
+    loop :: m ()
+    loop = do
+      ATransmissionOrError corrId entity cmdOrErr <- tGet SClient h
+      case cmdOrErr of
+        Right cmd -> write rcvQ $ ATransmission corrId entity cmd
+        Left e -> write sndQ $ ATransmission corrId entity $ ERR e
     write :: TBQueue (ATransmission p) -> ATransmission p -> m ()
     write q t = do
       logClient c "-->" t
@@ -123,15 +126,18 @@ send h c@AgentClient {sndQ} = forever $ do
   logClient c "<--" t
 
 logClient :: MonadUnliftIO m => AgentClient -> ByteString -> ATransmission a -> m ()
-logClient AgentClient {clientId} dir (CorrId corrId, cAlias, cmd) = do
-  logInfo . decodeUtf8 $ B.unwords [bshow clientId, dir, "A :", corrId, cAlias, B.takeWhile (/= ' ') $ serializeCommand cmd]
+logClient AgentClient {clientId} dir (ATransmission corrId entity cmd) = do
+  logInfo . decodeUtf8 $ B.unwords [bshow clientId, dir, "A :", corrId, serializeEntity entity, B.takeWhile (/= ' ') $ serializeCommand cmd]
 
-client :: (MonadUnliftIO m, MonadReader Env m) => AgentClient -> SQLiteStore -> m ()
-client c@AgentClient {rcvQ, sndQ} st = forever $ do
-  t@(corrId, cAlias, _) <- atomically $ readTBQueue rcvQ
-  runExceptT (processCommand c st t) >>= \case
-    Left e -> atomically $ writeTBQueue sndQ (corrId, cAlias, ERR e)
-    Right _ -> return ()
+client :: forall m. (MonadUnliftIO m, MonadReader Env m) => AgentClient -> SQLiteStore -> m ()
+client c@AgentClient {rcvQ, sndQ} st = forever loop
+  where
+    loop :: m ()
+    loop = do
+      t@(ATransmission corrId entity _) <- atomically $ readTBQueue rcvQ
+      runExceptT (processCommand c st t) >>= \case
+        Left e -> atomically . writeTBQueue sndQ $ ATransmission corrId entity (ERR e)
+        Right _ -> pure ()
 
 withStore ::
   AgentMonad m =>
@@ -151,24 +157,26 @@ withStore action = do
       e -> INTERNAL $ show e
 
 processCommand :: forall m. AgentMonad m => AgentClient -> SQLiteStore -> ATransmission 'Client -> m ()
-processCommand c@AgentClient {sndQ} st (corrId, connAlias, cmd) =
-  case cmd of
-    NEW -> createNewConnection
-    JOIN smpQueueInfo replyMode -> joinConnection smpQueueInfo replyMode
-    SUB -> subscribeConnection connAlias
-    SUBALL -> subscribeAll
-    SEND msgBody -> sendMessage msgBody
-    OFF -> suspendConnection
-    DEL -> deleteConnection
+processCommand c@AgentClient {sndQ} st (ATransmission corrId entity cmd) =
+  case entity of
+    Conn cId -> case cmd of
+      NEW -> createNewConnection cId
+      JOIN smpQueueInfo replyMode -> joinConnection cId smpQueueInfo replyMode
+      SUB -> subscribeConnection cId
+      SUBALL -> subscribeAll
+      SEND msgBody -> sendMessage cId msgBody
+      OFF -> suspendConnection cId
+      DEL -> deleteConnection cId
+    _ -> atomically . writeTBQueue sndQ . ATransmission corrId entity . ERR $ CMD ENTITY
   where
-    createNewConnection :: m ()
-    createNewConnection = do
+    createNewConnection :: ByteString -> m ()
+    createNewConnection cId = do
       -- TODO create connection alias if not passed
       -- make connAlias Maybe?
       srv <- getSMPServer
-      (rq, qInfo) <- newReceiveQueue c srv connAlias
+      (rq, qInfo) <- newReceiveQueue c srv cId
       withStore $ createRcvConn st rq
-      respond $ INV qInfo
+      respond (Conn cId) $ INV qInfo
 
     getSMPServer :: m SMPServer
     getSMPServer =
@@ -179,33 +187,33 @@ processCommand c@AgentClient {sndQ} st (corrId, connAlias, cmd) =
           i <- atomically . stateTVar gen $ randomR (0, L.length servers - 1)
           pure $ servers L.!! i
 
-    joinConnection :: SMPQueueInfo -> ReplyMode -> m ()
-    joinConnection qInfo (ReplyMode replyMode) = do
+    joinConnection :: ByteString -> SMPQueueInfo -> ReplyMode -> m ()
+    joinConnection cId qInfo (ReplyMode replyMode) = do
       -- TODO create connection alias if not passed
       -- make connAlias Maybe?
-      (sq, senderKey, verifyKey) <- newSendQueue qInfo connAlias
+      (sq, senderKey, verifyKey) <- newSendQueue qInfo cId
       withStore $ createSndConn st sq
       connectToSendQueue c st sq senderKey verifyKey
-      when (replyMode == On) $ createReplyQueue sq
+      when (replyMode == On) $ createReplyQueue cId sq
     -- TODO this response is disabled to avoid two responses in terminal client (OK + CON),
     -- respond OK
 
-    subscribeConnection :: ConnAlias -> m ()
-    subscribeConnection cAlias =
-      withStore (getConn st cAlias) >>= \case
+    subscribeConnection :: ByteString -> m ()
+    subscribeConnection cId =
+      withStore (getConn st cId) >>= \case
         SomeConn _ (DuplexConnection _ rq _) -> subscribe rq
         SomeConn _ (RcvConnection _ rq) -> subscribe rq
         _ -> throwError $ CONN SIMPLEX
       where
-        subscribe rq = subscribeQueue c rq cAlias >> respond' cAlias OK
+        subscribe rq = subscribeQueue c rq cId >> respond (Conn cId) OK
 
     -- TODO remove - hack for subscribing to all; respond' and parameterization of subscribeConnection are byproduct
     subscribeAll :: m ()
     subscribeAll = withStore (getAllConnAliases st) >>= mapM_ subscribeConnection
 
-    sendMessage :: MsgBody -> m ()
-    sendMessage msgBody =
-      withStore (getConn st connAlias) >>= \case
+    sendMessage :: ByteString -> MsgBody -> m ()
+    sendMessage cId msgBody =
+      withStore (getConn st cId) >>= \case
         SomeConn _ (DuplexConnection _ _ sq) -> sendMsg sq
         SomeConn _ (SndConnection _ sq) -> sendMsg sq
         _ -> throwError $ CONN SIMPLEX
@@ -226,35 +234,35 @@ processCommand c@AgentClient {sndQ} st (corrId, connAlias, cmd) =
             createSndMsg st sq $
               SndMsgData {internalId, internalSndId, internalTs, msgBody, internalHash = msgHash}
           sendAgentMessage c sq msgStr
-          respond $ SENT (unId internalId)
+          respond (Conn cId) $ SENT (unId internalId)
 
-    suspendConnection :: m ()
-    suspendConnection =
-      withStore (getConn st connAlias) >>= \case
+    suspendConnection :: ByteString -> m ()
+    suspendConnection cId =
+      withStore (getConn st cId) >>= \case
         SomeConn _ (DuplexConnection _ rq _) -> suspend rq
         SomeConn _ (RcvConnection _ rq) -> suspend rq
         _ -> throwError $ CONN SIMPLEX
       where
-        suspend rq = suspendQueue c rq >> respond OK
+        suspend rq = suspendQueue c rq >> respond (Conn cId) OK
 
-    deleteConnection :: m ()
-    deleteConnection =
-      withStore (getConn st connAlias) >>= \case
+    deleteConnection :: ByteString -> m ()
+    deleteConnection cId =
+      withStore (getConn st cId) >>= \case
         SomeConn _ (DuplexConnection _ rq _) -> delete rq
         SomeConn _ (RcvConnection _ rq) -> delete rq
         _ -> delConn
       where
-        delConn = withStore (deleteConn st connAlias) >> respond OK
+        delConn = withStore (deleteConn st cId) >> respond (Conn cId) OK
         delete rq = do
           deleteQueue c rq
-          removeSubscription c connAlias
+          removeSubscription c cId
           delConn
 
-    createReplyQueue :: SndQueue -> m ()
-    createReplyQueue sq = do
+    createReplyQueue :: ByteString -> SndQueue -> m ()
+    createReplyQueue cId sq = do
       srv <- getSMPServer
-      (rq, qInfo) <- newReceiveQueue c srv connAlias
-      withStore $ upgradeSndConnToDuplex st connAlias rq
+      (rq, qInfo) <- newReceiveQueue c srv cId
+      withStore $ upgradeSndConnToDuplex st cId rq
       senderTimestamp <- liftIO getCurrentTime
       sendAgentMessage c sq . serializeSMPMessage $
         SMPMessage
@@ -264,11 +272,8 @@ processCommand c@AgentClient {sndQ} st (corrId, connAlias, cmd) =
             agentMessage = REPLY qInfo
           }
 
-    respond :: ACommand 'Agent -> m ()
-    respond = respond' connAlias
-
-    respond' :: ConnAlias -> ACommand 'Agent -> m ()
-    respond' cAlias resp = atomically $ writeTBQueue sndQ (corrId, cAlias, resp)
+    respond :: EntityCommand t c => Entity t -> ACommand 'Agent c -> m ()
+    respond ent resp = atomically . writeTBQueue sndQ $ ATransmission corrId ent resp
 
 subscriber :: (MonadUnliftIO m, MonadReader Env m) => AgentClient -> SQLiteStore -> m ()
 subscriber c@AgentClient {msgQ} st = forever $ do
@@ -283,7 +288,7 @@ processSMPTransmission c@AgentClient {sndQ} st (srv, rId, cmd) = do
   withStore (getRcvConn st srv rId) >>= \case
     SomeConn SCDuplex (DuplexConnection _ rq _) -> processSMP SCDuplex rq
     SomeConn SCRcv (RcvConnection _ rq) -> processSMP SCRcv rq
-    _ -> atomically $ writeTBQueue sndQ ("", "", ERR $ CONN SIMPLEX)
+    _ -> atomically . writeTBQueue sndQ $ ATransmission "" (Conn "") (ERR $ CONN SIMPLEX)
   where
     processSMP :: SConnType c -> RcvQueue -> m ()
     processSMP cType rq@RcvQueue {connAlias, status} =
@@ -310,8 +315,8 @@ processSMPTransmission c@AgentClient {sndQ} st (srv, rId, cmd) = do
           logServer "<--" c srv rId $ "unexpected: " <> bshow cmd
           notify . ERR $ BROKER UNEXPECTED
       where
-        notify :: ACommand 'Agent -> m ()
-        notify msg = atomically $ writeTBQueue sndQ ("", connAlias, msg)
+        notify :: EntityCommand 'Conn_ c => ACommand 'Agent c -> m ()
+        notify msg = atomically . writeTBQueue sndQ $ ATransmission "" (Conn connAlias) msg
 
         prohibited :: m ()
         prohibited = notify . ERR $ AGENT A_PROHIBITED

src/Simplex/Messaging/Agent/Client.hs

@@ -118,7 +118,7 @@ getSMPServerClient c@AgentClient {smpClients, msgQ} srv =
         deleteKeys ks m = S.foldr' M.delete m ks
 
     notifySub :: ConnAlias -> IO ()
-    notifySub connAlias = atomically $ writeTBQueue (sndQ c) ("", connAlias, END)
+    notifySub connAlias = atomically . writeTBQueue (sndQ c) $ ATransmission "" (Conn connAlias) END
 
 closeSMPServerClients :: MonadUnliftIO m => AgentClient -> m ()
 closeSMPServerClients c = liftIO $ readTVarIO (smpClients c) >>= mapM_ closeSMPClient

src/Simplex/Messaging/Agent/Protocol.hs

@@ -1,14 +1,18 @@
 {-# LANGUAGE DataKinds #-}
 {-# LANGUAGE DeriveAnyClass #-}
 {-# LANGUAGE DeriveGeneric #-}
+{-# LANGUAGE FlexibleInstances #-}
 {-# LANGUAGE GADTs #-}
 {-# LANGUAGE LambdaCase #-}
 {-# LANGUAGE NamedFieldPuns #-}
 {-# LANGUAGE OverloadedStrings #-}
+{-# LANGUAGE PolyKinds #-}
 {-# LANGUAGE RankNTypes #-}
 {-# LANGUAGE ScopedTypeVariables #-}
 {-# LANGUAGE StandaloneDeriving #-}
+{-# LANGUAGE TypeApplications #-}
 {-# LANGUAGE TypeFamilies #-}
+{-# LANGUAGE TypeOperators #-}
 {-# OPTIONS_GHC -fno-warn-unticked-promoted-constructors #-}
 
 -- |
@@ -25,8 +29,14 @@
 -- See https://github.com/simplex-chat/simplexmq/blob/master/protocol/agent-protocol.md
 module Simplex.Messaging.Agent.Protocol
   ( -- * SMP agent protocol types
+    Entity (..),
+    EntityTag (..),
+    AnEntity (..),
+    EntityCommand,
+    entityCommand,
     ACommand (..),
     AParty (..),
+    APartyCmd (..),
     SAParty (..),
     SMPMessage (..),
     AMessage (..),
@@ -37,8 +47,8 @@ module Simplex.Messaging.Agent.Protocol
     ConnectionErrorType (..),
     BrokerErrorType (..),
     SMPAgentError (..),
-    ATransmission,
-    ATransmissionOrError,
+    ATransmission (..),
+    ATransmissionOrError (..),
     ARawTransmission,
     ConnAlias,
     ReplyMode (..),
@@ -51,15 +61,19 @@ module Simplex.Messaging.Agent.Protocol
     VerificationKey,
     EncryptionKey,
     DecryptionKey,
+    ACorrId,
+    AgentMsgId,
 
     -- * Parse and serialize
     serializeCommand,
+    serializeEntity,
     serializeSMPMessage,
     serializeMsgIntegrity,
     serializeServer,
     serializeSmpQueueInfo,
     serializeAgentError,
     commandP,
+    entityP,
     parseSMPMessage,
     smpServerP,
     smpQueueInfoP,
@@ -81,9 +95,11 @@ import qualified Data.Attoparsec.ByteString.Char8 as A
 import Data.ByteString.Base64
 import Data.ByteString.Char8 (ByteString)
 import qualified Data.ByteString.Char8 as B
+import Data.Constraint (Dict (..))
 import Data.Functor (($>))
 import Data.Int (Int64)
-import Data.Kind (Type)
+import Data.Kind (Constraint, Type)
+import Data.Maybe (isJust)
 import Data.String (IsString (..))
 import Data.Time.Clock (UTCTime)
 import Data.Time.ISO8601
@@ -91,12 +107,11 @@ import Data.Type.Equality
 import Data.Typeable ()
 import GHC.Generics (Generic)
 import Generic.Random (genericArbitraryU)
-import Network.Socket
+import Network.Socket (HostName, ServiceName)
 import qualified Simplex.Messaging.Crypto as C
 import Simplex.Messaging.Parsers
 import Simplex.Messaging.Protocol
-  ( CorrId (..),
-    ErrorType,
+  ( ErrorType,
     MsgBody,
     MsgId,
     SenderPublicKey,
@@ -112,10 +127,14 @@ import UnliftIO.Exception
 type ARawTransmission = (ByteString, ByteString, ByteString)
 
 -- | Parsed SMP agent protocol transmission.
-type ATransmission p = (CorrId, ConnAlias, ACommand p)
+data ATransmission p = forall t c. EntityCommand t c => ATransmission ACorrId (Entity t) (ACommand p c)
 
 -- | SMP agent protocol transmission or transmission error.
-type ATransmissionOrError p = (CorrId, ConnAlias, Either AgentErrorType (ACommand p))
+data ATransmissionOrError p = forall t c. EntityCommand t c => ATransmissionOrError ACorrId (Entity t) (Either AgentErrorType (ACommand p c))
+
+deriving instance Show (ATransmissionOrError p)
+
+type ACorrId = ByteString
 
 -- | SMP agent protocol participants.
 data AParty = Agent | Client
@@ -135,26 +154,107 @@ instance TestEquality SAParty where
   testEquality SClient SClient = Just Refl
   testEquality _ _ = Nothing
 
-data ACmd = forall p. ACmd (SAParty p) (ACommand p)
+-- | SMP agent protocol entity types
+data EntityTag = Conn_ | OpenConn_ | Broadcast_ | AGroup_
+
+data Entity :: EntityTag -> Type where
+  Conn :: ByteString -> Entity Conn_
+  OpenConn :: ByteString -> Entity OpenConn_
+  BroadCast :: ByteString -> Entity Broadcast_
+  AGroup :: ByteString -> Entity AGroup_
+
+deriving instance Show (Entity t)
+
+entityId :: Entity t -> ByteString
+entityId = \case
+  Conn bs -> bs
+  OpenConn bs -> bs
+  BroadCast bs -> bs
+  AGroup bs -> bs
+
+data AnEntity = forall t. AE (Entity t)
+
+data ACmd = forall (p :: AParty) (c :: ACmdTag). ACmd (SAParty p) (ACommand p c)
 
 deriving instance Show ACmd
 
+data APartyCmd (p :: AParty) = forall c. APartyCmd (ACommand p c)
+
+instance Eq (APartyCmd p) where
+  APartyCmd c1 == APartyCmd c2 = isJust $ testEquality c1 c2
+
+deriving instance Show (APartyCmd p)
+
+type family EntityCommand (t :: EntityTag) (c :: ACmdTag) :: Constraint where
+  EntityCommand Conn_ NEW_ = ()
+  EntityCommand Conn_ INV_ = ()
+  EntityCommand Conn_ JOIN_ = ()
+  EntityCommand Conn_ CON_ = ()
+  EntityCommand Conn_ SUB_ = ()
+  EntityCommand Conn_ SUBALL_ = ()
+  EntityCommand Conn_ END_ = ()
+  EntityCommand Conn_ SEND_ = ()
+  EntityCommand Conn_ SENT_ = ()
+  EntityCommand Conn_ MSG_ = ()
+  EntityCommand Conn_ OFF_ = ()
+  EntityCommand Conn_ DEL_ = ()
+  EntityCommand Conn_ OK_ = ()
+  EntityCommand Conn_ ERR_ = ()
+  EntityCommand _ ERR_ = ()
+
+entityCommand :: Entity t -> ACommand p c -> Maybe (Dict (EntityCommand t c))
+entityCommand = \case
+  Conn _ -> \case
+    NEW -> Just Dict
+    INV _ -> Just Dict
+    JOIN {} -> Just Dict
+    CON -> Just Dict
+    SUB -> Just Dict
+    SUBALL -> Just Dict
+    END -> Just Dict
+    SEND _ -> Just Dict
+    SENT _ -> Just Dict
+    MSG {} -> Just Dict
+    OFF -> Just Dict
+    DEL -> Just Dict
+    OK -> Just Dict
+    ERR _ -> Just Dict
+  _ -> \case
+    ERR _ -> Just Dict
+    _ -> Nothing
+
+data ACmdTag
+  = NEW_
+  | INV_
+  | JOIN_
+  | CON_
+  | SUB_
+  | SUBALL_
+  | END_
+  | SEND_
+  | SENT_
+  | MSG_
+  | OFF_
+  | DEL_
+  | OK_
+  | ERR_
+
 -- | Parameterized type for SMP agent protocol commands and responses from all participants.
-data ACommand (p :: AParty) where
-  NEW :: ACommand Client -- response INV
-  INV :: SMPQueueInfo -> ACommand Agent
-  JOIN :: SMPQueueInfo -> ReplyMode -> ACommand Client -- response OK
-  CON :: ACommand Agent -- notification that connection is established
+data ACommand (p :: AParty) (c :: ACmdTag) where
+  NEW :: ACommand Client NEW_ -- response INV
+  INV :: SMPQueueInfo -> ACommand Agent INV_
+  JOIN :: SMPQueueInfo -> ReplyMode -> ACommand Client JOIN_ -- response OK
+  CON :: ACommand Agent CON_ -- notification that connection is established
   -- TODO currently it automatically allows whoever sends the confirmation
   -- CONF :: OtherPartyId -> ACommand Agent
   -- LET :: OtherPartyId -> ACommand Client
-  SUB :: ACommand Client
-  SUBALL :: ACommand Client -- TODO should be moved to chat protocol - hack for subscribing to all
-  END :: ACommand Agent
+  SUB :: ACommand Client SUB_
+  SUBALL :: ACommand Client SUBALL_ -- TODO should be moved to chat protocol - hack for subscribing to all
+  END :: ACommand Agent END_
   -- QST :: QueueDirection -> ACommand Client
   -- STAT :: QueueDirection -> Maybe QueueStatus -> Maybe SubMode -> ACommand Agent
-  SEND :: MsgBody -> ACommand Client
-  SENT :: AgentMsgId -> ACommand Agent
+  SEND :: MsgBody -> ACommand Client SEND_
+  SENT :: AgentMsgId -> ACommand Agent SENT_
   MSG ::
     { recipientMeta :: (AgentMsgId, UTCTime),
       brokerMeta :: (MsgId, UTCTime),
@@ -162,17 +262,37 @@ data ACommand (p :: AParty) where
       msgIntegrity :: MsgIntegrity,
       msgBody :: MsgBody
     } ->
-    ACommand Agent
+    ACommand Agent MSG_
   -- ACK :: AgentMsgId -> ACommand Client
   -- RCVD :: AgentMsgId -> ACommand Agent
-  OFF :: ACommand Client
-  DEL :: ACommand Client
-  OK :: ACommand Agent
-  ERR :: AgentErrorType -> ACommand Agent
+  OFF :: ACommand Client MSG_
+  DEL :: ACommand Client DEL_
+  OK :: ACommand Agent OK_
+  ERR :: AgentErrorType -> ACommand Agent ERR_
 
-deriving instance Eq (ACommand p)
+deriving instance Eq (ACommand p c)
 
-deriving instance Show (ACommand p)
+deriving instance Show (ACommand p c)
+
+instance TestEquality (ACommand p) where
+  testEquality NEW NEW = Just Refl
+  testEquality c@INV {} c'@INV {} = refl c c'
+  testEquality c@JOIN {} c'@JOIN {} = refl c c'
+  testEquality CON CON = Just Refl
+  testEquality SUB SUB = Just Refl
+  testEquality SUBALL SUBALL = Just Refl
+  testEquality END END = Just Refl
+  testEquality c@SEND {} c'@SEND {} = refl c c'
+  testEquality c@SENT {} c'@SENT {} = refl c c'
+  testEquality c@MSG {} c'@MSG {} = refl c c'
+  testEquality OFF OFF = Just Refl
+  testEquality DEL DEL = Just Refl
+  testEquality OK OK = Just Refl
+  testEquality c@ERR {} c'@ERR {} = refl c c'
+  testEquality _ _ = Nothing
+
+refl :: Eq (f a) => f a -> f a -> Maybe (a :~: a)
+refl x x' = if x == x' then Just Refl else Nothing
 
 -- | SMP message formats.
 data SMPMessage
@@ -369,12 +489,16 @@ data AgentErrorType
 
 -- | SMP agent protocol command or response error.
 data CommandErrorType
-  = -- | command is prohibited
+  = -- | command is prohibited in this context
     PROHIBITED
+  | -- | command is not supported by this entity
+    ENTITY
   | -- | command syntax is invalid
     SYNTAX
-  | -- | connection alias is required with this command
-    NO_CONN
+  | -- | cannot parse entity
+    BAD_ENTITY
+  | -- | entity ID is required with this command
+    NO_ENTITY
   | -- | message size is not correct (no terminating space)
     SIZE
   | -- | message does not fit in SMP block
@@ -427,7 +551,24 @@ instance Arbitrary BrokerErrorType where arbitrary = genericArbitraryU
 
 instance Arbitrary SMPAgentError where arbitrary = genericArbitraryU
 
--- | AMP agent command and response parser
+entityP :: Parser AnEntity
+entityP =
+  ($)
+    <$> ( "C:" $> AE . Conn
+            <|> "O:" $> AE . OpenConn
+            <|> "B:" $> AE . BroadCast
+            <|> "G:" $> AE . AGroup
+        )
+    <*> A.takeTill (== ' ')
+
+serializeEntity :: Entity t -> ByteString
+serializeEntity = \case
+  Conn s -> "C:" <> s
+  OpenConn s -> "O:" <> s
+  BroadCast s -> "B:" <> s
+  AGroup s -> "G:" <> s
+
+-- | SMP agent command and response parser
 commandP :: Parser ACmd
 commandP =
   "NEW" $> ACmd SClient NEW
@@ -474,7 +615,7 @@ parseCommand :: ByteString -> Either AgentErrorType ACmd
 parseCommand = parse commandP $ CMD SYNTAX
 
 -- | Serialize SMP agent command.
-serializeCommand :: ACommand p -> ByteString
+serializeCommand :: ACommand p c -> ByteString
 serializeCommand = \case
   NEW -> "NEW"
   INV qInfo -> "INV " <> serializeSmpQueueInfo qInfo
@@ -540,9 +681,9 @@ serializeMsg body = bshow (B.length body) <> "\n" <> body
 
 -- | Send raw (unparsed) SMP agent protocol transmission to TCP connection.
 tPutRaw :: Transport c => c -> ARawTransmission -> IO ()
-tPutRaw h (corrId, connAlias, command) = do
+tPutRaw h (corrId, entity, command) = do
   putLn h corrId
-  putLn h connAlias
+  putLn h entity
   putLn h command
 
 -- | Receive raw (unparsed) SMP agent protocol transmission from TCP connection.
@@ -551,41 +692,54 @@ tGetRaw h = (,,) <$> getLn h <*> getLn h <*> getLn h
 
 -- | Send SMP agent protocol command (or response) to TCP connection.
 tPut :: (Transport c, MonadIO m) => c -> ATransmission p -> m ()
-tPut h (CorrId corrId, connAlias, command) =
-  liftIO $ tPutRaw h (corrId, connAlias, serializeCommand command)
+tPut h (ATransmission corrId ent cmd) =
+  liftIO $ tPutRaw h (corrId, serializeEntity ent, serializeCommand cmd)
 
 -- | Receive client and agent transmissions from TCP connection.
 tGet :: forall c m p. (Transport c, MonadIO m) => SAParty p -> c -> m (ATransmissionOrError p)
 tGet party h = liftIO (tGetRaw h) >>= tParseLoadBody
   where
     tParseLoadBody :: ARawTransmission -> m (ATransmissionOrError p)
-    tParseLoadBody t@(corrId, connAlias, command) = do
-      let cmd = parseCommand command >>= fromParty >>= tConnAlias t
-      fullCmd <- either (return . Left) cmdWithMsgBody cmd
-      return (CorrId corrId, connAlias, fullCmd)
+    tParseLoadBody (corrId, entityStr, command) =
+      case parseAll entityP entityStr of
+        Left _ -> pure $ ATransmissionOrError @_ @_ @ERR_ corrId (Conn "") $ Left $ CMD BAD_ENTITY
+        Right entity -> do
+          let cmd = parseCommand command >>= fromParty >>= hasEntityId entity
+          makeTransmission corrId entity <$> either (pure . Left) cmdWithMsgBody cmd
 
-    fromParty :: ACmd -> Either AgentErrorType (ACommand p)
+    fromParty :: ACmd -> Either AgentErrorType (APartyCmd p)
     fromParty (ACmd (p :: p1) cmd) = case testEquality party p of
-      Just Refl -> Right cmd
+      Just Refl -> Right $ APartyCmd cmd
       _ -> Left $ CMD PROHIBITED
 
-    tConnAlias :: ARawTransmission -> ACommand p -> Either AgentErrorType (ACommand p)
-    tConnAlias (_, connAlias, _) cmd = case cmd of
-      -- NEW and JOIN have optional connAlias
-      NEW -> Right cmd
-      JOIN _ _ -> Right cmd
-      -- ERROR response does not always have connAlias
-      ERR _ -> Right cmd
-      -- other responses must have connAlias
-      _
-        | B.null connAlias -> Left $ CMD NO_CONN
-        | otherwise -> Right cmd
+    hasEntityId :: AnEntity -> APartyCmd p -> Either AgentErrorType (APartyCmd p)
+    hasEntityId (AE entity) (APartyCmd cmd) =
+      APartyCmd <$> case cmd of
+        -- NEW and JOIN have optional entity
+        NEW -> Right cmd
+        JOIN _ _ -> Right cmd
+        -- ERROR response does not always have entity
+        ERR _ -> Right cmd
+        -- other responses must have entity
+        _
+          | B.null (entityId entity) -> Left $ CMD NO_ENTITY
+          | otherwise -> Right cmd
 
-    cmdWithMsgBody :: ACommand p -> m (Either AgentErrorType (ACommand p))
-    cmdWithMsgBody = \case
-      SEND body -> SEND <$$> getMsgBody body
-      MSG agentMsgId srvTS agentTS integrity body -> MSG agentMsgId srvTS agentTS integrity <$$> getMsgBody body
-      cmd -> return $ Right cmd
+    makeTransmission :: ACorrId -> AnEntity -> Either AgentErrorType (APartyCmd p) -> ATransmissionOrError p
+    makeTransmission corrId (AE entity) = \case
+      Left e -> err e
+      Right (APartyCmd cmd) -> case entityCommand entity cmd of
+        Just Dict -> ATransmissionOrError corrId entity $ Right cmd
+        _ -> err $ CMD ENTITY
+      where
+        err e = ATransmissionOrError @_ @_ @ERR_ corrId entity $ Left e
+
+    cmdWithMsgBody :: APartyCmd p -> m (Either AgentErrorType (APartyCmd p))
+    cmdWithMsgBody (APartyCmd cmd) =
+      APartyCmd <$$> case cmd of
+        SEND body -> SEND <$$> getMsgBody body
+        MSG agentMsgId srvTS agentTS integrity body -> MSG agentMsgId srvTS agentTS integrity <$$> getMsgBody body
+        _ -> pure $ Right cmd
 
     -- TODO refactor with server
     getMsgBody :: MsgBody -> m (Either AgentErrorType MsgBody)

tests/AgentTests.hs

@@ -5,7 +5,9 @@
 {-# LANGUAGE NamedFieldPuns #-}
 {-# LANGUAGE OverloadedStrings #-}
 {-# LANGUAGE PatternSynonyms #-}
+{-# LANGUAGE RankNTypes #-}
 {-# LANGUAGE ScopedTypeVariables #-}
+{-# LANGUAGE TupleSections #-}
 
 module AgentTests where
 
@@ -37,32 +39,44 @@ agentTests (ATransport t) = do
     it "should send notifications to client when server disconnects" $
       smpAgentServerTest $ testSubscrNotification t
 
+type TestTransmission p = (ACorrId, ByteString, APartyCmd p)
+
+type TestTransmission' p c = (ACorrId, ByteString, ACommand p c)
+
+type TestTransmissionOrError p = (ACorrId, ByteString, Either AgentErrorType (APartyCmd p))
+
+testTE :: ATransmissionOrError p -> TestTransmissionOrError p
+testTE (ATransmissionOrError corrId entity cmdOrErr) =
+  (corrId,serializeEntity entity,) $ case cmdOrErr of
+    Right cmd -> Right $ APartyCmd cmd
+    Left e -> Left e
+
 -- | send transmission `t` to handle `h` and get response
-(#:) :: Transport c => c -> (ByteString, ByteString, ByteString) -> IO (ATransmissionOrError 'Agent)
-h #: t = tPutRaw h t >> tGet SAgent h
+(#:) :: Transport c => c -> (ByteString, ByteString, ByteString) -> IO (TestTransmissionOrError 'Agent)
+h #: t = tPutRaw h t >> testTE <$> tGet SAgent h
 
 -- | action and expected response
 -- `h #:t #> r` is the test that sends `t` to `h` and validates that the response is `r`
-(#>) :: IO (ATransmissionOrError 'Agent) -> ATransmission 'Agent -> Expectation
-action #> (corrId, cAlias, cmd) = action `shouldReturn` (corrId, cAlias, Right cmd)
+(#>) :: IO (TestTransmissionOrError 'Agent) -> TestTransmission' 'Agent c -> Expectation
+action #> (corrId, cAlias, cmd) = action `shouldReturn` (corrId, cAlias, Right (APartyCmd cmd))
 
 -- | action and predicate for the response
 -- `h #:t =#> p` is the test that sends `t` to `h` and validates the response using `p`
-(=#>) :: IO (ATransmissionOrError 'Agent) -> (ATransmission 'Agent -> Bool) -> Expectation
+(=#>) :: IO (TestTransmissionOrError 'Agent) -> (TestTransmission 'Agent -> Bool) -> Expectation
 action =#> p = action >>= (`shouldSatisfy` p . correctTransmission)
 
-correctTransmission :: ATransmissionOrError a -> ATransmission a
+correctTransmission :: TestTransmissionOrError p -> TestTransmission p
 correctTransmission (corrId, cAlias, cmdOrErr) = case cmdOrErr of
   Right cmd -> (corrId, cAlias, cmd)
   Left e -> error $ show e
 
 -- | receive message to handle `h` and validate that it is the expected one
-(<#) :: Transport c => c -> ATransmission 'Agent -> Expectation
-h <# (corrId, cAlias, cmd) = tGet SAgent h `shouldReturn` (corrId, cAlias, Right cmd)
+(<#) :: Transport c => c -> TestTransmission' 'Agent c' -> Expectation
+h <# (corrId, cAlias, cmd) = tGet SAgent h >>= (`shouldBe` (corrId, cAlias, Right (APartyCmd cmd))) . testTE
 
 -- | receive message to handle `h` and validate it using predicate `p`
-(<#=) :: Transport c => c -> (ATransmission 'Agent -> Bool) -> Expectation
-h <#= p = tGet SAgent h >>= (`shouldSatisfy` p . correctTransmission)
+(<#=) :: Transport c => c -> (TestTransmission 'Agent -> Bool) -> Expectation
+h <#= p = tGet SAgent h >>= (`shouldSatisfy` p . correctTransmission . testTE)
 
 -- | test that nothing is delivered to handle `h` during 10ms
 (#:#) :: Transport c => c -> String -> Expectation
@@ -73,75 +87,81 @@ h #:# err = tryGet `shouldReturn` ()
         Just _ -> error err
         _ -> return ()
 
-pattern Msg :: MsgBody -> ACommand 'Agent
-pattern Msg msgBody <- MSG {msgBody, msgIntegrity = MsgOk}
+pattern Msg :: MsgBody -> APartyCmd 'Agent
+pattern Msg msgBody <- APartyCmd MSG {msgBody, msgIntegrity = MsgOk}
+
+pattern Sent :: AgentMsgId -> APartyCmd 'Agent
+pattern Sent msgId <- APartyCmd (SENT msgId)
+
+pattern Inv :: SMPQueueInfo -> APartyCmd 'Agent
+pattern Inv invitation <- APartyCmd (INV invitation)
 
 testDuplexConnection :: Transport c => TProxy c -> c -> c -> IO ()
 testDuplexConnection _ alice bob = do
-  ("1", "bob", Right (INV qInfo)) <- alice #: ("1", "bob", "NEW")
+  ("1", "C:bob", Right (Inv qInfo)) <- alice #: ("1", "C:bob", "NEW")
   let qInfo' = serializeSmpQueueInfo qInfo
-  bob #: ("11", "alice", "JOIN " <> qInfo') #> ("", "alice", CON)
-  alice <# ("", "bob", CON)
-  alice #: ("2", "bob", "SEND :hello") =#> \case ("2", "bob", SENT 1) -> True; _ -> False
-  alice #: ("3", "bob", "SEND :how are you?") =#> \case ("3", "bob", SENT 2) -> True; _ -> False
-  bob <#= \case ("", "alice", Msg "hello") -> True; _ -> False
-  bob <#= \case ("", "alice", Msg "how are you?") -> True; _ -> False
-  bob #: ("14", "alice", "SEND 9\nhello too") =#> \case ("14", "alice", SENT 3) -> True; _ -> False
-  alice <#= \case ("", "bob", Msg "hello too") -> True; _ -> False
-  bob #: ("15", "alice", "SEND 9\nmessage 1") =#> \case ("15", "alice", SENT 4) -> True; _ -> False
-  alice <#= \case ("", "bob", Msg "message 1") -> True; _ -> False
-  alice #: ("5", "bob", "OFF") #> ("5", "bob", OK)
-  bob #: ("17", "alice", "SEND 9\nmessage 3") #> ("17", "alice", ERR (SMP AUTH))
-  alice #: ("6", "bob", "DEL") #> ("6", "bob", OK)
+  bob #: ("11", "C:alice", "JOIN " <> qInfo') #> ("", "C:alice", CON)
+  alice <# ("", "C:bob", CON)
+  alice #: ("2", "C:bob", "SEND :hello") =#> \case ("2", "C:bob", Sent 1) -> True; _ -> False
+  alice #: ("3", "C:bob", "SEND :how are you?") =#> \case ("3", "C:bob", Sent 2) -> True; _ -> False
+  bob <#= \case ("", "C:alice", Msg "hello") -> True; _ -> False
+  bob <#= \case ("", "C:alice", Msg "how are you?") -> True; _ -> False
+  bob #: ("14", "C:alice", "SEND 9\nhello too") =#> \case ("14", "C:alice", Sent 3) -> True; _ -> False
+  alice <#= \case ("", "C:bob", Msg "hello too") -> True; _ -> False
+  bob #: ("15", "C:alice", "SEND 9\nmessage 1") =#> \case ("15", "C:alice", Sent 4) -> True; _ -> False
+  alice <#= \case ("", "C:bob", Msg "message 1") -> True; _ -> False
+  alice #: ("5", "C:bob", "OFF") #> ("5", "C:bob", OK)
+  bob #: ("17", "C:alice", "SEND 9\nmessage 3") #> ("17", "C:alice", ERR (SMP AUTH))
+  alice #: ("6", "C:bob", "DEL") #> ("6", "C:bob", OK)
   alice #:# "nothing else should be delivered to alice"
 
 testSubscription :: Transport c => TProxy c -> c -> c -> c -> IO ()
 testSubscription _ alice1 alice2 bob = do
-  ("1", "bob", Right (INV qInfo)) <- alice1 #: ("1", "bob", "NEW")
+  ("1", "C:bob", Right (Inv qInfo)) <- alice1 #: ("1", "C:bob", "NEW")
   let qInfo' = serializeSmpQueueInfo qInfo
-  bob #: ("11", "alice", "JOIN " <> qInfo') #> ("", "alice", CON)
-  bob #: ("12", "alice", "SEND 5\nhello") =#> \case ("12", "alice", SENT _) -> True; _ -> False
-  bob #: ("13", "alice", "SEND 11\nhello again") =#> \case ("13", "alice", SENT _) -> True; _ -> False
-  alice1 <# ("", "bob", CON)
-  alice1 <#= \case ("", "bob", Msg "hello") -> True; _ -> False
-  alice1 <#= \case ("", "bob", Msg "hello again") -> True; _ -> False
-  alice2 #: ("21", "bob", "SUB") #> ("21", "bob", OK)
-  alice1 <# ("", "bob", END)
-  bob #: ("14", "alice", "SEND 2\nhi") =#> \case ("14", "alice", SENT _) -> True; _ -> False
-  alice2 <#= \case ("", "bob", Msg "hi") -> True; _ -> False
+  bob #: ("11", "C:alice", "JOIN " <> qInfo') #> ("", "C:alice", CON)
+  bob #: ("12", "C:alice", "SEND 5\nhello") =#> \case ("12", "C:alice", Sent _) -> True; _ -> False
+  bob #: ("13", "C:alice", "SEND 11\nhello again") =#> \case ("13", "C:alice", Sent _) -> True; _ -> False
+  alice1 <# ("", "C:bob", CON)
+  alice1 <#= \case ("", "C:bob", Msg "hello") -> True; _ -> False
+  alice1 <#= \case ("", "C:bob", Msg "hello again") -> True; _ -> False
+  alice2 #: ("21", "C:bob", "SUB") #> ("21", "C:bob", OK)
+  alice1 <# ("", "C:bob", END)
+  bob #: ("14", "C:alice", "SEND 2\nhi") =#> \case ("14", "C:alice", Sent _) -> True; _ -> False
+  alice2 <#= \case ("", "C:bob", Msg "hi") -> True; _ -> False
   alice1 #:# "nothing else should be delivered to alice1"
 
 testSubscrNotification :: Transport c => TProxy c -> (ThreadId, ThreadId) -> c -> IO ()
 testSubscrNotification _ (server, _) client = do
-  client #: ("1", "conn1", "NEW") =#> \case ("1", "conn1", INV _) -> True; _ -> False
+  client #: ("1", "C:conn1", "NEW") =#> \case ("1", "C:conn1", Inv _) -> True; _ -> False
   client #:# "nothing should be delivered to client before the server is killed"
   killThread server
-  client <# ("", "conn1", END)
+  client <# ("", "C:conn1", END)
 
 samplePublicKey :: ByteString
 samplePublicKey = "rsa:MIIBoDANBgkqhkiG9w0BAQEFAAOCAY0AMIIBiAKCAQEAtn1NI2tPoOGSGfad0aUg0tJ0kG2nzrIPGLiz8wb3dQSJC9xkRHyzHhEE8Kmy2cM4q7rNZIlLcm4M7oXOTe7SC4x59bLQG9bteZPKqXu9wk41hNamV25PWQ4zIcIRmZKETVGbwN7jFMpH7wxLdI1zzMArAPKXCDCJ5ctWh4OWDI6OR6AcCtEj+toCI6N6pjxxn5VigJtwiKhxYpoUJSdNM60wVEDCSUrZYBAuDH8pOxPfP+Tm4sokaFDTIG3QJFzOjC+/9nW4MUjAOFll9PCp9kaEFHJ/YmOYKMWNOCCPvLS6lxA83i0UaardkNLNoFS5paWfTlroxRwOC2T6PwO2ywKBgDjtXcSED61zK1seocQMyGRINnlWdhceD669kIHju/f6kAayvYKW3/lbJNXCmyinAccBosO08/0sUxvtuniIo18kfYJE0UmP1ReCjhMP+O+yOmwZJini/QelJk/Pez8IIDDWnY1qYQsN/q7ocjakOYrpGG7mig6JMFpDJtD6istR"
 
 syntaxTests :: forall c. Transport c => TProxy c -> Spec
 syntaxTests t = do
-  it "unknown command" $ ("1", "5678", "HELLO") >#> ("1", "5678", "ERR CMD SYNTAX")
+  it "unknown command" $ ("1", "C:5678", "HELLO") >#> ("1", "C:5678", "ERR CMD SYNTAX")
   describe "NEW" do
     describe "valid" do
       -- TODO: ERROR no connection alias in the response (it does not generate it yet if not provided)
       -- TODO: add tests with defined connection alias
-      xit "without parameters" $ ("211", "", "NEW") >#>= \case ("211", "", "INV" : _) -> True; _ -> False
+      xit "without parameters" $ ("211", "C:", "NEW") >#>= \case ("211", "C:", "INV" : _) -> True; _ -> False
     describe "invalid" do
       -- TODO: add tests with defined connection alias
-      it "with parameters" $ ("222", "", "NEW hi") >#> ("222", "", "ERR CMD SYNTAX")
+      it "with parameters" $ ("222", "C:", "NEW hi") >#> ("222", "C:", "ERR CMD SYNTAX")
 
   describe "JOIN" do
     describe "valid" do
       -- TODO: ERROR no connection alias in the response (it does not generate it yet if not provided)
       -- TODO: add tests with defined connection alias
       it "using same server as in invitation" $
-        ("311", "", "JOIN smp::localhost:5000::1234::" <> samplePublicKey) >#> ("311", "", "ERR SMP AUTH")
+        ("311", "C:", "JOIN smp::localhost:5000::1234::" <> samplePublicKey) >#> ("311", "C:", "ERR SMP AUTH")
     describe "invalid" do
       -- TODO: JOIN is not merged yet - to be added
-      it "no parameters" $ ("321", "", "JOIN") >#> ("321", "", "ERR CMD SYNTAX")
+      it "no parameters" $ ("321", "C:", "JOIN") >#> ("321", "C:", "ERR CMD SYNTAX")
   where
     -- simple test for one command with the expected response
     (>#>) :: ARawTransmission -> ARawTransmission -> Expectation