From 1d56ef3a60553d9b9812b26be833fdda4122d676 Mon Sep 17 00:00:00 2001
From: Evgeny Poberezkin <2769109+epoberezkin@users.noreply.github.com>
Date: Thu, 16 Jun 2022 12:35:17 +0100
Subject: [PATCH] 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>
---
 src/Simplex/Messaging/Agent.hs        | 32 +++++++++++++++++++++++----
 src/Simplex/Messaging/Agent/Client.hs | 23 ++++++++++++++++++-
 src/Simplex/Messaging/Client.hs       | 13 +++++++++++
 src/Simplex/Messaging/Server.hs       |  2 +-
 tests/ServerTests.hs                  |  3 ++-
 5 files changed, 66 insertions(+), 7 deletions(-)

diff --git a/src/Simplex/Messaging/Agent.hs b/src/Simplex/Messaging/Agent.hs
index f958792ac..069b2fee2 100644
--- a/src/Simplex/Messaging/Agent.hs
+++ b/src/Simplex/Messaging/Agent.hs
@@ -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 ()
diff --git a/src/Simplex/Messaging/Agent/Client.hs b/src/Simplex/Messaging/Agent/Client.hs
index a68cec7b9..c626ed43b 100644
--- a/src/Simplex/Messaging/Agent/Client.hs
+++ b/src/Simplex/Messaging/Agent/Client.hs
@@ -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 ->
diff --git a/src/Simplex/Messaging/Client.hs b/src/Simplex/Messaging/Client.hs
index 6c37008af..92d749091 100644
--- a/src/Simplex/Messaging/Client.hs
+++ b/src/Simplex/Messaging/Client.hs
@@ -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
diff --git a/src/Simplex/Messaging/Server.hs b/src/Simplex/Messaging/Server.hs
index daeb49601..3f0ba8340 100644
--- a/src/Simplex/Messaging/Server.hs
+++ b/src/Simplex/Messaging/Server.hs
@@ -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
diff --git a/tests/ServerTests.hs b/tests/ServerTests.hs
index 706ac938a..6e561772b 100644
--- a/tests/ServerTests.hs
+++ b/tests/ServerTests.hs
@@ -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