From 1104b4413691df648d486190dc55fc15388c10b8 Mon Sep 17 00:00:00 2001 From: Alexander Bondarenko <486682+dpwiz@users.noreply.github.com> Date: Sat, 4 May 2024 01:29:19 +0300 Subject: [PATCH] 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 --- src/Simplex/Messaging/Agent.hs | 10 +++-- src/Simplex/Messaging/Agent/Client.hs | 59 +++++++++++++++++--------- tests/AgentTests/FunctionalAPITests.hs | 4 +- 3 files changed, 46 insertions(+), 27 deletions(-) diff --git a/src/Simplex/Messaging/Agent.hs b/src/Simplex/Messaging/Agent.hs index 57187036d..5ef3ab9b8 100644 --- a/src/Simplex/Messaging/Agent.hs +++ b/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 diff --git a/src/Simplex/Messaging/Agent/Client.hs b/src/Simplex/Messaging/Agent/Client.hs index baee9b974..22b49a719 100644 --- a/src/Simplex/Messaging/Agent/Client.hs +++ b/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)] diff --git a/tests/AgentTests/FunctionalAPITests.hs b/tests/AgentTests/FunctionalAPITests.hs index 79742efab..cb4f9bb55 100644 --- a/tests/AgentTests/FunctionalAPITests.hs +++ b/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