service subscriptions in the client

Author: epoberezkin

src/Simplex/Messaging/Agent.hs

@@ -594,6 +594,7 @@ testProtocolServer c nm userId srv = withAgentEnv' c $ case protocolTypeI @p of
   SPNTF -> runNTFServerTest c nm userId srv
 
 -- | set SOCKS5 proxy on/off and optionally set TCP timeouts for fast network
+-- TODO [certs rcv] should fail if any user is enabled to use services and per-connection isolation is chosen
 setNetworkConfig :: AgentClient -> NetworkConfig -> IO ()
 setNetworkConfig c@AgentClient {useNetworkConfig, proxySessTs} cfg' = do
   ts <- getCurrentTime
@@ -771,6 +772,7 @@ deleteUser' c@AgentClient {smpServersStats, xftpServersStats} userId delSMPQueue
       whenM (withStore' c (`deleteUserWithoutConns` userId)) . atomically $
         writeTBQueue (subQ c) ("", "", AEvt SAENone $ DEL_USER userId)
 
+-- TODO [certs rcv] should fail enabling if per-connection isolation is set
 setUserService' :: AgentClient -> UserId -> Bool -> AM ()
 setUserService' c userId enable = do
   wasEnabled <- liftIO $ fromMaybe False <$> TM.lookupIO userId (useClientServices c)
@@ -2870,7 +2872,7 @@ processSMPTransmissions c@AgentClient {subQ} (tSess@(userId, srv, _), _v, sessId
     processSubOk :: RcvQueue -> TVar [ConnId] -> IO ()
     processSubOk rq@RcvQueue {connId} upConnIds =
       atomically . whenM (isPendingSub rq) $ do
-        SS.addActiveSub tSess sessId (rcvQueueSub rq) $ currentSubs c
+        SS.addActiveSub tSess sessId rq $ currentSubs c
         modifyTVar' upConnIds (connId :)
     processSubErr :: RcvQueue -> SMPClientError -> AM' ()
     processSubErr rq@RcvQueue {connId} e = do

src/Simplex/Messaging/Agent/Client.hs

@@ -241,7 +241,7 @@ import Simplex.Messaging.Agent.RetryInterval
 import Simplex.Messaging.Agent.Stats
 import Simplex.Messaging.Agent.Store
 import Simplex.Messaging.Agent.Store.AgentStore
-import Simplex.Messaging.Agent.Store.Common (DBStore, withTransaction)
+import Simplex.Messaging.Agent.Store.Common (DBStore)
 import qualified Simplex.Messaging.Agent.Store.DB as DB
 import Simplex.Messaging.Agent.Store.Entity
 import Simplex.Messaging.Agent.TSessionSubs (TSessionSubs)
@@ -279,7 +279,7 @@ import Simplex.Messaging.Protocol
     RcvNtfPublicDhKey,
     SMPMsgMeta (..),
     SProtocolType (..),
-    ServiceSub,
+    ServiceSub (..),
     SndPublicAuthKey,
     SubscriptionMode (..),
     NewNtfCreds (..),
@@ -500,6 +500,7 @@ data UserNetworkType = UNNone | UNCellular | UNWifi | UNEthernet | UNOther
   deriving (Eq, Show)
 
 -- | Creates an SMP agent client instance that receives commands and sends responses via 'TBQueue's.
+-- TODO [certs rcv] should fail if both per-connection isolation is set and any users use services
 newAgentClient :: Int -> InitialAgentServers -> UTCTime -> Map (Maybe SMPServer) (Maybe SystemSeconds) -> Env -> IO AgentClient
 newAgentClient clientId InitialAgentServers {smp, ntf, xftp, netCfg, useServices, presetDomains, presetServers} currentTs notices agentEnv = do
   let cfg = config agentEnv
@@ -745,9 +746,11 @@ smpConnectClient c@AgentClient {smpClients, msgQ, proxySessTs, presetDomains} nm
       smp <- liftError (protocolClientError SMP $ B.unpack $ strEncode srv) $ do
         ts <- readTVarIO proxySessTs
         ExceptT $ getProtocolClient g nm tSess cfg' presetDomains (Just msgQ) ts $ smpClientDisconnected c tSess env v' prs
+      -- TODO [certs rcv] add service to SS, possibly combine with SS.setSessionId
       atomically $ SS.setSessionId tSess (sessionId $ thParams smp) $ currentSubs c
       updateClientService service smp
       pure SMPConnectedClient {connectedClient = smp, proxiedRelays = prs}
+    -- TODO [certs rcv] this should differentiate between service ID just set and service ID changed, and in the latter case disassociate the queue
     updateClientService service smp = case (service, smpClientService smp) of
       (Just (_, serviceId_), Just THClientService {serviceId})
         | serviceId_ /= Just serviceId -> withStore' c $ \db -> setClientServiceId db userId srv serviceId
@@ -764,32 +767,34 @@ smpClientDisconnected c@AgentClient {active, smpClients, smpProxiedRelays} tSess
     -- we make active subscriptions pending only if the client for tSess was current (in the map) and active,
     -- because we can have a race condition when a new current client could have already
     -- made subscriptions active, and the old client would be processing diconnection later.
-    removeClientAndSubs :: IO ([RcvQueueSub], [ConnId])
+    removeClientAndSubs :: IO ([RcvQueueSub], [ConnId], Maybe ServiceSub)
     removeClientAndSubs = atomically $ do
       removeSessVar v tSess smpClients
-      ifM (readTVar active) removeSubs (pure ([], []))
+      ifM (readTVar active) removeSubs (pure ([], [], Nothing))
       where
         sessId = sessionId $ thParams client
         removeSubs = do
           mode <- getSessionMode c
-          subs <- SS.setSubsPending mode tSess sessId $ currentSubs c
+          (subs, serviceSub_) <- SS.setSubsPending mode tSess sessId $ currentSubs c
           let qs = M.elems subs
               cs = nubOrd $ map qConnId qs
           -- this removes proxied relays that this client created sessions to
           destSrvs <- M.keys <$> readTVar prs
           forM_ destSrvs $ \destSrv -> TM.delete (userId, destSrv, cId) smpProxiedRelays
-          pure (qs, cs)
+          pure (qs, cs, serviceSub_)
 
-    serverDown :: ([RcvQueueSub], [ConnId]) -> IO ()
-    serverDown (qs, conns) = whenM (readTVarIO active) $ do
+    serverDown :: ([RcvQueueSub], [ConnId], Maybe ServiceSub) -> IO ()
+    serverDown (qs, conns, serviceSub_) = whenM (readTVarIO active) $ do
       notifySub c $ hostEvent' DISCONNECT client
       unless (null conns) $ notifySub c $ DOWN srv conns
-      unless (null qs) $ do
+      unless (null qs && isNothing serviceSub_) $ do
         releaseGetLocksIO c qs
         mode <- getSessionModeIO c
         let resubscribe
               | (mode == TSMEntity) == isJust cId = resubscribeSMPSession c tSess
-              | otherwise = void $ subscribeQueues c True qs
+              | otherwise = do
+                  mapM_ (runExceptT . resubscribeClientService c tSess) serviceSub_
+                  unless (null qs) $ void $ subscribeQueues c True qs
         runReaderT resubscribe env
 
 resubscribeSMPSession :: AgentClient -> SMPTransportSession -> AM' ()
@@ -808,11 +813,12 @@ resubscribeSMPSession c@AgentClient {smpSubWorkers, workerSeq} tSess = do
     runSubWorker = do
       ri <- asks $ reconnectInterval . config
       withRetryForeground ri isForeground (isNetworkOnline c) $ \_ loop -> do
-        pending <- atomically $ SS.getPendingSubs tSess $ currentSubs c
-        unless (M.null pending) $ do
+        (pendingSubs, pendingSS) <- atomically $ SS.getPendingSubs tSess $ currentSubs c
+        unless (M.null pendingSubs && isNothing pendingSS) $ do
           liftIO $ waitUntilForeground c
           liftIO $ waitForUserNetwork c
-          handleNotify $ resubscribeSessQueues c tSess $ M.elems pending
+          mapM_ (handleNotify . void . runExceptT . resubscribeClientService c tSess) pendingSS
+          unless (M.null pendingSubs) $ handleNotify $ resubscribeSessQueues c tSess $ M.elems pendingSubs
           loop
     isForeground = (ASForeground ==) <$> readTVar (agentState c)
     cleanup :: SessionVar (Async ()) -> STM ()
@@ -1509,25 +1515,25 @@ newRcvQueue_ c nm userId connId (ProtoServerWithAuth srv auth) vRange cqrd enabl
         newErr :: String -> AM (Maybe ShortLinkCreds)
         newErr = throwE . BROKER (B.unpack $ strEncode srv) . UNEXPECTED . ("Create queue: " <>)
 
-processSubResults :: AgentClient -> SMPTransportSession -> SessionId -> NonEmpty (RcvQueueSub, Either SMPClientError (Maybe ServiceId)) -> STM [(RcvQueueSub, Maybe ClientNotice)]
-processSubResults c tSess@(userId, srv, _) sessId rs = do
-  pendingSubs <- SS.getPendingSubs tSess $ currentSubs c
-  let (failed, subscribed, notices, ignored) = foldr (partitionResults pendingSubs) (M.empty, [], [], 0) rs
+processSubResults :: AgentClient -> SMPTransportSession -> SessionId -> Maybe ServiceId -> NonEmpty (RcvQueueSub, Either SMPClientError (Maybe ServiceId)) -> STM ([RcvQueueSub], [(RcvQueueSub, Maybe ClientNotice)])
+processSubResults c tSess@(userId, srv, _) sessId smpServiceId rs = do
+  pending <- SS.getPendingSubs tSess $ currentSubs c
+  let (failed, subscribed@(qs, sQs), notices, ignored) = foldr (partitionResults pending) (M.empty, ([], []), [], 0) rs
   unless (M.null failed) $ do
     incSMPServerStat' c userId srv connSubErrs $ M.size failed
     failSubscriptions c tSess failed
-  unless (null subscribed) $ do
-    incSMPServerStat' c userId srv connSubscribed $ length subscribed
+  unless (null qs && null sQs) $ do
+    incSMPServerStat' c userId srv connSubscribed $ length qs + length sQs
     SS.batchAddActiveSubs tSess sessId subscribed $ currentSubs c
   unless (ignored == 0) $ incSMPServerStat' c userId srv connSubIgnored ignored
-  pure notices
+  pure (sQs, notices)
   where
     partitionResults ::
-      Map SMP.RecipientId RcvQueueSub ->
+      (Map SMP.RecipientId RcvQueueSub, Maybe ServiceSub) ->
       (RcvQueueSub, Either SMPClientError (Maybe ServiceId)) ->
-      (Map SMP.RecipientId SMPClientError, [RcvQueueSub], [(RcvQueueSub, Maybe ClientNotice)], Int) ->
-      (Map SMP.RecipientId SMPClientError, [RcvQueueSub], [(RcvQueueSub, Maybe ClientNotice)], Int)
-    partitionResults pendingSubs (rq@RcvQueueSub {rcvId, clientNoticeId}, r) acc@(failed, subscribed, notices, ignored) = case r of
+      (Map SMP.RecipientId SMPClientError, ([RcvQueueSub], [RcvQueueSub]), [(RcvQueueSub, Maybe ClientNotice)], Int) ->
+      (Map SMP.RecipientId SMPClientError, ([RcvQueueSub], [RcvQueueSub]), [(RcvQueueSub, Maybe ClientNotice)], Int)
+    partitionResults (pendingSubs, pendingSS) (rq@RcvQueueSub {rcvId, clientNoticeId}, r) acc@(failed, subscribed@(qs, sQs), notices, ignored) = case r of
       Left e -> case smpErrorClientNotice e of
         Just notice_ -> (failed', subscribed, (rq, notice_) : notices, ignored)
           where
@@ -1537,8 +1543,12 @@ processSubResults c tSess@(userId, srv, _) sessId rs = do
           | otherwise -> (failed', subscribed, notices, ignored)
         where
           failed' = M.insert rcvId e failed
-      Right _serviceId -- TODO [certs rcv] store association with the service
-        | rcvId `M.member` pendingSubs -> (failed, rq : subscribed, notices', ignored)
+      Right serviceId_
+        | rcvId `M.member` pendingSubs ->
+            let subscribed' = case (smpServiceId, serviceId_, pendingSS) of
+                  (Just sId, Just sId', Just ServiceSub {serviceId}) | sId == sId' && sId == serviceId -> (qs, rq : sQs)
+                  _ -> (rq : qs, sQs)
+             in (failed, subscribed', notices', ignored)
         | otherwise -> (failed, subscribed, notices', ignored + 1)
         where
           notices' = if isJust clientNoticeId then (rq, Nothing) : notices else notices
@@ -1577,6 +1587,7 @@ serverHostError = \case
 
 -- | Batch by transport session and subscribe queues. The list of results can have a different order.
 subscribeQueues :: AgentClient -> Bool -> [RcvQueueSub] -> AM' [(RcvQueueSub, Either AgentErrorType (Maybe ServiceId))]
+subscribeQueues _ _ [] = pure []
 subscribeQueues c withEvents qs = do
   (errs, qs') <- checkQueues c qs
   atomically $ modifyTVar' (subscrConns c) (`S.union` S.fromList (map qConnId qs'))
@@ -1633,6 +1644,7 @@ checkQueues c = fmap partitionEithers . mapM checkQueue
 -- This function expects that all queues belong to one transport session,
 -- and that they are already added to pending subscriptions.
 resubscribeSessQueues :: AgentClient -> SMPTransportSession -> [RcvQueueSub] -> AM' ()
+resubscribeSessQueues _ _ [] = pure ()
 resubscribeSessQueues c tSess qs = do
   (errs, qs_) <- checkQueues c qs
   forM_ (L.nonEmpty qs_) $ \qs' -> void $ subscribeSessQueues_ c True (tSess, qs')
@@ -1651,13 +1663,15 @@ subscribeSessQueues_ c withEvents qs = sendClientBatch_ "SUB" False subscribe_ c
           then Just . S.fromList . map qConnId . M.elems <$> atomically (SS.getActiveSubs tSess $ currentSubs c)
           else pure Nothing
       active <- E.uninterruptibleMask_ $ do
-        (active, notices) <- atomically $ do
-          r@(_, notices) <- ifM
+        (active, (serviceQs, notices)) <- atomically $ do
+          r@(_, (_, notices)) <- ifM
             (activeClientSession c tSess sessId)
-            ((True,) <$> processSubResults c tSess sessId rs)
-            ((False, []) <$ incSMPServerStat' c userId srv connSubIgnored (length rs))
+            ((True,) <$> processSubResults c tSess sessId smpServiceId rs)
+            ((False, ([], [])) <$ incSMPServerStat' c userId srv connSubIgnored (length rs))
           unless (null notices) $ takeTMVar $ clientNoticesLock c
           pure r
+        unless (null serviceQs) $ void $
+          processRcvServiceAssocs c serviceQs `runReaderT` agentEnv c
         unless (null notices) $ void $
           (processClientNotices c tSess notices `runReaderT` agentEnv c)
             `E.finally` atomically (putTMVar (clientNoticesLock c) ())
@@ -1678,6 +1692,13 @@ subscribeSessQueues_ c withEvents qs = sendClientBatch_ "SUB" False subscribe_ c
       where
         tSess = transportSession' smp
         sessId = sessionId $ thParams smp
+        smpServiceId = (\THClientService {serviceId} -> serviceId) <$> smpClientService smp
+
+processRcvServiceAssocs :: AgentClient -> [RcvQueueSub] -> AM' ()
+processRcvServiceAssocs c serviceQs =
+  withStore' c (`setRcvServiceAssocs` serviceQs) `catchAllErrors'` \e -> do
+    logError $ "processClientNotices error: " <> tshow e
+    notifySub' c "" $ ERR e
 
 processClientNotices :: AgentClient -> SMPTransportSession -> [(RcvQueueSub, Maybe ClientNotice)] -> AM' ()
 processClientNotices c@AgentClient {presetServers} tSess notices = do
@@ -1690,10 +1711,35 @@ processClientNotices c@AgentClient {presetServers} tSess notices = do
       logError $ "processClientNotices error: " <> tshow e
       notifySub' c "" $ ERR e
 
+resubscribeClientService :: AgentClient -> SMPTransportSession -> ServiceSub -> AM ServiceSub
+resubscribeClientService c tSess (ServiceSub _ n idsHash) =
+  withServiceClient c tSess $ \smp _ -> do
+    subscribeClientService_ c tSess smp n idsHash
+
 subscribeClientService :: AgentClient -> UserId -> SMPServer -> Int64 -> IdsHash -> AM ServiceSub
 subscribeClientService c userId srv n idsHash =
-  withLogClient c NRMBackground (userId, srv, Nothing) B.empty "SUBS" $ \(SMPConnectedClient smp _) ->
-    subscribeService smp SMP.SRecipientService n idsHash
+  withServiceClient c tSess $ \smp smpServiceId -> do
+    let serviceSub = ServiceSub smpServiceId n idsHash
+    atomically $ SS.setPendingServiceSub tSess serviceSub $ currentSubs c
+    subscribeClientService_ c tSess smp n idsHash
+  where
+    tSess = (userId, srv, Nothing)
+
+withServiceClient :: AgentClient -> SMPTransportSession -> (SMPClient -> ServiceId -> ExceptT SMPClientError IO a) -> AM a
+withServiceClient c tSess action =
+  withLogClient c NRMBackground tSess B.empty "SUBS" $ \(SMPConnectedClient smp _) ->
+    case (\THClientService {serviceId} -> serviceId) <$> smpClientService smp of
+      Just smpServiceId -> action smp smpServiceId
+      Nothing -> throwE PCEServiceUnavailable
+
+subscribeClientService_ :: AgentClient -> SMPTransportSession -> SMPClient -> Int64 -> IdsHash -> ExceptT SMPClientError IO ServiceSub
+subscribeClientService_ c tSess smp n idsHash = do
+  -- TODO [certs rcv] handle error
+  serviceSub' <- subscribeService smp SMP.SRecipientService n idsHash
+  let sessId = sessionId $ thParams smp
+  atomically $ whenM (activeClientSession c tSess sessId) $
+    SS.setActiveServiceSub tSess sessId serviceSub' $ currentSubs c
+  pure serviceSub'
 
 activeClientSession :: AgentClient -> SMPTransportSession -> SessionId -> STM Bool
 activeClientSession c tSess sessId = sameSess <$> tryReadSessVar tSess (smpClients c)
@@ -1763,7 +1809,7 @@ addNewQueueSubscription c rq' tSess sessId = do
     modifyTVar' (subscrConns c) $ S.insert $ qConnId rq
     active <- activeClientSession c tSess sessId
     if active
-      then SS.addActiveSub tSess sessId rq $ currentSubs c
+      then SS.addActiveSub tSess sessId rq' $ currentSubs c
       else SS.addPendingSub tSess rq $ currentSubs c
     pure active
   unless same $ resubscribeSMPSession c tSess
@@ -1952,6 +1998,7 @@ releaseGetLock c rq =
 {-# INLINE releaseGetLock #-}
 
 releaseGetLocksIO :: SomeRcvQueue q => AgentClient -> [q] -> IO ()
+releaseGetLocksIO _ [] = pure ()
 releaseGetLocksIO c rqs = do
   locks <- readTVarIO $ getMsgLocks c
   forM_ rqs $ \rq ->

src/Simplex/Messaging/Agent/Store/AgentStore.hs

@@ -53,6 +53,7 @@ module Simplex.Messaging.Agent.Store.AgentStore
     getSubscriptionServers,
     getUserServerRcvQueueSubs,
     unsetQueuesToSubscribe,
+    setRcvServiceAssocs,
     getConnIds,
     getConn,
     getDeletedConn,
@@ -2249,6 +2250,14 @@ getUserServerRcvQueueSubs db userId srv onlyNeeded =
 unsetQueuesToSubscribe :: DB.Connection -> IO ()
 unsetQueuesToSubscribe db = DB.execute_ db "UPDATE rcv_queues SET to_subscribe = 0 WHERE to_subscribe = 1"
 
+setRcvServiceAssocs :: DB.Connection -> [RcvQueueSub] -> IO ()
+setRcvServiceAssocs db rqs =
+#if defined(dbPostgres)
+  DB.execute db "UPDATE rcv_queues SET rcv_service_assoc = 1 WHERE rcv_id IN " $ Only $ In (map queueId rqs)
+#else
+  DB.executeMany db "UPDATE rcv_queues SET rcv_service_assoc = 1 WHERE rcv_id = " $ map (Only . queueId) rqs
+#endif
+
 -- * getConn helpers
 
 getConnIds :: DB.Connection -> IO [ConnId]