try tracking service subs in the agent (WIP, does not compile)

This commit is contained in:
Evgeny Poberezkin
2025-09-03 08:33:01 +01:00
parent 7fba70acc1
commit 59e908100d
6 changed files with 92 additions and 55 deletions
+6 -5
View File
@@ -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'')
+3 -4
View File
@@ -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
+80 -43
View File
@@ -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
@@ -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
@@ -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
+1 -1
View File
@@ -1465,7 +1465,7 @@ type MsgId = ByteString
type MsgBody = ByteString
data ServiceSub = ServiceSub
{ serviceId :: ServiceId,
{ smpServiceId :: ServiceId,
smpQueueCount :: Int64,
smpQueueIdsHash :: IdsHash
}