client: prevent sub actions from zombie sessions (#1122)

* client: prevent sub actions from zombie sessions

* error handling

* add AERR to pass background errors to client

* switch to activeClientSession

* put closeClient under activeClientSession

* rename

* remove AERR, do not skip processing

* move check and state update to one transaction

* catch extra UPs

* fix

* check queue is still pending before making it active

---------

Co-authored-by: Evgeny Poberezkin <evgeny@poberezkin.com>

src/Simplex/Messaging/Agent.hs

@@ -890,7 +890,7 @@ subscribeConnections' c connIds = do
       (subRs, rcvQs) = M.mapEither rcvQueueOrResult cs
   mapM_ (mapM_ (\(cData, sqs) -> mapM_ (lift . resumeMsgDelivery c cData) sqs) . sndQueue) cs
   mapM_ (resumeConnCmds c) $ M.keys cs
-  rcvRs <- lift $ connResults <$> subscribeQueues c (concat $ M.elems rcvQs)
+  rcvRs <- lift $ connResults . fst <$> subscribeQueues c (concat $ M.elems rcvQs)
   ns <- asks ntfSupervisor
   tkn <- readTVarIO (ntfTkn ns)
   when (instantNotifications tkn) . void . lift . forkIO . void . runExceptT $ sendNtfCreate ns rcvRs conns
@@ -1957,11 +1957,13 @@ getSMPServer c userId = withUserServers c userId pickServer
 {-# INLINE getSMPServer #-}
 
 subscriber :: AgentClient -> AM' ()
-subscriber c@AgentClient {msgQ} = forever $ do
+subscriber c@AgentClient {subQ, msgQ} = forever $ do
   t <- atomically $ readTBQueue msgQ
   agentOperationBracket c AORcvNetwork waitUntilActive $
     runExceptT (processSMPTransmission c t) >>= \case
-      Left e -> liftIO $ print e
+      Left e -> do
+        logError $ tshow e
+        atomically $ writeTBQueue subQ ("", "", APC SAEConn $ ERR e)
       Right _ -> return ()
 
 cleanupManager :: AgentClient -> AM' ()
@@ -2238,7 +2240,7 @@ processSMPTransmission c@AgentClient {smpClients, subQ} (tSess@(_, srv, _), _v,
                   ((&&) <$> hasPendingSubscription c connId <*> activeClientSession c tSess sessId)
                   (True <$ addSubscription c rq)
                   (pure False)
-            when added . notify $ UP srv [connId]
+            when (added && isResponse) $ notify $ UP srv [connId]
 
           decryptClientMessage :: C.DhSecretX25519 -> SMP.ClientMsgEnvelope -> AM (SMP.PrivHeader, AgentMsgEnvelope)
           decryptClientMessage e2eDh SMP.ClientMsgEnvelope {cmNonce, cmEncBody} = do

src/Simplex/Messaging/Agent/Client.hs

@@ -155,7 +155,7 @@ import Data.Bifunctor (bimap, first, second)
 import Data.ByteString.Base64
 import Data.ByteString.Char8 (ByteString)
 import qualified Data.ByteString.Char8 as B
-import Data.Either (lefts, partitionEithers)
+import Data.Either (partitionEithers)
 import Data.Functor (($>))
 import Data.Int (Int64)
 import Data.List (deleteFirstsBy, foldl', partition, (\\))
@@ -643,7 +643,7 @@ reconnectSMPClient tc c tSess@(_, srv, _) qs = do
     resubscribe :: AM ()
     resubscribe = do
       cs <- readTVarIO $ RQ.getConnections $ activeSubs c
-      rs <- lift . subscribeQueues c $ L.toList qs
+      (rs, sessId_) <- lift . subscribeQueues c $ L.toList qs
       let (errs, okConns) = partitionEithers $ map (\(RcvQueue {connId}, r) -> bimap (connId,) (const connId) r) rs
       liftIO $ do
         let conns = filter (`M.notMember` cs) okConns
@@ -652,7 +652,10 @@ reconnectSMPClient tc c tSess@(_, srv, _) qs = do
       liftIO $ mapM_ (\(connId, e) -> notifySub connId $ ERR e) finalErrs
       forM_ (listToMaybe tempErrs) $ \(_, err) -> do
         when (null okConns && M.null cs && null finalErrs) . liftIO $
-          closeClient c smpClients tSess
+          forM_ sessId_ $ \sessId -> do
+            -- We only close the client session that was used to subscribe.
+            v_ <- atomically $ ifM (activeClientSession c tSess sessId) (TM.lookupDelete tSess $ smpClients c) (pure Nothing)
+            mapM_ (closeClient_ c) v_
         throwError err
     notifySub :: forall e. AEntityI e => ConnId -> ACommand 'Agent e -> IO ()
     notifySub connId cmd = atomically $ writeTBQueue (subQ c) ("", connId, APC (sAEntity @e) cmd)
@@ -1133,15 +1136,15 @@ newRcvQueue c userId connId (ProtoServerWithAuth srv auth) vRange subMode = do
       qUri = SMPQueueUri vRange $ SMPQueueAddress srv sndId e2eDhKey
   pure (rq, qUri, tSess, sessId)
 
-processSubResult :: AgentClient -> RcvQueue -> Either SMPClientError () -> IO (Either SMPClientError ())
-processSubResult c rq r = do
-  case r of
-    Left e ->
-      unless (temporaryClientError e) . atomically $ do
-        RQ.deleteQueue rq (pendingSubs c)
-        TM.insert (RQ.qKey rq) e (removedSubs c)
-    _ -> atomically $ addSubscription c rq
-  pure r
+processSubResult :: AgentClient -> RcvQueue -> Either SMPClientError () -> STM ()
+processSubResult c rq@RcvQueue {connId} = \case
+  Left e ->
+    unless (temporaryClientError e) $ do
+      RQ.deleteQueue rq (pendingSubs c)
+      TM.insert (RQ.qKey rq) e (removedSubs c)
+  Right () ->
+    whenM (hasPendingSubscription c connId) $
+      addSubscription c rq
 
 temporaryAgentError :: AgentErrorType -> Bool
 temporaryAgentError = \case
@@ -1158,7 +1161,7 @@ temporaryOrHostError = \case
 {-# INLINE temporaryOrHostError #-}
 
 -- | Subscribe to queues. The list of results can have a different order.
-subscribeQueues :: AgentClient -> [RcvQueue] -> AM' [(RcvQueue, Either AgentErrorType ())]
+subscribeQueues :: AgentClient -> [RcvQueue] -> AM' ([(RcvQueue, Either AgentErrorType ())], Maybe SessionId)
 subscribeQueues c qs = do
   (errs, qs') <- partitionEithers <$> mapM checkQueue qs
   atomically $ do
@@ -1166,18 +1169,32 @@ subscribeQueues c qs = do
     RQ.batchAddQueues (pendingSubs c) qs'
   env <- ask
   -- only "checked" queues are subscribed
-  (errs <>) <$> sendTSessionBatches "SUB" 90 id (subscribeQueues_ env) c qs'
+  session <- newTVarIO Nothing
+  rs <- sendTSessionBatches "SUB" 90 id (subscribeQueues_ env session) c qs'
+  (errs <> rs,) <$> readTVarIO session
   where
     checkQueue rq = do
       prohibited <- atomically $ hasGetLock c rq
       pure $ if prohibited then Left (rq, Left $ CMD PROHIBITED) else Right rq
-    subscribeQueues_ :: Env -> SMPClient -> NonEmpty RcvQueue -> IO (BatchResponses SMPClientError ())
-    subscribeQueues_ env smp qs' = do
+    subscribeQueues_ :: Env -> TVar (Maybe SessionId) -> SMPClient -> NonEmpty RcvQueue -> IO (BatchResponses SMPClientError ())
+    subscribeQueues_ env session smp qs' = do
       rs <- sendBatch subscribeSMPQueues smp qs'
-      mapM_ (uncurry $ processSubResult c) rs
-      when (any temporaryClientError . lefts . map snd $ L.toList rs) $
-        runReaderT (resubscribeSMPSession c $ transportSession' smp) env
-      pure rs
+      active <-
+        atomically $
+          ifM
+            (activeClientSession c tSess sessId)
+            (writeTVar session (Just sessId) >> mapM_ (uncurry $ processSubResult c) rs $> True)
+            (pure False)
+      if active
+        then when (hasTempErrors rs) resubscribe $> rs
+        else do
+          unless active $ logWarn "subcription batch result for replaced SMP client, resubscribing"
+          resubscribe $> L.map (second $ \_ -> Left PCENetworkError) rs
+      where
+        tSess = transportSession' smp
+        sessId = sessionId $ thParams smp
+        hasTempErrors = any (either temporaryClientError (const False) . snd)
+        resubscribe = resubscribeSMPSession c tSess `runReaderT` env
 
 activeClientSession :: AgentClient -> SMPTransportSession -> SessionId -> STM Bool
 activeClientSession c tSess sessId = sameSess <$> tryReadSessVar tSess (smpClients c)
@@ -1186,7 +1203,7 @@ activeClientSession c tSess sessId = sameSess <$> tryReadSessVar tSess (smpClien
       Just (Right smp) -> sessId == sessionId (thParams smp)
       _ -> False
 
-type BatchResponses e r = (NonEmpty (RcvQueue, Either e r))
+type BatchResponses e r = NonEmpty (RcvQueue, Either e r)
 
 -- statBatchSize is not used to batch the commands, only for traffic statistics
 sendTSessionBatches :: forall q r. ByteString -> Int -> (q -> RcvQueue) -> (SMPClient -> NonEmpty q -> IO (BatchResponses SMPClientError r)) -> AgentClient -> [q] -> AM' [(RcvQueue, Either AgentErrorType r)]

tests/AgentTests/FunctionalAPITests.hs

@@ -1393,10 +1393,10 @@ makeConnection = makeConnection_ PQSupportOn
 makeConnection_ :: PQSupport -> AgentClient -> AgentClient -> ExceptT AgentErrorType IO (ConnId, ConnId)
 makeConnection_ pqEnc alice bob = makeConnectionForUsers_ pqEnc alice 1 bob 1
 
-makeConnectionForUsers :: AgentClient -> UserId -> AgentClient -> UserId -> ExceptT AgentErrorType IO (ConnId, ConnId)
+makeConnectionForUsers :: HasCallStack => AgentClient -> UserId -> AgentClient -> UserId -> ExceptT AgentErrorType IO (ConnId, ConnId)
 makeConnectionForUsers = makeConnectionForUsers_ PQSupportOn
 
-makeConnectionForUsers_ :: PQSupport -> AgentClient -> UserId -> AgentClient -> UserId -> ExceptT AgentErrorType IO (ConnId, ConnId)
+makeConnectionForUsers_ :: HasCallStack => PQSupport -> AgentClient -> UserId -> AgentClient -> UserId -> ExceptT AgentErrorType IO (ConnId, ConnId)
 makeConnectionForUsers_ pqSupport alice aliceUserId bob bobUserId = do
   (bobId, qInfo) <- A.createConnection alice aliceUserId True SCMInvitation Nothing (CR.IKNoPQ pqSupport) SMSubscribe
   aliceId <- A.joinConnection bob bobUserId True qInfo "bob's connInfo" pqSupport SMSubscribe