agent method to GET message (#403)

* agent method to GET message * correction Co-authored-by: JRoberts <8711996+jr-simplex@users.noreply.github.com> Co-authored-by: JRoberts <8711996+jr-simplex@users.noreply.github.com>
2026-08-28 20:08:16 +00:00 · 2022-06-16 12:35:17 +01:00
parent 0c45e0bf76
commit 1d56ef3a60
5 changed files with 66 additions and 7 deletions
@@ -44,6 +44,8 @@ module Simplex.Messaging.Agent
    acceptContact,
    rejectContact,
    subscribeConnection,
+    getConnectionMessage,
+    getNotificationMessage,
    resubscribeConnection,
    sendMessage,
    ackMessage,
@@ -96,7 +98,7 @@ import Simplex.Messaging.Parsers (parse)
 import Simplex.Messaging.Protocol (BrokerMsg, ErrorType (AUTH), MsgBody, MsgFlags)
 import qualified Simplex.Messaging.Protocol as SMP
 import qualified Simplex.Messaging.TMap as TM
-import Simplex.Messaging.Util (bshow, liftError, tryError, unlessM, ($>>=))
+import Simplex.Messaging.Util (bshow, liftError, tryError, unlessM, whenM, ($>>=))
 import Simplex.Messaging.Version
 import System.Random (randomR)
 import UnliftIO.Async (async, race_)
@@ -149,6 +151,14 @@ rejectContact c = withAgentEnv c .: rejectContact' c
 subscribeConnection :: AgentErrorMonad m => AgentClient -> ConnId -> m ()
 subscribeConnection c = withAgentEnv c . subscribeConnection' c

+-- | Get connection message (GET command)
+getConnectionMessage :: AgentErrorMonad m => AgentClient -> ConnId -> m (Maybe (SMP.MsgId, MsgFlags))
+getConnectionMessage c = withAgentEnv c . getConnectionMessage' c
+
+-- | Get connection message for received notification
+getNotificationMessage :: AgentErrorMonad m => AgentClient -> ByteString -> m (Maybe (SMP.MsgId, MsgFlags))
+getNotificationMessage c = withAgentEnv c . getNotificationMessage' c
+
 resubscribeConnection :: AgentErrorMonad m => AgentClient -> ConnId -> m ()
 resubscribeConnection c = withAgentEnv c . resubscribeConnection' c

@@ -365,12 +375,24 @@ subscribeConnection' c connId =
      ns <- asks ntfSupervisor
      atomically $ sendNtfSubCommand ns (connId, NSCCreate)

-resubscribeConnection' :: forall m. AgentMonad m => AgentClient -> ConnId -> m ()
+resubscribeConnection' :: AgentMonad m => AgentClient -> ConnId -> m ()
 resubscribeConnection' c connId =
  unlessM
    (atomically $ hasActiveSubscription c connId)
    (subscribeConnection' c connId)

+getConnectionMessage' :: AgentMonad m => AgentClient -> ConnId -> m (Maybe (SMP.MsgId, MsgFlags))
+getConnectionMessage' c connId = do
+  whenM (atomically $ hasActiveSubscription c connId) . throwError $ CMD PROHIBITED
+  withStore c (`getConn` connId) >>= \case
+    SomeConn _ (DuplexConnection _ rq _) -> getQueueMessage c rq connId
+    SomeConn _ (RcvConnection _ rq) -> getQueueMessage c rq connId
+    SomeConn _ (ContactConnection {}) -> throwError $ CMD PROHIBITED
+    SomeConn _ (SndConnection {}) -> throwError $ CONN SIMPLEX
+
+getNotificationMessage' :: AgentErrorMonad m => AgentClient -> ByteString -> m (Maybe (SMP.MsgId, MsgFlags))
+getNotificationMessage' c _encMessageInfo = throwError $ CMD PROHIBITED
+
 -- | Send message to the connection (SEND command) in Reader monad
 sendMessage' :: forall m. AgentMonad m => AgentClient -> ConnId -> MsgFlags -> MsgBody -> m AgentMsgId
 sendMessage' c connId msgFlags msg =
@@ -532,7 +554,8 @@ ackMessage' c connId msgId = do
  withStore c (`getConn` connId) >>= \case
    SomeConn _ (DuplexConnection _ rq _) -> ack rq
    SomeConn _ (RcvConnection _ rq) -> ack rq
-    _ -> throwError $ CONN SIMPLEX
+    SomeConn _ (SndConnection _ _) -> throwError $ CONN SIMPLEX
+    SomeConn _ (ContactConnection _ _) -> throwError $ CMD PROHIBITED
  where
    ack :: RcvQueue -> m ()
    ack rq = do
@@ -549,7 +572,8 @@ suspendConnection' c connId =
  withStore c (`getConn` connId) >>= \case
    SomeConn _ (DuplexConnection _ rq _) -> suspendQueue c rq
    SomeConn _ (RcvConnection _ rq) -> suspendQueue c rq
-    _ -> throwError $ CONN SIMPLEX
+    SomeConn _ (ContactConnection _ rq) -> suspendQueue c rq
+    SomeConn _ (SndConnection _ _) -> throwError $ CONN SIMPLEX

 -- | Delete SMP agent connection (DEL command) in Reader monad
 deleteConnection' :: forall m. AgentMonad m => AgentClient -> ConnId -> m ()
@@ -20,6 +20,7 @@ module Simplex.Messaging.Agent.Client
    closeAgentClient,
    newRcvQueue,
    subscribeQueue,
+    getQueueMessage,
    addSubscription,
    getSubscriptions,
    sendConfirmation,
@@ -113,6 +114,7 @@ data AgentClient = AgentClient
    connMsgsQueued :: TMap ConnId Bool,
    smpQueueMsgQueues :: TMap (ConnId, SMPServer, SMP.SenderId) (TQueue InternalId),
    smpQueueMsgDeliveries :: TMap (ConnId, SMPServer, SMP.SenderId) (Async ()),
+    getMsgLocks :: TMap (ConnId, SMPServer, SMP.RecipientId) (TMVar ()),
    reconnections :: TVar [Async ()],
    asyncClients :: TVar [Async ()],
    clientId :: Int,
@@ -137,11 +139,12 @@ newAgentClient InitialAgentServers {smp, ntf} agentEnv = do
  connMsgsQueued <- TM.empty
  smpQueueMsgQueues <- TM.empty
  smpQueueMsgDeliveries <- TM.empty
+  getMsgLocks <- TM.empty
  reconnections <- newTVar []
  asyncClients <- newTVar []
  clientId <- stateTVar (clientCounter agentEnv) $ \i -> (i + 1, i + 1)
  lock <- newTMVar ()
-  return AgentClient {active, rcvQ, subQ, msgQ, smpServers, smpClients, ntfServers, ntfClients, subscrSrvrs, pendingSubscrSrvrs, subscrConns, connMsgsQueued, smpQueueMsgQueues, smpQueueMsgDeliveries, reconnections, asyncClients, clientId, agentEnv, lock}
+  return AgentClient {active, rcvQ, subQ, msgQ, smpServers, smpClients, ntfServers, ntfClients, subscrSrvrs, pendingSubscrSrvrs, subscrConns, connMsgsQueued, smpQueueMsgQueues, smpQueueMsgDeliveries, getMsgLocks, reconnections, asyncClients, clientId, agentEnv, lock}

 agentDbPath :: AgentClient -> FilePath
 agentDbPath AgentClient {agentEnv = Env {store = SQLiteStore {dbFilePath}}} = dbFilePath
@@ -423,6 +426,7 @@ newRcvQueue_ a c srv = do

 subscribeQueue :: AgentMonad m => AgentClient -> RcvQueue -> ConnId -> m ()
 subscribeQueue c rq@RcvQueue {server, rcvPrivateKey, rcvId} connId = do
+  whenM (atomically . TM.member (connId, server, rcvId) $ getMsgLocks c) . throwError $ CMD PROHIBITED
  atomically $ addPendingSubscription c rq connId
  withLogClient c server rcvId "SUB" $ \smp -> do
    liftIO (runExceptT $ subscribeSMPQueue smp rcvPrivateKey rcvId) >>= \case
@@ -501,6 +505,23 @@ sendInvitation c (Compatible SMPQueueInfo {smpServer, senderId, dhPublicKey}) co
      agentCbEncryptOnce dhPublicKey . smpEncode $
        SMP.ClientMessage SMP.PHEmpty $ smpEncode agentEnvelope

+getQueueMessage :: AgentMonad m => AgentClient -> RcvQueue -> ConnId -> m (Maybe (MsgId, MsgFlags))
+getQueueMessage c@AgentClient {getMsgLocks} RcvQueue {server, rcvId, rcvPrivateKey} connId =
+  E.bracket (atomically createTakeLock) (atomically . (`putTMVar` ())) $ \_ ->
+    withLogClient c server rcvId "GET" $ \smp ->
+      getSMPMessage smp rcvPrivateKey rcvId
+  where
+    k = (connId, server, rcvId)
+    createTakeLock = do
+      l <- TM.lookup k getMsgLocks >>= maybe newLock pure
+      takeTMVar l
+      pure l
+      where
+        newLock = do
+          l <- newTMVar ()
+          TM.insert k l getMsgLocks
+          pure l
+
 secureQueue :: AgentMonad m => AgentClient -> RcvQueue -> SndPublicVerifyKey -> m ()
 secureQueue c RcvQueue {server, rcvId, rcvPrivateKey} senderKey =
  withLogClient c server rcvId "KEY <key>" $ \smp ->
@@ -32,6 +32,7 @@ module Simplex.Messaging.Client
    -- * SMP protocol command functions
    createSMPQueue,
    subscribeSMPQueue,
+    getSMPMessage,
    subscribeSMPQueueNotifications,
    secureSMPQueue,
    enableSMPQueueNotifications,
@@ -287,6 +288,18 @@ subscribeSMPQueue c@ProtocolClient {protocolServer, sessionId, msgQ} rpKey rId =
      lift . atomically $ mapM_ (`writeTBQueue` (protocolServer, sessionId, rId, cmd)) msgQ
    _ -> throwE PCEUnexpectedResponse

+-- | Get message from SMP queue. The server returns ERR PROHIBITED if a client uses SUB and GET via the same transport connection for the same queue
+--
+-- https://github.covm/simplex-chat/simplexmq/blob/master/protocol/simplex-messaging.md#receive-a-message-from-the-queue
+getSMPMessage :: SMPClient -> RcvPrivateSignKey -> RecipientId -> ExceptT ProtocolClientError IO (Maybe (MsgId, MsgFlags))
+getSMPMessage c@ProtocolClient {protocolServer, sessionId, msgQ} rpKey rId =
+  sendSMPCommand c (Just rpKey) rId GET >>= \case
+    OK -> pure Nothing
+    cmd@(MSG msgId _ msgFlags _) -> do
+      lift . atomically $ mapM_ (`writeTBQueue` (protocolServer, sessionId, rId, cmd)) msgQ
+      pure $ Just (msgId, msgFlags)
+    _ -> throwE PCEUnexpectedResponse
+
 -- | Subscribe to the SMP queue notifications.
 --
 -- https://github.com/simplex-chat/simplexmq/blob/master/protocol/simplex-messaging.md#subscribe-to-queue-notifications
@@ -442,7 +442,7 @@ client clnt@Client {subscriptions, ntfSubscriptions, rcvQ, sndQ} Server {subscri
              atomically $
                tryPeekMsg q >>= \case
                  Just msg -> setDelivered s msg $> (corrId, queueId, msgCmd msg)
-                  _ -> pure (corrId, queueId, ERR NO_MSG)
+                  _ -> pure (corrId, queueId, OK)

        subscribeNotifications :: m (Transmission BrokerMsg)
        subscribeNotifications = atomically $ do
@@ -339,6 +339,7 @@ testGetCommand t =
        Resp "2" _ (MSG mId1 _ _ msg1) <- signSendRecv rh rKey ("2", rId, GET)
        (dec mId1 msg1, Right "hello") #== "retrieved from queue"
        Resp "3" _ OK <- signSendRecv rh rKey ("3", rId, ACK mId1)
+        Resp "4" _ OK <- signSendRecv rh rKey ("4", rId, GET)
        pure ()

 testGetSubCommands :: forall c. Transport c => TProxy c -> Spec
@@ -387,7 +388,7 @@ testGetSubCommands t =
      (dec mId4 msg4, Right "hello 4") #== "retrieved from queue with GET"
      Resp "11" _ OK <- signSendRecv rh1 rKey ("11", rId, ACK mId4)
      -- no more messages for getter too
-      Resp "12" _ (ERR NO_MSG) <- signSendRecv rh2 rKey ("12", rId, GET)
+      Resp "12" _ OK <- signSendRecv rh2 rKey ("12", rId, GET)
      pure ()

 testWithStoreLog :: ATransport -> Spec