ntf server: use different client certificates for each SMP server

Author: epoberezkin

src/Simplex/FileTransfer/Client.hs

@@ -11,6 +11,7 @@
 
 module Simplex.FileTransfer.Client where
 
+import qualified Control.Exception as E
 import Control.Logger.Simple
 import Control.Monad
 import Control.Monad.Except
@@ -264,7 +265,7 @@ downloadXFTPChunk g c@XFTPClient {config} rpKey fId chunkSpec@XFTPRcvChunkSpec {
         where
           errors =
             [ Handler $ \(e :: H.HTTP2Error) -> pure $ Left $ PCENetworkError $ NEConnectError $ displayException e,
-              Handler $ \(e :: IOException) -> pure $ Left $ PCEIOError e,
+              Handler $ \(e :: IOException) -> pure $ Left $ PCEIOError $ E.displayException e,
               Handler $ \(e :: SomeException) -> pure $ Left $ PCENetworkError $ toNetworkError e
             ]
           download cbState =

src/Simplex/Messaging/Agent/Client.hs

@@ -751,8 +751,8 @@ smpConnectClient c@AgentClient {smpClients, msgQ, proxySessTs, presetDomains} nm
       atomically $ SS.setSessionId tSess (sessionId $ thParams smp) $ currentSubs c
       updateClientService service smp
       pure SMPConnectedClient {connectedClient = smp, proxiedRelays = prs}
-    updateClientService service smp = case (service, smpClientService smp) of
-      (Just (_, serviceId_), Just THClientService {serviceId}) -> withStore' c $ \db -> do
+    updateClientService service smp = case (service, smpClientServiceId smp) of
+      (Just (_, serviceId_), Just serviceId) -> withStore' c $ \db -> do
         setClientServiceId db userId srv serviceId
         forM_ serviceId_ $ \sId -> when (sId /= serviceId) $ removeRcvServiceAssocs db userId srv
       (Just _, Nothing) -> withStore' c $ \db -> deleteClientService db userId srv -- e.g., server version downgrade
@@ -1255,7 +1255,7 @@ protocolClientError protocolError_ host = \case
   PCETransportError e -> BROKER host $ TRANSPORT e
   e@PCECryptoError {} -> INTERNAL $ show e
   PCEServiceUnavailable {} -> BROKER host NO_SERVICE
-  PCEIOError e -> BROKER host $ NETWORK $ NEConnectError $ E.displayException e
+  PCEIOError e -> BROKER host $ NETWORK $ NEConnectError e
 
 -- it is consistent with smpClientServiceError
 clientServiceError :: AgentErrorType -> Bool
@@ -1546,6 +1546,7 @@ processSubResults c tSess@(userId, srv, _) sessId serviceId_ rs = do
       Left e -> case smpErrorClientNotice e of
         Just notice_ -> (failed', subscribed, (rq, notice_) : notices, ignored)
           where
+            -- TODO [certs rcv] not used?
             notices' = if isJust notice_ || isJust clientNoticeId then (rq, notice_) : notices else notices
         Nothing
           | temporaryClientError e -> acc
@@ -1678,7 +1679,7 @@ subscribeSessQueues_ c withEvents qs = sendClientBatch_ "SUB" False subscribe_ c
         (active, (serviceQs, notices)) <- atomically $ do
           r@(_, (_, notices)) <- ifM
             (activeClientSession c tSess sessId)
-            ((True,) <$> processSubResults c tSess sessId smpServiceId rs)
+            ((True,) <$> processSubResults c tSess sessId (smpClientServiceId smp) rs)
             ((False, ([], [])) <$ incSMPServerStat' c userId srv connSubIgnored (length rs))
           unless (null notices) $ takeTMVar $ clientNoticesLock c
           pure r
@@ -1704,7 +1705,6 @@ subscribeSessQueues_ c withEvents qs = sendClientBatch_ "SUB" False subscribe_ c
       where
         tSess = transportSession' smp
         sessId = sessionId $ thParams smp
-        smpServiceId = (\THClientService {serviceId} -> serviceId) <$> smpClientService smp
 
 processRcvServiceAssocs :: SMPQueue q => AgentClient -> [q] -> AM' ()
 processRcvServiceAssocs _ [] = pure ()
@@ -1752,7 +1752,7 @@ subscribeClientService c withEvent userId srv (ServiceSub _ n idsHash) =
 withServiceClient :: AgentClient -> SMPTransportSession -> (SMPClient -> ServiceId -> ExceptT SMPClientError IO a) -> AM a
 withServiceClient c tSess subscribe =
   withLogClient c NRMBackground tSess B.empty "SUBS" $ \(SMPConnectedClient smp _) ->
-    case (\THClientService {serviceId} -> serviceId) <$> smpClientService smp of
+    case smpClientServiceId smp of
       Just smpServiceId -> subscribe smp smpServiceId
       Nothing -> throwE PCEServiceUnavailable
 

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

@@ -472,15 +472,21 @@ toServerService (host, port, kh, serviceId, n, Binary idsHash) =
   (SMPServer host port kh, ServiceSub serviceId n (IdsHash idsHash))
 
 setClientServiceId :: DB.Connection -> UserId -> SMPServer -> ServiceId -> IO ()
-setClientServiceId db userId srv serviceId =
+setClientServiceId db userId (SMPServer h p kh) serviceId =
   DB.execute
     db
     [sql|
       UPDATE client_services
       SET service_id = ?
-      WHERE user_id = ? AND host = ? AND port = ?
+      FROM servers s
+      WHERE client_services.user_id = ?
+        AND client_services.host = ?
+        AND client_services.port = ?
+        AND s.host = client_services.host
+        AND s.port = client_services.port
+        AND COALESCE(client_services.server_key_hash, s.key_hash) = ?
     |]
-    (serviceId, userId, host srv, port srv)
+    (serviceId, userId, h, p, kh)
 
 deleteClientService :: DB.Connection -> UserId -> SMPServer -> IO ()
 deleteClientService db userId (SMPServer h p kh) =
@@ -2307,7 +2313,7 @@ unsetQueuesToSubscribe db = DB.execute_ db "UPDATE rcv_queues SET to_subscribe =
 setRcvServiceAssocs :: SMPQueue q => DB.Connection -> [q] -> 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)
+  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

src/Simplex/Messaging/Client.hs

@@ -52,6 +52,7 @@ module Simplex.Messaging.Client
     subscribeSMPQueuesNtfs,
     subscribeService,
     smpClientService,
+    smpClientServiceId,
     secureSMPQueue,
     secureSndSMPQueue,
     proxySecureSndSMPQueue,
@@ -128,7 +129,8 @@ import Control.Applicative ((<|>))
 import Control.Concurrent (ThreadId, forkFinally, forkIO, killThread, mkWeakThreadId)
 import Control.Concurrent.Async
 import Control.Concurrent.STM
-import Control.Exception
+import Control.Exception (Exception, SomeException)
+import qualified Control.Exception as E
 import Control.Logger.Simple
 import Control.Monad
 import Control.Monad.Except
@@ -565,7 +567,7 @@ getProtocolClient g nm transportSession@(_, srv, _) cfg@ProtocolClientConfig {qS
   case chooseTransportHost networkConfig (host srv) of
     Right useHost ->
       (getCurrentTime >>= mkProtocolClient useHost >>= runClient useTransport useHost)
-        `catch` \(e :: IOException) -> pure . Left $ PCEIOError e
+        `E.catch` \(e :: SomeException) -> pure $ Left $ PCEIOError $ E.displayException e
     Left e -> pure $ Left e
   where
     NetworkConfig {tcpConnectTimeout, tcpTimeout, smpPingInterval} = networkConfig
@@ -638,7 +640,7 @@ getProtocolClient g nm transportSession@(_, srv, _) cfg@ProtocolClientConfig {qS
             writeTVar (connected c) True
             putTMVar cVar $ Right c'
           raceAny_ ([send c' th, process c', receive c' th] <> [monitor c' | smpPingInterval > 0])
-            `finally` disconnected c'
+            `E.finally` disconnected c'
 
     send :: Transport c => ProtocolClient v err msg -> THandle v c 'TClient -> IO ()
     send ProtocolClient {client_ = PClient {sndQ}} h = forever $ atomically (readTBQueue sndQ) >>= sendPending
@@ -765,7 +767,7 @@ data ProtocolClientError err
   | -- | Error when cryptographically "signing" the command or when initializing crypto_box.
     PCECryptoError C.CryptoError
   | -- | IO Error
-    PCEIOError IOException
+    PCEIOError String
   deriving (Eq, Show, Exception)
 
 type SMPClientError = ProtocolClientError ErrorType
@@ -926,6 +928,10 @@ smpClientService :: SMPClient -> Maybe THClientService
 smpClientService = thAuth . thParams >=> clientService
 {-# INLINE smpClientService #-}
 
+smpClientServiceId :: SMPClient -> Maybe ServiceId
+smpClientServiceId = fmap (\THClientService {serviceId} -> serviceId) . smpClientService
+{-# INLINE smpClientServiceId #-}
+
 enablePings :: SMPClient -> IO ()
 enablePings ProtocolClient {client_ = PClient {sendPings}} = atomically $ writeTVar sendPings True
 {-# INLINE enablePings #-}

src/Simplex/Messaging/Client/Agent.hs

@@ -15,6 +15,7 @@ module Simplex.Messaging.Client.Agent
   ( SMPClientAgent (..),
     SMPClientAgentConfig (..),
     SMPClientAgentEvent (..),
+    DBService (..),
     OwnServer,
     defaultSMPClientAgentConfig,
     newSMPClientAgent,
@@ -133,6 +134,7 @@ defaultSMPClientAgentConfig =
 data SMPClientAgent p = SMPClientAgent
   { agentCfg :: SMPClientAgentConfig,
     agentParty :: SParty p,
+    dbService :: Maybe DBService,
     active :: TVar Bool,
     startedAt :: UTCTime,
     msgQ :: TBQueue (ServerTransmissionBatch SMPVersion ErrorType BrokerMsg),
@@ -155,8 +157,8 @@ data SMPClientAgent p = SMPClientAgent
 
 type OwnServer = Bool
 
-newSMPClientAgent :: SParty p -> SMPClientAgentConfig -> TVar ChaChaDRG -> IO (SMPClientAgent p)
-newSMPClientAgent agentParty agentCfg@SMPClientAgentConfig {msgQSize, agentQSize} randomDrg = do
+newSMPClientAgent :: SParty p -> SMPClientAgentConfig -> Maybe DBService -> TVar ChaChaDRG -> IO (SMPClientAgent p)
+newSMPClientAgent agentParty agentCfg@SMPClientAgentConfig {msgQSize, agentQSize} dbService randomDrg = do
   active <- newTVarIO True
   startedAt <- getCurrentTime
   msgQ <- newTBQueueIO msgQSize
@@ -173,6 +175,7 @@ newSMPClientAgent agentParty agentCfg@SMPClientAgentConfig {msgQSize, agentQSize
     SMPClientAgent
       { agentCfg,
         agentParty,
+        dbService,
         active,
         startedAt,
         msgQ,
@@ -188,6 +191,11 @@ newSMPClientAgent agentParty agentCfg@SMPClientAgentConfig {msgQSize, agentQSize
         workerSeq
       }
 
+data DBService = DBService
+  { getCredentials :: SMPServer -> IO (Either SMPClientError ServiceCredentials),
+    updateServiceId :: SMPServer -> Maybe ServiceId -> IO (Either SMPClientError ())
+  }
+
 -- | Get or create SMP client for SMPServer
 getSMPServerClient' :: SMPClientAgent p -> SMPServer -> ExceptT SMPClientError IO SMPClient
 getSMPServerClient' ca srv = snd <$> getSMPServerClient'' ca srv
@@ -218,7 +226,7 @@ getSMPServerClient'' ca@SMPClientAgent {agentCfg, smpClients, smpSessions, worke
 
     newSMPClient :: SMPClientVar -> IO (Either SMPClientError (OwnServer, SMPClient))
     newSMPClient v = do
-      r <- connectClient ca srv v `E.catch` (pure . Left . PCEIOError)
+      r <- connectClient ca srv v `E.catch` \(e :: E.SomeException) -> pure $ Left $ PCEIOError $ E.displayException e
       case r of
         Right smp -> do
           logInfo . decodeUtf8 $ "Agent connected to " <> showServer srv
@@ -227,8 +235,7 @@ getSMPServerClient'' ca@SMPClientAgent {agentCfg, smpClients, smpSessions, worke
           atomically $ do
             putTMVar (sessionVar v) (Right c)
             TM.insert (sessionId $ thParams smp) c smpSessions
-          let serviceId_ = (\THClientService {serviceId} -> serviceId) <$> smpClientService smp
-          notify ca $ CAConnected srv serviceId_
+          notify ca $ CAConnected srv $ smpClientServiceId smp
           pure $ Right c
         Left e -> do
           let ei = persistErrorInterval agentCfg
@@ -249,9 +256,18 @@ isOwnServer SMPClientAgent {agentCfg} ProtocolServer {host} =
 
 -- | Run an SMP client for SMPClientVar
 connectClient :: SMPClientAgent p -> SMPServer -> SMPClientVar -> IO (Either SMPClientError SMPClient)
-connectClient ca@SMPClientAgent {agentCfg, smpClients, smpSessions, msgQ, randomDrg, startedAt} srv v =
-  getProtocolClient randomDrg NRMBackground (1, srv, Nothing) (smpCfg agentCfg) [] (Just msgQ) startedAt clientDisconnected
+connectClient ca@SMPClientAgent {agentCfg, dbService, smpClients, smpSessions, msgQ, randomDrg, startedAt} srv v = case dbService of
+  Just dbs -> runExceptT $ do
+    creds <- ExceptT $ getCredentials dbs srv
+    smp <- ExceptT $ getClient cfg {serviceCredentials = Just creds}
+    whenM (atomically $ activeClientSession ca smp srv) $
+      ExceptT $ updateServiceId dbs srv $ smpClientServiceId smp
+    pure smp
+  Nothing -> getClient cfg
   where
+    cfg = smpCfg agentCfg
+    getClient cfg' = getProtocolClient randomDrg NRMBackground (1, srv, Nothing) cfg' [] (Just msgQ) startedAt clientDisconnected
+
     clientDisconnected :: SMPClient -> IO ()
     clientDisconnected smp = do
       removeClientAndSubs smp >>= serverDown
@@ -435,7 +451,7 @@ smpSubscribeQueues ca smp srv subs = do
       unless (null notPending) $ removePendingSubs ca srv notPending
       pure acc
     sessId = sessionId $ thParams smp
-    smpServiceId = (\THClientService {serviceId} -> serviceId) <$> smpClientService smp
+    smpServiceId = smpClientServiceId smp
     groupSub ::
       Map QueueId C.APrivateAuthKey ->
       ((QueueId, C.APrivateAuthKey), Either SMPClientError (Maybe ServiceId)) ->

src/Simplex/Messaging/Notifications/Server.hs

@@ -588,7 +588,7 @@ ntfSubscriber NtfSubscriber {smpAgent = ca@SMPClientAgent {msgQ, agentQ}} =
             logError $ "SMP server service subscription error " <> showService srv serviceSub <> ": " <> tshow e
           CAServiceUnavailable srv serviceSub -> do
             logError $ "SMP server service unavailable: " <> showService srv serviceSub
-            removeServiceAssociation st srv >>= \case
+            removeServiceAndAssociations st srv >>= \case
               Right (srvId, updated) -> do
                 logSubStatus srv "removed service association" updated updated
                 void $ subscribeSrvSubs ca st batchSize (srv, srvId, Nothing)

src/Simplex/Messaging/Notifications/Server/Env.hs

@@ -4,11 +4,14 @@
 {-# LANGUAGE LambdaCase #-}
 {-# LANGUAGE KindSignatures #-}
 {-# LANGUAGE NamedFieldPuns #-}
+{-# LANGUAGE OverloadedLists #-}
 {-# LANGUAGE OverloadedStrings #-}
 
 module Simplex.Messaging.Notifications.Server.Env where
 
 import Control.Concurrent (ThreadId)
+import Control.Monad.Except
+import Control.Monad.Trans.Except
 import Crypto.Random
 import Data.Int (Int64)
 import Data.List.NonEmpty (NonEmpty)
@@ -19,7 +22,7 @@ import qualified Data.X509.Validation as XV
 import Network.Socket
 import qualified Network.TLS as TLS
 import Numeric.Natural
-import Simplex.Messaging.Client (ProtocolClientConfig (..))
+import Simplex.Messaging.Client (ProtocolClientError (..), SMPClientError)
 import Simplex.Messaging.Client.Agent
 import qualified Simplex.Messaging.Crypto as C
 import Simplex.Messaging.Notifications.Protocol
@@ -28,16 +31,17 @@ import Simplex.Messaging.Notifications.Server.Stats
 import Simplex.Messaging.Notifications.Server.Store.Postgres
 import Simplex.Messaging.Notifications.Server.Store.Types
 import Simplex.Messaging.Notifications.Transport (NTFVersion, VersionRangeNTF)
-import Simplex.Messaging.Protocol (BasicAuth, CorrId, Party (..), SMPServer, SParty (..), Transmission)
+import Simplex.Messaging.Protocol (BasicAuth, CorrId, Party (..), SMPServer, SParty (..), ServiceId, Transmission)
 import Simplex.Messaging.Server.Env.STM (StartOptions (..))
 import Simplex.Messaging.Server.Expiration
 import Simplex.Messaging.Server.QueueStore.Postgres.Config (PostgresStoreCfg (..))
 import Simplex.Messaging.Session
 import Simplex.Messaging.TMap (TMap)
 import qualified Simplex.Messaging.TMap as TM
 import Simplex.Messaging.Transport (ASrvTransport, SMPServiceRole (..), ServiceCredentials (..), THandleParams, TransportPeer (..))
+import Simplex.Messaging.Transport.Credentials (genCredentials, tlsCredentials)
 import Simplex.Messaging.Transport.Server (AddHTTP, ServerCredentials, TransportServerConfig, loadFingerprint, loadServerCredential)
-import System.Exit (exitFailure)
+import Simplex.Messaging.Util (liftEitherWith)
 import System.Mem.Weak (Weak)
 import UnliftIO.STM
 
@@ -95,20 +99,35 @@ newNtfServerEnv config@NtfServerConfig {pushQSize, smpAgentCfg, apnsConfig, dbSt
   random <- C.newRandom
   store <- newNtfDbStore dbStoreConfig
   tlsServerCreds <- loadServerCredential ntfCredentials
-  serviceCertHash@(XV.Fingerprint fp) <- loadFingerprint ntfCredentials
-  smpAgentCfg' <-
-    if useServiceCreds
-      then do
-        serviceSignKey <- case C.x509ToPrivate' $ snd tlsServerCreds of
-          Right pk -> pure pk
-          Left e -> putStrLn ("Server has no valid key: " <> show e) >> exitFailure
-        let service = ServiceCredentials {serviceRole = SRNotifier, serviceCreds = tlsServerCreds, serviceCertHash, serviceSignKey}
-        pure smpAgentCfg {smpCfg = (smpCfg smpAgentCfg) {serviceCredentials = Just service}}
-      else pure smpAgentCfg
-  subscriber <- newNtfSubscriber smpAgentCfg' random
+  XV.Fingerprint fp <- loadFingerprint ntfCredentials
+  let dbService = if useServiceCreds then Just $ mkDbService random store else Nothing
+  subscriber <- newNtfSubscriber smpAgentCfg dbService random
   pushServer <- newNtfPushServer pushQSize apnsConfig
   serverStats <- newNtfServerStats =<< getCurrentTime
   pure NtfEnv {config, subscriber, pushServer, store, random, tlsServerCreds, serverIdentity = C.KeyHash fp, serverStats}
+  where
+    mkDbService g store =
+      DBService
+        { getCredentials = getCredentials g store,
+          updateServiceId = updateServiceId store
+        }
+    getCredentials :: TVar ChaChaDRG -> NtfPostgresStore -> SMPServer -> IO (Either SMPClientError ServiceCredentials)
+    getCredentials g st srv = runExceptT $ do
+      ExceptT (withClientDB "" st $ \db -> getNtfServiceCredentials db srv >>= mapM (mkServiceCreds db)) >>= \case
+        Just (C.KeyHash kh, serviceCreds) -> do
+          serviceSignKey <- liftEitherWith PCEIOError $ C.x509ToPrivate' $ snd serviceCreds
+          pure ServiceCredentials {serviceRole = SRNotifier, serviceCreds, serviceCertHash = XV.Fingerprint kh, serviceSignKey}
+        Nothing -> throwE PCEServiceUnavailable -- this error cannot happen, as clients never connect to unknown servers
+      where
+        mkServiceCreds db = \case
+          (_, Just tlsCreds) -> pure tlsCreds
+          (srvId, Nothing) -> do
+            cred <- genCredentials g Nothing (25, 24 * 999999) "simplex"
+            let tlsCreds = tlsCredentials [cred]
+            setNtfServiceCredentials db srvId tlsCreds
+            pure tlsCreds
+    updateServiceId :: NtfPostgresStore -> SMPServer -> Maybe ServiceId -> IO (Either SMPClientError ())
+    updateServiceId st srv serviceId_ = withClientDB "" st $ \db -> updateNtfServiceId db srv serviceId_
 
 data NtfSubscriber = NtfSubscriber
   { smpSubscribers :: TMap SMPServer SMPSubscriberVar,
@@ -118,11 +137,11 @@ data NtfSubscriber = NtfSubscriber
 
 type SMPSubscriberVar = SessionVar SMPSubscriber
 
-newNtfSubscriber :: SMPClientAgentConfig -> TVar ChaChaDRG -> IO NtfSubscriber
-newNtfSubscriber smpAgentCfg random = do
+newNtfSubscriber :: SMPClientAgentConfig -> Maybe DBService -> TVar ChaChaDRG -> IO NtfSubscriber
+newNtfSubscriber smpAgentCfg dbService random = do
   smpSubscribers <- TM.emptyIO
   subscriberSeq <- newTVarIO 0
-  smpAgent <- newSMPClientAgent SNotifierService smpAgentCfg random
+  smpAgent <- newSMPClientAgent SNotifierService smpAgentCfg dbService random
   pure NtfSubscriber {smpSubscribers, subscriberSeq, smpAgent}
 
 data SMPSubscriber = SMPSubscriber