dandri/simplexmq
1
0
Fork 0
mirror of https://github.com/simplex-chat/simplexmq.git synced 2026-03-30 18:35:59 +00:00
Code Issues Packages Projects Releases Wiki Activity
Files
40875e319985b6dcec43a4a646aaaff7775ea751
simplexmq/src/Simplex/Messaging/Agent/Client.hs
Evgeny Poberezkin ac825b0ff3 Merge branch 'master' into rcv-services
2026-01-24 14:10:51 +00:00

2869 lines
131 KiB
Haskell
Raw Blame History

{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE CPP #-}
{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE DeriveAnyClass #-}
{-# LANGUAGE DuplicateRecordFields #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE OverloadedLists #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE StrictData #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeFamilyDependencies #-}
module Simplex.Messaging.Agent.Client
( AgentClient (..),
ProtocolTestFailure (..),
ProtocolTestStep (..),
ClntQueueReqData (..),
CQRData (..),
newAgentClient,
withConnLock,
withConnLocks,
withInvLock,
withLockMap,
withLocksMap,
getMapLock,
ipAddressProtected,
closeAgentClient,
closeProtocolServerClients,
reconnectServerClients,
reconnectSMPServer,
closeXFTPServerClient,
runSMPServerTest,
runXFTPServerTest,
runNTFServerTest,
getXFTPWorkPath,
newRcvQueue,
newRcvQueue_,
subscribeQueues,
subscribeUserServerQueues,
subscribeClientService,
processRcvServiceAssocs,
processClientNotices,
getQueueMessage,
decryptSMPMessage,
failSubscription,
addNewQueueSubscription,
getSubscriptions,
sendConfirmation,
sendInvitation,
temporaryAgentError,
temporaryOrHostError,
serverHostError,
secureQueue,
secureSndQueue,
addQueueLink,
deleteQueueLink,
secureGetQueueLink,
getQueueLink,
enableQueueNotifications,
EnableQueueNtfReq (..),
enableQueuesNtfs,
disableQueueNotifications,
DisableQueueNtfReq,
disableQueuesNtfs,
sendAgentMessage,
getQueueInfo,
agentNtfRegisterToken,
agentNtfVerifyToken,
agentNtfCheckToken,
agentNtfReplaceToken,
agentNtfDeleteToken,
agentNtfSetCronInterval,
agentNtfCreateSubscription,
agentNtfCreateSubscriptions,
agentNtfCheckSubscription,
agentNtfCheckSubscriptions,
agentNtfDeleteSubscription,
agentXFTPDownloadChunk,
agentXFTPNewChunk,
agentXFTPUploadChunk,
agentXFTPAddRecipients,
agentXFTPDeleteChunk,
agentCbDecrypt,
cryptoError,
sendAck,
suspendQueue,
deleteQueue,
deleteQueues,
logServer,
logSecret,
logSecret',
removeSubscription,
removeSubscriptions,
hasActiveSubscription,
hasPendingSubscription,
hasRemovedSubscription,
hasGetLock,
releaseGetLock,
activeClientSession,
agentClientStore,
agentDRG,
ServerQueueInfo (..),
AgentServersSummary (..),
ServerSessions (..),
SMPServerSubs (..),
getAgentSubsTotal,
getAgentServersSummary,
getAgentSubscriptions,
slowNetworkConfig,
protocolClientError,
clientServiceError,
Worker (..),
SessionVar (..),
SubscriptionsInfo (..),
SubInfo (..),
AgentOperation (..),
AgentOpState (..),
AgentState (..),
AgentLocks (..),
getAgentWorker,
getAgentWorker',
cancelWorker,
waitForWork,
hasWorkToDo,
hasWorkToDo',
withWork,
withWork_,
withWorkItems,
agentOperations,
agentOperationBracket,
waitUntilActive,
UserNetworkInfo (..),
UserNetworkType (..),
getFastNetworkConfig,
waitForUserNetwork,
isNetworkOnline,
isOnline,
throwWhenInactive,
throwWhenNoDelivery,
beginAgentOperation,
endAgentOperation,
waitUntilForeground,
waitWhileSuspended,
suspendSendingAndDatabase,
suspendOperation,
notifySuspended,
whenSuspending,
withStore,
withStore',
withStoreBatch,
withStoreBatch',
unsafeWithStore,
storeError,
notifySub,
notifySub',
userServers,
pickServer,
getNextServer,
withNextSrv,
incSMPServerStat,
incSMPServerStat',
incXFTPServerStat,
incXFTPServerStat',
incXFTPServerSizeStat,
incNtfServerStat,
incNtfServerStat',
AgentWorkersDetails (..),
getAgentWorkersDetails,
AgentWorkersSummary (..),
getAgentWorkersSummary,
AgentQueuesInfo (..),
getAgentQueuesInfo,
SMPTransportSession,
NtfTransportSession,
XFTPTransportSession,
ProxiedRelay (..),
SMPConnectedClient (..),
)
where
import Control.Applicative ((<|>))
import Control.Concurrent (ThreadId, killThread)
import Control.Concurrent.Async (Async, uninterruptibleCancel)
import Control.Concurrent.STM (retry)
import Control.Exception (AsyncException (..), BlockedIndefinitelyOnSTM (..))
import Control.Logger.Simple
import Control.Monad
import Control.Monad.Except
import Control.Monad.IO.Unlift
import Control.Monad.Reader
import Control.Monad.Trans.Except
import Crypto.Random (ChaChaDRG)
import qualified Data.Aeson as J
import qualified Data.Aeson.TH as J
import Data.Bifunctor (bimap, first, second)
import qualified Data.ByteString.Base64 as B64
import Data.ByteString.Char8 (ByteString)
import qualified Data.ByteString.Char8 as B
import Data.Composition ((.:), (.:.))
import Data.Containers.ListUtils (nubOrd)
import Data.Either (isRight, partitionEithers)
import Data.Functor (($>))
import Data.Int (Int64)
import Data.List (find, foldl')
import Data.List.NonEmpty (NonEmpty (..), (<|))
import qualified Data.List.NonEmpty as L
import Data.Map.Strict (Map)
import qualified Data.Map.Strict as M
import Data.Maybe (catMaybes, fromMaybe, isJust, isNothing, mapMaybe)
import Data.Set (Set)
import qualified Data.Set as S
import Data.Text (Text)
import Data.Text.Encoding
import Data.Time (UTCTime, addUTCTime, defaultTimeLocale, formatTime, getCurrentTime)
import Data.Time.Clock.System (getSystemTime)
import Data.Word (Word16)
import qualified Data.X509.Validation as XV
import Network.Socket (HostName)
import Simplex.FileTransfer.Client (XFTPChunkSpec (..), XFTPClient, XFTPClientConfig (..), XFTPClientError)
import qualified Simplex.FileTransfer.Client as X
import Simplex.FileTransfer.Description (ChunkReplicaId (..), FileDigest (..), kb)
import Simplex.FileTransfer.Protocol (FileInfo (..), FileResponse)
import Simplex.FileTransfer.Transport (XFTPErrorType (DIGEST), XFTPRcvChunkSpec (..), XFTPVersion)
import qualified Simplex.FileTransfer.Transport as XFTP
import Simplex.FileTransfer.Types (DeletedSndChunkReplica (..), NewSndChunkReplica (..), RcvFileChunkReplica (..), SndFileChunk (..), SndFileChunkReplica (..))
import Simplex.FileTransfer.Util (uniqueCombine)
import Simplex.Messaging.Agent.Env.SQLite
import Simplex.Messaging.Agent.Lock
import Simplex.Messaging.Agent.Protocol
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)
import Simplex.Messaging.Agent.Store.DB (SQLError)
import qualified Simplex.Messaging.Agent.Store.DB as DB
import Simplex.Messaging.Agent.Store.Entity
import Simplex.Messaging.Agent.TSessionSubs (TSessionSubs)
import qualified Simplex.Messaging.Agent.TSessionSubs as SS
import Simplex.Messaging.Client
import qualified Simplex.Messaging.Crypto as C
import Simplex.Messaging.Encoding
import Simplex.Messaging.Encoding.String
import Simplex.Messaging.Notifications.Client
import Simplex.Messaging.Notifications.Protocol
import Simplex.Messaging.Notifications.Transport (NTFVersion)
import Simplex.Messaging.Notifications.Types
import Simplex.Messaging.Parsers (defaultJSON, dropPrefix, enumJSON, parse, sumTypeJSON)
import Simplex.Messaging.Protocol
( AProtocolType (..),
BrokerMsg,
EntityId (..),
ServiceId,
ErrorType,
NetworkError (..),
MsgFlags (..),
MsgId,
NtfServer,
NtfServerWithAuth,
ProtoServer,
ProtoServerWithAuth (..),
Protocol (..),
ProtocolServer (..),
ProtocolType (..),
ProtocolTypeI (..),
QueueIdsKeys (..),
ServerNtfCreds (..),
RcvMessage (..),
RcvNtfPublicDhKey,
SMPMsgMeta (..),
SProtocolType (..),
ServiceSub (..),
ServiceSubResult (..),
ServiceSubError (..),
SndPublicAuthKey,
SubscriptionMode (..),
NewNtfCreds (..),
QueueReqData (..),
QueueLinkData,
UserProtocol,
VersionRangeSMPC,
VersionSMPC,
XFTPServer,
XFTPServerWithAuth,
pattern NoEntity,
senderCanSecure,
serviceSubResult,
)
import qualified Simplex.Messaging.Protocol as SMP
import Simplex.Messaging.Protocol.Types
import Simplex.Messaging.Server.QueueStore.QueueInfo
import Simplex.Messaging.Session
import Simplex.Messaging.SystemTime
import Simplex.Messaging.TMap (TMap)
import qualified Simplex.Messaging.TMap as TM
import Simplex.Messaging.Transport (HandshakeError (..), SMPServiceRole (..), SMPVersion, ServiceCredentials (..), SessionId, THClientService' (..), THandleAuth (..), THandleParams (sessionId, thAuth, thVersion), TransportError (..), TransportPeer (..), sndAuthKeySMPVersion, shortLinksSMPVersion, newNtfCredsSMPVersion)
import Simplex.Messaging.Transport.Client (TransportHost (..))
import Simplex.Messaging.Transport.Credentials
import Simplex.Messaging.Util
import Simplex.Messaging.Version
import System.Mem.Weak (Weak, deRefWeak)
import System.Random (randomR)
import UnliftIO (mapConcurrently, timeout)
import UnliftIO.Async (async)
import UnliftIO.Concurrent (forkIO, mkWeakThreadId)
import UnliftIO.Directory (doesFileExist, getTemporaryDirectory, removeFile)
import qualified UnliftIO.Exception as E
import UnliftIO.STM
#if !defined(dbPostgres)
import qualified Database.SQLite.Simple as SQL
#endif
type ClientVar msg = SessionVar (Either (AgentErrorType, Maybe UTCTime) (Client msg))
type SMPClientVar = ClientVar SMP.BrokerMsg
type NtfClientVar = ClientVar NtfResponse
type XFTPClientVar = ClientVar FileResponse
type NtfTransportSession = TransportSession NtfResponse
type XFTPTransportSession = TransportSession FileResponse
data AgentClient = AgentClient
{ acThread :: TVar (Maybe (Weak ThreadId)),
active :: TVar Bool,
subQ :: TBQueue ATransmission,
msgQ :: TBQueue (ServerTransmissionBatch SMPVersion ErrorType BrokerMsg),
smpServers :: TMap UserId (UserServers 'PSMP),
smpClients :: TMap SMPTransportSession SMPClientVar,
useClientServices :: TMap UserId Bool,
-- smpProxiedRelays:
-- SMPTransportSession defines connection from proxy to relay,
-- SMPServerWithAuth defines client connected to SMP proxy (with the same userId and entityId in TransportSession)
smpProxiedRelays :: TMap SMPTransportSession SMPServerWithAuth,
ntfServers :: TVar [NtfServer],
ntfClients :: TMap NtfTransportSession NtfClientVar,
xftpServers :: TMap UserId (UserServers 'PXFTP),
xftpClients :: TMap XFTPTransportSession XFTPClientVar,
useNetworkConfig :: TVar (NetworkConfig, NetworkConfig), -- (slow, fast) networks
presetDomains :: [HostName],
presetServers :: [SMPServer],
userNetworkInfo :: TVar UserNetworkInfo,
userNetworkUpdated :: TVar (Maybe UTCTime),
subscrConns :: TVar (Set ConnId),
currentSubs :: TSessionSubs,
removedSubs :: TMap (UserId, SMPServer) (TMap SMP.RecipientId SMPClientError),
clientNotices :: TMap (Maybe SMPServer) (Maybe SystemSeconds),
clientNoticesLock :: TMVar (),
workerSeq :: TVar Int,
smpDeliveryWorkers :: TMap SndQAddr (Worker, TMVar ()),
asyncCmdWorkers :: TMap (ConnId, Maybe SMPServer) Worker,
ntfNetworkOp :: TVar AgentOpState,
rcvNetworkOp :: TVar AgentOpState,
msgDeliveryOp :: TVar AgentOpState,
sndNetworkOp :: TVar AgentOpState,
databaseOp :: TVar AgentOpState,
agentState :: TVar AgentState,
getMsgLocks :: TMap (SMPServer, SMP.RecipientId) (TMVar ()),
-- locks to prevent concurrent operations with connection
connLocks :: TMap ConnId Lock,
-- locks to prevent concurrent operations with connection request invitations
invLocks :: TMap ByteString Lock,
-- lock to prevent concurrency between periodic and async connection deletions
deleteLock :: Lock,
-- smpSubWorkers for SMP servers sessions
smpSubWorkers :: TMap SMPTransportSession (SessionVar (Async ())),
clientId :: Int,
agentEnv :: Env,
proxySessTs :: TVar UTCTime,
smpServersStats :: TMap (UserId, SMPServer) AgentSMPServerStats,
xftpServersStats :: TMap (UserId, XFTPServer) AgentXFTPServerStats,
ntfServersStats :: TMap (UserId, NtfServer) AgentNtfServerStats,
srvStatsStartedAt :: TVar UTCTime
}
data SMPConnectedClient = SMPConnectedClient
{ connectedClient :: SMPClient,
proxiedRelays :: TMap SMPServer ProxiedRelayVar
}
type ProxiedRelayVar = SessionVar (Either AgentErrorType ProxiedRelay)
getAgentWorker :: (Ord k, Show k, AnyError e, MonadUnliftIO m) => String -> Bool -> AgentClient -> k -> TMap k Worker -> (Worker -> ExceptT e m ()) -> m Worker
getAgentWorker = getAgentWorker' id pure
{-# INLINE getAgentWorker #-}
getAgentWorker' :: forall a k e m. (Ord k, Show k, AnyError e, MonadUnliftIO m) => (a -> Worker) -> (Worker -> STM a) -> String -> Bool -> AgentClient -> k -> TMap k a -> (a -> ExceptT e m ()) -> m a
getAgentWorker' toW fromW name hasWork c@AgentClient {agentEnv} key ws work = do
atomically (getWorker >>= maybe createWorker whenExists) >>= \w -> runWorker w $> w
where
getWorker = TM.lookup key ws
createWorker = do
w <- fromW =<< newWorker c
TM.insert key w ws
pure w
whenExists w
| hasWork = hasWorkToDo (toW w) $> w
| otherwise = pure w
runWorker w = runWorkerAsync (toW w) runWork
where
runWork :: m ()
runWork = tryAllErrors' (work w) >>= restartOrDelete
restartOrDelete :: Either e () -> m ()
restartOrDelete e_ = do
t <- liftIO getSystemTime
let maxRestarts = maxWorkerRestartsPerMin $ config agentEnv
-- worker may terminate because it was deleted from the map (getWorker returns Nothing), then it won't restart
restart <- atomically $ getWorker >>= maybe (pure False) (shouldRestart e_ (toW w) t maxRestarts)
when restart runWork
shouldRestart e_ Worker {workerId = wId, doWork, action, restarts} t maxRestarts w'
| wId == workerId (toW w') = do
rc <- readTVar restarts
isActive <- readTVar $ active c
checkRestarts isActive $ updateRestartCount t rc
| otherwise =
pure False -- there is a new worker in the map, no action
where
checkRestarts isActive rc
| isActive && restartCount rc < maxRestarts = do
writeTVar restarts rc
hasWorkToDo' doWork
void $ tryPutTMVar action Nothing
notifyErr INTERNAL
pure True
| otherwise = do
TM.delete key ws
when isActive $ notifyErr $ CRITICAL True
pure False
where
notifyErr err = do
let e = either ((", error: " <>) . show) (\_ -> ", no error") e_
msg = "Worker " <> name <> " for " <> show key <> " terminated " <> show (restartCount rc) <> " times" <> e
writeTBQueue (subQ c) ("", "", AEvt SAEConn $ ERR $ err msg)
newWorker :: AgentClient -> STM Worker
newWorker c = do
workerId <- stateTVar (workerSeq c) $ \next -> (next, next + 1)
doWork <- newTMVar () -- new worker is created with "some work to do" (indicated by () in TMVar)
action <- newTMVar Nothing
restarts <- newTVar $ RestartCount 0 0
pure Worker {workerId, doWork, action, restarts}
runWorkerAsync :: MonadUnliftIO m => Worker -> m () -> m ()
runWorkerAsync Worker {action} work =
E.bracket
(atomically $ takeTMVar action) -- get current action, locking to avoid race conditions
(atomically . tryPutTMVar action) -- if it was running (or if start crashes), put it back and unlock (don't lock if it was just started)
(\a -> when (isNothing a) start) -- start worker if it's not running
where
start = atomically . putTMVar action . Just =<< mkWeakThreadId =<< forkIO work
data AgentOperation = AONtfNetwork | AORcvNetwork | AOMsgDelivery | AOSndNetwork | AODatabase
deriving (Eq, Show)
agentOpSel :: AgentOperation -> (AgentClient -> TVar AgentOpState)
agentOpSel = \case
AONtfNetwork -> ntfNetworkOp
AORcvNetwork -> rcvNetworkOp
AOMsgDelivery -> msgDeliveryOp
AOSndNetwork -> sndNetworkOp
AODatabase -> databaseOp
agentOperations :: [AgentClient -> TVar AgentOpState]
agentOperations = [ntfNetworkOp, rcvNetworkOp, msgDeliveryOp, sndNetworkOp, databaseOp]
data AgentOpState = AgentOpState {opSuspended :: !Bool, opsInProgress :: !Int}
data AgentState = ASForeground | ASSuspending | ASSuspended
deriving (Eq, Show)
data AgentLocks = AgentLocks
{ connLocks :: Map Text Text,
invLocks :: Map Text Text,
delLock :: Maybe Text
}
deriving (Show)
data UserNetworkInfo = UserNetworkInfo
{ networkType :: UserNetworkType,
online :: Bool
}
deriving (Show)
isNetworkOnline :: AgentClient -> STM Bool
isNetworkOnline c = isOnline <$> readTVar (userNetworkInfo c)
isOnline :: UserNetworkInfo -> Bool
isOnline UserNetworkInfo {networkType, online} = networkType /= UNNone && online
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.
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
qSize = tbqSize cfg
proxySessTs <- newTVarIO =<< getCurrentTime
acThread <- newTVarIO Nothing
active <- newTVarIO True
subQ <- newTBQueueIO qSize
msgQ <- newTBQueueIO qSize
smpServers <- newTVarIO $ M.map mkUserServers smp
smpClients <- TM.emptyIO
useClientServices <- newTVarIO useServices
smpProxiedRelays <- TM.emptyIO
ntfServers <- newTVarIO ntf
ntfClients <- TM.emptyIO
xftpServers <- newTVarIO $ M.map mkUserServers xftp
xftpClients <- TM.emptyIO
useNetworkConfig <- newTVarIO (slowNetworkConfig netCfg, netCfg)
userNetworkInfo <- newTVarIO $ UserNetworkInfo UNOther True
userNetworkUpdated <- newTVarIO Nothing
subscrConns <- newTVarIO S.empty
currentSubs <- SS.emptyIO
removedSubs <- TM.emptyIO
clientNotices <- newTVarIO notices
clientNoticesLock <- newTMVarIO ()
workerSeq <- newTVarIO 0
smpDeliveryWorkers <- TM.emptyIO
asyncCmdWorkers <- TM.emptyIO
ntfNetworkOp <- newTVarIO $ AgentOpState False 0
rcvNetworkOp <- newTVarIO $ AgentOpState False 0
msgDeliveryOp <- newTVarIO $ AgentOpState False 0
sndNetworkOp <- newTVarIO $ AgentOpState False 0
databaseOp <- newTVarIO $ AgentOpState False 0
agentState <- newTVarIO ASForeground
getMsgLocks <- TM.emptyIO
connLocks <- TM.emptyIO
invLocks <- TM.emptyIO
deleteLock <- createLockIO
smpSubWorkers <- TM.emptyIO
smpServersStats <- TM.emptyIO
xftpServersStats <- TM.emptyIO
ntfServersStats <- TM.emptyIO
srvStatsStartedAt <- newTVarIO currentTs
return
AgentClient
{ acThread,
active,
subQ,
msgQ,
smpServers,
smpClients,
useClientServices,
smpProxiedRelays,
ntfServers,
ntfClients,
xftpServers,
xftpClients,
useNetworkConfig,
presetDomains,
presetServers,
userNetworkInfo,
userNetworkUpdated,
subscrConns,
currentSubs,
removedSubs,
clientNotices,
clientNoticesLock,
workerSeq,
smpDeliveryWorkers,
asyncCmdWorkers,
ntfNetworkOp,
rcvNetworkOp,
msgDeliveryOp,
sndNetworkOp,
databaseOp,
agentState,
getMsgLocks,
connLocks,
invLocks,
deleteLock,
smpSubWorkers,
clientId,
agentEnv,
proxySessTs,
smpServersStats,
xftpServersStats,
ntfServersStats,
srvStatsStartedAt
}
slowNetworkConfig :: NetworkConfig -> NetworkConfig
slowNetworkConfig cfg@NetworkConfig {tcpConnectTimeout, tcpTimeout, tcpTimeoutPerKb} =
cfg {tcpConnectTimeout = slowTimeout tcpConnectTimeout, tcpTimeout = slowTimeout tcpTimeout, tcpTimeoutPerKb = slow tcpTimeoutPerKb}
where
slowTimeout (NetworkTimeout t1 t2) = NetworkTimeout (slow t1) (slow t2)
slow :: Integral a => a -> a
slow t = (t * 3) `div` 2
agentClientStore :: AgentClient -> DBStore
agentClientStore AgentClient {agentEnv = Env {store}} = store
{-# INLINE agentClientStore #-}
agentDRG :: AgentClient -> TVar ChaChaDRG
agentDRG AgentClient {agentEnv = Env {random}} = random
{-# INLINE agentDRG #-}
getServiceCredentials :: AgentClient -> UserId -> SMPServer -> AM (Maybe (ServiceCredentials, Maybe ServiceId))
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_) <-
withStore' c $ \db ->
getClientServiceCredentials db userId srv >>= \case
Just service -> pure service
Nothing -> do
cred <- genCredentials g Nothing (25, 24 * 999999) "simplex"
let tlsCreds = tlsCredentials [cred]
createClientService db userId srv tlsCreds
pure (tlsCreds, Nothing)
serviceSignKey <- liftEitherWith INTERNAL $ C.x509ToPrivate' $ snd serviceCreds
let creds = ServiceCredentials {serviceRole = SRMessaging, serviceCreds, serviceCertHash = XV.Fingerprint kh, serviceSignKey}
pure (creds, serviceId_)
class (Encoding err, Show err) => ProtocolServerClient v err msg | msg -> v, msg -> err where
type Client msg = c | c -> msg
getProtocolServerClient :: AgentClient -> NetworkRequestMode -> TransportSession msg -> AM (Client msg)
type ProtoClient msg = c | c -> msg
protocolClient :: Client msg -> ProtoClient msg
clientProtocolError :: HostName -> err -> AgentErrorType
closeProtocolServerClient :: ProtoClient msg -> IO ()
clientServer :: ProtoClient msg -> String
clientTransportHost :: ProtoClient msg -> TransportHost
instance ProtocolServerClient SMPVersion ErrorType BrokerMsg where
type Client BrokerMsg = SMPConnectedClient
getProtocolServerClient = getSMPServerClient
type ProtoClient BrokerMsg = ProtocolClient SMPVersion ErrorType BrokerMsg
protocolClient = connectedClient
clientProtocolError = SMP
closeProtocolServerClient = closeProtocolClient
clientServer = protocolClientServer
clientTransportHost = transportHost'
instance ProtocolServerClient NTFVersion ErrorType NtfResponse where
type Client NtfResponse = ProtocolClient NTFVersion ErrorType NtfResponse
getProtocolServerClient = getNtfServerClient
type ProtoClient NtfResponse = ProtocolClient NTFVersion ErrorType NtfResponse
protocolClient = id
clientProtocolError = NTF
closeProtocolServerClient = closeProtocolClient
clientServer = protocolClientServer
clientTransportHost = transportHost'
instance ProtocolServerClient XFTPVersion XFTPErrorType FileResponse where
type Client FileResponse = XFTPClient
getProtocolServerClient c _ = getXFTPServerClient c
type ProtoClient FileResponse = XFTPClient
protocolClient = id
clientProtocolError = XFTP
closeProtocolServerClient = X.closeXFTPClient
clientServer = X.xftpClientServer
clientTransportHost = X.xftpTransportHost
getSMPServerClient :: AgentClient -> NetworkRequestMode -> SMPTransportSession -> AM SMPConnectedClient
getSMPServerClient c@AgentClient {active, smpClients, workerSeq} nm tSess = do
unlessM (readTVarIO active) $ throwE INACTIVE
ts <- liftIO getCurrentTime
atomically (getSessVar workerSeq tSess smpClients ts)
>>= either newClient (waitForProtocolClient c nm tSess smpClients)
where
newClient v = do
prs <- liftIO TM.emptyIO
smpConnectClient c nm tSess prs v
getSMPProxyClient :: AgentClient -> NetworkRequestMode -> Maybe SMPServerWithAuth -> SMPTransportSession -> AM (SMPConnectedClient, Either AgentErrorType ProxiedRelay)
getSMPProxyClient c@AgentClient {active, smpClients, smpProxiedRelays, workerSeq} nm proxySrv_ destSess@(userId, destSrv, qId) = do
unlessM (readTVarIO active) $ throwE INACTIVE
proxySrv <- maybe (getNextServer c userId proxySrvs [destSrv]) pure proxySrv_
ts <- liftIO getCurrentTime
atomically (getClientVar proxySrv ts) >>= \(tSess, auth, v) ->
either (newProxyClient tSess auth ts) (waitForProxyClient tSess auth) v
where
getClientVar :: SMPServerWithAuth -> UTCTime -> STM (SMPTransportSession, Maybe SMP.BasicAuth, Either SMPClientVar SMPClientVar)
getClientVar proxySrv ts = do
ProtoServerWithAuth srv auth <- TM.lookup destSess smpProxiedRelays >>= maybe (TM.insert destSess proxySrv smpProxiedRelays $> proxySrv) pure
let tSess = (userId, srv, qId)
(tSess,auth,) <$> getSessVar workerSeq tSess smpClients ts
newProxyClient :: SMPTransportSession -> Maybe SMP.BasicAuth -> UTCTime -> SMPClientVar -> AM (SMPConnectedClient, Either AgentErrorType ProxiedRelay)
newProxyClient tSess auth ts v = do
prs <- liftIO TM.emptyIO
-- we do not need to check if it is a new proxied relay session,
-- as the client is just created and there are no sessions yet
rv <- atomically $ either id id <$> getSessVar workerSeq destSrv prs ts
clnt <- smpConnectClient c nm tSess prs v
(clnt,) <$> newProxiedRelay clnt auth rv
waitForProxyClient :: SMPTransportSession -> Maybe SMP.BasicAuth -> SMPClientVar -> AM (SMPConnectedClient, Either AgentErrorType ProxiedRelay)
waitForProxyClient tSess auth v = do
clnt@(SMPConnectedClient _ prs) <- waitForProtocolClient c nm tSess smpClients v
ts <- liftIO getCurrentTime
sess <-
atomically (getSessVar workerSeq destSrv prs ts)
>>= either (newProxiedRelay clnt auth) (waitForProxiedRelay tSess)
pure (clnt, sess)
newProxiedRelay :: SMPConnectedClient -> Maybe SMP.BasicAuth -> ProxiedRelayVar -> AM (Either AgentErrorType ProxiedRelay)
newProxiedRelay (SMPConnectedClient smp prs) proxyAuth rv =
tryAllErrors (liftClient SMP (clientServer smp) $ connectSMPProxiedRelay smp nm destSrv proxyAuth) >>= \case
Right sess -> do
atomically $ putTMVar (sessionVar rv) (Right sess)
pure $ Right sess
Left e -> do
atomically $ do
unless (serverHostError e) $ do
removeSessVar rv destSrv prs
TM.delete destSess smpProxiedRelays
putTMVar (sessionVar rv) (Left e)
pure $ Left e
waitForProxiedRelay :: SMPTransportSession -> ProxiedRelayVar -> AM (Either AgentErrorType ProxiedRelay)
waitForProxiedRelay (_, srv, _) rv = do
NetworkConfig {tcpConnectTimeout} <- getNetworkConfig c
sess_ <- liftIO $ netTimeoutInt tcpConnectTimeout nm `timeout` atomically (readTMVar $ sessionVar rv)
pure $ case sess_ of
Just (Right sess) -> Right sess
Just (Left e) -> Left e
Nothing -> Left $ BROKER (B.unpack $ strEncode srv) TIMEOUT
smpConnectClient :: AgentClient -> NetworkRequestMode -> SMPTransportSession -> TMap SMPServer ProxiedRelayVar -> SMPClientVar -> AM SMPConnectedClient
smpConnectClient c@AgentClient {smpClients, msgQ, proxySessTs, presetDomains} nm tSess@(userId, srv, _) prs v =
newProtocolClient c tSess smpClients connectClient v
`catchAllErrors` \e -> lift (resubscribeSMPSession c tSess) >> throwE e
where
connectClient :: SMPClientVar -> AM SMPConnectedClient
connectClient v' = do
cfg <- lift $ getClientConfig c smpCfg
g <- asks random
service <- getServiceCredentials c userId srv
let cfg' = cfg {serviceCredentials = fst <$> service}
env <- ask
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
atomically $ SS.setSessionId tSess (sessionId $ thParams smp) $ currentSubs c
updateClientService service smp
pure SMPConnectedClient {connectedClient = smp, proxiedRelays = prs}
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
(Nothing, Just _) -> logError "server returned serviceId without service credentials in request"
(Nothing, Nothing) -> pure ()
smpClientDisconnected :: AgentClient -> SMPTransportSession -> Env -> SMPClientVar -> TMap SMPServer ProxiedRelayVar -> SMPClient -> IO ()
smpClientDisconnected c@AgentClient {active, smpClients, smpProxiedRelays} tSess@(userId, srv, cId) env v prs client = do
removeClientAndSubs >>= serverDown
logInfo . decodeUtf8 $ "Agent disconnected from " <> showServer srv
where
-- 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], Maybe ServiceSub)
removeClientAndSubs = atomically $ do
removeSessVar v tSess smpClients
ifM (readTVar active) removeSubs (pure ([], [], Nothing))
where
sessId = sessionId $ thParams client
removeSubs = do
mode <- getSessionMode 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, serviceSub_)
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
mapM_ (notifySub c . SERVICE_DOWN srv) serviceSub_
unless (null qs && isNothing serviceSub_) $ do
releaseGetLocksIO c qs
mode <- getSessionModeIO c
let resubscribe
| (mode == TSMEntity) == isJust cId = resubscribeSMPSession c tSess
| otherwise = do
mapM_ (runExceptT . resubscribeClientService c tSess) serviceSub_
unless (null qs) $ void $ subscribeQueues c True qs
runReaderT resubscribe env
resubscribeSMPSession :: AgentClient -> SMPTransportSession -> AM' ()
resubscribeSMPSession c@AgentClient {smpSubWorkers, workerSeq} tSess = do
ts <- liftIO getCurrentTime
atomically (getWorkerVar ts) >>= mapM_ (either newSubWorker (\_ -> pure ()))
where
getWorkerVar ts =
ifM
(not <$> SS.hasPendingSubs tSess (currentSubs c))
(pure Nothing) -- prevent race with cleanup and adding pending queues in another call
(Just <$> getSessVar workerSeq tSess smpSubWorkers ts)
newSubWorker v = do
a <- async $ void (E.tryAny runSubWorker) >> atomically (cleanup v)
atomically $ putTMVar (sessionVar v) a
runSubWorker = do
ri <- asks $ reconnectInterval . config
withRetryForeground ri isForeground (isNetworkOnline c) $ \_ loop -> do
(pendingSubs, pendingSS) <- atomically $ SS.getPendingSubs tSess $ currentSubs c
unless (M.null pendingSubs && isNothing pendingSS) $ do
liftIO $ waitUntilForeground c
liftIO $ waitForUserNetwork c
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 ()
cleanup v = do
-- Here we wait until TMVar is not empty to prevent worker cleanup happening before worker is added to TMVar.
-- Not waiting may result in terminated worker remaining in the map.
whenM (isEmptyTMVar $ sessionVar v) retry
removeSessVar v tSess smpSubWorkers
handleNotify :: AM' () -> AM' ()
handleNotify = E.handleAny $ notifySub' c "" . ERR . INTERNAL . show
notifySub' :: forall e m. (AEntityI e, MonadIO m) => AgentClient -> ConnId -> AEvent e -> m ()
notifySub' c connId cmd = liftIO $ nonBlockingWriteTBQueue (subQ c) (B.empty, connId, AEvt (sAEntity @e) cmd)
{-# INLINE notifySub' #-}
notifySub :: MonadIO m => AgentClient -> AEvent 'AENone -> m ()
notifySub c = notifySub' c ""
{-# INLINE notifySub #-}
getNtfServerClient :: AgentClient -> NetworkRequestMode -> NtfTransportSession -> AM NtfClient
getNtfServerClient c@AgentClient {active, ntfClients, workerSeq, proxySessTs, presetDomains} nm tSess@(_, srv, _) = do
unlessM (readTVarIO active) $ throwE INACTIVE
ts <- liftIO getCurrentTime
atomically (getSessVar workerSeq tSess ntfClients ts)
>>= either
(newProtocolClient c tSess ntfClients connectClient)
(waitForProtocolClient c nm tSess ntfClients)
where
connectClient :: NtfClientVar -> AM NtfClient
connectClient v = do
cfg <- lift $ getClientConfig c ntfCfg
g <- asks random
ts <- readTVarIO proxySessTs
liftError' (protocolClientError NTF $ B.unpack $ strEncode srv) $
getProtocolClient g nm tSess cfg presetDomains Nothing ts $
clientDisconnected v
clientDisconnected :: NtfClientVar -> NtfClient -> IO ()
clientDisconnected v client = do
atomically $ removeSessVar v tSess ntfClients
atomically $ writeTBQueue (subQ c) ("", "", AEvt SAENone $ hostEvent DISCONNECT client)
logInfo . decodeUtf8 $ "Agent disconnected from " <> showServer srv
getXFTPServerClient :: AgentClient -> XFTPTransportSession -> AM XFTPClient
getXFTPServerClient c@AgentClient {active, xftpClients, workerSeq, proxySessTs, presetDomains} tSess@(_, srv, _) = do
unlessM (readTVarIO active) $ throwE INACTIVE
ts <- liftIO getCurrentTime
atomically (getSessVar workerSeq tSess xftpClients ts)
>>= either
(newProtocolClient c tSess xftpClients connectClient)
(waitForProtocolClient c NRMBackground tSess xftpClients)
where
connectClient :: XFTPClientVar -> AM XFTPClient
connectClient v = do
cfg <- asks $ xftpCfg . config
xftpNetworkConfig <- getNetworkConfig c
ts <- readTVarIO proxySessTs
liftError' (protocolClientError XFTP $ B.unpack $ strEncode srv) $
X.getXFTPClient tSess cfg {xftpNetworkConfig} presetDomains ts $
clientDisconnected v
clientDisconnected :: XFTPClientVar -> XFTPClient -> IO ()
clientDisconnected v client = do
atomically $ removeSessVar v tSess xftpClients
atomically $ writeTBQueue (subQ c) ("", "", AEvt SAENone $ hostEvent DISCONNECT client)
logInfo . decodeUtf8 $ "Agent disconnected from " <> showServer srv
waitForProtocolClient ::
(ProtocolTypeI (ProtoType msg), ProtocolServerClient v err msg) =>
AgentClient ->
NetworkRequestMode ->
TransportSession msg ->
TMap (TransportSession msg) (ClientVar msg) ->
ClientVar msg ->
AM (Client msg)
waitForProtocolClient c nm tSess@(_, srv, _) clients v = do
NetworkConfig {tcpConnectTimeout} <- getNetworkConfig c
client_ <- liftIO $ netTimeoutInt tcpConnectTimeout nm `timeout` atomically (readTMVar $ sessionVar v)
case client_ of
Just (Right smpClient) -> pure smpClient
Just (Left (e, ts_)) -> case ts_ of
Nothing -> throwE e
Just ts ->
ifM
((ts <) <$> liftIO getCurrentTime)
(atomically (removeSessVar v tSess clients) >> getProtocolServerClient c nm tSess)
(throwE e)
Nothing -> throwE $ BROKER (B.unpack $ strEncode srv) TIMEOUT
newProtocolClient ::
forall v err msg.
(ProtocolTypeI (ProtoType msg), ProtocolServerClient v err msg) =>
AgentClient ->
TransportSession msg ->
TMap (TransportSession msg) (ClientVar msg) ->
(ClientVar msg -> AM (Client msg)) ->
ClientVar msg ->
AM (Client msg)
newProtocolClient c tSess@(userId, srv, entityId_) clients connectClient v =
tryAllErrors (connectClient v) >>= \case
Right client -> do
logInfo . decodeUtf8 $ "Agent connected to " <> showServer srv <> " (user " <> bshow userId <> maybe "" (" for entity " <>) entityId_ <> ")"
atomically $ putTMVar (sessionVar v) (Right client)
liftIO $ nonBlockingWriteTBQueue (subQ c) ("", "", AEvt SAENone $ hostEvent CONNECT client)
pure client
Left e -> do
ei <- asks $ persistErrorInterval . config
if ei == 0
then atomically $ do
removeSessVar v tSess clients
putTMVar (sessionVar v) (Left (e, Nothing))
else do
ts <- addUTCTime ei <$> liftIO getCurrentTime
atomically $ putTMVar (sessionVar v) (Left (e, Just ts))
throwE e -- signal error to caller
hostEvent :: forall v err msg. (ProtocolTypeI (ProtoType msg), ProtocolServerClient v err msg) => (AProtocolType -> TransportHost -> AEvent 'AENone) -> Client msg -> AEvent 'AENone
hostEvent event = hostEvent' event . protocolClient
{-# INLINE hostEvent #-}
hostEvent' :: forall v err msg. (ProtocolTypeI (ProtoType msg), ProtocolServerClient v err msg) => (AProtocolType -> TransportHost -> AEvent 'AENone) -> ProtoClient msg -> AEvent 'AENone
hostEvent' event = event (AProtocolType $ protocolTypeI @(ProtoType msg)) . clientTransportHost
getClientConfig :: AgentClient -> (AgentConfig -> ProtocolClientConfig v) -> AM' (ProtocolClientConfig v)
getClientConfig c cfgSel = do
cfg <- asks $ cfgSel . config
networkConfig <- getNetworkConfig c
pure cfg {networkConfig}
getNetworkConfig :: MonadIO m => AgentClient -> m NetworkConfig
getNetworkConfig c = do
(slowCfg, fastCfg) <- readTVarIO $ useNetworkConfig c
UserNetworkInfo {networkType} <- readTVarIO $ userNetworkInfo c
pure $ case networkType of
UNCellular -> slowCfg
UNNone -> slowCfg
_ -> fastCfg
-- returns fast network config
getFastNetworkConfig :: AgentClient -> IO NetworkConfig
getFastNetworkConfig = fmap snd . readTVarIO . useNetworkConfig
{-# INLINE getFastNetworkConfig #-}
waitForUserNetwork :: AgentClient -> IO ()
waitForUserNetwork c =
unlessM (isOnline <$> readTVarIO (userNetworkInfo c)) $ do
delay <- registerDelay $ userNetworkInterval $ config $ agentEnv c
atomically $ unlessM (isNetworkOnline c) $ unlessM (readTVar delay) retry
closeAgentClient :: AgentClient -> IO ()
closeAgentClient c = do
atomically $ writeTVar (active c) False
closeProtocolServerClients c smpClients
closeProtocolServerClients c ntfClients
closeProtocolServerClients c xftpClients
atomically $ writeTVar (smpProxiedRelays c) M.empty
atomically (swapTVar (smpSubWorkers c) M.empty) >>= mapM_ cancelReconnect
clearWorkers smpDeliveryWorkers >>= mapM_ (cancelWorker . fst)
clearWorkers asyncCmdWorkers >>= mapM_ cancelWorker
atomically $ SS.clear $ currentSubs c
clear subscrConns
clear getMsgLocks
where
clearWorkers :: Ord k => (AgentClient -> TMap k a) -> IO (Map k a)
clearWorkers workers = atomically $ swapTVar (workers c) mempty
clear :: Monoid m => (AgentClient -> TVar m) -> IO ()
clear sel = atomically $ writeTVar (sel c) mempty
cancelReconnect :: SessionVar (Async ()) -> IO ()
cancelReconnect v = void . forkIO $ atomically (readTMVar $ sessionVar v) >>= uninterruptibleCancel
cancelWorker :: Worker -> IO ()
cancelWorker Worker {doWork, action} = do
noWorkToDo doWork
atomically (tryTakeTMVar action) >>= mapM_ (mapM_ $ deRefWeak >=> mapM_ killThread)
waitUntilActive :: AgentClient -> IO ()
waitUntilActive AgentClient {active} = unlessM (readTVarIO active) $ atomically $ unlessM (readTVar active) retry
throwWhenInactive :: AgentClient -> IO ()
throwWhenInactive c = unlessM (readTVarIO $ active c) $ E.throwIO ThreadKilled
{-# INLINE throwWhenInactive #-}
-- this function is used to remove workers once delivery is complete, not when it is removed from the map
throwWhenNoDelivery :: AgentClient -> SndQueue -> IO ()
throwWhenNoDelivery c sq =
unlessM (TM.memberIO (qAddress sq) $ smpDeliveryWorkers c) $
E.throwIO ThreadKilled
closeProtocolServerClients :: ProtocolServerClient v err msg => AgentClient -> (AgentClient -> TMap (TransportSession msg) (ClientVar msg)) -> IO ()
closeProtocolServerClients c clientsSel =
atomically (clientsSel c `swapTVar` M.empty) >>= mapM_ (forkIO . closeClient_ c)
reconnectServerClients :: ProtocolServerClient v err msg => AgentClient -> (AgentClient -> TMap (TransportSession msg) (ClientVar msg)) -> IO ()
reconnectServerClients c clientsSel =
readTVarIO (clientsSel c) >>= mapM_ (forkIO . closeClient_ c)
reconnectSMPServer :: AgentClient -> UserId -> SMPServer -> IO ()
reconnectSMPServer c userId srv = do
cs <- readTVarIO $ smpClients c
let vs = M.foldrWithKey srvClient [] cs
mapM_ (forkIO . closeClient_ c) vs
where
srvClient (userId', srv', _) v
| userId == userId' && srv == srv' = (v :)
| otherwise = id
closeClient :: ProtocolServerClient v err msg => AgentClient -> (AgentClient -> TMap (TransportSession msg) (ClientVar msg)) -> TransportSession msg -> IO ()
closeClient c clientSel tSess =
atomically (TM.lookupDelete tSess $ clientSel c) >>= mapM_ (closeClient_ c)
closeClient_ :: ProtocolServerClient v err msg => AgentClient -> ClientVar msg -> IO ()
closeClient_ c v = do
NetworkConfig {tcpConnectTimeout} <- getNetworkConfig c
E.handle (\BlockedIndefinitelyOnSTM -> pure ()) $
netTimeoutInt tcpConnectTimeout NRMBackground `timeout` atomically (readTMVar $ sessionVar v) >>= \case
Just (Right client) -> closeProtocolServerClient (protocolClient client) `catchAll_` pure ()
_ -> pure ()
closeXFTPServerClient :: AgentClient -> UserId -> XFTPServer -> FileDigest -> IO ()
closeXFTPServerClient c userId server (FileDigest chunkDigest) =
mkTransportSession c userId server chunkDigest >>= closeClient c xftpClients
withConnLock :: AgentClient -> ConnId -> Text -> AM a -> AM a
withConnLock c connId name = ExceptT . withConnLock' c connId name . runExceptT
{-# INLINE withConnLock #-}
withConnLock' :: AgentClient -> ConnId -> Text -> AM' a -> AM' a
withConnLock' _ "" _ = id
withConnLock' AgentClient {connLocks} connId name = withLockMap connLocks connId name
{-# INLINE withConnLock' #-}
withInvLock :: AgentClient -> ByteString -> Text -> AM a -> AM a
withInvLock c key name = ExceptT . withInvLock' c key name . runExceptT
{-# INLINE withInvLock #-}
withInvLock' :: AgentClient -> ByteString -> Text -> AM' a -> AM' a
withInvLock' AgentClient {invLocks} = withLockMap invLocks
{-# INLINE withInvLock' #-}
withConnLocks :: AgentClient -> Set ConnId -> Text -> AM' a -> AM' a
withConnLocks AgentClient {connLocks} = withLocksMap connLocks
{-# INLINE withConnLocks #-}
withLockMap :: (Ord k, MonadUnliftIO m) => TMap k Lock -> k -> Text -> m a -> m a
withLockMap = withGetLock . getMapLock
{-# INLINE withLockMap #-}
withLocksMap :: (Ord k, MonadUnliftIO m) => TMap k Lock -> Set k -> Text -> m a -> m a
withLocksMap = withGetLocks . getMapLock
{-# INLINE withLocksMap #-}
getMapLock :: Ord k => TMap k Lock -> k -> STM Lock
getMapLock locks key = TM.lookup key locks >>= maybe newLock pure
where
newLock = createLock >>= \l -> TM.insert key l locks $> l
withClient_ :: forall a v err msg. ProtocolServerClient v err msg => AgentClient -> NetworkRequestMode -> TransportSession msg -> (Client msg -> AM a) -> AM a
withClient_ c nm tSess@(_, srv, _) action = do
cl <- getProtocolServerClient c nm tSess
action cl `catchAllErrors` logServerError
where
logServerError :: AgentErrorType -> AM a
logServerError e = do
logServer "<--" c srv NoEntity $ bshow e
throwE e
withProxySession :: AgentClient -> NetworkRequestMode -> Maybe SMPServerWithAuth -> SMPTransportSession -> SMP.SenderId -> ByteString -> ((SMPConnectedClient, ProxiedRelay) -> AM a) -> AM a
withProxySession c nm proxySrv_ destSess@(_, destSrv, _) entId cmdStr action = do
(cl, sess_) <- getSMPProxyClient c nm proxySrv_ destSess
logServer ("--> " <> proxySrv cl <> " >") c destSrv entId cmdStr
case sess_ of
Right sess -> do
r <- action (cl, sess) `catchAllErrors` logServerError cl
logServer ("<-- " <> proxySrv cl <> " <") c destSrv entId "OK"
pure r
Left e -> logServerError cl e
where
proxySrv = showServer . protocolClientServer' . protocolClient
logServerError :: SMPConnectedClient -> AgentErrorType -> AM a
logServerError cl e = do
logServer ("<-- " <> proxySrv cl <> " <") c destSrv NoEntity $ bshow e
throwE e
withLogClient_ :: ProtocolServerClient v err msg => AgentClient -> NetworkRequestMode -> TransportSession msg -> ByteString -> ByteString -> (Client msg -> AM a) -> AM a
withLogClient_ c nm tSess@(_, srv, _) entId cmdStr action = do
logServer' "-->" c srv entId cmdStr
res <- withClient_ c nm tSess action
logServer' "<--" c srv entId "OK"
return res
withClient :: forall v err msg a. ProtocolServerClient v err msg => AgentClient -> NetworkRequestMode -> TransportSession msg -> (Client msg -> ExceptT (ProtocolClientError err) IO a) -> AM a
withClient c nm tSess action = withClient_ c nm tSess $ \client -> liftClient (clientProtocolError @v @err @msg) (clientServer $ protocolClient client) $ action client
{-# INLINE withClient #-}
withLogClient :: forall v err msg a. ProtocolServerClient v err msg => AgentClient -> NetworkRequestMode -> TransportSession msg -> ByteString -> ByteString -> (Client msg -> ExceptT (ProtocolClientError err) IO a) -> AM a
withLogClient c nm tSess entId cmdStr action = withLogClient_ c nm tSess entId cmdStr $ \client -> liftClient (clientProtocolError @v @err @msg) (clientServer $ protocolClient client) $ action client
{-# INLINE withLogClient #-}
withSMPClient :: SMPQueueRec q => AgentClient -> NetworkRequestMode -> q -> ByteString -> (SMPClient -> ExceptT SMPClientError IO a) -> AM a
withSMPClient c nm q cmdStr action = do
tSess <- mkSMPTransportSessionIO c q
withLogClient c nm tSess (unEntityId $ queueId q) cmdStr $ action . connectedClient
sendOrProxySMPMessage :: AgentClient -> NetworkRequestMode -> UserId -> SMPServer -> ConnId -> ByteString -> Maybe SMP.SndPrivateAuthKey -> SMP.SenderId -> MsgFlags -> SMP.MsgBody -> AM (Maybe SMPServer)
sendOrProxySMPMessage c nm userId destSrv connId cmdStr spKey_ senderId msgFlags msg =
fst <$> sendOrProxySMPCommand c nm userId destSrv connId cmdStr senderId sendViaProxy sendDirectly
where
sendViaProxy smp proxySess = do
atomically $ incSMPServerStat c userId destSrv sentViaProxyAttempts
atomically $ incSMPServerStat c userId (protocolClientServer' smp) sentProxiedAttempts
proxySMPMessage smp nm proxySess spKey_ senderId msgFlags msg
sendDirectly smp = do
atomically $ incSMPServerStat c userId destSrv sentDirectAttempts
sendSMPMessage smp nm spKey_ senderId msgFlags msg
sendOrProxySMPCommand ::
forall a.
AgentClient ->
NetworkRequestMode ->
UserId ->
SMPServer ->
ConnId -> -- session entity ID, for short links LinkId is used
ByteString ->
SMP.EntityId -> -- sender or link ID
(SMPClient -> ProxiedRelay -> ExceptT SMPClientError IO (Either ProxyClientError a)) ->
(SMPClient -> ExceptT SMPClientError IO a) ->
AM (Maybe SMPServer, a)
sendOrProxySMPCommand c nm userId destSrv@ProtocolServer {host = destHosts} connId cmdStr entId sendCmdViaProxy sendCmdDirectly = do
tSess <- mkTransportSession c userId destSrv connId
ifM shouldUseProxy (sendViaProxy Nothing tSess) ((Nothing,) <$> sendDirectly tSess)
where
shouldUseProxy = do
cfg <- getNetworkConfig c
case smpProxyMode cfg of
SPMAlways -> pure True
SPMUnknown -> unknownServer
SPMUnprotected
| ipAddressProtected cfg destSrv -> pure False
| otherwise -> unknownServer
SPMNever -> pure False
directAllowed = do
cfg <- getNetworkConfig c
pure $ case smpProxyFallback cfg of
SPFAllow -> True
SPFAllowProtected -> ipAddressProtected cfg destSrv
SPFProhibit -> False
unknownServer = liftIO $ maybe True (\srvs -> all (`S.notMember` knownHosts srvs) destHosts) <$> TM.lookupIO userId (smpServers c)
sendViaProxy :: Maybe SMPServerWithAuth -> SMPTransportSession -> AM (Maybe SMPServer, a)
sendViaProxy proxySrv_ destSess@(_, _, connId_) = do
r <- tryAllErrors . withProxySession c nm proxySrv_ destSess entId ("PFWD " <> cmdStr) $ \(SMPConnectedClient smp _, proxySess@ProxiedRelay {prBasicAuth}) -> do
r' <- liftClient SMP (clientServer smp) $ sendCmdViaProxy smp proxySess
let proxySrv = protocolClientServer' smp
case r' of
Right r -> pure (Just proxySrv, r)
Left proxyErr -> do
case proxyErr of
ProxyProtocolError (SMP.PROXY SMP.NO_SESSION) -> do
atomically deleteRelaySession
case proxySrv_ of
Just _ -> proxyError
-- sendViaProxy is called recursively here to re-create the session via the same server
-- to avoid failure in interactive calls that don't retry after the session disconnection.
Nothing -> sendViaProxy (Just $ ProtoServerWithAuth proxySrv prBasicAuth) destSess
_ -> proxyError
where
proxyError =
throwE
PROXY
{ proxyServer = protocolClientServer smp,
relayServer = B.unpack $ strEncode destSrv,
proxyErr
}
-- checks that the current proxied relay session is the same one that was used to send the message and removes it
deleteRelaySession =
( TM.lookup destSess (smpProxiedRelays c)
$>>= \(ProtoServerWithAuth srv _) -> tryReadSessVar (userId, srv, connId_) (smpClients c)
)
>>= \case
Just (Right (SMPConnectedClient smp' prs))
| sameClient smp' ->
tryReadSessVar destSrv prs >>= \case
Just (Right proxySess') | sameProxiedRelay proxySess' -> TM.delete destSrv prs
_ -> pure ()
_ -> pure ()
sameClient smp' = sessionId (thParams smp) == sessionId (thParams smp')
sameProxiedRelay proxySess' = prSessionId proxySess == prSessionId proxySess'
case r of
Right r'@(srv_, _) -> do
atomically $ incSMPServerStat c userId destSrv sentViaProxy
forM_ srv_ $ \proxySrv -> atomically $ incSMPServerStat c userId proxySrv sentProxied
pure r'
Left e
| serverHostError e -> ifM directAllowed ((Nothing,) <$> sendDirectly destSess) (throwE e)
| otherwise -> throwE e
sendDirectly tSess =
withLogClient_ c nm tSess (unEntityId entId) ("SEND " <> cmdStr) $ \(SMPConnectedClient smp _) -> do
tryAllErrors (liftClient SMP (clientServer smp) $ sendCmdDirectly smp) >>= \case
Right r -> r <$ atomically (incSMPServerStat c userId destSrv sentDirect)
Left e -> throwE e
ipAddressProtected :: NetworkConfig -> ProtocolServer p -> Bool
ipAddressProtected NetworkConfig {socksProxy, hostMode} (ProtocolServer _ hosts _ _) = do
isJust socksProxy || (hostMode == HMOnion && any isOnionHost hosts)
where
isOnionHost = \case THOnionHost _ -> True; _ -> False
withNtfClient :: AgentClient -> NetworkRequestMode -> NtfServer -> EntityId -> ByteString -> (NtfClient -> ExceptT NtfClientError IO a) -> AM a
withNtfClient c nm srv (EntityId entId) = withLogClient c nm (0, srv, Nothing) entId
withXFTPClient ::
ProtocolServerClient v err msg =>
AgentClient ->
(UserId, ProtoServer msg, ByteString) ->
ByteString ->
(Client msg -> ExceptT (ProtocolClientError err) IO b) ->
AM b
withXFTPClient c (userId, srv, sessEntId) cmdStr action = do
tSess <- mkTransportSession c userId srv sessEntId
withLogClient c NRMBackground tSess sessEntId cmdStr action
liftClient :: (Show err, Encoding err) => (HostName -> err -> AgentErrorType) -> HostName -> ExceptT (ProtocolClientError err) IO a -> AM a
liftClient protocolError_ = liftError . protocolClientError protocolError_
{-# INLINE liftClient #-}
protocolClientError :: (Show err, Encoding err) => (HostName -> err -> AgentErrorType) -> HostName -> ProtocolClientError err -> AgentErrorType
protocolClientError protocolError_ host = \case
PCEProtocolError e -> protocolError_ host e
PCEResponseError e -> BROKER host $ RESPONSE $ B.unpack $ smpEncode e
PCEUnexpectedResponse e -> BROKER host $ UNEXPECTED $ B.unpack e
PCEResponseTimeout -> BROKER host TIMEOUT
PCENetworkError e -> BROKER host $ NETWORK e
PCEIncompatibleHost -> BROKER host HOST
PCETransportError e -> BROKER host $ TRANSPORT e
e@PCECryptoError {} -> INTERNAL $ show e
PCEServiceUnavailable {} -> BROKER host NO_SERVICE
PCEIOError e -> BROKER host $ NETWORK $ NEConnectError e
-- it is consistent with smpClientServiceError
clientServiceError :: AgentErrorType -> Bool
clientServiceError = \case
BROKER _ NO_SERVICE -> True
SMP _ SMP.SERVICE -> True
SMP _ (SMP.PROXY (SMP.BROKER NO_SERVICE)) -> True -- for completeness, it cannot happen.
_ -> False
data ProtocolTestStep
= TSConnect
| TSDisconnect
| TSCreateQueue
| TSSecureQueue
| TSDeleteQueue
| TSCreateFile
| TSUploadFile
| TSDownloadFile
| TSCompareFile
| TSDeleteFile
| TSCreateNtfToken
| TSDeleteNtfToken
deriving (Eq, Show)
data ProtocolTestFailure = ProtocolTestFailure
{ testStep :: ProtocolTestStep,
testError :: AgentErrorType
}
deriving (Eq, Show)
runSMPServerTest :: AgentClient -> NetworkRequestMode -> UserId -> SMPServerWithAuth -> AM' (Maybe ProtocolTestFailure)
runSMPServerTest c@AgentClient {presetDomains} nm userId (ProtoServerWithAuth srv auth) = do
cfg <- getClientConfig c smpCfg
C.AuthAlg ra <- asks $ rcvAuthAlg . config
C.AuthAlg sa <- asks $ sndAuthAlg . config
g <- asks random
liftIO $ do
let tSess = (userId, srv, Nothing)
ts <- readTVarIO $ proxySessTs c
getProtocolClient g nm tSess cfg presetDomains Nothing ts (\_ -> pure ()) >>= \case
Right smp -> do
rKeys@(_, rpKey) <- atomically $ C.generateAuthKeyPair ra g
(sKey, spKey) <- atomically $ C.generateAuthKeyPair sa g
(dhKey, _) <- atomically $ C.generateKeyPair g
r <- runExceptT $ do
SMP.QIK {rcvId, sndId, queueMode} <- liftError (testErr TSCreateQueue) $ createSMPQueue smp nm Nothing rKeys dhKey auth SMSubscribe (QRMessaging Nothing) Nothing
liftError (testErr TSSecureQueue) $
case queueMode of
Just QMMessaging -> secureSndSMPQueue smp nm spKey sndId
_ -> secureSMPQueue smp nm rpKey rcvId sKey
liftError (testErr TSDeleteQueue) $ deleteSMPQueue smp nm rpKey rcvId
ok <- netTimeoutInt (tcpTimeout $ networkConfig cfg) nm `timeout` closeProtocolClient smp
pure $ either Just (const Nothing) r <|> maybe (Just (ProtocolTestFailure TSDisconnect $ BROKER addr TIMEOUT)) (const Nothing) ok
Left e -> pure (Just $ testErr TSConnect e)
where
addr = B.unpack $ strEncode srv
testErr :: ProtocolTestStep -> SMPClientError -> ProtocolTestFailure
testErr step = ProtocolTestFailure step . protocolClientError SMP addr
runXFTPServerTest :: AgentClient -> NetworkRequestMode -> UserId -> XFTPServerWithAuth -> AM' (Maybe ProtocolTestFailure)
runXFTPServerTest c@AgentClient {presetDomains} nm userId (ProtoServerWithAuth srv auth) = do
cfg <- asks $ xftpCfg . config
g <- asks random
xftpNetworkConfig <- getNetworkConfig c
workDir <- getXFTPWorkPath
filePath <- getTempFilePath workDir
rcvPath <- getTempFilePath workDir
liftIO $ do
let tSess = (userId, srv, Nothing)
ts <- readTVarIO $ proxySessTs c
X.getXFTPClient tSess cfg {xftpNetworkConfig} presetDomains ts (\_ -> pure ()) >>= \case
Right xftp -> withTestChunk filePath $ do
(sndKey, spKey) <- atomically $ C.generateAuthKeyPair C.SEd25519 g
(rcvKey, rpKey) <- atomically $ C.generateAuthKeyPair C.SEd25519 g
digest <- liftIO $ C.sha256Hash <$> B.readFile filePath
let file = FileInfo {sndKey, size = chSize, digest}
chunkSpec = X.XFTPChunkSpec {filePath, chunkOffset = 0, chunkSize = chSize}
r <- runExceptT $ do
(sId, [rId]) <- liftError (testErr TSCreateFile) $ X.createXFTPChunk xftp spKey file [rcvKey] auth
liftError (testErr TSUploadFile) $ X.uploadXFTPChunk xftp spKey sId chunkSpec
liftError (testErr TSDownloadFile) $ X.downloadXFTPChunk g xftp rpKey rId $ XFTPRcvChunkSpec rcvPath chSize digest
rcvDigest <- liftIO $ C.sha256Hash <$> B.readFile rcvPath
unless (digest == rcvDigest) $ throwE $ ProtocolTestFailure TSCompareFile $ XFTP (B.unpack $ strEncode srv) DIGEST
liftError (testErr TSDeleteFile) $ X.deleteXFTPChunk xftp spKey sId
ok <- netTimeoutInt (tcpTimeout xftpNetworkConfig) nm `timeout` X.closeXFTPClient xftp
pure $ either Just (const Nothing) r <|> maybe (Just (ProtocolTestFailure TSDisconnect $ BROKER addr TIMEOUT)) (const Nothing) ok
Left e -> pure (Just $ testErr TSConnect e)
where
addr = B.unpack $ strEncode srv
testErr :: ProtocolTestStep -> XFTPClientError -> ProtocolTestFailure
testErr step = ProtocolTestFailure step . protocolClientError XFTP addr
chSize :: Integral a => a
chSize = kb 64
getTempFilePath :: FilePath -> AM' FilePath
getTempFilePath workPath = do
ts <- liftIO getCurrentTime
let isoTime = formatTime defaultTimeLocale "%Y-%m-%dT%H%M%S.%6q" ts
uniqueCombine workPath isoTime
withTestChunk :: FilePath -> IO a -> IO a
withTestChunk fp =
E.bracket_
(createTestChunk fp)
(whenM (doesFileExist fp) $ removeFile fp `catchAll_` pure ())
-- this creates a new DRG on purpose to avoid blocking the one used in the agent
createTestChunk :: FilePath -> IO ()
createTestChunk fp = B.writeFile fp =<< atomically . C.randomBytes chSize =<< C.newRandom
runNTFServerTest :: AgentClient -> NetworkRequestMode -> UserId -> NtfServerWithAuth -> AM' (Maybe ProtocolTestFailure)
runNTFServerTest c@AgentClient {presetDomains} nm userId (ProtoServerWithAuth srv _) = do
cfg <- getClientConfig c ntfCfg
C.AuthAlg a <- asks $ rcvAuthAlg . config
g <- asks random
liftIO $ do
let tSess = (userId, srv, Nothing)
ts <- readTVarIO $ proxySessTs c
getProtocolClient g nm tSess cfg presetDomains Nothing ts (\_ -> pure ()) >>= \case
Right ntf -> do
(nKey, npKey) <- atomically $ C.generateAuthKeyPair a g
(dhKey, _) <- atomically $ C.generateKeyPair g
r <- runExceptT $ do
let deviceToken = DeviceToken PPApnsNull "test_ntf_token"
(tknId, _) <- liftError (testErr TSCreateNtfToken) $ ntfRegisterToken ntf nm npKey (NewNtfTkn deviceToken nKey dhKey)
liftError (testErr TSDeleteNtfToken) $ ntfDeleteToken ntf nm npKey tknId
ok <- netTimeoutInt (tcpTimeout $ networkConfig cfg) nm `timeout` closeProtocolClient ntf
pure $ either Just (const Nothing) r <|> maybe (Just (ProtocolTestFailure TSDisconnect $ BROKER addr TIMEOUT)) (const Nothing) ok
Left e -> pure (Just $ testErr TSConnect e)
where
addr = B.unpack $ strEncode srv
testErr :: ProtocolTestStep -> SMPClientError -> ProtocolTestFailure
testErr step = ProtocolTestFailure step . protocolClientError NTF addr
getXFTPWorkPath :: AM' FilePath
getXFTPWorkPath = do
workDir <- readTVarIO =<< asks (xftpWorkDir . xftpAgent)
maybe getTemporaryDirectory pure workDir
mkTransportSession :: MonadIO m => AgentClient -> UserId -> ProtoServer msg -> ByteString -> m (TransportSession msg)
mkTransportSession c userId srv sessEntId = mkTSession userId srv sessEntId <$> getSessionModeIO c
{-# INLINE mkTransportSession #-}
mkTSession :: UserId -> ProtoServer msg -> ByteString -> TransportSessionMode -> TransportSession msg
mkTSession userId srv sessEntId mode = (userId, srv, if mode == TSMEntity then Just sessEntId else Nothing)
{-# INLINE mkTSession #-}
mkSMPTransportSessionIO :: (SMPQueueRec q, MonadIO m) => AgentClient -> q -> m SMPTransportSession
mkSMPTransportSessionIO c q = mkSMPTSession q <$> getSessionModeIO c
{-# INLINE mkSMPTransportSessionIO #-}
mkSMPTransportSession :: SMPQueueRec q => AgentClient -> q -> STM SMPTransportSession
mkSMPTransportSession c q = mkSMPTSession q <$> getSessionMode c
{-# INLINE mkSMPTransportSession #-}
mkSMPTSession :: SMPQueueRec q => q -> TransportSessionMode -> SMPTransportSession
mkSMPTSession q = mkTSession (qUserId q) (qServer q) (qConnId q)
{-# INLINE mkSMPTSession #-}
getSessionModeIO :: MonadIO m => AgentClient -> m TransportSessionMode
getSessionModeIO = fmap (sessionMode . snd) . readTVarIO . useNetworkConfig
{-# INLINE getSessionModeIO #-}
getSessionMode :: AgentClient -> STM TransportSessionMode
getSessionMode = fmap (sessionMode . snd) . readTVar . useNetworkConfig
{-# INLINE getSessionMode #-}
newRcvQueue :: AgentClient -> NetworkRequestMode -> UserId -> ConnId -> SMPServerWithAuth -> VersionRangeSMPC -> SConnectionMode c -> Bool -> SubscriptionMode -> AM (NewRcvQueue, SMPQueueUri, SMPTransportSession, SessionId, Maybe ServiceId)
newRcvQueue c nm userId connId srv vRange cMode enableNtfs subMode = do
let qrd = case cMode of SCMInvitation -> CQRMessaging Nothing; SCMContact -> CQRContact Nothing
e2eKeys <- atomically . C.generateKeyPair =<< asks random
newRcvQueue_ c nm userId connId srv vRange qrd enableNtfs subMode Nothing e2eKeys
data ClntQueueReqData
= CQRMessaging (Maybe (CQRData (SMP.SenderId, QueueLinkData)))
| CQRContact (Maybe (CQRData (SMP.LinkId, (SMP.SenderId, QueueLinkData))))
data CQRData r = CQRData
{ linkKey :: LinkKey,
privSigKey :: C.PrivateKeyEd25519,
srvReq :: r
}
queueReqData :: ClntQueueReqData -> QueueReqData
queueReqData = \case
CQRMessaging d -> QRMessaging $ srvReq <$> d
CQRContact d -> QRContact $ srvReq <$> d
newRcvQueue_ :: AgentClient -> NetworkRequestMode -> UserId -> ConnId -> SMPServerWithAuth -> VersionRangeSMPC -> ClntQueueReqData -> Bool -> SubscriptionMode -> Maybe C.CbNonce -> C.KeyPairX25519 -> AM (NewRcvQueue, SMPQueueUri, SMPTransportSession, SessionId, Maybe ServiceId)
newRcvQueue_ c nm userId connId (ProtoServerWithAuth srv auth) vRange cqrd enableNtfs subMode nonce_ (e2eDhKey, e2ePrivKey) = do
C.AuthAlg a <- asks (rcvAuthAlg . config)
g <- asks random
rKeys@(_, rcvPrivateKey) <- atomically $ C.generateAuthKeyPair a g
(dhKey, privDhKey) <- atomically $ C.generateKeyPair g
logServer "-->" c srv NoEntity "NEW"
tSess <- mkTransportSession c userId srv connId
(thParams', ntfKeys, qik@QIK {rcvId, sndId, rcvPublicDhKey, queueMode, serviceId, serverNtfCreds}) <-
withClient c nm tSess $ \(SMPConnectedClient smp _) -> do
(ntfKeys, ntfCreds) <- liftIO $ mkNtfCreds a g smp
(thParams smp,ntfKeys,) <$> createSMPQueue smp nm nonce_ rKeys dhKey auth subMode (queueReqData cqrd) ntfCreds
let sessServiceId = (\THClientService {serviceId = sId} -> sId) <$> (clientService =<< thAuth thParams')
when (isJust serviceId && serviceId /= sessServiceId) $ logError "incorrect service ID in NEW response"
liftIO . logServer "<--" c srv NoEntity $ B.unwords ["IDS", logSecret rcvId, logSecret sndId]
shortLink <- mkShortLinkCreds thParams' qik
let rq =
RcvQueue
{ userId,
connId,
server = srv,
rcvId,
rcvPrivateKey,
rcvDhSecret = C.dh' rcvPublicDhKey privDhKey,
e2ePrivKey,
e2eDhSecret = Nothing,
sndId,
queueMode,
shortLink,
rcvServiceAssoc = isJust serviceId && serviceId == sessServiceId,
status = New,
enableNtfs,
clientNoticeId = Nothing,
dbQueueId = DBNewEntity,
primary = True,
dbReplaceQueueId = Nothing,
rcvSwchStatus = Nothing,
smpClientVersion = maxVersion vRange,
clientNtfCreds = mkClientNtfCreds ntfKeys serverNtfCreds,
deleteErrors = 0
}
qUri = SMPQueueUri vRange $ SMPQueueAddress srv sndId e2eDhKey queueMode
pure (rq, qUri, tSess, sessionId thParams', sessServiceId)
where
mkNtfCreds :: (C.AlgorithmI a, C.AuthAlgorithm a) => C.SAlgorithm a -> TVar ChaChaDRG -> SMPClient -> IO (Maybe (C.AAuthKeyPair, C.PrivateKeyX25519), Maybe NewNtfCreds)
mkNtfCreds a g smp
| enableNtfs && thVersion (thParams smp) >= newNtfCredsSMPVersion = do
authKeys@(k, _) <- atomically $ C.generateAuthKeyPair a g
(dhk, dhpk) <- atomically $ C.generateKeyPair g
pure (Just (authKeys, dhpk), Just $ NewNtfCreds k dhk)
| otherwise = pure (Nothing, Nothing)
mkClientNtfCreds :: Maybe (C.AAuthKeyPair, C.PrivateKeyX25519) -> Maybe ServerNtfCreds -> Maybe ClientNtfCreds
mkClientNtfCreds ntfKeys serverNtfCreds = case (ntfKeys, serverNtfCreds) of
(Just ((ntfPublicKey, ntfPrivateKey), dhpk), Just (ServerNtfCreds notifierId dhk')) ->
Just ClientNtfCreds {ntfPublicKey, ntfPrivateKey, notifierId, rcvNtfDhSecret = C.dh' dhk' dhpk}
_ -> Nothing
mkShortLinkCreds :: THandleParams SMPVersion 'TClient -> QueueIdsKeys -> AM (Maybe ShortLinkCreds)
mkShortLinkCreds thParams' QIK {sndId, queueMode, linkId} = case (cqrd, queueMode) of
(CQRMessaging ld, Just QMMessaging) ->
withLinkData ld $ \lnkId CQRData {linkKey, privSigKey, srvReq = (sndId', d)} ->
if sndId == sndId'
then pure $ Just $ ShortLinkCreds lnkId linkKey privSigKey Nothing (fst d)
else newErr "different sender ID"
(CQRContact ld, Just QMContact) ->
withLinkData ld $ \lnkId CQRData {linkKey, privSigKey, srvReq = (lnkId', (sndId', d))} ->
if sndId == sndId' && lnkId == lnkId'
then pure $ Just $ ShortLinkCreds lnkId linkKey privSigKey Nothing (fst d)
else newErr "different sender or link IDs"
(_, Nothing) -> case linkId of
Nothing | v < sndAuthKeySMPVersion -> pure Nothing
_ -> newErr "unexpected link ID"
_ -> newErr "unexpected queue mode"
where
v = thVersion thParams'
withLinkData :: Maybe d -> (SMP.LinkId -> d -> AM (Maybe ShortLinkCreds)) -> AM (Maybe ShortLinkCreds)
withLinkData ld_ mkLink = case (ld_, linkId) of
(Just ld, Just lnkId) -> mkLink lnkId ld
(Just _, Nothing) | v < shortLinksSMPVersion -> pure Nothing
(Nothing, Nothing) -> pure Nothing
_ -> newErr "unexpected or absent link ID"
newErr :: String -> AM (Maybe ShortLinkCreds)
newErr = throwE . BROKER (B.unpack $ strEncode srv) . UNEXPECTED . ("Create queue: " <>)
processSubResults :: AgentClient -> SMPTransportSession -> SessionId -> Maybe ServiceId -> NonEmpty (RcvQueueSub, Either SMPClientError (Maybe ServiceId)) -> STM ([RcvQueueSub], [(RcvQueueSub, Maybe ClientNotice)])
processSubResults c tSess@(userId, srv, _) sessId serviceId_ rs = do
pendingSubs <- SS.getPendingQueueSubs tSess $ currentSubs c
let (failed, subscribed@(qs, sQs), notices, ignored) = foldr (partitionResults pendingSubs) (M.empty, ([], []), [], 0) rs
unless (M.null failed) $ do
incSMPServerStat' c userId srv connSubErrs $ M.size failed
failSubscriptions c tSess failed
unless (null qs && null sQs) $ do
incSMPServerStat' c userId srv connSubscribed $ length qs + length sQs
SS.batchAddActiveSubs tSess sessId serviceId_ subscribed $ currentSubs c
unless (ignored == 0) $ incSMPServerStat' c userId srv connSubIgnored ignored
pure (sQs, notices)
where
partitionResults ::
Map SMP.RecipientId RcvQueueSub ->
(RcvQueueSub, Either SMPClientError (Maybe ServiceId)) ->
(Map SMP.RecipientId SMPClientError, ([RcvQueueSub], [RcvQueueSub]), [(RcvQueueSub, Maybe ClientNotice)], Int) ->
(Map SMP.RecipientId SMPClientError, ([RcvQueueSub], [RcvQueueSub]), [(RcvQueueSub, Maybe ClientNotice)], Int)
partitionResults pendingSubs (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, notices', ignored)
where
notices' = if isJust notice_ || isJust clientNoticeId then (rq, notice_) : notices else notices
Nothing
| temporaryClientError e -> acc
| otherwise -> (failed', subscribed, notices, ignored)
where
failed' = M.insert rcvId e failed
Right serviceId_'
| rcvId `M.member` pendingSubs ->
let subscribed' = case (serviceId_, serviceId_') of
(Just sId, Just sId') | sId == sId' -> (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
temporaryAgentError :: AgentErrorType -> Bool
temporaryAgentError = \case
BROKER _ e -> tempBrokerError e
SMP _ (SMP.PROXY (SMP.BROKER e)) -> tempBrokerError e
SMP _ (SMP.STORE _) -> True
NTF _ (SMP.STORE _) -> True
XFTP _ XFTP.TIMEOUT -> True
PROXY _ _ (ProxyProtocolError (SMP.PROXY (SMP.BROKER e))) -> tempBrokerError e
PROXY _ _ (ProxyProtocolError (SMP.PROXY SMP.NO_SESSION)) -> True
INACTIVE -> True
CRITICAL True _ -> True -- critical errors that do not show restart button are likely to be permanent
_ -> False
where
tempBrokerError = \case
NETWORK _ -> True
TIMEOUT -> True
TRANSPORT (TEHandshake BAD_SERVICE) -> True -- this error is considered temporary because it is DB error
_ -> False
temporaryOrHostError :: AgentErrorType -> Bool
temporaryOrHostError e = temporaryAgentError e || serverHostError e
{-# INLINE temporaryOrHostError #-}
serverHostError :: AgentErrorType -> Bool
serverHostError = \case
BROKER _ e -> brokerHostError e
SMP _ (SMP.PROXY (SMP.BROKER e)) -> brokerHostError e
PROXY _ _ (ProxyProtocolError (SMP.PROXY (SMP.BROKER e))) -> brokerHostError e
_ -> False
where
brokerHostError = \case
HOST -> True
SMP.TRANSPORT TEVersion -> True
_ -> False
-- | 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'))
qss <- batchQueues mkSMPTSession qs' <$> getSessionModeIO c
mapM_ (addPendingSubs c) qss
rs <- mapConcurrently (subscribeQueues_ c withEvents) qss
when withEvents $ forM_ (L.nonEmpty errs) $ notifySub c . ERRS . L.map (first qConnId)
pure $ map (second Left) errs <> concatMap L.toList rs
addPendingSubs :: AgentClient -> (SMPTransportSession, NonEmpty RcvQueueSub) -> AM' ()
addPendingSubs c (tSess, qs') = atomically $ SS.batchAddPendingSubs tSess (L.toList qs') $ currentSubs c
subscribeQueues_ :: AgentClient -> Bool -> (SMPTransportSession, NonEmpty RcvQueueSub) -> AM' (BatchResponses RcvQueueSub AgentErrorType (Maybe ServiceId))
subscribeQueues_ c withEvents qs'@(tSess@(_, srv, _), _) = do
(rs, active) <- subscribeSessQueues_ c withEvents qs'
if active
then when (hasTempErrors rs) resubscribe $> rs
else do
logWarn "subcription batch result for replaced SMP client, resubscribing"
-- we use BROKER NETWORK error here instead of the original error, so it becomes temporary.
resubscribe $> L.map (second $ Left . toNESubscribeError) rs
where
-- treating host errors as temporary here as well
hasTempErrors = any (either temporaryOrHostError (const False) . snd)
toNESubscribeError = BROKER (B.unpack $ strEncode srv) . NETWORK . NESubscribeError . show
resubscribe = resubscribeSMPSession c tSess
subscribeUserServerQueues :: AgentClient -> UserId -> SMPServer -> [RcvQueueSub] -> AM' [(RcvQueueSub, Either AgentErrorType (Maybe ServiceId))]
subscribeUserServerQueues c userId srv qs = do
mode <- getSessionModeIO c
if mode == TSMEntity
then subscribeQueues c True qs
else do
let tSess = (userId, srv, Nothing)
(errs, qs_) <- checkQueues c qs
forM_ (L.nonEmpty errs) $ notifySub c . ERRS . L.map (first qConnId)
let errs' = map (second Left) errs
case L.nonEmpty qs_ of
Just qs' -> do
atomically $ modifyTVar' (subscrConns c) (`S.union` S.fromList (map qConnId $ L.toList qs'))
addPendingSubs c (tSess, qs')
rs <- subscribeQueues_ c True (tSess, qs')
pure $ errs' <> L.toList rs
Nothing -> pure errs'
-- only "checked" queues are subscribed
checkQueues :: AgentClient -> [RcvQueueSub] -> AM' ([(RcvQueueSub, AgentErrorType)], [RcvQueueSub])
checkQueues c = fmap partitionEithers . mapM checkQueue
where
checkQueue rq = do
prohibited <- liftIO $ hasGetLock c rq
pure $ if prohibited then Left (rq, CMD PROHIBITED "checkQueues") else Right rq
-- 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')
forM_ (L.nonEmpty errs) $ notifySub c . ERRS . L.map (first qConnId)
subscribeSessQueues_ :: AgentClient -> Bool -> (SMPTransportSession, NonEmpty RcvQueueSub) -> AM' (BatchResponses RcvQueueSub AgentErrorType (Maybe ServiceId), Bool)
subscribeSessQueues_ c withEvents qs = sendClientBatch_ "SUB" False subscribe_ c NRMBackground qs
where
subscribe_ :: SMPClient -> NonEmpty RcvQueueSub -> IO (BatchResponses RcvQueueSub SMPClientError (Maybe ServiceId), Bool)
subscribe_ smp qs' = do
let (userId, srv, _) = tSess
atomically $ incSMPServerStat' c userId srv connSubAttempts $ length qs'
rs <- sendBatch (\smp' _ -> subscribeSMPQueues smp') smp NRMBackground qs'
cs_ <-
if withEvents
then Just . S.fromList . map qConnId . M.elems <$> atomically (SS.getActiveSubs tSess $ currentSubs c)
else pure Nothing
active <- E.uninterruptibleMask_ $ do
(active, (serviceQs, notices)) <- atomically $ do
r@(_, (_, notices)) <- ifM
(activeClientSession c tSess sessId)
((True,) <$> processSubResults c tSess sessId (smpClientServiceId smp) 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) ())
pure active
forM_ cs_ $ \cs -> do
let (errs, okConns) = partitionEithers $ map (\(RcvQueueSub {connId}, r) -> bimap (connId,) (const connId) r) $ L.toList rs
conns = filter (`S.notMember` cs) okConns
unless (null conns) $ notifySub c $ UP srv conns
forM_ (L.nonEmpty errs) $ \errs' -> do
let noFinalErrs = all (temporaryClientError . snd) errs'
addr = B.unpack $ strEncode srv
notifySub c $ ERRS $ L.map (second $ protocolClientError SMP addr) errs'
when (null okConns && S.null cs && noFinalErrs && active) $ liftIO $ do
-- We only close the client session that was used to subscribe.
v_ <- atomically $ ifM (activeClientSession c tSess sessId) (TM.lookupDelete tSess $ smpClients c) (pure Nothing)
mapM_ (closeClient_ c) v_
pure (rs, active)
where
tSess = transportSession' smp
sessId = sessionId $ thParams smp
processRcvServiceAssocs :: SMPQueue q => AgentClient -> [q] -> AM' ()
processRcvServiceAssocs _ [] = pure ()
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
now <- liftIO getSystemSeconds
tryAllErrors' (withStore' c $ \db -> (,) <$> updateClientNotices db tSess now notices <*> getClientNotices db presetServers) >>= \case
Right (noticeIds, clntNotices) -> atomically $ do
SS.updateClientNotices tSess noticeIds $ currentSubs c
writeTVar (clientNotices c) clntNotices
Left e -> do
logError $ "processClientNotices error: " <> tshow e
notifySub' c "" $ ERR e
resubscribeClientService :: AgentClient -> SMPTransportSession -> ServiceSub -> AM ServiceSubResult
resubscribeClientService c tSess@(userId, srv, _) serviceSub =
tryAllErrors (withServiceClient c tSess $ \smp _ -> subscribeClientService_ c True tSess smp serviceSub) >>= \case
Right r@(ServiceSubResult e _) -> case e of
Just SSErrorServiceId {} -> unassocSubscribeQueues $> r
_ -> pure r
Left e -> do
when (clientServiceError e) $ unassocSubscribeQueues
atomically $ writeTBQueue (subQ c) ("", "", AEvt SAEConn $ ERR e)
throwE e
where
unassocSubscribeQueues = do
qs <- withStore' c $ \db -> unassocUserServerRcvQueueSubs db userId srv
void $ lift $ subscribeUserServerQueues c userId srv qs
-- TODO [certs rcv] update service in the database if it has different ID and re-associate queues, and send event
subscribeClientService :: AgentClient -> Bool -> UserId -> SMPServer -> ServiceSub -> AM ServiceSubResult
subscribeClientService c withEvent userId srv (ServiceSub _ n idsHash) =
withServiceClient c tSess $ \smp smpServiceId -> do
let serviceSub = ServiceSub smpServiceId n idsHash
atomically $ SS.setPendingServiceSub tSess serviceSub $ currentSubs c
subscribeClientService_ c withEvent tSess smp serviceSub
where
tSess = (userId, srv, Nothing)
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 smpClientServiceId smp of
Just smpServiceId -> subscribe smp smpServiceId
Nothing -> throwE PCEServiceUnavailable
-- TODO [certs rcv] send subscription error event?
subscribeClientService_ :: AgentClient -> Bool -> SMPTransportSession -> SMPClient -> ServiceSub -> ExceptT SMPClientError IO ServiceSubResult
subscribeClientService_ c withEvent tSess@(_, srv, _) smp expected@(ServiceSub _ n idsHash) = do
subscribed <- subscribeService smp SMP.SRecipientService n idsHash
let sessId = sessionId $ thParams smp
r = serviceSubResult expected subscribed
atomically $ whenM (activeClientSession c tSess sessId) $
SS.setActiveServiceSub tSess sessId subscribed $ currentSubs c
when withEvent $ notifySub c $ SERVICE_UP srv r
pure r
activeClientSession :: AgentClient -> SMPTransportSession -> SessionId -> STM Bool
activeClientSession c tSess sessId = sameSess <$> tryReadSessVar tSess (smpClients c)
where
sameSess = \case
Just (Right (SMPConnectedClient smp _)) -> sessId == sessionId (thParams smp)
_ -> False
type BatchResponses q e r = NonEmpty (q, Either e r)
-- Please note: this function does not preserve order of results to be the same as the order of arguments,
-- it includes arguments in the results instead.
sendTSessionBatches :: forall q r. ByteString -> (q -> TransportSessionMode -> SMPTransportSession) -> (SMPClient -> NonEmpty q -> IO (BatchResponses q SMPClientError r)) -> AgentClient -> NetworkRequestMode -> [q] -> AM' [(q, Either AgentErrorType r)]
sendTSessionBatches statCmd mkSession action c nm qs = do
qs' <- batchQueues mkSession qs <$> getSessionModeIO c
concatMap L.toList <$> mapConcurrently (sendClientBatch statCmd action c nm) qs'
batchQueues :: (q -> TransportSessionMode -> SMPTransportSession) -> [q] -> TransportSessionMode -> [(SMPTransportSession, NonEmpty q)]
batchQueues mkSession qs mode = M.assocs $ foldr batch M.empty qs
where
batch q m =
let tSess = mkSession q mode
in M.alter (Just . maybe [q] (q <|)) tSess m
sendClientBatch :: ByteString -> (SMPClient -> NonEmpty q -> IO (BatchResponses q SMPClientError r)) -> AgentClient -> NetworkRequestMode -> (SMPTransportSession, NonEmpty q) -> AM' (BatchResponses q AgentErrorType r)
sendClientBatch statCmd action = fmap fst .:. sendClientBatch_ statCmd () (fmap (,()) .: action)
{-# INLINE sendClientBatch #-}
sendClientBatch_ :: ByteString -> res -> (SMPClient -> NonEmpty q -> IO (BatchResponses q SMPClientError r, res)) -> AgentClient -> NetworkRequestMode -> (SMPTransportSession, NonEmpty q) -> AM' (BatchResponses q AgentErrorType r, res)
sendClientBatch_ statCmd errRes action c nm (tSess@(_, srv, _), qs') =
tryAllErrors' (getSMPServerClient c nm tSess) >>= \case
Left e -> pure (L.map (,Left e) qs', errRes)
Right (SMPConnectedClient smp _) -> liftIO $ do
logServer' "-->" c srv (bshow (length qs') <> " queues") statCmd
first (L.map agentError) <$> action smp qs'
where
agentError = second . first $ protocolClientError SMP $ B.unpack (strEncode srv)
sendBatch :: SomeRcvQueue q => (SMPClient -> NetworkRequestMode -> NonEmpty (SMP.RecipientId, SMP.RcvPrivateAuthKey) -> IO (NonEmpty (Either SMPClientError a))) -> SMPClient -> NetworkRequestMode -> NonEmpty q -> IO (BatchResponses q SMPClientError a)
sendBatch smpCmdFunc smp nm qs = L.zip qs <$> smpCmdFunc smp nm (L.map queueCreds qs)
where
queueCreds q = (queueId q, rcvAuthKey q)
failSubscription :: SomeRcvQueue q => AgentClient -> SMPTransportSession -> q -> SMPClientError -> STM ()
failSubscription c tSess rq e = do
let rId = queueId rq
TM.insert rId e =<< getRemovedSubs c (qUserId rq, qServer rq)
SS.deletePendingSub tSess rId $ currentSubs c
failSubscriptions :: AgentClient -> SMPTransportSession -> Map SMP.RecipientId SMPClientError -> STM ()
failSubscriptions c tSess@(uId, srv, _) qs = do
TM.union qs =<< getRemovedSubs c (uId, srv)
SS.batchDeletePendingSubs tSess (M.keysSet qs) $ currentSubs c
getRemovedSubs :: AgentClient -> (UserId, SMPServer) -> STM (TMap SMP.RecipientId SMPClientError)
getRemovedSubs AgentClient {removedSubs} k = TM.lookup k removedSubs >>= maybe new pure
where
new = do
s <- newTVar M.empty
TM.insert k s removedSubs
pure s
addNewQueueSubscription :: AgentClient -> RcvQueue -> SMPTransportSession -> SessionId -> Maybe ServiceId -> AM' ()
addNewQueueSubscription c rq' tSess sessId serviceId_ = do
let rq = rcvQueueSub rq'
same <- atomically $ do
modifyTVar' (subscrConns c) $ S.insert $ qConnId rq
active <- activeClientSession c tSess sessId
if active
then SS.addActiveSub tSess sessId serviceId_ rq' $ currentSubs c
else SS.addPendingSub tSess rq $ currentSubs c
pure active
unless same $ resubscribeSMPSession c tSess
hasActiveSubscription :: SomeRcvQueue q => AgentClient -> q -> STM Bool
hasActiveSubscription c rq = do
tSess <- mkSMPTransportSession c rq
SS.hasActiveSub tSess (queueId rq) $ currentSubs c
{-# INLINE hasActiveSubscription #-}
hasPendingSubscription :: SomeRcvQueue q => AgentClient -> q -> STM Bool
hasPendingSubscription c rq = do
tSess <- mkSMPTransportSession c rq
SS.hasPendingSub tSess (queueId rq) $ currentSubs c
{-# INLINE hasPendingSubscription #-}
hasRemovedSubscription :: SomeRcvQueue q => AgentClient -> q -> STM (Maybe SMPClientError)
hasRemovedSubscription c rq = do
TM.lookup (qUserId rq, qServer rq) (removedSubs c) $>>= TM.lookup (queueId rq)
removeSubscription :: SomeRcvQueue q => AgentClient -> SMPTransportSession -> ConnId -> q -> STM ()
removeSubscription c tSess connId rq = do
modifyTVar' (subscrConns c) $ S.delete connId
SS.deleteSub tSess (queueId rq) $ currentSubs c
removeSubscriptions :: SomeRcvQueue q => AgentClient -> [ConnId] -> [q] -> STM ()
removeSubscriptions c connIds qs = do
unless (null connIds) $ modifyTVar' (subscrConns c) (`S.difference` (S.fromList connIds))
qss <- batchQueues mkSMPTSession qs <$> getSessionMode c
forM_ qss $ \(tSess, qs') -> SS.batchDeleteSubs tSess (L.toList qs') $ currentSubs c
getSubscriptions :: AgentClient -> IO (Set ConnId)
getSubscriptions = readTVarIO . subscrConns
{-# INLINE getSubscriptions #-}
logServer :: MonadIO m => ByteString -> AgentClient -> ProtocolServer s -> EntityId -> ByteString -> m ()
logServer dir c srv = logServer' dir c srv . unEntityId
{-# INLINE logServer #-}
logServer' :: MonadIO m => ByteString -> AgentClient -> ProtocolServer s -> ByteString -> ByteString -> m ()
logServer' dir AgentClient {clientId} srv qStr cmdStr =
logInfo . decodeUtf8 $ B.unwords ["A", "(" <> bshow clientId <> ")", dir, showServer srv, ":", logSecret' qStr, cmdStr]
showServer :: ProtocolServer s -> ByteString
showServer ProtocolServer {host, port} =
strEncode host <> B.pack (if null port then "" else ':' : port)
{-# INLINE showServer #-}
logSecret :: EntityId -> ByteString
logSecret = logSecret' . unEntityId
{-# INLINE logSecret #-}
logSecret' :: ByteString -> ByteString
logSecret' = B64.encode . B.take 3
{-# INLINE logSecret' #-}
sendConfirmation :: AgentClient -> NetworkRequestMode -> SndQueue -> ByteString -> AM (Maybe SMPServer)
sendConfirmation c nm sq@SndQueue {userId, server, connId, sndId, queueMode, sndPrivateKey, e2ePubKey = e2ePubKey@Just {}} agentConfirmation = do
let (privHdr, spKey) = if senderCanSecure queueMode then (SMP.PHEmpty, Just sndPrivateKey) else (SMP.PHConfirmation (C.toPublic sndPrivateKey), Nothing)
clientMsg = SMP.ClientMessage privHdr agentConfirmation
msg <- agentCbEncrypt sq e2ePubKey $ smpEncode clientMsg
sendOrProxySMPMessage c nm userId server connId "<CONF>" spKey sndId (MsgFlags {notification = True}) msg
sendConfirmation _ _ _ _ = throwE $ INTERNAL "sendConfirmation called without snd_queue public key(s) in the database"
sendInvitation :: AgentClient -> NetworkRequestMode -> UserId -> ConnId -> Compatible SMPQueueInfo -> Compatible VersionSMPA -> ConnectionRequestUri 'CMInvitation -> ConnInfo -> AM (Maybe SMPServer)
sendInvitation c nm userId connId (Compatible (SMPQueueInfo v SMPQueueAddress {smpServer, senderId, dhPublicKey})) (Compatible agentVersion) connReq connInfo = do
msg <- mkInvitation
sendOrProxySMPMessage c nm userId smpServer connId "<INV>" Nothing senderId (MsgFlags {notification = True}) msg
where
mkInvitation :: AM ByteString
-- this is only encrypted with per-queue E2E, not with double ratchet
mkInvitation = do
let agentEnvelope = AgentInvitation {agentVersion, connReq, connInfo}
agentCbEncryptOnce v dhPublicKey . smpEncode $
SMP.ClientMessage SMP.PHEmpty (smpEncode agentEnvelope)
getQueueMessage :: AgentClient -> RcvQueue -> AM (Maybe SMPMsgMeta)
getQueueMessage c rq@RcvQueue {server, rcvId, rcvPrivateKey} = do
atomically createTakeGetLock
msg_ <- withSMPClient c NRMInteractive rq "GET" $ \smp ->
getSMPMessage smp rcvPrivateKey rcvId
mapM decryptMeta msg_
where
decryptMeta msg@SMP.RcvMessage {msgId} = SMP.rcvMessageMeta msgId <$> decryptSMPMessage rq msg
createTakeGetLock = TM.alterF takeLock (server, rcvId) $ getMsgLocks c
where
takeLock l_ = do
l <- maybe (newTMVar ()) pure l_
takeTMVar l
pure $ Just l
{-# INLINE getQueueMessage #-}
decryptSMPMessage :: RcvQueue -> SMP.RcvMessage -> AM SMP.ClientRcvMsgBody
decryptSMPMessage rq SMP.RcvMessage {msgId, msgBody = SMP.EncRcvMsgBody body} =
liftEither $ parse SMP.clientRcvMsgBodyP (AGENT A_MESSAGE) =<< decrypt body
where
decrypt = agentCbDecrypt (rcvDhSecret rq) (C.cbNonce msgId)
secureQueue :: AgentClient -> NetworkRequestMode -> RcvQueue -> SndPublicAuthKey -> AM ()
secureQueue c nm rq@RcvQueue {rcvId, rcvPrivateKey} senderKey =
withSMPClient c nm rq "KEY <key>" $ \smp ->
secureSMPQueue smp nm rcvPrivateKey rcvId senderKey
secureSndQueue :: AgentClient -> NetworkRequestMode -> SndQueue -> AM ()
secureSndQueue c nm SndQueue {userId, connId, server, sndId, sndPrivateKey} =
void $ sendOrProxySMPCommand c nm userId server connId "SKEY <key>" sndId secureViaProxy secureDirectly
where
-- TODO track statistics
secureViaProxy smp proxySess = proxySecureSndSMPQueue smp nm proxySess sndPrivateKey sndId
secureDirectly smp = secureSndSMPQueue smp nm sndPrivateKey sndId
addQueueLink :: AgentClient -> NetworkRequestMode -> RcvQueue -> SMP.LinkId -> QueueLinkData -> AM ()
addQueueLink c nm rq@RcvQueue {rcvId, rcvPrivateKey} lnkId d =
withSMPClient c nm rq "LSET" $ \smp -> addSMPQueueLink smp nm rcvPrivateKey rcvId lnkId d
deleteQueueLink :: AgentClient -> NetworkRequestMode -> RcvQueue -> AM ()
deleteQueueLink c nm rq@RcvQueue {rcvId, rcvPrivateKey} =
withSMPClient c nm rq "LDEL" $ \smp -> deleteSMPQueueLink smp nm rcvPrivateKey rcvId
secureGetQueueLink :: AgentClient -> NetworkRequestMode -> UserId -> InvShortLink -> AM (SMP.SenderId, QueueLinkData)
secureGetQueueLink c nm userId InvShortLink {server, linkId, sndPrivateKey} =
snd <$> sendOrProxySMPCommand c nm userId server (unEntityId linkId) "LKEY <key>" linkId secureGetViaProxy secureGetDirectly
where
secureGetViaProxy smp proxySess = proxySecureGetSMPQueueLink smp nm proxySess sndPrivateKey linkId
secureGetDirectly smp = secureGetSMPQueueLink smp nm sndPrivateKey linkId
getQueueLink :: AgentClient -> NetworkRequestMode -> UserId -> SMPServer -> SMP.LinkId -> AM (SMP.SenderId, QueueLinkData)
getQueueLink c nm userId server lnkId =
snd <$> sendOrProxySMPCommand c nm userId server (unEntityId lnkId) "LGET" lnkId getViaProxy getDirectly
where
getViaProxy smp proxySess = proxyGetSMPQueueLink smp nm proxySess lnkId
getDirectly smp = getSMPQueueLink smp nm lnkId
enableQueueNotifications :: AgentClient -> RcvQueue -> SMP.NtfPublicAuthKey -> SMP.RcvNtfPublicDhKey -> AM (SMP.NotifierId, SMP.RcvNtfPublicDhKey)
enableQueueNotifications c rq@RcvQueue {rcvId, rcvPrivateKey} notifierKey rcvNtfPublicDhKey =
withSMPClient c NRMBackground rq "NKEY <nkey>" $ \smp ->
enableSMPQueueNotifications smp rcvPrivateKey rcvId notifierKey rcvNtfPublicDhKey
data EnableQueueNtfReq = EnableQueueNtfReq
{ eqnrNtfSub :: NtfSubscription,
eqnrRq :: RcvQueue,
eqnrAuthKeyPair :: C.AAuthKeyPair,
eqnrRcvKeyPair :: C.KeyPairX25519
}
enableQueuesNtfs :: AgentClient -> [EnableQueueNtfReq] -> AM' [(EnableQueueNtfReq, Either AgentErrorType (SMP.NotifierId, SMP.RcvNtfPublicDhKey))]
enableQueuesNtfs c = sendTSessionBatches "NKEY" (mkSMPTSession . eqnrRq) enableQueues_ c NRMBackground
where
enableQueues_ :: SMPClient -> NonEmpty EnableQueueNtfReq -> IO (NonEmpty (EnableQueueNtfReq, Either (ProtocolClientError ErrorType) (SMP.NotifierId, RcvNtfPublicDhKey)))
enableQueues_ smp qs' = L.zip qs' <$> enableSMPQueuesNtfs smp (L.map queueCreds qs')
queueCreds :: EnableQueueNtfReq -> (SMP.RecipientId, SMP.RcvPrivateAuthKey, SMP.NtfPublicAuthKey, SMP.RcvNtfPublicDhKey)
queueCreds EnableQueueNtfReq {eqnrRq, eqnrAuthKeyPair, eqnrRcvKeyPair} =
let RcvQueue {rcvPrivateKey, rcvId} = eqnrRq
(ntfPublicKey, _) = eqnrAuthKeyPair
(rcvNtfPubDhKey, _) = eqnrRcvKeyPair
in (rcvId, rcvPrivateKey, ntfPublicKey, rcvNtfPubDhKey)
disableQueueNotifications :: AgentClient -> RcvQueue -> AM ()
disableQueueNotifications c rq@RcvQueue {rcvId, rcvPrivateKey} =
withSMPClient c NRMBackground rq "NDEL" $ \smp ->
disableSMPQueueNotifications smp rcvPrivateKey rcvId
type DisableQueueNtfReq = (NtfSubscription, RcvQueue)
disableQueuesNtfs :: AgentClient -> [DisableQueueNtfReq] -> AM' [(DisableQueueNtfReq, Either AgentErrorType ())]
disableQueuesNtfs c = sendTSessionBatches "NDEL" (mkSMPTSession . snd) disableQueues_ c NRMBackground
where
disableQueues_ :: SMPClient -> NonEmpty DisableQueueNtfReq -> IO (NonEmpty (DisableQueueNtfReq, Either (ProtocolClientError ErrorType) ()))
disableQueues_ smp qs' = L.zip qs' <$> disableSMPQueuesNtfs smp (L.map queueCreds qs')
queueCreds :: DisableQueueNtfReq -> (SMP.RecipientId, SMP.RcvPrivateAuthKey)
queueCreds (_, RcvQueue {rcvPrivateKey, rcvId}) = (rcvId, rcvPrivateKey)
sendAck :: AgentClient -> RcvQueue -> MsgId -> AM ()
sendAck c rq@RcvQueue {rcvId, rcvPrivateKey} msgId =
withSMPClient c NRMBackground rq ("ACK:" <> logSecret' msgId) $ \smp ->
ackSMPMessage smp rcvPrivateKey rcvId msgId
hasGetLock :: SomeRcvQueue q => AgentClient -> q -> IO Bool
hasGetLock c rq =
TM.memberIO (qServer rq, queueId rq) $ getMsgLocks c
{-# INLINE hasGetLock #-}
releaseGetLock :: SomeRcvQueue q => AgentClient -> q -> STM ()
releaseGetLock c rq =
TM.lookup (qServer rq, queueId rq) (getMsgLocks c) >>= mapM_ (`tryPutTMVar` ())
{-# INLINE releaseGetLock #-}
releaseGetLocksIO :: SomeRcvQueue q => AgentClient -> [q] -> IO ()
releaseGetLocksIO _ [] = pure ()
releaseGetLocksIO c rqs = do
locks <- readTVarIO $ getMsgLocks c
forM_ rqs $ \rq ->
forM_ (M.lookup ((qServer rq, queueId rq)) locks) $ \lock ->
atomically $ tryPutTMVar lock ()
suspendQueue :: AgentClient -> NetworkRequestMode -> RcvQueue -> AM ()
suspendQueue c nm rq@RcvQueue {rcvId, rcvPrivateKey} =
withSMPClient c nm rq "OFF" $ \smp ->
suspendSMPQueue smp nm rcvPrivateKey rcvId
deleteQueue :: AgentClient -> NetworkRequestMode -> RcvQueue -> AM ()
deleteQueue c nm rq@RcvQueue {rcvId, rcvPrivateKey} = do
withSMPClient c nm rq "DEL" $ \smp ->
deleteSMPQueue smp nm rcvPrivateKey rcvId
deleteQueues :: AgentClient -> NetworkRequestMode -> [RcvQueue] -> AM' [(RcvQueue, Either AgentErrorType ())]
deleteQueues c nm = sendTSessionBatches "DEL" mkSMPTSession deleteQueues_ c nm
where
deleteQueues_ smp rqs = do
let (userId, srv, _) = transportSession' smp
atomically $ incSMPServerStat' c userId srv connDelAttempts $ length rqs
rs <- sendBatch deleteSMPQueues smp nm rqs
let successes = foldl' (\n (_, r) -> if isRight r then n + 1 else n) 0 rs
atomically $ incSMPServerStat' c userId srv connDeleted successes
pure rs
-- This is only used in background
sendAgentMessage :: AgentClient -> SndQueue -> MsgFlags -> ByteString -> AM (Maybe SMPServer)
sendAgentMessage c sq@SndQueue {userId, server, connId, sndId, sndPrivateKey} msgFlags agentMsg = do
let clientMsg = SMP.ClientMessage SMP.PHEmpty agentMsg
msg <- agentCbEncrypt sq Nothing $ smpEncode clientMsg
sendOrProxySMPMessage c NRMBackground userId server connId "<MSG>" (Just sndPrivateKey) sndId msgFlags msg
data ServerQueueInfo = ServerQueueInfo
{ server :: SMPServer,
rcvId :: Text,
sndId :: Text,
ntfId :: Maybe Text,
status :: Text,
info :: QueueInfo
}
deriving (Show)
getQueueInfo :: AgentClient -> NetworkRequestMode -> RcvQueue -> AM ServerQueueInfo
getQueueInfo c nm rq@RcvQueue {server, rcvId, rcvPrivateKey, sndId, status, clientNtfCreds} =
withSMPClient c nm rq "QUE" $ \smp -> do
info <- getSMPQueueInfo smp nm rcvPrivateKey rcvId
let ntfId = enc . (\ClientNtfCreds {notifierId} -> notifierId) <$> clientNtfCreds
pure ServerQueueInfo {server, rcvId = enc rcvId, sndId = enc sndId, ntfId, status = serializeQueueStatus status, info}
where
enc = decodeLatin1 . B64.encode . unEntityId
agentNtfRegisterToken :: AgentClient -> NetworkRequestMode -> NtfToken -> SMP.NtfPublicAuthKey -> C.PublicKeyX25519 -> AM (NtfTokenId, C.PublicKeyX25519)
agentNtfRegisterToken c nm NtfToken {deviceToken, ntfServer, ntfPrivKey} ntfPubKey pubDhKey =
withClient c nm (0, ntfServer, Nothing) $ \ntf -> ntfRegisterToken ntf nm ntfPrivKey (NewNtfTkn deviceToken ntfPubKey pubDhKey)
agentNtfVerifyToken :: AgentClient -> NetworkRequestMode -> NtfTokenId -> NtfToken -> NtfRegCode -> AM ()
agentNtfVerifyToken c nm tknId NtfToken {ntfServer, ntfPrivKey} code =
withNtfClient c nm ntfServer tknId "TVFY" $ \ntf -> ntfVerifyToken ntf nm ntfPrivKey tknId code
agentNtfCheckToken :: AgentClient -> NetworkRequestMode -> NtfTokenId -> NtfToken -> AM NtfTknStatus
agentNtfCheckToken c nm tknId NtfToken {ntfServer, ntfPrivKey} =
withNtfClient c nm ntfServer tknId "TCHK" $ \ntf -> ntfCheckToken ntf nm ntfPrivKey tknId
agentNtfReplaceToken :: AgentClient -> NetworkRequestMode -> NtfTokenId -> NtfToken -> DeviceToken -> AM ()
agentNtfReplaceToken c nm tknId NtfToken {ntfServer, ntfPrivKey} token =
withNtfClient c nm ntfServer tknId "TRPL" $ \ntf -> ntfReplaceToken ntf nm ntfPrivKey tknId token
agentNtfDeleteToken :: AgentClient -> NetworkRequestMode -> NtfServer -> C.APrivateAuthKey -> NtfTokenId -> AM ()
agentNtfDeleteToken c nm ntfServer ntfPrivKey tknId =
withNtfClient c nm ntfServer tknId "TDEL" $ \ntf -> ntfDeleteToken ntf nm ntfPrivKey tknId
-- set to 0 to disable
agentNtfSetCronInterval :: AgentClient -> NetworkRequestMode -> NtfTokenId -> NtfToken -> Word16 -> AM ()
agentNtfSetCronInterval c nm tknId NtfToken {ntfServer, ntfPrivKey} interval =
withNtfClient c nm ntfServer tknId "TCRN" $ \ntf -> ntfSetCronInterval ntf nm ntfPrivKey tknId interval
agentNtfCreateSubscription :: AgentClient -> NtfTokenId -> NtfToken -> SMPQueueNtf -> SMP.NtfPrivateAuthKey -> AM NtfSubscriptionId
agentNtfCreateSubscription c tknId NtfToken {ntfServer, ntfPrivKey} smpQueue nKey =
withNtfClient c NRMBackground ntfServer tknId "SNEW" $ \ntf -> ntfCreateSubscription ntf ntfPrivKey (NewNtfSub tknId smpQueue nKey)
agentNtfCreateSubscriptions :: AgentClient -> NtfToken -> NonEmpty (NewNtfEntity 'Subscription) -> AM' (NonEmpty (Either AgentErrorType NtfSubscriptionId))
agentNtfCreateSubscriptions = withNtfBatch "SNEW" ntfCreateSubscriptions
agentNtfCheckSubscription :: AgentClient -> NtfToken -> NtfSubscriptionId -> AM NtfSubStatus
agentNtfCheckSubscription c NtfToken {ntfServer, ntfPrivKey} subId =
withNtfClient c NRMBackground ntfServer subId "SCHK" $ \ntf -> ntfCheckSubscription ntf ntfPrivKey subId
agentNtfCheckSubscriptions :: AgentClient -> NtfToken -> NonEmpty NtfSubscriptionId -> AM' (NonEmpty (Either AgentErrorType NtfSubStatus))
agentNtfCheckSubscriptions = withNtfBatch "SCHK" ntfCheckSubscriptions
-- This batch sends all commands to one ntf server (client can only use one server at a time)
withNtfBatch ::
ByteString ->
(NtfClient -> C.APrivateAuthKey -> NonEmpty a -> IO (NonEmpty (Either NtfClientError r))) ->
AgentClient ->
NtfToken ->
NonEmpty a ->
AM' (NonEmpty (Either AgentErrorType r))
withNtfBatch cmdStr action c NtfToken {ntfServer, ntfPrivKey} subs = do
let tSess = (0, ntfServer, Nothing)
tryAllErrors' (getNtfServerClient c NRMBackground tSess) >>= \case
Left e -> pure $ L.map (\_ -> Left e) subs
Right ntf -> liftIO $ do
logServer' "-->" c ntfServer (bshow (length subs) <> " subscriptions") cmdStr
L.map agentError <$> action ntf ntfPrivKey subs
where
agentError = first $ protocolClientError NTF $ B.unpack (strEncode ntfServer)
agentNtfDeleteSubscription :: AgentClient -> NtfSubscriptionId -> NtfToken -> AM ()
agentNtfDeleteSubscription c subId NtfToken {ntfServer, ntfPrivKey} =
withNtfClient c NRMBackground ntfServer subId "SDEL" $ \ntf -> ntfDeleteSubscription ntf ntfPrivKey subId
agentXFTPDownloadChunk :: AgentClient -> UserId -> FileDigest -> RcvFileChunkReplica -> XFTPRcvChunkSpec -> AM ()
agentXFTPDownloadChunk c userId (FileDigest chunkDigest) RcvFileChunkReplica {server, replicaId = ChunkReplicaId fId, replicaKey} chunkSpec = do
g <- asks random
withXFTPClient c (userId, server, chunkDigest) "FGET" $ \xftp -> X.downloadXFTPChunk g xftp replicaKey fId chunkSpec
agentXFTPNewChunk :: AgentClient -> SndFileChunk -> Int -> XFTPServerWithAuth -> AM NewSndChunkReplica
agentXFTPNewChunk c SndFileChunk {userId, chunkSpec = XFTPChunkSpec {chunkSize}, digest = FileDigest chunkDigest} n (ProtoServerWithAuth srv auth) = do
rKeys <- xftpRcvKeys n
(sndKey, replicaKey) <- atomically . C.generateAuthKeyPair C.SEd25519 =<< asks random
let fileInfo = FileInfo {sndKey, size = chunkSize, digest = chunkDigest}
logServer "-->" c srv NoEntity "FNEW"
tSess <- mkTransportSession c userId srv chunkDigest
(sndId, rIds) <- withClient c NRMBackground tSess $ \xftp -> X.createXFTPChunk xftp replicaKey fileInfo (L.map fst rKeys) auth
logServer "<--" c srv NoEntity $ B.unwords ["SIDS", logSecret sndId]
pure NewSndChunkReplica {server = srv, replicaId = ChunkReplicaId sndId, replicaKey, rcvIdsKeys = L.toList $ xftpRcvIdsKeys rIds rKeys}
agentXFTPUploadChunk :: AgentClient -> UserId -> FileDigest -> SndFileChunkReplica -> XFTPChunkSpec -> AM ()
agentXFTPUploadChunk c userId (FileDigest chunkDigest) SndFileChunkReplica {server, replicaId = ChunkReplicaId fId, replicaKey} chunkSpec =
withXFTPClient c (userId, server, chunkDigest) "FPUT" $ \xftp -> X.uploadXFTPChunk xftp replicaKey fId chunkSpec
agentXFTPAddRecipients :: AgentClient -> UserId -> FileDigest -> SndFileChunkReplica -> Int -> AM (NonEmpty (ChunkReplicaId, C.APrivateAuthKey))
agentXFTPAddRecipients c userId (FileDigest chunkDigest) SndFileChunkReplica {server, replicaId = ChunkReplicaId fId, replicaKey} n = do
rKeys <- xftpRcvKeys n
rIds <- withXFTPClient c (userId, server, chunkDigest) "FADD" $ \xftp -> X.addXFTPRecipients xftp replicaKey fId (L.map fst rKeys)
pure $ xftpRcvIdsKeys rIds rKeys
agentXFTPDeleteChunk :: AgentClient -> UserId -> DeletedSndChunkReplica -> AM ()
agentXFTPDeleteChunk c userId DeletedSndChunkReplica {server, replicaId = ChunkReplicaId fId, replicaKey, chunkDigest = FileDigest chunkDigest} =
withXFTPClient c (userId, server, chunkDigest) "FDEL" $ \xftp -> X.deleteXFTPChunk xftp replicaKey fId
xftpRcvKeys :: Int -> AM (NonEmpty C.AAuthKeyPair)
xftpRcvKeys n = do
rKeys <- atomically . replicateM n . C.generateAuthKeyPair C.SEd25519 =<< asks random
case L.nonEmpty rKeys of
Just rKeys' -> pure rKeys'
_ -> throwE $ INTERNAL "non-positive number of recipients"
xftpRcvIdsKeys :: NonEmpty EntityId -> NonEmpty C.AAuthKeyPair -> NonEmpty (ChunkReplicaId, C.APrivateAuthKey)
xftpRcvIdsKeys rIds rKeys = L.map ChunkReplicaId rIds `L.zip` L.map snd rKeys
agentCbEncrypt :: SndQueue -> Maybe C.PublicKeyX25519 -> ByteString -> AM ByteString
agentCbEncrypt SndQueue {e2eDhSecret, smpClientVersion} e2ePubKey msg = do
cmNonce <- atomically . C.randomCbNonce =<< asks random
let paddedLen = maybe SMP.e2eEncMessageLength (const SMP.e2eEncConfirmationLength) e2ePubKey
cmEncBody <-
liftEither . first cryptoError $
C.cbEncrypt e2eDhSecret cmNonce msg paddedLen
let cmHeader = SMP.PubHeader smpClientVersion e2ePubKey
pure $ smpEncode SMP.ClientMsgEnvelope {cmHeader, cmNonce, cmEncBody}
-- add encoding as AgentInvitation'?
agentCbEncryptOnce :: VersionSMPC -> C.PublicKeyX25519 -> ByteString -> AM ByteString
agentCbEncryptOnce clientVersion dhRcvPubKey msg = do
g <- asks random
(dhSndPubKey, dhSndPrivKey) <- atomically $ C.generateKeyPair g
let e2eDhSecret = C.dh' dhRcvPubKey dhSndPrivKey
cmNonce <- atomically $ C.randomCbNonce g
cmEncBody <-
liftEither . first cryptoError $
C.cbEncrypt e2eDhSecret cmNonce msg SMP.e2eEncConfirmationLength
let cmHeader = SMP.PubHeader clientVersion (Just dhSndPubKey)
pure $ smpEncode SMP.ClientMsgEnvelope {cmHeader, cmNonce, cmEncBody}
-- | NaCl crypto-box decrypt - both for messages received from the server
-- and per-queue E2E encrypted messages from the sender that were inside.
agentCbDecrypt :: C.DhSecretX25519 -> C.CbNonce -> ByteString -> Either AgentErrorType ByteString
agentCbDecrypt dhSecret nonce msg =
first cryptoError $
C.cbDecrypt dhSecret nonce msg
cryptoError :: C.CryptoError -> AgentErrorType
cryptoError = \case
C.CryptoLargeMsgError -> CMD LARGE "CryptoLargeMsgError"
C.CryptoHeaderError _ -> AGENT A_MESSAGE -- parsing error
C.CERatchetDuplicateMessage -> AGENT A_DUPLICATE
C.AESDecryptError -> c DECRYPT_AES
C.CBDecryptError -> c DECRYPT_CB
C.CERatchetHeader -> c RATCHET_HEADER
C.CERatchetTooManySkipped n -> c $ RATCHET_SKIPPED n
C.CERatchetEarlierMessage n -> c $ RATCHET_EARLIER n
e -> INTERNAL $ show e
where
c = AGENT . A_CRYPTO
waitForWork :: MonadIO m => TMVar () -> m ()
waitForWork = void . atomically . readTMVar
{-# INLINE waitForWork #-}
withWork :: AgentClient -> TMVar () -> (DB.Connection -> IO (Either StoreError (Maybe a))) -> (a -> AM ()) -> AM ()
withWork c doWork = withWork_ c doWork . withStore' c
{-# INLINE withWork #-}
-- setting doWork flag to "no work" before getWork rather than after prevents race condition when flag is set to "has work" by another thread after getWork call.
withWork_ :: (AnyStoreError e', MonadIO m) => AgentClient -> TMVar () -> ExceptT e m (Either e' (Maybe a)) -> (a -> ExceptT e m ()) -> ExceptT e m ()
withWork_ c doWork getWork action =
noWork >> getWork >>= \case
Right (Just r) -> hasWork >> action r
Right Nothing -> pure ()
Left e
| isWorkItemError e -> notifyErr (CRITICAL False) e -- worker remains stopped here because the next iteration is likely to produce the same result
| otherwise -> hasWork >> notifyErr INTERNAL e
where
hasWork = atomically $ hasWorkToDo' doWork
noWork = liftIO $ noWorkToDo doWork
notifyErr err e = do
logError $ "withWork_ error: " <> tshow e
atomically $ writeTBQueue (subQ c) ("", "", AEvt SAEConn $ ERR $ err $ show e)
withWorkItems :: (AnyStoreError e', MonadIO m) => AgentClient -> TMVar () -> ExceptT e m (Either e' [Either e' a]) -> (NonEmpty a -> ExceptT e m ()) -> ExceptT e m ()
withWorkItems c doWork getWork action = do
noWork >> getWork >>= \case
Right [] -> pure ()
Right rs -> do
let (errs, items) = partitionEithers rs
case L.nonEmpty items of
Just items' -> hasWork >> action items'
Nothing -> do
case find isWorkItemError errs of
Nothing -> hasWork
Just err -> do
notifyErr (CRITICAL False) err
unless (all isWorkItemError errs) hasWork
forM_ (L.nonEmpty errs) $ notifySub c . ERRS . L.map (\e -> ("", INTERNAL $ show e))
Left e
| isWorkItemError e -> notifyErr (CRITICAL False) e
| otherwise -> hasWork >> notifyErr INTERNAL e
where
hasWork = atomically $ hasWorkToDo' doWork
noWork = liftIO $ noWorkToDo doWork
notifyErr err e = do
logError $ "withWorkItems error: " <> tshow e
atomically $ writeTBQueue (subQ c) ("", "", AEvt SAEConn $ ERR $ err $ show e)
noWorkToDo :: TMVar () -> IO ()
noWorkToDo = void . atomically . tryTakeTMVar
{-# INLINE noWorkToDo #-}
hasWorkToDo :: Worker -> STM ()
hasWorkToDo = hasWorkToDo' . doWork
{-# INLINE hasWorkToDo #-}
hasWorkToDo' :: TMVar () -> STM ()
hasWorkToDo' = void . (`tryPutTMVar` ())
{-# INLINE hasWorkToDo' #-}
endAgentOperation :: AgentClient -> AgentOperation -> STM ()
endAgentOperation c op = endOperation c op $ case op of
AONtfNetwork -> pure ()
AORcvNetwork ->
suspendOperation c AOMsgDelivery $
suspendSendingAndDatabase c
AOMsgDelivery ->
suspendSendingAndDatabase c
AOSndNetwork ->
suspendOperation c AODatabase $
notifySuspended c
AODatabase ->
notifySuspended c
suspendSendingAndDatabase :: AgentClient -> STM ()
suspendSendingAndDatabase c =
suspendOperation c AOSndNetwork $
suspendOperation c AODatabase $
notifySuspended c
suspendOperation :: AgentClient -> AgentOperation -> STM () -> STM ()
suspendOperation c op endedAction = do
n <- stateTVar (agentOpSel op c) $ \s -> (opsInProgress s, s {opSuspended = True})
-- unsafeIOToSTM $ putStrLn $ "suspendOperation_ " <> show op <> " " <> show n
when (n == 0) $ whenSuspending c endedAction
notifySuspended :: AgentClient -> STM ()
notifySuspended c = do
-- unsafeIOToSTM $ putStrLn "notifySuspended"
writeTBQueue (subQ c) ("", "", AEvt SAENone SUSPENDED)
writeTVar (agentState c) ASSuspended
endOperation :: AgentClient -> AgentOperation -> STM () -> STM ()
endOperation c op endedAction = do
(suspended, n) <- stateTVar (agentOpSel op c) $ \s ->
let n = max 0 (opsInProgress s - 1)
in ((opSuspended s, n), s {opsInProgress = n})
-- unsafeIOToSTM $ putStrLn $ "endOperation: " <> show op <> " " <> show suspended <> " " <> show n
when (suspended && n == 0) $ whenSuspending c endedAction
whenSuspending :: AgentClient -> STM () -> STM ()
whenSuspending c = whenM ((== ASSuspending) <$> readTVar (agentState c))
{-# INLINE whenSuspending #-}
beginAgentOperation :: AgentClient -> AgentOperation -> STM ()
beginAgentOperation c op = do
let opVar = agentOpSel op c
s <- readTVar opVar
-- unsafeIOToSTM $ putStrLn $ "beginOperation? " <> show op <> " " <> show (opsInProgress s)
when (opSuspended s) retry
-- unsafeIOToSTM $ putStrLn $ "beginOperation! " <> show op <> " " <> show (opsInProgress s + 1)
writeTVar opVar $! s {opsInProgress = opsInProgress s + 1}
agentOperationBracket :: MonadUnliftIO m => AgentClient -> AgentOperation -> (AgentClient -> IO ()) -> m a -> m a
agentOperationBracket c op check action =
E.bracket
(liftIO (check c) >> atomically (beginAgentOperation c op))
(\_ -> atomically $ endAgentOperation c op)
(const action)
waitUntilForeground :: AgentClient -> IO ()
waitUntilForeground c =
unlessM (foreground readTVarIO) $ atomically $ unlessM (foreground readTVar) retry
where
foreground :: Monad m => (TVar AgentState -> m AgentState) -> m Bool
foreground rd = (ASForeground ==) <$> rd (agentState c)
-- This function waits while agent is suspended, but will proceed while it is suspending,
-- to allow completing in-flight operations.
waitWhileSuspended :: AgentClient -> IO ()
waitWhileSuspended c =
whenM (suspended readTVarIO) $ atomically $ whenM (suspended readTVar) retry
where
suspended :: Monad m => (TVar AgentState -> m AgentState) -> m Bool
suspended rd = (ASSuspended ==) <$> rd (agentState c)
withStore' :: AgentClient -> (DB.Connection -> IO a) -> AM a
withStore' c action = withStore c $ fmap Right . action
{-# INLINE withStore' #-}
withStore :: AgentClient -> (DB.Connection -> IO (Either StoreError a)) -> AM a
withStore c action = do
st <- asks store
withExceptT storeError . ExceptT . liftIO . agentOperationBracket c AODatabase (\_ -> pure ()) $
withTransaction st action `E.catch` handleDBErrors
where
handleDBErrors :: E.SomeException -> IO (Either StoreError a)
handleDBErrors e = pure $ Left $ case E.fromException e of
Just (e' :: SQLError) ->
#if defined(dbPostgres)
SEInternal $ bshow e'
#else
let se = SQL.sqlError e'
busy = se == SQL.ErrorBusy || se == SQL.ErrorLocked
in (if busy then SEDatabaseBusy else SEInternal) $ bshow e'
#endif
Nothing -> SEInternal $ bshow e
unsafeWithStore :: AgentClient -> (DB.Connection -> IO a) -> AM' a
unsafeWithStore c action = do
st <- asks store
liftIO $ agentOperationBracket c AODatabase (\_ -> pure ()) $ withTransaction st action
withStoreBatch :: Traversable t => AgentClient -> (DB.Connection -> t (IO (Either AgentErrorType a))) -> AM' (t (Either AgentErrorType a))
withStoreBatch c actions = do
st <- asks store
liftIO . agentOperationBracket c AODatabase (\_ -> pure ()) $
withTransaction st $
mapM (`E.catch` handleInternal) . actions
where
handleInternal :: E.SomeException -> IO (Either AgentErrorType a)
handleInternal = pure . Left . INTERNAL . show
withStoreBatch' :: Traversable t => AgentClient -> (DB.Connection -> t (IO a)) -> AM' (t (Either AgentErrorType a))
withStoreBatch' c actions = withStoreBatch c (fmap (fmap Right) . actions)
{-# INLINE withStoreBatch' #-}
storeError :: StoreError -> AgentErrorType
storeError = \case
SEConnNotFound -> CONN NOT_FOUND ""
SEUserNotFound -> NO_USER
SERatchetNotFound -> CONN NOT_FOUND ""
SEConnDuplicate -> CONN DUPLICATE ""
SEBadConnType cxt CRcv -> CONN SIMPLEX cxt
SEBadConnType cxt CSnd -> CONN SIMPLEX cxt
SEInvitationNotFound cxt invId -> CMD PROHIBITED $ "SEInvitationNotFound " <> cxt <> ", invitationId = " <> show invId
-- this error is never reported as store error,
-- it is used to wrap agent operations when "transaction-like" store access is needed
-- NOTE: network IO should NOT be used inside AgentStoreMonad
SEAgentError e -> e
SEDatabaseBusy e -> CRITICAL True $ B.unpack e
e -> INTERNAL $ show e
userServers :: forall p. (ProtocolTypeI p, UserProtocol p) => AgentClient -> TMap UserId (UserServers p)
userServers c = case protocolTypeI @p of
SPSMP -> smpServers c
SPXFTP -> xftpServers c
{-# INLINE userServers #-}
pickServer :: NonEmpty (Maybe OperatorId, ProtoServerWithAuth p) -> AM (ProtoServerWithAuth p)
pickServer = \case
(_, srv) :| [] -> pure srv
servers -> do
gen <- asks randomServer
atomically $ snd . (servers L.!!) <$> stateTVar gen (randomR (0, L.length servers - 1))
getNextServer ::
(ProtocolTypeI p, UserProtocol p) =>
AgentClient ->
UserId ->
(UserServers p -> NonEmpty (Maybe OperatorId, ProtoServerWithAuth p)) ->
[ProtocolServer p] ->
AM (ProtoServerWithAuth p)
getNextServer c userId srvsSel usedSrvs = do
srvs <- getUserServers_ c userId srvsSel
snd <$> getNextServer_ srvs (usedOperatorsHosts srvs usedSrvs)
usedOperatorsHosts :: NonEmpty (Maybe OperatorId, ProtoServerWithAuth p) -> [ProtocolServer p] -> (Set (Maybe OperatorId), Set TransportHost)
usedOperatorsHosts srvs usedSrvs = (usedOperators, usedHosts)
where
usedHosts = S.unions $ map serverHosts usedSrvs
usedOperators = S.fromList $ mapMaybe usedOp $ L.toList srvs
usedOp (op, srv) = if hasUsedHost srv then Just op else Nothing
hasUsedHost (ProtoServerWithAuth srv _) = any (`S.member` usedHosts) $ serverHosts srv
getNextServer_ ::
(ProtocolTypeI p, UserProtocol p) =>
NonEmpty (Maybe OperatorId, ProtoServerWithAuth p) ->
(Set (Maybe OperatorId), Set TransportHost) ->
AM (NonEmpty (Maybe OperatorId, ProtoServerWithAuth p), ProtoServerWithAuth p)
getNextServer_ servers (usedOperators, usedHosts) = do
-- choose from servers of unused operators, when possible
let otherOpsSrvs = filterOrAll ((`S.notMember` usedOperators) . fst) servers
-- choose from servers with unused hosts when possible
unusedSrvs = filterOrAll (isUnusedServer usedHosts) otherOpsSrvs
(otherOpsSrvs,) <$> pickServer unusedSrvs
where
filterOrAll p srvs = fromMaybe srvs $ L.nonEmpty $ L.filter p srvs
isUnusedServer :: Set TransportHost -> (Maybe OperatorId, ProtoServerWithAuth p) -> Bool
isUnusedServer usedHosts (_, ProtoServerWithAuth ProtocolServer {host} _) = all (`S.notMember` usedHosts) host
getUserServers_ ::
(ProtocolTypeI p, UserProtocol p) =>
AgentClient ->
UserId ->
(UserServers p -> NonEmpty (Maybe OperatorId, ProtoServerWithAuth p)) ->
AM (NonEmpty (Maybe OperatorId, ProtoServerWithAuth p))
getUserServers_ c userId srvsSel =
liftIO (TM.lookupIO userId $ userServers c) >>= \case
Just srvs -> pure $ srvsSel srvs
_ -> throwE $ INTERNAL "unknown userId - no user servers"
-- This function checks used servers and operators every time to allow
-- changing configuration while retry look is executing.
-- This function is not thread safe.
withNextSrv ::
(ProtocolTypeI p, UserProtocol p) =>
AgentClient ->
UserId ->
(UserServers p -> NonEmpty (Maybe OperatorId, ProtoServerWithAuth p)) ->
TVar (Set TransportHost) ->
[ProtocolServer p] ->
(ProtoServerWithAuth p -> AM a) ->
AM a
withNextSrv c userId srvsSel triedHosts usedSrvs action = do
srvs <- getUserServers_ c userId srvsSel
let (usedOperators, usedHosts) = usedOperatorsHosts srvs usedSrvs
tried <- readTVarIO triedHosts
let triedOrUsed = S.union tried usedHosts
(otherOpsSrvs, srvAuth@(ProtoServerWithAuth srv _)) <- getNextServer_ srvs (usedOperators, triedOrUsed)
let newHosts = serverHosts srv
unusedSrvs = L.filter (isUnusedServer $ S.union triedOrUsed newHosts) otherOpsSrvs
!tried' = if null unusedSrvs then S.empty else S.union tried newHosts
atomically $ writeTVar triedHosts tried'
action srvAuth
incSMPServerStat :: AgentClient -> UserId -> SMPServer -> (AgentSMPServerStats -> TVar Int) -> STM ()
incSMPServerStat c userId srv sel = incSMPServerStat' c userId srv sel 1
incSMPServerStat' :: AgentClient -> UserId -> SMPServer -> (AgentSMPServerStats -> TVar Int) -> Int -> STM ()
incSMPServerStat' = incServerStat (\AgentClient {smpServersStats = s} -> s) newAgentSMPServerStats
incXFTPServerStat :: AgentClient -> UserId -> XFTPServer -> (AgentXFTPServerStats -> TVar Int) -> STM ()
incXFTPServerStat c userId srv sel = incXFTPServerStat_ c userId srv sel 1
{-# INLINE incXFTPServerStat #-}
incXFTPServerStat' :: AgentClient -> UserId -> XFTPServer -> (AgentXFTPServerStats -> TVar Int) -> Int -> STM ()
incXFTPServerStat' = incXFTPServerStat_
{-# INLINE incXFTPServerStat' #-}
incXFTPServerSizeStat :: AgentClient -> UserId -> XFTPServer -> (AgentXFTPServerStats -> TVar Int64) -> Int64 -> STM ()
incXFTPServerSizeStat = incXFTPServerStat_
{-# INLINE incXFTPServerSizeStat #-}
incXFTPServerStat_ :: Num n => AgentClient -> UserId -> XFTPServer -> (AgentXFTPServerStats -> TVar n) -> n -> STM ()
incXFTPServerStat_ = incServerStat (\AgentClient {xftpServersStats = s} -> s) newAgentXFTPServerStats
{-# INLINE incXFTPServerStat_ #-}
incNtfServerStat :: AgentClient -> UserId -> NtfServer -> (AgentNtfServerStats -> TVar Int) -> STM ()
incNtfServerStat c userId srv sel = incNtfServerStat' c userId srv sel 1
{-# INLINE incNtfServerStat #-}
incNtfServerStat' :: AgentClient -> UserId -> NtfServer -> (AgentNtfServerStats -> TVar Int) -> Int -> STM ()
incNtfServerStat' = incServerStat (\AgentClient {ntfServersStats = s} -> s) newAgentNtfServerStats
{-# INLINE incNtfServerStat' #-}
incServerStat :: Num n => (AgentClient -> TMap (UserId, ProtocolServer p) s) -> STM s -> AgentClient -> UserId -> ProtocolServer p -> (s -> TVar n) -> n -> STM ()
incServerStat statsSel mkNewStats c userId srv sel n = do
TM.lookup (userId, srv) (statsSel c) >>= \case
Just v -> modifyTVar' (sel v) (+ n)
Nothing -> do
newStats <- mkNewStats
modifyTVar' (sel newStats) (+ n)
TM.insert (userId, srv) newStats (statsSel c)
data AgentServersSummary = AgentServersSummary
{ smpServersStats :: Map (UserId, SMPServer) AgentSMPServerStatsData,
xftpServersStats :: Map (UserId, XFTPServer) AgentXFTPServerStatsData,
ntfServersStats :: Map (UserId, NtfServer) AgentNtfServerStatsData,
statsStartedAt :: UTCTime,
smpServersSessions :: Map (UserId, SMPServer) ServerSessions,
smpServersSubs :: Map (UserId, SMPServer) SMPServerSubs,
xftpServersSessions :: Map (UserId, XFTPServer) ServerSessions,
xftpRcvInProgress :: [XFTPServer],
xftpSndInProgress :: [XFTPServer],
xftpDelInProgress :: [XFTPServer],
ntfServersSessions :: Map (UserId, NtfServer) ServerSessions
}
deriving (Show)
data SMPServerSubs = SMPServerSubs
{ ssActive :: Int, -- based on activeSubs
ssPending :: Int -- based on pendingSubs
}
deriving (Show)
data ServerSessions = ServerSessions
{ ssConnected :: Int,
ssErrors :: Int,
ssConnecting :: Int
}
deriving (Show)
getAgentSubsTotal :: AgentClient -> [UserId] -> IO (SMPServerSubs, Bool)
getAgentSubsTotal c userIds = do
(ssActive, ssPending) <- SS.foldSessionSubs addSub (0, 0) $ currentSubs c
sess <- hasSession . M.toList =<< readTVarIO (smpClients c)
pure (SMPServerSubs {ssActive, ssPending}, sess)
where
addSub :: (Int, Int) -> (SMPTransportSession, SS.SessSubs) -> IO (Int, Int)
addSub acc@(!ssActive, !ssPending) ((userId, _, _), s)
| userId `elem` userIds = do
(active, pending) <- SS.mapSubs M.size s
pure (ssActive + active, ssPending + pending)
| otherwise = pure acc
hasSession :: [(SMPTransportSession, SMPClientVar)] -> IO Bool
hasSession = \case
[] -> pure False
(s : ss) -> ifM (isConnected s) (pure True) (hasSession ss)
isConnected ((userId, _, _), SessionVar {sessionVar})
| userId `elem` userIds = atomically $ maybe False isRight <$> tryReadTMVar sessionVar
| otherwise = pure False
getAgentServersSummary :: AgentClient -> IO AgentServersSummary
getAgentServersSummary c@AgentClient {smpServersStats, xftpServersStats, ntfServersStats, srvStatsStartedAt, agentEnv} = do
sss <- mapM getAgentSMPServerStats =<< readTVarIO smpServersStats
xss <- mapM getAgentXFTPServerStats =<< readTVarIO xftpServersStats
nss <- mapM getAgentNtfServerStats =<< readTVarIO ntfServersStats
statsStartedAt <- readTVarIO srvStatsStartedAt
smpServersSessions <- countSessions =<< readTVarIO (smpClients c)
smpServersSubs <- getServerSubs
xftpServersSessions <- countSessions =<< readTVarIO (xftpClients c)
xftpRcvInProgress <- catMaybes <$> getXFTPWorkerSrvs xftpRcvWorkers
xftpSndInProgress <- catMaybes <$> getXFTPWorkerSrvs xftpSndWorkers
xftpDelInProgress <- getXFTPWorkerSrvs xftpDelWorkers
ntfServersSessions <- countSessions =<< readTVarIO (ntfClients c)
pure
AgentServersSummary
{ smpServersStats = sss,
xftpServersStats = xss,
ntfServersStats = nss,
statsStartedAt,
smpServersSessions,
smpServersSubs,
xftpServersSessions,
xftpRcvInProgress,
xftpSndInProgress,
xftpDelInProgress,
ntfServersSessions
}
where
getServerSubs = SS.foldSessionSubs addSub M.empty $ currentSubs c
where
addSub subs ((userId, srv, _), s) = do
(active, pending) <- SS.mapSubs M.size s
let add ss = ss {ssActive = ssActive ss + active, ssPending = ssPending ss + pending}
pure $ M.alter (Just . add . fromMaybe (SMPServerSubs 0 0)) (userId, srv) subs
Env {xftpAgent = XFTPAgent {xftpRcvWorkers, xftpSndWorkers, xftpDelWorkers}} = agentEnv
getXFTPWorkerSrvs workers = foldM addSrv [] . M.toList =<< readTVarIO workers
where
addSrv acc (srv, Worker {doWork}) = do
hasWork <- atomically $ not <$> isEmptyTMVar doWork
pure $ if hasWork then srv : acc else acc
countSessions :: Map (TransportSession msg) (ClientVar msg) -> IO (Map (UserId, ProtoServer msg) ServerSessions)
countSessions = foldM addClient M.empty . M.toList
where
addClient !acc ((userId, srv, _), SessionVar {sessionVar}) = do
c_ <- atomically $ tryReadTMVar sessionVar
pure $ M.alter (Just . add c_) (userId, srv) acc
where
add c_ = modifySessions c_ . fromMaybe ServerSessions {ssConnected = 0, ssErrors = 0, ssConnecting = 0}
modifySessions c_ ss = case c_ of
Just (Right _) -> ss {ssConnected = ssConnected ss + 1}
Just (Left _) -> ss {ssErrors = ssErrors ss + 1}
Nothing -> ss {ssConnecting = ssConnecting ss + 1}
data SubInfo = SubInfo {userId :: UserId, server :: Text, rcvId :: Text, subError :: Maybe String}
deriving (Show)
data SubscriptionsInfo = SubscriptionsInfo
{ activeSubscriptions :: [SubInfo],
pendingSubscriptions :: [SubInfo],
removedSubscriptions :: [SubInfo]
}
deriving (Show)
getAgentSubscriptions :: AgentClient -> IO SubscriptionsInfo
getAgentSubscriptions c = do
(activeSubscriptions, pendingSubscriptions) <- SS.foldSessionSubs addSubs ([], []) $ currentSubs c
removedSubscriptions <- getRemoved
pure $ SubscriptionsInfo {activeSubscriptions, pendingSubscriptions, removedSubscriptions}
where
addSubs :: ([SubInfo], [SubInfo]) -> (SMPTransportSession, SS.SessSubs) -> IO ([SubInfo], [SubInfo])
addSubs (active, pending) ((userId, srv, _), s) = do
(active', pending') <- SS.mapSubs (map (\rId -> subInfo (userId, srv, rId) Nothing) . M.keys) s
pure (active' ++ active, pending' ++ pending)
getRemoved :: IO [SubInfo]
getRemoved = foldM addSubInfo [] . M.assocs =<< readTVarIO (removedSubs c)
where
addSubInfo :: [SubInfo] -> ((UserId, SMPServer), TMap SMP.RecipientId SMPClientError) -> IO [SubInfo]
addSubInfo ss ((uId, srv), errs) = do
ss' <- map (\(rId, e) -> subInfo (uId, srv, rId) (Just e)) . M.assocs <$> readTVarIO errs
pure $ ss' ++ ss
subInfo :: (UserId, SMPServer, SMP.RecipientId) -> Maybe SMPClientError -> SubInfo
subInfo (uId, srv, rId) err = SubInfo {userId = uId, server = enc srv, rcvId = enc rId, subError = show <$> err}
enc :: StrEncoding a => a -> Text
enc = decodeLatin1 . strEncode
data AgentWorkersDetails = AgentWorkersDetails
{ smpClients_ :: [Text],
ntfClients_ :: [Text],
xftpClients_ :: [Text],
smpDeliveryWorkers_ :: Map Text WorkersDetails,
asyncCmdWorkers_ :: Map Text WorkersDetails,
smpSubWorkers_ :: [Text],
ntfWorkers_ :: Map Text WorkersDetails,
ntfSMPWorkers_ :: Map Text WorkersDetails,
xftpRcvWorkers_ :: Map Text WorkersDetails,
xftpSndWorkers_ :: Map Text WorkersDetails,
xftpDelWorkers_ :: Map Text WorkersDetails
}
deriving (Show)
data WorkersDetails = WorkersDetails
{ restarts :: Int,
hasWork :: Bool,
hasAction :: Bool
}
deriving (Show)
getAgentWorkersDetails :: AgentClient -> IO AgentWorkersDetails
getAgentWorkersDetails AgentClient {smpClients, ntfClients, xftpClients, smpDeliveryWorkers, asyncCmdWorkers, smpSubWorkers, agentEnv} = do
smpClients_ <- textKeys <$> readTVarIO smpClients
ntfClients_ <- textKeys <$> readTVarIO ntfClients
xftpClients_ <- textKeys <$> readTVarIO xftpClients
smpDeliveryWorkers_ <- workerStats . fmap fst =<< readTVarIO smpDeliveryWorkers
asyncCmdWorkers_ <- workerStats =<< readTVarIO asyncCmdWorkers
smpSubWorkers_ <- textKeys <$> readTVarIO smpSubWorkers
ntfWorkers_ <- workerStats =<< readTVarIO ntfWorkers
ntfSMPWorkers_ <- workerStats =<< readTVarIO ntfSMPWorkers
xftpRcvWorkers_ <- workerStats =<< readTVarIO xftpRcvWorkers
xftpSndWorkers_ <- workerStats =<< readTVarIO xftpSndWorkers
xftpDelWorkers_ <- workerStats =<< readTVarIO xftpDelWorkers
pure
AgentWorkersDetails
{ smpClients_,
ntfClients_,
xftpClients_,
smpDeliveryWorkers_,
asyncCmdWorkers_,
smpSubWorkers_,
ntfWorkers_,
ntfSMPWorkers_,
xftpRcvWorkers_,
xftpSndWorkers_,
xftpDelWorkers_
}
where
textKeys :: StrEncoding k => Map k v -> [Text]
textKeys = map textKey . M.keys
textKey :: StrEncoding k => k -> Text
textKey = decodeASCII . strEncode
workerStats :: StrEncoding k => Map k Worker -> IO (Map Text WorkersDetails)
workerStats ws = fmap M.fromList . forM (M.toList ws) $ \(qa, Worker {restarts, doWork, action}) -> do
RestartCount {restartCount} <- readTVarIO restarts
hasWork <- atomically $ not <$> isEmptyTMVar doWork
hasAction <- atomically $ not <$> isEmptyTMVar action
pure (textKey qa, WorkersDetails {restarts = restartCount, hasWork, hasAction})
Env {ntfSupervisor, xftpAgent} = agentEnv
NtfSupervisor {ntfWorkers, ntfSMPWorkers} = ntfSupervisor
XFTPAgent {xftpRcvWorkers, xftpSndWorkers, xftpDelWorkers} = xftpAgent
data AgentWorkersSummary = AgentWorkersSummary
{ smpClientsCount :: Int,
ntfClientsCount :: Int,
xftpClientsCount :: Int,
smpDeliveryWorkersCount :: WorkersSummary,
asyncCmdWorkersCount :: WorkersSummary,
smpSubWorkersCount :: Int,
ntfWorkersCount :: WorkersSummary,
ntfSMPWorkersCount :: WorkersSummary,
xftpRcvWorkersCount :: WorkersSummary,
xftpSndWorkersCount :: WorkersSummary,
xftpDelWorkersCount :: WorkersSummary
}
deriving (Show)
data WorkersSummary = WorkersSummary
{ numActive :: Int,
numIdle :: Int,
totalRestarts :: Int
}
deriving (Show)
getAgentWorkersSummary :: AgentClient -> IO AgentWorkersSummary
getAgentWorkersSummary AgentClient {smpClients, ntfClients, xftpClients, smpDeliveryWorkers, asyncCmdWorkers, smpSubWorkers, agentEnv} = do
smpClientsCount <- M.size <$> readTVarIO smpClients
ntfClientsCount <- M.size <$> readTVarIO ntfClients
xftpClientsCount <- M.size <$> readTVarIO xftpClients
smpDeliveryWorkersCount <- readTVarIO smpDeliveryWorkers >>= workerSummary . fmap fst
asyncCmdWorkersCount <- readTVarIO asyncCmdWorkers >>= workerSummary
smpSubWorkersCount <- M.size <$> readTVarIO smpSubWorkers
ntfWorkersCount <- readTVarIO ntfWorkers >>= workerSummary
ntfSMPWorkersCount <- readTVarIO ntfSMPWorkers >>= workerSummary
xftpRcvWorkersCount <- readTVarIO xftpRcvWorkers >>= workerSummary
xftpSndWorkersCount <- readTVarIO xftpSndWorkers >>= workerSummary
xftpDelWorkersCount <- readTVarIO xftpDelWorkers >>= workerSummary
pure
AgentWorkersSummary
{ smpClientsCount,
ntfClientsCount,
xftpClientsCount,
smpDeliveryWorkersCount,
asyncCmdWorkersCount,
smpSubWorkersCount,
ntfWorkersCount,
ntfSMPWorkersCount,
xftpRcvWorkersCount,
xftpSndWorkersCount,
xftpDelWorkersCount
}
where
Env {ntfSupervisor, xftpAgent} = agentEnv
NtfSupervisor {ntfWorkers, ntfSMPWorkers} = ntfSupervisor
XFTPAgent {xftpRcvWorkers, xftpSndWorkers, xftpDelWorkers} = xftpAgent
workerSummary :: M.Map k Worker -> IO WorkersSummary
workerSummary = liftIO . foldM byWork WorkersSummary {numActive = 0, numIdle = 0, totalRestarts = 0}
where
byWork WorkersSummary {numActive, numIdle, totalRestarts} Worker {action, restarts} = do
RestartCount {restartCount} <- readTVarIO restarts
ifM
(atomically $ isJust <$> tryReadTMVar action)
(pure WorkersSummary {numActive, numIdle = numIdle + 1, totalRestarts = totalRestarts + restartCount})
(pure WorkersSummary {numActive = numActive + 1, numIdle, totalRestarts = totalRestarts + restartCount})
data AgentQueuesInfo = AgentQueuesInfo
{ msgQInfo :: TBQueueInfo,
subQInfo :: TBQueueInfo,
smpClientsQueues :: Map Text (Int, UTCTime, ClientInfo)
}
deriving (Show)
data ClientInfo
= ClientInfoQueues {sndQInfo :: TBQueueInfo, rcvQInfo :: TBQueueInfo}
| ClientInfoError {clientError :: (AgentErrorType, Maybe UTCTime)}
| ClientInfoConnecting
deriving (Show)
getAgentQueuesInfo :: AgentClient -> IO AgentQueuesInfo
getAgentQueuesInfo AgentClient {msgQ, subQ, smpClients} = do
msgQInfo <- atomically $ getTBQueueInfo msgQ
subQInfo <- atomically $ getTBQueueInfo subQ
smpClientsMap <- readTVarIO smpClients
let smpClientsMap' = M.mapKeys (decodeLatin1 . strEncode) smpClientsMap
smpClientsQueues <- mapM getClientQueuesInfo smpClientsMap'
pure AgentQueuesInfo {msgQInfo, subQInfo, smpClientsQueues}
where
getClientQueuesInfo :: SMPClientVar -> IO (Int, UTCTime, ClientInfo)
getClientQueuesInfo SessionVar {sessionVar, sessionVarId, sessionVarTs} = do
clientInfo <-
atomically (tryReadTMVar sessionVar) >>= \case
Just (Right c) -> do
(sndQInfo, rcvQInfo) <- getProtocolClientQueuesInfo $ protocolClient c
pure ClientInfoQueues {sndQInfo, rcvQInfo}
Just (Left e) -> pure $ ClientInfoError e
Nothing -> pure ClientInfoConnecting
pure (sessionVarId, sessionVarTs, clientInfo)
$(J.deriveJSON defaultJSON ''AgentLocks)
$(J.deriveJSON (enumJSON $ dropPrefix "TS") ''ProtocolTestStep)
$(J.deriveJSON defaultJSON ''ProtocolTestFailure)
$(J.deriveJSON defaultJSON ''ServerSessions)
$(J.deriveJSON defaultJSON ''SMPServerSubs)
$(J.deriveJSON defaultJSON ''AgentServersSummary)
$(J.deriveJSON defaultJSON ''SubInfo)
$(J.deriveJSON defaultJSON ''SubscriptionsInfo)
$(J.deriveJSON defaultJSON ''WorkersDetails)
$(J.deriveJSON defaultJSON ''WorkersSummary)
$(J.deriveJSON defaultJSON {J.fieldLabelModifier = takeWhile (/= '_')} ''AgentWorkersDetails)
$(J.deriveJSON defaultJSON ''AgentWorkersSummary)
$(J.deriveJSON (sumTypeJSON $ dropPrefix "ClientInfo") ''ClientInfo)
$(J.deriveJSON defaultJSON ''AgentQueuesInfo)
$(J.deriveJSON (enumJSON $ dropPrefix "UN") ''UserNetworkType)
$(J.deriveJSON defaultJSON ''UserNetworkInfo)
$(J.deriveJSON defaultJSON ''ServerQueueInfo)
Reference in New Issue View Git Blame Copy Permalink