SMP protocol extension to debug subscribed SMP queues (#1181)

* SMP protocol extension to debug subscribed SMP queues

* fix, test

* corrections

Co-authored-by: spaced4ndy <8711996+spaced4ndy@users.noreply.github.com>

* increase delays

* increase timeout

* delay

* delay

* enable all tests

---------

Co-authored-by: spaced4ndy <8711996+spaced4ndy@users.noreply.github.com>

src/Simplex/Messaging/Agent.hs

@@ -70,6 +70,7 @@ module Simplex.Messaging.Agent
     sendMessages,
     sendMessagesB,
     ackMessage,
+    getConnectionQueueInfo,
     switchConnection,
     abortConnectionSwitch,
     synchronizeRatchet,
@@ -176,6 +177,7 @@ import Simplex.Messaging.Notifications.Types
 import Simplex.Messaging.Parsers (parse)
 import Simplex.Messaging.Protocol (BrokerMsg, Cmd (..), EntityId, ErrorType (AUTH), MsgBody, MsgFlags (..), NtfServer, ProtoServerWithAuth, ProtocolTypeI (..), SMPMsgMeta, SParty (..), SProtocolType (..), SndPublicAuthKey, SubscriptionMode (..), UserProtocol, VersionSMPC, XFTPServerWithAuth)
 import qualified Simplex.Messaging.Protocol as SMP
+import Simplex.Messaging.Server.QueueStore.QueueInfo
 import Simplex.Messaging.ServiceScheme (ServiceScheme (..))
 import qualified Simplex.Messaging.TMap as TM
 import Simplex.Messaging.Transport (SMPVersion, THandleParams (sessionId))
@@ -371,6 +373,10 @@ ackMessage :: AgentClient -> ConnId -> AgentMsgId -> Maybe MsgReceiptInfo -> AE
 ackMessage c = withAgentEnv c .:. ackMessage' c
 {-# INLINE ackMessage #-}
 
+getConnectionQueueInfo :: AgentClient -> ConnId -> AE QueueInfo
+getConnectionQueueInfo c = withAgentEnv c . getConnectionQueueInfo' c
+{-# INLINE getConnectionQueueInfo #-}
+
 -- | Switch connection to the new receive queue
 switchConnection :: AgentClient -> ConnId -> AE ConnectionStats
 switchConnection c = withAgentEnv c . switchConnection' c
@@ -1510,6 +1516,16 @@ ackMessage' c connId msgId rcptInfo_ = withConnLock c connId "ackMessage" $ do
             withStore' c $ \db -> deleteDeliveredSndMsg db connId $ InternalId sndMsgId
           _ -> pure ()
 
+getConnectionQueueInfo' :: AgentClient -> ConnId -> AM QueueInfo
+getConnectionQueueInfo' c connId = do
+  SomeConn _ conn <- withStore c (`getConn` connId)
+  case conn of
+    DuplexConnection _ (rq :| _) _ -> getQueueInfo c rq
+    RcvConnection _ rq -> getQueueInfo c rq
+    ContactConnection _ rq -> getQueueInfo c rq
+    SndConnection {} -> throwE $ CONN SIMPLEX
+    NewConnection _ -> throwE $ CMD PROHIBITED "getConnectionQueueInfo': NewConnection"
+
 switchConnection' :: AgentClient -> ConnId -> AM ConnectionStats
 switchConnection' c connId =
   withConnLock c connId "switchConnection" $

src/Simplex/Messaging/Agent/Client.hs

@@ -56,6 +56,7 @@ module Simplex.Messaging.Agent.Client
     disableQueueNotifications,
     disableQueuesNtfs,
     sendAgentMessage,
+    getQueueInfo,
     agentNtfRegisterToken,
     agentNtfVerifyToken,
     agentNtfCheckToken,
@@ -238,6 +239,7 @@ import Simplex.Messaging.Protocol
     sameSrvAddr',
   )
 import qualified Simplex.Messaging.Protocol as SMP
+import Simplex.Messaging.Server.QueueStore.QueueInfo
 import Simplex.Messaging.Session
 import Simplex.Messaging.TMap (TMap)
 import qualified Simplex.Messaging.TMap as TM
@@ -1577,6 +1579,11 @@ sendAgentMessage c sq@SndQueue {userId, server, sndId, sndPrivateKey} msgFlags a
   msg <- agentCbEncrypt sq Nothing $ smpEncode clientMsg
   sendOrProxySMPMessage c userId server "<MSG>" (Just sndPrivateKey) sndId msgFlags msg
 
+getQueueInfo :: AgentClient -> RcvQueue -> AM QueueInfo
+getQueueInfo c rq@RcvQueue {rcvId, rcvPrivateKey} =
+  withSMPClient c rq "QUE" $ \smp ->
+    getSMPQueueInfo smp rcvPrivateKey rcvId
+
 agentNtfRegisterToken :: AgentClient -> NtfToken -> NtfPublicAuthKey -> C.PublicKeyX25519 -> AM (NtfTokenId, C.PublicKeyX25519)
 agentNtfRegisterToken c NtfToken {deviceToken, ntfServer, ntfPrivKey} ntfPubKey pubDhKey =
   withClient c (0, ntfServer, Nothing) "TNEW" $ \ntf -> ntfRegisterToken ntf ntfPrivKey (NewNtfTkn deviceToken ntfPubKey pubDhKey)

src/Simplex/Messaging/Client.hs

@@ -59,6 +59,7 @@ module Simplex.Messaging.Client
     connectSMPProxiedRelay,
     proxySMPMessage,
     forwardSMPMessage,
+    getSMPQueueInfo,
     sendProtocolCommand,
 
     -- * Supporting types and client configuration
@@ -128,6 +129,7 @@ import Simplex.Messaging.Encoding
 import Simplex.Messaging.Encoding.String
 import Simplex.Messaging.Parsers (defaultJSON, dropPrefix, enumJSON)
 import Simplex.Messaging.Protocol
+import Simplex.Messaging.Server.QueueStore.QueueInfo
 import Simplex.Messaging.TMap (TMap)
 import qualified Simplex.Messaging.TMap as TM
 import Simplex.Messaging.Transport
@@ -920,6 +922,12 @@ forwardSMPMessage c@ProtocolClient {thParams, client_ = PClient {clientCorrId =
       pure fwdResponse
     r -> throwE $ unexpectedResponse r
 
+getSMPQueueInfo :: SMPClient -> C.APrivateAuthKey -> QueueId -> ExceptT SMPClientError IO QueueInfo
+getSMPQueueInfo c pKey qId =
+  sendSMPCommand c (Just pKey) qId QUE >>= \case
+    INFO info -> pure info
+    r -> throwE $ unexpectedResponse r
+
 okSMPCommand :: PartyI p => Command p -> SMPClient -> C.APrivateAuthKey -> QueueId -> ExceptT SMPClientError IO ()
 okSMPCommand cmd c pKey qId =
   sendSMPCommand c (Just pKey) qId cmd >>= \case

src/Simplex/Messaging/Protocol.hs

@@ -156,6 +156,7 @@ module Simplex.Messaging.Protocol
     sameSrvAddr,
     sameSrvAddr',
     noAuthSrv,
+    toMsgInfo,
 
     -- * TCP transport functions
     TransportBatch (..),
@@ -197,8 +198,8 @@ import qualified Data.List.NonEmpty as L
 import Data.Maybe (isJust, isNothing)
 import Data.String
 import qualified Data.Text as T
-import Data.Text.Encoding (encodeUtf8)
-import Data.Time.Clock.System (SystemTime (..))
+import Data.Text.Encoding (decodeLatin1, encodeUtf8)
+import Data.Time.Clock.System (SystemTime (..), systemToUTCTime)
 import Data.Type.Equality
 import Data.Word (Word16)
 import qualified Data.X509 as X
@@ -210,6 +211,7 @@ import qualified Simplex.Messaging.Crypto as C
 import Simplex.Messaging.Encoding
 import Simplex.Messaging.Encoding.String
 import Simplex.Messaging.Parsers
+import Simplex.Messaging.Server.QueueStore.QueueInfo
 import Simplex.Messaging.ServiceScheme
 import Simplex.Messaging.Transport
 import Simplex.Messaging.Transport.Client (TransportHost, TransportHosts (..))
@@ -386,6 +388,7 @@ data Command (p :: Party) where
   ACK :: MsgId -> Command Recipient
   OFF :: Command Recipient
   DEL :: Command Recipient
+  QUE :: Command Recipient
   -- SMP sender commands
   -- SEND v1 has to be supported for encoding/decoding
   -- SEND :: MsgBody -> Command Sender
@@ -463,6 +466,7 @@ data BrokerMsg where
   RRES :: EncFwdResponse -> BrokerMsg -- relay to proxy
   PRES :: EncResponse -> BrokerMsg -- proxy to client
   END :: BrokerMsg
+  INFO :: QueueInfo -> BrokerMsg
   OK :: BrokerMsg
   ERR :: ErrorType -> BrokerMsg
   PONG :: BrokerMsg
@@ -505,6 +509,13 @@ data Message
         msgTs :: SystemTime
       }
 
+toMsgInfo :: Message -> MsgInfo
+toMsgInfo = \case
+  Message {msgId, msgTs} -> msgInfo msgId msgTs MTMessage
+  MessageQuota {msgId, msgTs} -> msgInfo msgId msgTs MTQuota
+  where
+    msgInfo msgId msgTs msgType = MsgInfo {msgId = decodeLatin1 $ B64.encode msgId, msgTs = systemToUTCTime msgTs, msgType}
+
 messageId :: Message -> MsgId
 messageId = \case
   Message {msgId} -> msgId
@@ -652,6 +663,7 @@ data CommandTag (p :: Party) where
   ACK_ :: CommandTag Recipient
   OFF_ :: CommandTag Recipient
   DEL_ :: CommandTag Recipient
+  QUE_ :: CommandTag Recipient
   SEND_ :: CommandTag Sender
   PING_ :: CommandTag Sender
   PRXY_ :: CommandTag ProxiedClient
@@ -674,6 +686,7 @@ data BrokerMsgTag
   | RRES_
   | PRES_
   | END_
+  | INFO_
   | OK_
   | ERR_
   | PONG_
@@ -698,6 +711,7 @@ instance PartyI p => Encoding (CommandTag p) where
     ACK_ -> "ACK"
     OFF_ -> "OFF"
     DEL_ -> "DEL"
+    QUE_ -> "QUE"
     SEND_ -> "SEND"
     PING_ -> "PING"
     PRXY_ -> "PRXY"
@@ -717,6 +731,7 @@ instance ProtocolMsgTag CmdTag where
     "ACK" -> Just $ CT SRecipient ACK_
     "OFF" -> Just $ CT SRecipient OFF_
     "DEL" -> Just $ CT SRecipient DEL_
+    "QUE" -> Just $ CT SRecipient QUE_
     "SEND" -> Just $ CT SSender SEND_
     "PING" -> Just $ CT SSender PING_
     "PRXY" -> Just $ CT SProxiedClient PRXY_
@@ -742,6 +757,7 @@ instance Encoding BrokerMsgTag where
     RRES_ -> "RRES"
     PRES_ -> "PRES"
     END_ -> "END"
+    INFO_ -> "INFO"
     OK_ -> "OK"
     ERR_ -> "ERR"
     PONG_ -> "PONG"
@@ -757,6 +773,7 @@ instance ProtocolMsgTag BrokerMsgTag where
     "RRES" -> Just RRES_
     "PRES" -> Just PRES_
     "END" -> Just END_
+    "INFO" -> Just INFO_
     "OK" -> Just OK_
     "ERR" -> Just ERR_
     "PONG" -> Just PONG_
@@ -1275,6 +1292,7 @@ instance PartyI p => ProtocolEncoding SMPVersion ErrorType (Command p) where
     ACK msgId -> e (ACK_, ' ', msgId)
     OFF -> e OFF_
     DEL -> e DEL_
+    QUE -> e QUE_
     SEND flags msg -> e (SEND_, ' ', flags, ' ', Tail msg)
     PING -> e PING_
     NSUB -> e NSUB_
@@ -1340,6 +1358,7 @@ instance ProtocolEncoding SMPVersion ErrorType Cmd where
         ACK_ -> ACK <$> _smpP
         OFF_ -> pure OFF
         DEL_ -> pure DEL
+        QUE_ -> pure QUE
     CT SSender tag ->
       Cmd SSender <$> case tag of
         SEND_ -> SEND <$> _smpP <*> (unTail <$> _smpP)
@@ -1368,6 +1387,7 @@ instance ProtocolEncoding SMPVersion ErrorType BrokerMsg where
     RRES (EncFwdResponse encBlock) -> e (RRES_, ' ', Tail encBlock)
     PRES (EncResponse encBlock) -> e (PRES_, ' ', Tail encBlock)
     END -> e END_
+    INFO info -> e (INFO_, ' ', info)
     OK -> e OK_
     ERR err -> e (ERR_, ' ', err)
     PONG -> e PONG_
@@ -1388,6 +1408,7 @@ instance ProtocolEncoding SMPVersion ErrorType BrokerMsg where
     RRES_ -> RRES <$> (EncFwdResponse . unTail <$> _smpP)
     PRES_ -> PRES <$> (EncResponse . unTail <$> _smpP)
     END_ -> pure END
+    INFO_ -> INFO <$> _smpP
     OK_ -> pure OK
     ERR_ -> ERR <$> _smpP
     PONG_ -> pure PONG

src/Simplex/Messaging/Server.hs

@@ -59,7 +59,7 @@ import Data.List (intercalate, mapAccumR)
 import Data.List.NonEmpty (NonEmpty (..))
 import qualified Data.List.NonEmpty as L
 import qualified Data.Map.Strict as M
-import Data.Maybe (catMaybes, fromMaybe, isNothing)
+import Data.Maybe (catMaybes, fromMaybe, isJust, isNothing)
 import qualified Data.Text as T
 import Data.Text.Encoding (decodeLatin1)
 import Data.Time.Clock (UTCTime (..), diffTimeToPicoseconds, getCurrentTime)
@@ -82,6 +82,7 @@ import Simplex.Messaging.Server.Expiration
 import Simplex.Messaging.Server.MsgStore
 import Simplex.Messaging.Server.MsgStore.STM
 import Simplex.Messaging.Server.QueueStore
+import Simplex.Messaging.Server.QueueStore.QueueInfo
 import Simplex.Messaging.Server.QueueStore.STM as QS
 import Simplex.Messaging.Server.Stats
 import Simplex.Messaging.Server.StoreLog
@@ -791,6 +792,7 @@ client thParams' clnt@Client {subscriptions, ntfSubscriptions, rcvQ, sndQ, sessi
           NDEL -> deleteQueueNotifier_ st
           OFF -> suspendQueue_ st
           DEL -> delQueueAndMsgs st
+          QUE -> withQueue getQueueInfo
       where
         createQueue :: QueueStore -> RcvPublicAuthKey -> RcvPublicDhKey -> SubscriptionMode -> M (Transmission BrokerMsg)
         createQueue st recipientKey dhKey subMode = time "NEW" $ do
@@ -1162,6 +1164,26 @@ client thParams' clnt@Client {subscriptions, ntfSubscriptions, rcvQ, sndQ, sessi
             Right q -> updateDeletedStats q $> ok
             Left e -> pure $ err e
 
+        getQueueInfo :: QueueRec -> M (Transmission BrokerMsg)
+        getQueueInfo QueueRec {senderKey, notifier} = do
+          q@MsgQueue {size} <- getStoreMsgQueue "getQueueInfo" queueId
+          info <- atomically $ do
+            qiSub <- TM.lookup queueId subscriptions >>= mapM mkQSub
+            qiSize <- readTVar size
+            qiMsg <- toMsgInfo <$$> tryPeekMsg q
+            pure QueueInfo {qiSnd = isJust senderKey, qiNtf = isJust notifier, qiSub, qiSize, qiMsg}
+          pure (corrId, queueId, INFO info)
+          where
+            mkQSub sub = do
+              Sub {subThread, delivered} <- readTVar sub
+              let qSubThread = case subThread of
+                    NoSub -> QNoSub
+                    SubPending -> QSubPending
+                    SubThread _ -> QSubThread
+                    ProhibitSub -> QProhibitSub
+              qDelivered <- decodeLatin1 . encode <$$> tryReadTMVar delivered
+              pure QSub {qSubThread, qDelivered}
+          
         ok :: Transmission BrokerMsg
         ok = (corrId, queueId, OK)
 

src/Simplex/Messaging/Server/QueueStore/QueueInfo.hs

@@ -0,0 +1,55 @@
+{-# LANGUAGE TemplateHaskell #-}
+
+module Simplex.Messaging.Server.QueueStore.QueueInfo where
+
+import qualified Data.Aeson as J
+import qualified Data.Aeson.TH as JQ
+import qualified Data.Attoparsec.ByteString.Char8 as A
+import qualified Data.ByteString.Lazy.Char8 as LB
+import Data.Text (Text)
+import Data.Time.Clock (UTCTime)
+import Simplex.Messaging.Encoding
+import Simplex.Messaging.Parsers (defaultJSON, dropPrefix, enumJSON)
+import Simplex.Messaging.Util ((<$?>))
+
+data QueueInfo = QueueInfo
+  { qiSnd :: Bool,
+    qiNtf :: Bool,
+    qiSub :: Maybe QSub,
+    qiSize :: Int,
+    qiMsg :: Maybe MsgInfo
+  }
+  deriving (Eq, Show)
+
+data QSub = QSub
+  { qSubThread :: QSubThread,
+    qDelivered :: Maybe Text
+  }
+  deriving (Eq, Show)
+
+data QSubThread = QNoSub | QSubPending | QSubThread | QProhibitSub
+  deriving (Eq, Show)
+
+data MsgInfo = MsgInfo
+  { msgId :: Text,
+    msgTs :: UTCTime,
+    msgType :: MsgType
+  }
+  deriving (Eq, Show)
+
+data MsgType = MTMessage | MTQuota
+  deriving (Eq, Show)
+
+$(JQ.deriveJSON (enumJSON $ dropPrefix "Q") ''QSubThread)
+
+$(JQ.deriveJSON defaultJSON ''QSub)
+
+$(JQ.deriveJSON (enumJSON $ dropPrefix "MT") ''MsgType)
+
+$(JQ.deriveJSON defaultJSON ''MsgInfo)
+
+$(JQ.deriveJSON defaultJSON ''QueueInfo)
+
+instance Encoding QueueInfo where
+  smpEncode = LB.toStrict . J.encode
+  smpP = J.eitherDecodeStrict <$?> A.takeByteString

src/Simplex/RemoteControl/Invitation.hs

@@ -123,7 +123,7 @@ instance StrEncoding RCSignedInvitation where
     idsig <- requiredP sigs "idsig" $ parseAll strP
     pure RCSignedInvitation {invitation, ssig, idsig}
 
-signInvitation :: C.PrivateKey C.Ed25519 -> C.PrivateKey C.Ed25519 -> RCInvitation -> RCSignedInvitation
+signInvitation :: C.PrivateKey 'C.Ed25519 -> C.PrivateKey 'C.Ed25519 -> RCInvitation -> RCSignedInvitation
 signInvitation sKey idKey invitation = RCSignedInvitation {invitation, ssig, idsig}
   where
     uri = strEncode invitation