diff --git a/src/Simplex/Messaging/Agent.hs b/src/Simplex/Messaging/Agent.hs index f705f6666..3e4967e05 100644 --- a/src/Simplex/Messaging/Agent.hs +++ b/src/Simplex/Messaging/Agent.hs @@ -85,6 +85,7 @@ import Data.Bifunctor (bimap, first, second) import Data.ByteString.Char8 (ByteString) import Data.Composition ((.:), (.:.)) import Data.Functor (($>)) +import Data.List (deleteFirstsBy) import Data.List.NonEmpty (NonEmpty (..)) import qualified Data.List.NonEmpty as L import Data.Map.Strict (Map) @@ -359,8 +360,11 @@ ackMessageAsync' c connId msgId = enqueueCommand c connId (Just server) $ ACK msgId newConn :: AgentMonad m => AgentClient -> ConnId -> Bool -> Bool -> SConnectionMode c -> m (ConnId, ConnectionRequestUri c) -newConn c connId asyncMode enableNtfs cMode = do - srv <- getSMPServer c +newConn c connId asyncMode enableNtfs cMode = + getSMPServer c >>= newConnSrv c connId asyncMode enableNtfs cMode + +newConnSrv :: AgentMonad m => AgentClient -> ConnId -> Bool -> Bool -> SConnectionMode c -> SMPServer -> m (ConnId, ConnectionRequestUri c) +newConnSrv c connId asyncMode enableNtfs cMode srv = do clientVRange <- asks $ smpClientVRange . config (rq, qUri) <- newRcvQueue c srv clientVRange connId' <- setUpConn asyncMode rq @@ -387,7 +391,11 @@ newConn c connId asyncMode enableNtfs cMode = do withStore c $ \db -> createRcvConn db g cData rq cMode joinConn :: AgentMonad m => AgentClient -> ConnId -> Bool -> Bool -> ConnectionRequestUri c -> ConnInfo -> m ConnId -joinConn c connId asyncMode enableNtfs (CRInvitationUri (ConnReqUriData _ agentVRange (qUri :| _)) e2eRcvParamsUri) cInfo = do +joinConn c connId asyncMode enableNtfs connReq cInfo = + getSMPServer c >>= joinConnSrv c connId asyncMode enableNtfs connReq cInfo + +joinConnSrv :: AgentMonad m => AgentClient -> ConnId -> Bool -> Bool -> ConnectionRequestUri c -> ConnInfo -> SMPServer -> m ConnId +joinConnSrv c connId asyncMode enableNtfs (CRInvitationUri (ConnReqUriData _ agentVRange (qUri :| _)) e2eRcvParamsUri) cInfo srv = do aVRange <- asks $ smpAgentVRange . config clientVRange <- asks $ smpClientVRange . config case ( qUri `compatibleVersion` clientVRange, @@ -403,7 +411,7 @@ joinConn c connId asyncMode enableNtfs (CRInvitationUri (ConnReqUriData _ agentV cData = ConnData {connId, connAgentVersion, enableNtfs, duplexHandshake = Just duplexHS} connId' <- setUpConn asyncMode cData sq rc let cData' = (cData :: ConnData) {connId = connId'} - tryError (confirmQueue aVersion c cData' sq cInfo $ Just e2eSndParams) >>= \case + tryError (confirmQueue aVersion c cData' sq srv cInfo $ Just e2eSndParams) >>= \case Right _ -> do unless duplexHS . void $ enqueueMessage c cData' sq SMP.noMsgFlags HELLO pure connId' @@ -424,23 +432,22 @@ joinConn c connId asyncMode enableNtfs (CRInvitationUri (ConnReqUriData _ agentV liftIO $ createRatchet db connId' rc pure connId' _ -> throwError $ AGENT A_VERSION -joinConn c connId False enableNtfs (CRContactUri (ConnReqUriData _ agentVRange (qUri :| _))) cInfo = do +joinConnSrv c connId False enableNtfs (CRContactUri (ConnReqUriData _ agentVRange (qUri :| _))) cInfo srv = do aVRange <- asks $ smpAgentVRange . config clientVRange <- asks $ smpClientVRange . config case ( qUri `compatibleVersion` clientVRange, agentVRange `compatibleVersion` aVRange ) of (Just qInfo, Just vrsn) -> do - (connId', cReq) <- newConn c connId False enableNtfs SCMInvitation + (connId', cReq) <- newConnSrv c connId False enableNtfs SCMInvitation srv sendInvitation c qInfo vrsn cReq cInfo pure connId' _ -> throwError $ AGENT A_VERSION -joinConn _c _connId True _enableNtfs (CRContactUri _) _cInfo = do +joinConnSrv _c _connId True _enableNtfs (CRContactUri _) _cInfo _srv = do throwError $ CMD PROHIBITED -createReplyQueue :: AgentMonad m => AgentClient -> ConnData -> SndQueue -> m SMPQueueInfo -createReplyQueue c ConnData {connId, enableNtfs} SndQueue {smpClientVersion} = do - srv <- getSMPServer c +createReplyQueue :: AgentMonad m => AgentClient -> ConnData -> SndQueue -> SMPServer -> m SMPQueueInfo +createReplyQueue c ConnData {connId, enableNtfs} SndQueue {smpClientVersion} srv = do (rq, qUri) <- newRcvQueue c srv $ versionToRange smpClientVersion let qInfo = toVersionT qUri smpClientVersion addSubscription c rq connId @@ -490,21 +497,21 @@ processConfirmation c rq@RcvQueue {e2ePrivKey, smpClientVersion = v} SMPConfirma -- | Subscribe to receive connection messages (SUB command) in Reader monad subscribeConnection' :: forall m. AgentMonad m => AgentClient -> ConnId -> m () subscribeConnection' c connId = - withStore c (`getConn` connId) >>= \case - SomeConn _ (DuplexConnection cData rq sq) -> do - resumeMsgDelivery c cData sq - subscribe rq - resumeConnCmds c connId - SomeConn _ (SndConnection cData sq) -> do - resumeMsgDelivery c cData sq - case status (sq :: SndQueue) of - Confirmed -> pure () - Active -> throwError $ CONN SIMPLEX - _ -> throwError $ INTERNAL "unexpected queue status" - resumeConnCmds c connId - SomeConn _ (RcvConnection _ rq) -> subscribe rq >> resumeConnCmds c connId - SomeConn _ (ContactConnection _ rq) -> subscribe rq >> resumeConnCmds c connId - SomeConn _ (NewConnection _) -> resumeConnCmds c connId + withStore c (`getConn` connId) >>= \conn -> do + resumeConnCmds c connId + case conn of + SomeConn _ (DuplexConnection cData rq sq) -> do + resumeMsgDelivery c cData sq + subscribe rq + SomeConn _ (SndConnection cData sq) -> do + resumeMsgDelivery c cData sq + case status (sq :: SndQueue) of + Confirmed -> pure () + Active -> throwError $ CONN SIMPLEX + _ -> throwError $ INTERNAL "unexpected queue status" + SomeConn _ (RcvConnection _ rq) -> subscribe rq + SomeConn _ (ContactConnection _ rq) -> subscribe rq + SomeConn _ (NewConnection _) -> pure () where subscribe :: RcvQueue -> m () subscribe rq = do @@ -521,7 +528,7 @@ subscribeConnections' c connIds = do (subRs, rcvQs) = M.mapEither rcvQueueOrResult cs srvRcvQs :: Map SMPServer (Map ConnId (RcvQueue, ConnData)) = M.foldlWithKey' addRcvQueue M.empty rcvQs mapM_ (mapM_ (uncurry $ resumeMsgDelivery c) . sndQueue) cs - forM_ (M.keys cs) $ resumeConnCmds c + mapM_ (resumeConnCmds c) $ M.keys cs rcvRs <- mapConcurrently subscribe (M.assocs srvRcvQs) ns <- asks ntfSupervisor tkn <- readTVarIO (ntfTkn ns) @@ -679,26 +686,35 @@ runCommandProcessing c@AgentClient {subQ} server = do E.try (withStore c $ \db -> getPendingCommand db cmdId) >>= \case Left (e :: E.SomeException) -> notify "" $ ERR (INTERNAL $ show e) - Right (connId, ACmd _ cmd) -> + Right (connId, ACmd _ cmd) -> do + usedSrvs <- newTVarIO ([] :: [SMPServer]) withRetryInterval ri $ \loop -> do resp <- tryError $ case cmd of - NEW enableNtfs (ACM cMode) -> do - (_, cReq) <- newConn c connId True enableNtfs cMode - notify connId $ INV (ACR cMode cReq) - JOIN enableNtfs (ACR _ cReq) connInfo -> void $ joinConn c connId True enableNtfs cReq connInfo + NEW enableNtfs (ACM cMode) -> + withNextSrv usedSrvs $ \srv -> do + (_, cReq) <- newConnSrv c connId True enableNtfs cMode srv + notify connId $ INV (ACR cMode cReq) + JOIN enableNtfs (ACR _ cReq) connInfo -> + withNextSrv usedSrvs $ \srv -> + void $ joinConnSrv c connId True enableNtfs cReq connInfo srv LET confId ownCInfo -> allowConnection' c connId confId ownCInfo ACK msgId -> ackMessage' c connId msgId - _ -> notify "" $ ERR (INTERNAL "") + _ -> notify connId $ ERR $ INTERNAL $ "unsupported async command " <> show cmd case resp of - Left _ -> - -- TODO retry NEW and JOIN on different server - -- TODO depending on command, some errors shouldn't be retried - retryCommand loop - Right () -> do - delCmd cmdId + Left e + | temporaryAgentError e || e == BROKER HOST -> retryCommand loop + | otherwise -> notify connId $ ERR e + Right () -> withStore' c (`deleteCommand` cmdId) where - delCmd :: AsyncCmdId -> m () - delCmd cmdId = withStore' c $ \db -> deleteCommand db cmdId + withNextSrv :: TVar [SMPServer] -> (SMPServer -> m ()) -> m () + withNextSrv usedSrvs action = do + used <- readTVarIO usedSrvs + srv <- getNextSMPServer c used + atomically $ do + srvs <- readTVar $ smpServers c + let used' = if length used + 1 >= L.length srvs then [] else srv : used + writeTVar usedSrvs used' + action srv notify :: ConnId -> ACommand 'Agent -> m () notify connId cmd = atomically $ writeTBQueue subQ ("", connId, cmd) retryCommand loop = do @@ -844,7 +860,8 @@ runSmpQueueMsgDelivery c@AgentClient {subQ} cData@ConnData {connId, duplexHandsh -- and this branch should never be reached as receive is created before the confirmation, -- so the condition is not necessary here, strictly speaking. _ -> unless (duplexHandshake == Just True) $ do - qInfo <- createReplyQueue c cData sq + srv <- getSMPServer c + qInfo <- createReplyQueue c cData sq srv void . enqueueMessage c cData sq SMP.noMsgFlags $ REPLY [qInfo] AM_A_MSG_ -> notify $ SENT mId _ -> pure () @@ -1155,14 +1172,23 @@ suspendAgent' c@AgentClient {agentState = as} maxDelay = do suspendSendingAndDatabase c getSMPServer :: AgentMonad m => AgentClient -> m SMPServer -getSMPServer c = do - smpServers <- readTVarIO $ smpServers c - case smpServers of - srv :| [] -> pure srv - servers -> do - gen <- asks randomServer - atomically . stateTVar gen $ - first (servers L.!!) . randomR (0, L.length servers - 1) +getSMPServer c = readTVarIO (smpServers c) >>= pickServer + +pickServer :: AgentMonad m => NonEmpty SMPServer -> m SMPServer +pickServer = \case + srv :| [] -> pure srv + servers -> do + gen <- asks randomServer + atomically $ (servers L.!!) <$> stateTVar gen (randomR (0, L.length servers - 1)) + +getNextSMPServer :: AgentMonad m => AgentClient -> [SMPServer] -> m SMPServer +getNextSMPServer c usedSrvs = do + srvs <- readTVarIO $ smpServers c + case L.nonEmpty $ deleteFirstsBy different (L.toList srvs) usedSrvs of + Just srvs' -> pickServer srvs' + _ -> pickServer srvs + where + different (SMPServer host port _) (SMPServer host' port' _) = host /= host' || port /= port' subscriber :: (MonadUnliftIO m, MonadReader Env m) => AgentClient -> m () subscriber c@AgentClient {msgQ} = forever $ do @@ -1400,8 +1426,8 @@ connectReplyQueues c cData@ConnData {connId} ownConnInfo (qInfo :| _) = do withStore c $ \db -> upgradeRcvConnToDuplex db connId sq enqueueConfirmation c cData sq ownConnInfo Nothing -confirmQueue :: forall m. AgentMonad m => Compatible Version -> AgentClient -> ConnData -> SndQueue -> ConnInfo -> Maybe (CR.E2ERatchetParams 'C.X448) -> m () -confirmQueue (Compatible agentVersion) c cData@ConnData {connId} sq connInfo e2eEncryption = do +confirmQueue :: forall m. AgentMonad m => Compatible Version -> AgentClient -> ConnData -> SndQueue -> SMPServer -> ConnInfo -> Maybe (CR.E2ERatchetParams 'C.X448) -> m () +confirmQueue (Compatible agentVersion) c cData@ConnData {connId} sq srv connInfo e2eEncryption = do aMessage <- mkAgentMessage agentVersion msg <- mkConfirmation aMessage sendConfirmation c sq msg @@ -1415,7 +1441,7 @@ confirmQueue (Compatible agentVersion) c cData@ConnData {connId} sq connInfo e2e mkAgentMessage :: Version -> m AgentMessage mkAgentMessage 1 = pure $ AgentConnInfo connInfo mkAgentMessage _ = do - qInfo <- createReplyQueue c cData sq + qInfo <- createReplyQueue c cData sq srv pure $ AgentConnInfoReply (qInfo :| []) connInfo enqueueConfirmation :: forall m. AgentMonad m => AgentClient -> ConnData -> SndQueue -> ConnInfo -> Maybe (CR.E2ERatchetParams 'C.X448) -> m ()