From 59e908100d21ddb6eb95c75d49821d2349fc4d6c Mon Sep 17 00:00:00 2001 From: Evgeny Poberezkin Date: Wed, 3 Sep 2025 08:33:01 +0100 Subject: [PATCH] try tracking service subs in the agent (WIP, does not compile) --- src/Simplex/Messaging/Agent.hs | 11 +- src/Simplex/Messaging/Agent/Client.hs | 7 +- .../Messaging/Agent/Store/AgentStore.hs | 123 ++++++++++++------ .../Migrations/M20250815_service_certs.hs | 2 +- .../Migrations/M20250815_service_certs.hs | 2 +- src/Simplex/Messaging/Protocol.hs | 2 +- 6 files changed, 92 insertions(+), 55 deletions(-) diff --git a/src/Simplex/Messaging/Agent.hs b/src/Simplex/Messaging/Agent.hs index d7270efeb..1b2388fb8 100644 --- a/src/Simplex/Messaging/Agent.hs +++ b/src/Simplex/Messaging/Agent.hs @@ -850,7 +850,7 @@ newConn c nm userId enableNtfs cMode userData_ clientData pqInitKeys subMode = d srv <- getSMPServer c userId connId <- newConnNoQueues c userId enableNtfs cMode (CR.connPQEncryption pqInitKeys) (connId,) <$> newRcvConnSrv c nm userId connId enableNtfs cMode userData_ clientData pqInitKeys subMode srv - `catchE` \e -> withStore' c (`deleteConnRecord` connId) >> throwE e + `catchE` \e -> withStore' c (\db -> deleteConnRecord db userId connId) >> throwE e setConnShortLink' :: AgentClient -> NetworkRequestMode -> ConnId -> SConnectionMode c -> UserLinkData -> Maybe CRClientData -> AM (ConnShortLink c) setConnShortLink' c nm connId cMode userData clientData = @@ -2160,7 +2160,8 @@ prepareDeleteConnections_ getConnections c waitDelivery connIds = do -- ! if it was used to notify about the result, it might be necessary to differentiate -- ! between completed deletions of connections, and deletions delayed due to wait for delivery (see deleteConn) deliveryTimeout <- if waitDelivery then asks (Just . connDeleteDeliveryTimeout . config) else pure Nothing - cIds_ <- lift $ L.nonEmpty . catMaybes . rights <$> withStoreBatch' c (\db -> map (deleteConn db deliveryTimeout) (M.keys delRs)) + let delConns = map (\(connId, RcvQueue {userId}) -> (userId, connId)) $ M.toList delRs + cIds_ <- lift $ L.nonEmpty . catMaybes . rights <$> withStoreBatch' c (\db -> map (uncurry $ deleteConn db deliveryTimeout) delConns) forM_ cIds_ $ \cIds -> notify ("", "", AEvt SAEConn $ DEL_CONNS cIds) pure (errs' <> delRs, rqs, connIds') where @@ -2179,9 +2180,9 @@ prepareDeleteConnections_ getConnections c waitDelivery connIds = do deleteConnQueues :: AgentClient -> NetworkRequestMode -> Bool -> Bool -> [RcvQueue] -> AM' (Map ConnId (Either AgentErrorType ())) deleteConnQueues c nm waitDelivery ntf rqs = do rs <- connResults <$> (deleteQueueRecs =<< deleteQueues c nm rqs) - let connIds = M.keys $ M.filter isRight rs + let conns = map (\(connId, RcvQueue {userId}) -> (userId, connId)) $ M.toList $ M.filter isRight rs deliveryTimeout <- if waitDelivery then asks (Just . connDeleteDeliveryTimeout . config) else pure Nothing - cIds_ <- L.nonEmpty . catMaybes . rights <$> withStoreBatch' c (\db -> map (deleteConn db deliveryTimeout) connIds) + cIds_ <- L.nonEmpty . catMaybes . rights <$> withStoreBatch' c (\db -> map (uncurry $ deleteConn db deliveryTimeout) conns) forM_ cIds_ $ \cIds -> notify ("", "", AEvt SAEConn $ DEL_CONNS cIds) pure rs where @@ -2474,7 +2475,7 @@ sendNtfConnCommands :: AgentClient -> NtfSupervisorCommand -> AM () sendNtfConnCommands c cmd = do ns <- asks ntfSupervisor connIds <- liftIO $ S.toList <$> getSubscriptions c - rs <- lift $ withStoreBatch' c (\db -> map (getConnData db) connIds) + rs <- lift $ withStoreBatch' c (\db -> map (getConnData False db) connIds) let (connIds', cErrs) = enabledNtfConns (zip connIds rs) forM_ (L.nonEmpty connIds') $ \connIds'' -> atomically $ writeTBQueue (ntfSubQ ns) (cmd, connIds'') diff --git a/src/Simplex/Messaging/Agent/Client.hs b/src/Simplex/Messaging/Agent/Client.hs index be5fe2321..c421b3887 100644 --- a/src/Simplex/Messaging/Agent/Client.hs +++ b/src/Simplex/Messaging/Agent/Client.hs @@ -271,7 +271,7 @@ import Simplex.Messaging.Protocol RcvNtfPublicDhKey, SMPMsgMeta (..), SProtocolType (..), - ServiceSub, + ServiceSub (..), SndPublicAuthKey, SubscriptionMode (..), NewNtfCreds (..), @@ -599,10 +599,9 @@ getServiceCredentials c userId srv = liftIO (TM.lookupIO userId $ useClientServices c) $>>= \useService -> if useService then Just <$> getService else pure Nothing where - getService :: AM (ServiceCredentials, Maybe ServiceId) getService = do let g = agentDRG c - ((C.KeyHash kh, serviceCreds), serviceId_) <- + ((C.KeyHash kh, serviceCreds), serviceSub_) <- withStore' c $ \db -> getClientService db userId srv >>= \case Just service -> pure service @@ -614,7 +613,7 @@ getServiceCredentials c userId srv = (_, pk) <- atomically $ C.generateKeyPair g let serviceSignKey = C.APrivateSignKey C.SEd25519 pk creds = ServiceCredentials {serviceRole = SRMessaging, serviceCreds, serviceCertHash = XV.Fingerprint kh, serviceSignKey} - pure (creds, serviceId_) + pure (creds, smpServiceId <$> serviceSub_) class (Encoding err, Show err) => ProtocolServerClient v err msg | msg -> v, msg -> err where type Client msg = c | c -> msg diff --git a/src/Simplex/Messaging/Agent/Store/AgentStore.hs b/src/Simplex/Messaging/Agent/Store/AgentStore.hs index cde6b56c8..aa8940367 100644 --- a/src/Simplex/Messaging/Agent/Store/AgentStore.hs +++ b/src/Simplex/Messaging/Agent/Store/AgentStore.hs @@ -395,19 +395,36 @@ createClientService db userId srv (kh, (cert, pk)) = |] (userId, host srv, port srv, kh, cert, pk) -getClientService :: DB.Connection -> UserId -> SMPServer -> IO (Maybe ((C.KeyHash, TLS.Credential), Maybe ServiceId)) +getClientService :: DB.Connection -> UserId -> SMPServer -> IO (Maybe ((C.KeyHash, TLS.Credential), Maybe ServiceSub)) getClientService db userId srv = maybeFirstRow toService $ DB.query db [sql| - SELECT service_cert_hash, service_cert, service_priv_key, rcv_service_id + SELECT service_cert_hash, service_cert, service_priv_key, rcv_service_id, FROM client_services WHERE user_id = ? AND host = ? AND port = ? |] (userId, host srv, port srv) where - toService (kh, cert, pk, serviceId_) = ((kh, (cert, pk)), serviceId_) + toService (kh, cert, pk, serviceId_, n, idsHash) = + let service_ = (\serviceId -> ServiceSub serviceId n idsHash) <$> serviceId_ + in ((kh, (cert, pk)), service_) + +updateServiceAggregates :: DB.Connection -> UserId -> SMPServer -> Int64 -> [RecipientId] -> IO () +updateServiceAggregates db userId server change rcvIds = + getClientService db userId server >>= \case + Just (_, Just (ServiceSub serviceId n idsHash)) -> + DB.execute + db + [sql| + UPDATE client_services + SET service_queue_count = ?, + service_queue_ids_hash = ? + WHERE user_id = ? AND rcv_service_id = ? + |] + (n + change, idsHash <> queueIdsHash rcvIds, userId, serviceId) + _ -> pure () getClientServiceServers :: DB.Connection -> UserId -> IO [(SMPServer, ServiceSub)] getClientServiceServers db userId = @@ -499,8 +516,13 @@ createConnRecord db connId ConnData {userId, connAgentVersion, enableNtfs, pqSup |] (userId, connId, cMode, connAgentVersion, BI enableNtfs, pqSupport, BI True) -deleteConnRecord :: DB.Connection -> ConnId -> IO () -deleteConnRecord db connId = DB.execute db "DELETE FROM connections WHERE conn_id = ?" (Only connId) +deleteConnRecord :: DB.Connection -> UserId -> ConnId -> IO () +deleteConnRecord db userId connId = do + -- TODO [certs rcv] it needs to be grouped per server here + rIds :: [RecipientId] <- + map fromOnly <$> DB.query db "SELECT rcv_id FROM rcv_queues WHERE conn_id = ? AND deleted = 0 AND rcv_service_assoc = 1" (Only connId) + unless (null rIds) $ updateServiceAggregates db userId server (- length rIds') rIds' + DB.execute db "DELETE FROM connections WHERE conn_id = ?" (Only connId) checkConfirmedSndQueueExists_ :: DB.Connection -> NewSndQueue -> IO Bool checkConfirmedSndQueueExists_ db SndQueue {server, sndId} = do @@ -521,8 +543,8 @@ getRcvConn db ProtocolServer {host, port} rcvId = runExceptT $ do (rq,) <$> ExceptT (getConn db connId) -- | Deletes connection, optionally checking for pending snd message deliveries; returns connection id if it was deleted -deleteConn :: DB.Connection -> Maybe NominalDiffTime -> ConnId -> IO (Maybe ConnId) -deleteConn db waitDeliveryTimeout_ connId = case waitDeliveryTimeout_ of +deleteConn :: DB.Connection -> Maybe NominalDiffTime -> UserId -> ConnId -> IO (Maybe ConnId) +deleteConn db waitDeliveryTimeout_ userId connId = case waitDeliveryTimeout_ of Nothing -> delete Just timeout -> ifM @@ -534,7 +556,7 @@ deleteConn db waitDeliveryTimeout_ connId = case waitDeliveryTimeout_ of (pure Nothing) ) where - delete = deleteConnRecord db connId $> Just connId + delete = deleteConnRecord db userId connId $> Just connId checkNoPendingDeliveries_ = do r :: (Maybe Int64) <- maybeFirstRow fromOnly $ @@ -609,15 +631,13 @@ setRcvSwitchStatus db rq@RcvQueue {rcvId, server = ProtocolServer {host, port}} pure rq {rcvSwchStatus} setRcvQueueDeleted :: DB.Connection -> RcvQueue -> IO () -setRcvQueueDeleted db RcvQueue {rcvId, server = ProtocolServer {host, port}} = do - DB.execute - db - [sql| - UPDATE rcv_queues - SET deleted = 1 - WHERE host = ? AND port = ? AND rcv_id = ? - |] - (host, port, rcvId) +setRcvQueueDeleted db RcvQueue {connId, dbQueueId, rcvId} = do + q_ :: Maybe BoolInt <- + maybeFirstRow fromOnly $ + DB.query db "SELECT rcv_service_assoc FROM rcv_queues WHERE conn_id = ? AND rcv_queue_id = ? AND deleted = 0" (connId, dbQueueId) + forM_ q_ $ \(BI rcvServiceAssoc) -> do + DB.execute db "UPDATE rcv_queues SET deleted = 1 WHERE host = ? AND port = ? AND rcv_id = ?" (host, port, rcvId) + when rcvServiceAssoc $ updateServiceAggregates db userId server (-1) [rcvId] setRcvQueueConfirmedE2E :: DB.Connection -> RcvQueue -> C.DhSecretX25519 -> VersionSMPC -> IO () setRcvQueueConfirmedE2E db RcvQueue {rcvId, server = ProtocolServer {host, port}} e2eDhSecret smpClientVersion = @@ -677,8 +697,13 @@ incRcvDeleteErrors db RcvQueue {connId, dbQueueId} = DB.execute db "UPDATE rcv_queues SET delete_errors = delete_errors + 1 WHERE conn_id = ? AND rcv_queue_id = ?" (connId, dbQueueId) deleteConnRcvQueue :: DB.Connection -> RcvQueue -> IO () -deleteConnRcvQueue db RcvQueue {connId, dbQueueId} = - DB.execute db "DELETE FROM rcv_queues WHERE conn_id = ? AND rcv_queue_id = ?" (connId, dbQueueId) +deleteConnRcvQueue db RcvQueue {connId, dbQueueId, server} = do + q_ :: Maybe (BoolInt, BoolInt) <- + maybeFirstRow id $ + DB.query db "SELECT deleted, rcv_service_assoc FROM rcv_queues WHERE conn_id = ? AND rcv_queue_id = ?" (connId, dbQueueId) + forM_ q_ $ \(BI deleted, BI rcvServiceAssoc) -> do + DB.execute db "DELETE FROM rcv_queues WHERE conn_id = ? AND rcv_queue_id = ?" (connId, dbQueueId) + when (not deleted && rcvServiceAssoc) $ updateServiceAggregates db userId server (-1) [rcvId] deleteConnSndQueue :: DB.Connection -> ConnId -> SndQueue -> IO () deleteConnSndQueue db connId SndQueue {dbQueueId} = do @@ -687,7 +712,7 @@ deleteConnSndQueue db connId SndQueue {dbQueueId} = do getPrimaryRcvQueue :: DB.Connection -> ConnId -> IO (Either StoreError RcvQueue) getPrimaryRcvQueue db connId = - maybe (Left SEConnNotFound) (Right . L.head) <$> getRcvQueuesByConnId_ db connId + maybe (Left SEConnNotFound) (Right . L.head) <$> getRcvQueuesByConnId_ db connId False getRcvQueue :: DB.Connection -> ConnId -> SMPServer -> SMP.RecipientId -> IO (Either StoreError RcvQueue) getRcvQueue db connId (SMPServer host port _) rcvId = @@ -1015,7 +1040,7 @@ getPendingQueueMsg db connId SndQueue {dbQueueId} = getMsgData :: InternalId -> IO (Either StoreError (Maybe RcvQueue, PendingMsgData)) getMsgData msgId = runExceptT $ do msg <- ExceptT $ firstRow' pendingMsgData err getMsgData_ - rq_ <- liftIO $ L.head <$$> getRcvQueuesByConnId_ db connId + rq_ <- liftIO $ L.head <$$> getRcvQueuesByConnId_ db connId False pure (rq_, msg) where getMsgData_ = @@ -2061,6 +2086,7 @@ insertRcvQueue_ db connId' rq@RcvQueue {..} serverKeyHash_ = do :. (shortLinkId <$> shortLink, shortLinkKey <$> shortLink, linkPrivSigKey <$> shortLink, linkEncFixedData <$> shortLink) :. ntfCredsFields ) + when rcvServiceAssoc $ updateServiceAggregates db userId server 1 [rcvId] pure (rq :: NewRcvQueue) {connId = connId', dbQueueId = qId} where ntfCredsFields = case clientNtfCreds of @@ -2119,21 +2145,19 @@ getDeletedConn = getAnyConn True {-# INLINE getDeletedConn #-} getAnyConn :: Bool -> DB.Connection -> ConnId -> IO (Either StoreError SomeConn) -getAnyConn deleted' dbConn connId = - getConnData dbConn connId >>= \case +getAnyConn deleted dbConn connId = + getConnData deleted dbConn connId >>= \case Nothing -> pure $ Left SEConnNotFound - Just (cData@ConnData {deleted}, cMode) - | deleted /= deleted' -> pure $ Left SEConnNotFound - | otherwise -> do - rQ <- getRcvQueuesByConnId_ dbConn connId - sQ <- getSndQueuesByConnId_ dbConn connId - pure $ case (rQ, sQ, cMode) of - (Just rqs, Just sqs, CMInvitation) -> Right $ SomeConn SCDuplex (DuplexConnection cData rqs sqs) - (Just (rq :| _), Nothing, CMInvitation) -> Right $ SomeConn SCRcv (RcvConnection cData rq) - (Nothing, Just (sq :| _), CMInvitation) -> Right $ SomeConn SCSnd (SndConnection cData sq) - (Just (rq :| _), Nothing, CMContact) -> Right $ SomeConn SCContact (ContactConnection cData rq) - (Nothing, Nothing, _) -> Right $ SomeConn SCNew (NewConnection cData) - _ -> Left SEConnNotFound + Just (cData, cMode) -> do + rQ <- getRcvQueuesByConnId_ dbConn connId deleted + sQ <- getSndQueuesByConnId_ dbConn connId + pure $ case (rQ, sQ, cMode) of + (Just rqs, Just sqs, CMInvitation) -> Right $ SomeConn SCDuplex (DuplexConnection cData rqs sqs) + (Just (rq :| _), Nothing, CMInvitation) -> Right $ SomeConn SCRcv (RcvConnection cData rq) + (Nothing, Just (sq :| _), CMInvitation) -> Right $ SomeConn SCSnd (SndConnection cData sq) + (Just (rq :| _), Nothing, CMContact) -> Right $ SomeConn SCContact (ContactConnection cData rq) + (Nothing, Nothing, _) -> Right $ SomeConn SCNew (NewConnection cData) + _ -> Left SEConnNotFound getConns :: DB.Connection -> [ConnId] -> IO [Either StoreError SomeConn] getConns = getAnyConns_ False @@ -2149,8 +2173,8 @@ getAnyConns_ deleted' db connIds = forM connIds $ E.handle handleDBError . getAn handleDBError :: E.SomeException -> IO (Either StoreError SomeConn) handleDBError = pure . Left . SEInternal . bshow -getConnData :: DB.Connection -> ConnId -> IO (Maybe (ConnData, ConnectionMode)) -getConnData db connId' = +getConnData :: Bool -> DB.Connection -> ConnId -> IO (Maybe (ConnData, ConnectionMode)) +getConnData deleted db connId' = maybeFirstRow cData $ DB.query db @@ -2159,9 +2183,9 @@ getConnData db connId' = user_id, conn_id, conn_mode, smp_agent_version, enable_ntfs, last_external_snd_msg_id, deleted, ratchet_sync_state, pq_support FROM connections - WHERE conn_id = ? + WHERE conn_id = ? AND deleted = ? |] - (Only connId') + (connId', BI deleted) where cData (userId, connId, cMode, connAgentVersion, enableNtfs_, lastExternalSndId, BI deleted, ratchetSyncState, pqSupport) = (ConnData {userId, connId, connAgentVersion, enableNtfs = maybe True unBI enableNtfs_, lastExternalSndId, deleted, ratchetSyncState, pqSupport}, cMode) @@ -2171,8 +2195,21 @@ setConnDeleted db waitDelivery connId | waitDelivery = do currentTs <- getCurrentTime DB.execute db "UPDATE connections SET deleted_at_wait_delivery = ? WHERE conn_id = ?" (currentTs, connId) - | otherwise = + | otherwise = do DB.execute db "UPDATE connections SET deleted = ? WHERE conn_id = ?" (BI True, connId) + -- TODO [certs rcv] it needs to be grouped per server + qs :: [(Int64, RecipientId, BoolInt)] <- + DB.query db "SELECT rcv_queue_id, rcv_id, rcv_service_assoc FROM rcv_queues WHERE conn_id = ? AND deleted = 0" (Only connId) + unless (null qs) $ do +#if defined(dbPostgres) + let dbQIds = In (map (\(dbQId, _, _) -> dbQId) qs) + DB.execute db "UPDATE rcv_queues SET deleted = 1 WHERE conn_id = ? AND rcv_queue_id IN ?" (connId, dbQIds) +#else + let ids = map (\(dbQId, _, _) -> (connId, dbQId)) qs + DB.executeMany db "UPDATE rcv_queues SET deleted = 1 WHERE conn_id = ? AND rcv_queue_id = ?" ids +#endif + let rIds' = mapMaybe (\case (_, rId, BI True) -> Just rId; _ -> Nothing) qs + unless (null rIds') $ updateServiceAggregates db userId server (- length rIds') rIds' setConnUserId :: DB.Connection -> UserId -> ConnId -> UserId -> IO () setConnUserId db oldUserId connId newUserId = @@ -2218,10 +2255,10 @@ deleteRatchetKeyHashesExpired db ttl = do DB.execute db "DELETE FROM processed_ratchet_key_hashes WHERE created_at < ?" (Only cutoffTs) -- | returns all connection queues, the first queue is the primary one -getRcvQueuesByConnId_ :: DB.Connection -> ConnId -> IO (Maybe (NonEmpty RcvQueue)) -getRcvQueuesByConnId_ db connId = +getRcvQueuesByConnId_ :: DB.Connection -> ConnId -> Bool -> IO (Maybe (NonEmpty RcvQueue)) +getRcvQueuesByConnId_ db connId deleted = L.nonEmpty . sortBy primaryFirst . map toRcvQueue - <$> DB.query db (rcvQueueQuery <> " WHERE q.conn_id = ? AND q.deleted = 0") (Only connId) + <$> DB.query db (rcvQueueQuery <> " WHERE q.conn_id = ? AND q.deleted = ?") (connId, deleted) where primaryFirst RcvQueue {primary = p, dbReplaceQueueId = i} RcvQueue {primary = p', dbReplaceQueueId = i'} = -- the current primary queue is ordered first, the next primary - second diff --git a/src/Simplex/Messaging/Agent/Store/Postgres/Migrations/M20250815_service_certs.hs b/src/Simplex/Messaging/Agent/Store/Postgres/Migrations/M20250815_service_certs.hs index 8bae5081d..7c50eac7c 100644 --- a/src/Simplex/Messaging/Agent/Store/Postgres/Migrations/M20250815_service_certs.hs +++ b/src/Simplex/Messaging/Agent/Store/Postgres/Migrations/M20250815_service_certs.hs @@ -17,7 +17,7 @@ CREATE TABLE client_services( service_cert BYTEA NOT NULL, service_cert_hash BYTEA NOT NULL, service_priv_key BYTEA NOT NULL, - service_id BYTEA, + rcv_service_id BYTEA, service_queue_count BIGINT NOT NULL DEFAULT 0, service_queue_ids_hash BYTEA NOT NULL DEFAULT '\x00000000000000000000000000000000', FOREIGN KEY(host, port) REFERENCES servers ON UPDATE CASCADE ON DELETE RESTRICT diff --git a/src/Simplex/Messaging/Agent/Store/SQLite/Migrations/M20250815_service_certs.hs b/src/Simplex/Messaging/Agent/Store/SQLite/Migrations/M20250815_service_certs.hs index c7801c328..557f0a958 100644 --- a/src/Simplex/Messaging/Agent/Store/SQLite/Migrations/M20250815_service_certs.hs +++ b/src/Simplex/Messaging/Agent/Store/SQLite/Migrations/M20250815_service_certs.hs @@ -15,7 +15,7 @@ CREATE TABLE client_services( service_cert BLOB NOT NULL, service_cert_hash BLOB NOT NULL, service_priv_key BLOB NOT NULL, - service_id BLOB, + rcv_service_id BLOB, service_queue_count INTEGER NOT NULL DEFAULT 0, service_queue_ids_hash BLOB NOT NULL DEFAULT x'00000000000000000000000000000000', FOREIGN KEY(host, port) REFERENCES servers ON UPDATE CASCADE ON DELETE RESTRICT diff --git a/src/Simplex/Messaging/Protocol.hs b/src/Simplex/Messaging/Protocol.hs index f4d4288b5..8ea9eac69 100644 --- a/src/Simplex/Messaging/Protocol.hs +++ b/src/Simplex/Messaging/Protocol.hs @@ -1465,7 +1465,7 @@ type MsgId = ByteString type MsgBody = ByteString data ServiceSub = ServiceSub - { serviceId :: ServiceId, + { smpServiceId :: ServiceId, smpQueueCount :: Int64, smpQueueIdsHash :: IdsHash }