From ff2b00a0299343e39bd8d56fc43d97f7a8fe214c Mon Sep 17 00:00:00 2001
From: spaced4ndy <8711996+spaced4ndy@users.noreply.github.com>
Date: Thu, 11 Jul 2024 19:27:04 +0400
Subject: [PATCH 01/10] agent: change ProxyClientError json encoding (#1226)

---
 src/Simplex/Messaging/Client.hs | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/src/Simplex/Messaging/Client.hs b/src/Simplex/Messaging/Client.hs
index e20b00039..1837a256b 100644
--- a/src/Simplex/Messaging/Client.hs
+++ b/src/Simplex/Messaging/Client.hs
@@ -130,7 +130,7 @@ import Numeric.Natural
 import qualified Simplex.Messaging.Crypto as C
 import Simplex.Messaging.Encoding
 import Simplex.Messaging.Encoding.String
-import Simplex.Messaging.Parsers (defaultJSON, dropPrefix, enumJSON)
+import Simplex.Messaging.Parsers (defaultJSON, dropPrefix, enumJSON, sumTypeJSON)
 import Simplex.Messaging.Protocol
 import Simplex.Messaging.Server.QueueStore.QueueInfo
 import Simplex.Messaging.TMap (TMap)
@@ -847,11 +847,11 @@ data ProxiedRelay = ProxiedRelay
 
 data ProxyClientError
   = -- | protocol error response from proxy
-    ProxyProtocolError ErrorType
+    ProxyProtocolError {protocolErr :: ErrorType}
   | -- | unexpexted response
-    ProxyUnexpectedResponse String
+    ProxyUnexpectedResponse {responseStr :: String}
   | -- | error between proxy and server
-    ProxyResponseError ErrorType
+    ProxyResponseError {responseErr :: ErrorType}
   deriving (Eq, Show, Exception)
 
 instance StrEncoding ProxyClientError where
@@ -1139,6 +1139,6 @@ $(J.deriveJSON (enumJSON $ dropPrefix "SPF") ''SMPProxyFallback)
 
 $(J.deriveJSON defaultJSON ''NetworkConfig)
 
-$(J.deriveJSON (enumJSON $ dropPrefix "Proxy") ''ProxyClientError)
+$(J.deriveJSON (sumTypeJSON $ dropPrefix "Proxy") ''ProxyClientError)
 
 $(J.deriveJSON defaultJSON ''TBQueueInfo)

From e56bd0b47b22781f9b5936e3b8a9d51cc3f9e8d5 Mon Sep 17 00:00:00 2001
From: Evgeny Poberezkin <evgeny@poberezkin.com>
Date: Fri, 12 Jul 2024 12:41:55 +0100
Subject: [PATCH 02/10] agent: add known servers (#1225)

* agent: add known servers

* test delays

* ServerCfg

* json encoding

* enabledServerCfg

* checkUserServers
---
 src/Simplex/Messaging/Agent.hs            | 35 ++++++++++-----
 src/Simplex/Messaging/Agent/Client.hs     | 17 ++++----
 src/Simplex/Messaging/Agent/Env/SQLite.hs | 52 +++++++++++++++++++++--
 tests/AgentTests/FunctionalAPITests.hs    | 24 +++++------
 tests/SMPAgentClient.hs                   | 19 ++++++---
 tests/SMPProxyTests.hs                    |  8 ++--
 6 files changed, 111 insertions(+), 44 deletions(-)

diff --git a/src/Simplex/Messaging/Agent.hs b/src/Simplex/Messaging/Agent.hs
index 2262a96e5..c08e04298 100644
--- a/src/Simplex/Messaging/Agent.hs
+++ b/src/Simplex/Messaging/Agent.hs
@@ -77,6 +77,7 @@ module Simplex.Messaging.Agent
     getConnectionServers,
     getConnectionRatchetAdHash,
     setProtocolServers,
+    checkUserServers,
     testProtocolServer,
     setNtfServers,
     setNetworkConfig,
@@ -172,7 +173,7 @@ import Simplex.Messaging.Notifications.Protocol (DeviceToken, NtfRegCode (NtfReg
 import Simplex.Messaging.Notifications.Server.Push.APNS (PNMessageData (..))
 import Simplex.Messaging.Notifications.Types
 import Simplex.Messaging.Parsers (parse)
-import Simplex.Messaging.Protocol (BrokerMsg, Cmd (..), EntityId, ErrorType (AUTH), MsgBody, MsgFlags (..), NtfServer, ProtoServerWithAuth, ProtocolTypeI (..), SMPMsgMeta, SParty (..), SProtocolType (..), SndPublicAuthKey, SubscriptionMode (..), UserProtocol, VersionSMPC, XFTPServerWithAuth, sndAuthKeySMPClientVersion)
+import Simplex.Messaging.Protocol (BrokerMsg, Cmd (..), EntityId, ErrorType (AUTH), MsgBody, MsgFlags (..), NtfServer, ProtoServerWithAuth, ProtocolType (..), ProtocolTypeI (..), SMPMsgMeta, SParty (..), SProtocolType (..), SndPublicAuthKey, SubscriptionMode (..), UserProtocol, VersionSMPC, sndAuthKeySMPClientVersion)
 import qualified Simplex.Messaging.Protocol as SMP
 import Simplex.Messaging.ServiceScheme (ServiceScheme (..))
 import qualified Simplex.Messaging.TMap as TM
@@ -197,15 +198,18 @@ getSMPAgentClient = getSMPAgentClient_ 1
 {-# INLINE getSMPAgentClient #-}
 
 getSMPAgentClient_ :: Int -> AgentConfig -> InitialAgentServers -> SQLiteStore -> Bool -> IO AgentClient
-getSMPAgentClient_ clientId cfg initServers store backgroundMode =
-  liftIO $ newSMPAgentEnv cfg store >>= runReaderT runAgent
+getSMPAgentClient_ clientId cfg initServers@InitialAgentServers {smp, xftp} store backgroundMode =
+  newSMPAgentEnv cfg store >>= runReaderT runAgent
   where
     runAgent = do
+      liftIO $ checkServers "SMP" smp >> checkServers "XFTP" xftp
       currentTs <- liftIO getCurrentTime
       c@AgentClient {acThread} <- atomically . newAgentClient clientId initServers currentTs =<< ask
       t <- runAgentThreads c `forkFinally` const (liftIO $ disconnectAgentClient c)
       atomically . writeTVar acThread . Just =<< mkWeakThreadId t
       pure c
+    checkServers protocol srvs =
+      forM_ (M.assocs srvs) $ \(userId, srvs') -> checkUserServers ("getSMPAgentClient " <> protocol <> " " <> tshow userId) srvs'
     runAgentThreads c
       | backgroundMode = run c "subscriber" $ subscriber c
       | otherwise = do
@@ -271,7 +275,7 @@ resumeAgentClient :: AgentClient -> IO ()
 resumeAgentClient c = atomically $ writeTVar (active c) True
 {-# INLINE resumeAgentClient #-}
 
-createUser :: AgentClient -> NonEmpty SMPServerWithAuth -> NonEmpty XFTPServerWithAuth -> AE UserId
+createUser :: AgentClient -> NonEmpty (ServerCfg 'PSMP) -> NonEmpty (ServerCfg 'PXFTP) -> AE UserId
 createUser c = withAgentEnv c .: createUser' c
 {-# INLINE createUser #-}
 
@@ -600,11 +604,13 @@ logConnection c connected =
   let event = if connected then "connected to" else "disconnected from"
    in logInfo $ T.unwords ["client", tshow (clientId c), event, "Agent"]
 
-createUser' :: AgentClient -> NonEmpty SMPServerWithAuth -> NonEmpty XFTPServerWithAuth -> AM UserId
+createUser' :: AgentClient -> NonEmpty (ServerCfg 'PSMP) -> NonEmpty (ServerCfg 'PXFTP) -> AM UserId
 createUser' c smp xftp = do
+  liftIO $ checkUserServers "createUser SMP" smp
+  liftIO $ checkUserServers "createUser XFTP" xftp
   userId <- withStore' c createUserRecord
-  atomically $ TM.insert userId smp $ smpServers c
-  atomically $ TM.insert userId xftp $ xftpServers c
+  atomically $ TM.insert userId (mkUserServers smp) $ smpServers c
+  atomically $ TM.insert userId (mkUserServers xftp) $ xftpServers c
   pure userId
 
 deleteUser' :: AgentClient -> UserId -> Bool -> AM ()
@@ -1815,10 +1821,17 @@ connectionStats = \case
           ratchetSyncSupported = connAgentVersion >= ratchetSyncSMPAgentVersion
         }
 
--- | Change servers to be used for creating new queues, in Reader monad
-setProtocolServers :: (ProtocolTypeI p, UserProtocol p) => AgentClient -> UserId -> NonEmpty (ProtoServerWithAuth p) -> IO ()
-setProtocolServers c userId srvs = atomically $ TM.insert userId srvs (userServers c)
-{-# INLINE setProtocolServers #-}
+-- | Change servers to be used for creating new queues.
+-- This function will set all servers as enabled in case all passed servers are disabled.
+setProtocolServers :: forall p. (ProtocolTypeI p, UserProtocol p) => AgentClient -> UserId -> NonEmpty (ServerCfg p) -> IO ()
+setProtocolServers c userId srvs = do
+  checkUserServers "setProtocolServers" srvs
+  atomically $ TM.insert userId (mkUserServers srvs) (userServers c)
+
+checkUserServers :: Text -> NonEmpty (ServerCfg p) -> IO ()
+checkUserServers name srvs =
+  unless (any (\ServerCfg {enabled} -> enabled) srvs) $
+    logWarn (name <> ": all passed servers are disabled, using all servers.")
 
 registerNtfToken' :: AgentClient -> DeviceToken -> NotificationsMode -> AM NtfTknStatus
 registerNtfToken' c suppliedDeviceToken suppliedNtfMode =
diff --git a/src/Simplex/Messaging/Agent/Client.hs b/src/Simplex/Messaging/Agent/Client.hs
index 0467c31f8..b8614c460 100644
--- a/src/Simplex/Messaging/Agent/Client.hs
+++ b/src/Simplex/Messaging/Agent/Client.hs
@@ -236,6 +236,7 @@ import Simplex.Messaging.Protocol
     ProtoServerWithAuth (..),
     Protocol (..),
     ProtocolServer (..),
+    ProtocolType (..),
     ProtocolTypeI (..),
     QueueId,
     QueueIdsKeys (..),
@@ -289,7 +290,7 @@ data AgentClient = AgentClient
     active :: TVar Bool,
     subQ :: TBQueue ATransmission,
     msgQ :: TBQueue (ServerTransmissionBatch SMPVersion ErrorType BrokerMsg),
-    smpServers :: TMap UserId (NonEmpty SMPServerWithAuth),
+    smpServers :: TMap UserId (UserServers 'PSMP),
     smpClients :: TMap SMPTransportSession SMPClientVar,
     -- smpProxiedRelays:
     -- SMPTransportSession defines connection from proxy to relay,
@@ -297,7 +298,7 @@ data AgentClient = AgentClient
     smpProxiedRelays :: TMap SMPTransportSession SMPServerWithAuth,
     ntfServers :: TVar [NtfServer],
     ntfClients :: TMap NtfTransportSession NtfClientVar,
-    xftpServers :: TMap UserId (NonEmpty XFTPServerWithAuth),
+    xftpServers :: TMap UserId (UserServers 'PXFTP),
     xftpClients :: TMap XFTPTransportSession XFTPClientVar,
     useNetworkConfig :: TVar (NetworkConfig, NetworkConfig), -- (slow, fast) networks
     userNetworkInfo :: TVar UserNetworkInfo,
@@ -456,12 +457,12 @@ newAgentClient clientId InitialAgentServers {smp, ntf, xftp, netCfg} currentTs a
   active <- newTVar True
   subQ <- newTBQueue qSize
   msgQ <- newTBQueue qSize
-  smpServers <- newTVar smp
+  smpServers <- newTVar $ M.map mkUserServers smp
   smpClients <- TM.empty
   smpProxiedRelays <- TM.empty
   ntfServers <- newTVar ntf
   ntfClients <- TM.empty
-  xftpServers <- newTVar xftp
+  xftpServers <- newTVar $ M.map mkUserServers xftp
   xftpClients <- TM.empty
   useNetworkConfig <- newTVar (slowNetworkConfig netCfg, netCfg)
   userNetworkInfo <- newTVar $ UserNetworkInfo UNOther True
@@ -1069,7 +1070,7 @@ sendOrProxySMPCommand c userId destSrv cmdStr senderId sendCmdViaProxy sendCmdDi
         SPFAllow -> True
         SPFAllowProtected -> ipAddressProtected cfg destSrv
         SPFProhibit -> False
-    unknownServer = maybe True (all ((destSrv /=) . protoServer)) <$> TM.lookup userId (userServers c)
+    unknownServer = maybe True (notElem destSrv . knownSrvs) <$> TM.lookup userId (smpServers c)
     sendViaProxy destSess@(_, _, qId) = do
       r <- tryAgentError . withProxySession c destSess senderId ("PFWD " <> cmdStr) $ \(SMPConnectedClient smp _, proxySess) -> do
         r' <- liftClient SMP (clientServer smp) $ sendCmdViaProxy smp proxySess
@@ -1904,7 +1905,7 @@ storeError = \case
   SEDatabaseBusy e -> CRITICAL True $ B.unpack e
   e -> INTERNAL $ show e
 
-userServers :: forall p. (ProtocolTypeI p, UserProtocol p) => AgentClient -> TMap UserId (NonEmpty (ProtoServerWithAuth p))
+userServers :: forall p. (ProtocolTypeI p, UserProtocol p) => AgentClient -> TMap UserId (UserServers p)
 userServers c = case protocolTypeI @p of
   SPSMP -> smpServers c
   SPXFTP -> xftpServers c
@@ -1926,7 +1927,7 @@ getNextServer c userId usedSrvs = withUserServers c userId $ \srvs ->
 withUserServers :: forall p a. (ProtocolTypeI p, UserProtocol p) => AgentClient -> UserId -> (NonEmpty (ProtoServerWithAuth p) -> AM a) -> AM a
 withUserServers c userId action =
   atomically (TM.lookup userId $ userServers c) >>= \case
-    Just srvs -> action srvs
+    Just srvs -> action $ enabledSrvs srvs
     _ -> throwE $ INTERNAL "unknown userId - no user servers"
 
 withNextSrv :: forall p a. (ProtocolTypeI p, UserProtocol p) => AgentClient -> UserId -> TVar [ProtocolServer p] -> [ProtocolServer p] -> (ProtoServerWithAuth p -> AM a) -> AM a
@@ -1935,7 +1936,7 @@ withNextSrv c userId usedSrvs initUsed action = do
   srvAuth@(ProtoServerWithAuth srv _) <- getNextServer c userId used
   atomically $ do
     srvs_ <- TM.lookup userId $ userServers c
-    let unused = maybe [] ((\\ used) . map protoServer . L.toList) srvs_
+    let unused = maybe [] ((\\ used) . map protoServer . L.toList . enabledSrvs) srvs_
         used' = if null unused then initUsed else srv : used
     writeTVar usedSrvs $! used'
   action srvAuth
diff --git a/src/Simplex/Messaging/Agent/Env/SQLite.hs b/src/Simplex/Messaging/Agent/Env/SQLite.hs
index 2ae2ad5c0..0f88508b9 100644
--- a/src/Simplex/Messaging/Agent/Env/SQLite.hs
+++ b/src/Simplex/Messaging/Agent/Env/SQLite.hs
@@ -7,6 +7,7 @@
 {-# LANGUAGE NumericUnderscores #-}
 {-# LANGUAGE RankNTypes #-}
 {-# LANGUAGE ScopedTypeVariables #-}
+{-# LANGUAGE TemplateHaskell #-}
 {-# LANGUAGE TypeApplications #-}
 {-# OPTIONS_GHC -fno-warn-unticked-promoted-constructors #-}
 
@@ -15,7 +16,12 @@ module Simplex.Messaging.Agent.Env.SQLite
     AM,
     AgentConfig (..),
     InitialAgentServers (..),
+    ServerCfg (..),
+    UserServers (..),
     NetworkConfig (..),
+    presetServerCfg,
+    enabledServerCfg,
+    mkUserServers,
     defaultAgentConfig,
     defaultReconnectInterval,
     tryAgentError,
@@ -39,10 +45,14 @@ import Control.Monad.Except
 import Control.Monad.IO.Unlift
 import Control.Monad.Reader
 import Crypto.Random
+import Data.Aeson (FromJSON (..), ToJSON (..))
+import qualified Data.Aeson.TH as JQ
 import Data.ByteArray (ScrubbedBytes)
 import Data.Int (Int64)
 import Data.List.NonEmpty (NonEmpty)
+import qualified Data.List.NonEmpty as L
 import Data.Map (Map)
+import Data.Maybe (fromMaybe)
 import Data.Time.Clock (NominalDiffTime, nominalDay)
 import Data.Time.Clock.System (SystemTime (..))
 import Data.Word (Word16)
@@ -59,7 +69,8 @@ import Simplex.Messaging.Crypto.Ratchet (VersionRangeE2E, supportedE2EEncryptVRa
 import Simplex.Messaging.Notifications.Client (defaultNTFClientConfig)
 import Simplex.Messaging.Notifications.Transport (NTFVersion)
 import Simplex.Messaging.Notifications.Types
-import Simplex.Messaging.Protocol (NtfServer, VersionRangeSMPC, XFTPServer, XFTPServerWithAuth, supportedSMPClientVRange)
+import Simplex.Messaging.Parsers (defaultJSON)
+import Simplex.Messaging.Protocol (NtfServer, ProtoServerWithAuth, ProtocolServer, ProtocolType (..), ProtocolTypeI, VersionRangeSMPC, XFTPServer, supportedSMPClientVRange)
 import Simplex.Messaging.TMap (TMap)
 import qualified Simplex.Messaging.TMap as TM
 import Simplex.Messaging.Transport (SMPVersion, TLS, Transport (..))
@@ -74,12 +85,38 @@ type AM' a = ReaderT Env IO a
 type AM a = ExceptT AgentErrorType (ReaderT Env IO) a
 
 data InitialAgentServers = InitialAgentServers
-  { smp :: Map UserId (NonEmpty SMPServerWithAuth),
+  { smp :: Map UserId (NonEmpty (ServerCfg 'PSMP)),
     ntf :: [NtfServer],
-    xftp :: Map UserId (NonEmpty XFTPServerWithAuth),
+    xftp :: Map UserId (NonEmpty (ServerCfg 'PXFTP)),
     netCfg :: NetworkConfig
   }
 
+data ServerCfg p = ServerCfg
+  { server :: ProtoServerWithAuth p,
+    preset :: Bool,
+    tested :: Maybe Bool,
+    enabled :: Bool
+  }
+  deriving (Show)
+
+enabledServerCfg :: ProtoServerWithAuth p -> ServerCfg p
+enabledServerCfg server = ServerCfg {server, preset = False, tested = Nothing, enabled = True}
+
+presetServerCfg :: Bool -> ProtoServerWithAuth p -> ServerCfg p
+presetServerCfg enabled server = ServerCfg {server, preset = True, tested = Nothing, enabled}
+
+data UserServers p = UserServers
+  { enabledSrvs :: NonEmpty (ProtoServerWithAuth p),
+    knownSrvs :: NonEmpty (ProtocolServer p)
+  }
+
+-- This function sets all servers as enabled in case all passed servers are disabled.
+mkUserServers :: NonEmpty (ServerCfg p) -> UserServers p
+mkUserServers srvs = UserServers {enabledSrvs, knownSrvs}
+  where
+    enabledSrvs = L.map (\ServerCfg {server} -> server) $ fromMaybe srvs $ L.nonEmpty $ L.filter (\ServerCfg {enabled} -> enabled) srvs
+    knownSrvs = L.map (\ServerCfg {server = ProtoServerWithAuth srv _} -> srv) srvs
+
 data AgentConfig = AgentConfig
   { tcpPort :: Maybe ServiceName,
     rcvAuthAlg :: C.AuthAlg,
@@ -294,3 +331,12 @@ updateRestartCount :: SystemTime -> RestartCount -> RestartCount
 updateRestartCount t (RestartCount minute count) = do
   let min' = systemSeconds t `div` 60
    in RestartCount min' $ if minute == min' then count + 1 else 1
+
+$(pure [])
+
+instance ProtocolTypeI p => ToJSON (ServerCfg p) where
+  toEncoding = $(JQ.mkToEncoding defaultJSON ''ServerCfg)
+  toJSON = $(JQ.mkToJSON defaultJSON ''ServerCfg)
+
+instance ProtocolTypeI p => FromJSON (ServerCfg p) where
+  parseJSON = $(JQ.mkParseJSON defaultJSON ''ServerCfg)
diff --git a/tests/AgentTests/FunctionalAPITests.hs b/tests/AgentTests/FunctionalAPITests.hs
index ab78d2ee9..1893ae105 100644
--- a/tests/AgentTests/FunctionalAPITests.hs
+++ b/tests/AgentTests/FunctionalAPITests.hs
@@ -980,7 +980,7 @@ testAsyncServerOffline t = withAgentClients2 $ \alice bob -> do
 
 testAllowConnectionClientRestart :: HasCallStack => ATransport -> IO ()
 testAllowConnectionClientRestart t = do
-  let initAgentServersSrv2 = initAgentServers {smp = userServers [noAuthSrv testSMPServer2]}
+  let initAgentServersSrv2 = initAgentServers {smp = userServers [testSMPServer2]}
   alice <- getSMPAgentClient' 1 agentCfg initAgentServers testDB
   bob <- getSMPAgentClient' 2 agentCfg initAgentServersSrv2 testDB2
   withSmpServerStoreLogOn t testPort $ \_ -> do
@@ -1335,7 +1335,7 @@ testExpireMessage t =
 
 testExpireManyMessages :: HasCallStack => ATransport -> IO ()
 testExpireManyMessages t =
-  withAgent 1 agentCfg {messageTimeout = 1.5, messageRetryInterval = fastMessageRetryInterval} initAgentServers testDB $ \a ->
+  withAgent 1 agentCfg {messageTimeout = 2, messageRetryInterval = fastMessageRetryInterval} initAgentServers testDB $ \a ->
     withAgent 2 agentCfg initAgentServers testDB2 $ \b -> do
       (aId, bId) <- withSmpServerStoreLogOn t testPort $ \_ -> runRight $ makeConnection a b
       runRight_ $ do
@@ -1344,7 +1344,7 @@ testExpireManyMessages t =
         2 <- sendMessage a bId SMP.noMsgFlags "1"
         3 <- sendMessage a bId SMP.noMsgFlags "2"
         4 <- sendMessage a bId SMP.noMsgFlags "3"
-        liftIO $ threadDelay 1500000
+        liftIO $ threadDelay 2000000
         5 <- sendMessage a bId SMP.noMsgFlags "4" -- this won't expire
         get a =##> \case ("", c, MERR 2 (BROKER _ e)) -> bId == c && (e == TIMEOUT || e == NETWORK); _ -> False
         -- get a =##> \case ("", c, MERRS [5, 6] (BROKER _ e)) -> bId == c && (e == TIMEOUT || e == NETWORK); _ -> False
@@ -1401,7 +1401,7 @@ testExpireMessageQuota t = withSmpServerConfigOn t cfg {msgQueueQuota = 1} testP
 
 testExpireManyMessagesQuota :: ATransport -> IO ()
 testExpireManyMessagesQuota t = withSmpServerConfigOn t cfg {msgQueueQuota = 1} testPort $ \_ -> do
-  a <- getSMPAgentClient' 1 agentCfg {quotaExceededTimeout = 1, messageRetryInterval = fastMessageRetryInterval} initAgentServers testDB
+  a <- getSMPAgentClient' 1 agentCfg {quotaExceededTimeout = 2, messageRetryInterval = fastMessageRetryInterval} initAgentServers testDB
   b <- getSMPAgentClient' 2 agentCfg initAgentServers testDB2
   (aId, bId) <- runRight $ do
     (aId, bId) <- makeConnection a b
@@ -1411,7 +1411,7 @@ testExpireManyMessagesQuota t = withSmpServerConfigOn t cfg {msgQueueQuota = 1}
     3 <- sendMessage a bId SMP.noMsgFlags "2"
     4 <- sendMessage a bId SMP.noMsgFlags "3"
     5 <- sendMessage a bId SMP.noMsgFlags "4"
-    liftIO $ threadDelay 1000000
+    liftIO $ threadDelay 2000000
     6 <- sendMessage a bId SMP.noMsgFlags "5" -- this won't expire
     get a =##> \case ("", c, MERR 3 (SMP _ QUOTA)) -> bId == c; _ -> False
     get a >>= \case
@@ -2226,7 +2226,7 @@ testWaitDeliveryTimeout2 t =
 
 testJoinConnectionAsyncReplyErrorV8 :: HasCallStack => ATransport -> IO ()
 testJoinConnectionAsyncReplyErrorV8 t = do
-  let initAgentServersSrv2 = initAgentServers {smp = userServers [noAuthSrv testSMPServer2]}
+  let initAgentServersSrv2 = initAgentServers {smp = userServers [testSMPServer2]}
   withAgent 1 agentCfgVPrevPQ initAgentServers testDB $ \a ->
     withAgent 2 agentCfgVPrevPQ initAgentServersSrv2 testDB2 $ \b -> do
       (aId, bId) <- withSmpServerStoreLogOn t testPort $ \_ -> runRight $ do
@@ -2265,7 +2265,7 @@ testJoinConnectionAsyncReplyErrorV8 t = do
 
 testJoinConnectionAsyncReplyError :: HasCallStack => ATransport -> IO ()
 testJoinConnectionAsyncReplyError t = do
-  let initAgentServersSrv2 = initAgentServers {smp = userServers [noAuthSrv testSMPServer2]}
+  let initAgentServersSrv2 = initAgentServers {smp = userServers [testSMPServer2]}
   withAgent 1 agentCfg initAgentServers testDB $ \a ->
     withAgent 2 agentCfg initAgentServersSrv2 testDB2 $ \b -> do
       (aId, bId) <- withSmpServerStoreLogOn t testPort $ \_ -> runRight $ do
@@ -2308,7 +2308,7 @@ testUsers =
   withAgentClients2 $ \a b -> runRight_ $ do
     (aId, bId) <- makeConnection a b
     exchangeGreetings a bId b aId
-    auId <- createUser a [noAuthSrv testSMPServer] [noAuthSrv testXFTPServer]
+    auId <- createUser a [noAuthSrvCfg testSMPServer] [noAuthSrvCfg testXFTPServer]
     (aId', bId') <- makeConnectionForUsers a auId b 1
     exchangeGreetings a bId' b aId'
     deleteUser a auId True
@@ -2323,7 +2323,7 @@ testDeleteUserQuietly =
   withAgentClients2 $ \a b -> runRight_ $ do
     (aId, bId) <- makeConnection a b
     exchangeGreetings a bId b aId
-    auId <- createUser a [noAuthSrv testSMPServer] [noAuthSrv testXFTPServer]
+    auId <- createUser a [noAuthSrvCfg testSMPServer] [noAuthSrvCfg testXFTPServer]
     (aId', bId') <- makeConnectionForUsers a auId b 1
     exchangeGreetings a bId' b aId'
     deleteUser a auId False
@@ -2335,7 +2335,7 @@ testUsersNoServer t = withAgentClientsCfg2 aCfg agentCfg $ \a b -> do
   (aId, bId, auId, _aId', bId') <- withSmpServerStoreLogOn t testPort $ \_ -> runRight $ do
     (aId, bId) <- makeConnection a b
     exchangeGreetings a bId b aId
-    auId <- createUser a [noAuthSrv testSMPServer] [noAuthSrv testXFTPServer]
+    auId <- createUser a [noAuthSrvCfg testSMPServer] [noAuthSrvCfg testXFTPServer]
     (aId', bId') <- makeConnectionForUsers a auId b 1
     exchangeGreetings a bId' b aId'
     pure (aId, bId, auId, aId', bId')
@@ -2759,7 +2759,7 @@ testCreateQueueAuth srvVersion clnt1 clnt2 baseId = do
   pure r
   where
     getClient clientId (clntAuth, clntVersion) db =
-      let servers = initAgentServers {smp = userServers [ProtoServerWithAuth testSMPServer clntAuth]}
+      let servers = initAgentServers {smp = userServers' [ProtoServerWithAuth testSMPServer clntAuth]}
           alpn_ = if clntVersion >= authCmdsSMPVersion then Just supportedSMPHandshakes else Nothing
           smpCfg = defaultClientConfig alpn_ $ V.mkVersionRange (prevVersion basicAuthSMPVersion) clntVersion
           sndAuthAlg = if srvVersion >= authCmdsSMPVersion && clntVersion >= authCmdsSMPVersion then C.AuthAlg C.SX25519 else C.AuthAlg C.SEd25519
@@ -2931,7 +2931,7 @@ testTwoUsers = withAgentClients2 $ \a b -> do
     ("", "", UP _ _) <- nGet a
     a `hasClients` 1
 
-    aUserId2 <- createUser a [noAuthSrv testSMPServer] [noAuthSrv testXFTPServer]
+    aUserId2 <- createUser a [noAuthSrvCfg testSMPServer] [noAuthSrvCfg testXFTPServer]
     (aId2, bId2) <- makeConnectionForUsers a aUserId2 b 1
     exchangeGreetings a bId2 b aId2
     (aId2', bId2') <- makeConnectionForUsers a aUserId2 b 1
diff --git a/tests/SMPAgentClient.hs b/tests/SMPAgentClient.hs
index 7cb2a88c5..e0de57466 100644
--- a/tests/SMPAgentClient.hs
+++ b/tests/SMPAgentClient.hs
@@ -11,6 +11,7 @@
 module SMPAgentClient where
 
 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 NtfClient (ntfTestPort)
@@ -20,7 +21,7 @@ import Simplex.Messaging.Agent.Protocol
 import Simplex.Messaging.Agent.RetryInterval
 import Simplex.Messaging.Client (ProtocolClientConfig (..), SMPProxyFallback, SMPProxyMode, defaultNetworkConfig, defaultSMPClientConfig)
 import Simplex.Messaging.Notifications.Client (defaultNTFClientConfig)
-import Simplex.Messaging.Protocol (NtfServer, ProtoServerWithAuth)
+import Simplex.Messaging.Protocol (NtfServer, ProtoServerWithAuth (..), ProtocolServer)
 import Simplex.Messaging.Transport
 import XFTPClient (testXFTPServer)
 
@@ -48,14 +49,14 @@ testNtfServer2 = "ntf://LcJUMfVhwD8yxjAiSaDzzGF3-kLG4Uh0Fl_ZIjrRwjI=@localhost:6
 initAgentServers :: InitialAgentServers
 initAgentServers =
   InitialAgentServers
-    { smp = userServers [noAuthSrv testSMPServer],
+    { smp = userServers [testSMPServer],
       ntf = [testNtfServer],
-      xftp = userServers [noAuthSrv testXFTPServer],
+      xftp = userServers [testXFTPServer],
       netCfg = defaultNetworkConfig {tcpTimeout = 500_000, tcpConnectTimeout = 500_000}
     }
 
 initAgentServers2 :: InitialAgentServers
-initAgentServers2 = initAgentServers {smp = userServers [noAuthSrv testSMPServer, noAuthSrv testSMPServer2]}
+initAgentServers2 = initAgentServers {smp = userServers [testSMPServer, testSMPServer2]}
 
 initAgentServersProxy :: SMPProxyMode -> SMPProxyFallback -> InitialAgentServers
 initAgentServersProxy smpProxyMode smpProxyFallback =
@@ -89,5 +90,11 @@ fastRetryInterval = defaultReconnectInterval {initialInterval = 50_000}
 fastMessageRetryInterval :: RetryInterval2
 fastMessageRetryInterval = RetryInterval2 {riFast = fastRetryInterval, riSlow = fastRetryInterval}
 
-userServers :: NonEmpty (ProtoServerWithAuth p) -> Map UserId (NonEmpty (ProtoServerWithAuth p))
-userServers srvs = M.fromList [(1, srvs)]
+userServers :: NonEmpty (ProtocolServer p) -> Map UserId (NonEmpty (ServerCfg p))
+userServers = userServers' . L.map noAuthSrv
+
+userServers' :: NonEmpty (ProtoServerWithAuth p) -> Map UserId (NonEmpty (ServerCfg p))
+userServers' srvs = M.fromList [(1, L.map (presetServerCfg True) srvs)]
+
+noAuthSrvCfg :: ProtocolServer p -> ServerCfg p
+noAuthSrvCfg = presetServerCfg True . noAuthSrv
diff --git a/tests/SMPProxyTests.hs b/tests/SMPProxyTests.hs
index 6452d2677..625dfbda7 100644
--- a/tests/SMPProxyTests.hs
+++ b/tests/SMPProxyTests.hs
@@ -234,7 +234,7 @@ agentDeliverMessageViaProxy aTestCfg@(aSrvs, _, aViaProxy) bTestCfg@(bSrvs, _, b
   where
     msgId = subtract baseId . fst
     aCfg = agentCfg {sndAuthAlg = C.AuthAlg alg, rcvAuthAlg = C.AuthAlg alg}
-    servers (srvs, smpProxyMode, _) = (initAgentServersProxy smpProxyMode SPFAllow) {smp = userServers $ L.map noAuthSrv srvs}
+    servers (srvs, smpProxyMode, _) = (initAgentServersProxy smpProxyMode SPFAllow) {smp = userServers srvs}
 
 agentDeliverMessagesViaProxyConc :: [NonEmpty SMPServer] -> [MsgBody] -> IO ()
 agentDeliverMessagesViaProxyConc agentServers msgs =
@@ -299,7 +299,7 @@ agentDeliverMessagesViaProxyConc agentServers msgs =
       logDebug "run finished"
     pqEnc = CR.PQEncOn
     aCfg = agentCfg {sndAuthAlg = C.AuthAlg C.SEd448, rcvAuthAlg = C.AuthAlg C.SEd448}
-    servers srvs = (initAgentServersProxy SPMAlways SPFAllow) {smp = userServers $ L.map noAuthSrv srvs}
+    servers srvs = (initAgentServersProxy SPMAlways SPFAllow) {smp = userServers srvs}
 
 agentViaProxyVersionError :: IO ()
 agentViaProxyVersionError =
@@ -310,7 +310,7 @@ agentViaProxyVersionError =
         A.joinConnection bob 1 Nothing True qInfo "bob's connInfo" PQSupportOn SMSubscribe
     pure ()
   where
-    servers srvs = (initAgentServersProxy SPMUnknown SPFProhibit) {smp = userServers $ L.map noAuthSrv srvs}
+    servers srvs = (initAgentServersProxy SPMUnknown SPFProhibit) {smp = userServers srvs}
 
 agentViaProxyRetryOffline :: IO ()
 agentViaProxyRetryOffline = do
@@ -372,7 +372,7 @@ agentViaProxyRetryOffline = do
     aCfg = agentCfg {messageRetryInterval = fastMessageRetryInterval}
     baseId = 1
     msgId = subtract baseId . fst
-    servers srv = (initAgentServersProxy SPMAlways SPFProhibit) {smp = userServers $ L.map noAuthSrv [srv]}
+    servers srv = (initAgentServersProxy SPMAlways SPFProhibit) {smp = userServers [srv]}
 
 testNoProxy :: IO ()
 testNoProxy = do

From 8dd54ced0eb344e04298acf4a6b2d71e54cede01 Mon Sep 17 00:00:00 2001
From: Evgeny Poberezkin <evgeny@poberezkin.com>
Date: Sat, 13 Jul 2024 10:06:48 +0100
Subject: [PATCH 03/10] agent: retry proxied command on NO_SESSION error, to
 prevent failure of proxied interactive commands (like joining connection)
 (#1227)

---
 src/Simplex/Messaging/Agent/Client.hs | 45 ++++++++++++++++-----------
 src/Simplex/Messaging/Client.hs       |  5 +--
 tests/SMPProxyTests.hs                | 31 +++++++++++++++---
 3 files changed, 57 insertions(+), 24 deletions(-)

diff --git a/src/Simplex/Messaging/Agent/Client.hs b/src/Simplex/Messaging/Agent/Client.hs
index b8614c460..43b3b8064 100644
--- a/src/Simplex/Messaging/Agent/Client.hs
+++ b/src/Simplex/Messaging/Agent/Client.hs
@@ -597,10 +597,10 @@ getSMPServerClient c@AgentClient {active, smpClients, workerSeq} tSess = do
       prs <- atomically TM.empty
       smpConnectClient c tSess prs v
 
-getSMPProxyClient :: AgentClient -> SMPTransportSession -> AM (SMPConnectedClient, Either AgentErrorType ProxiedRelay)
-getSMPProxyClient c@AgentClient {active, smpClients, smpProxiedRelays, workerSeq} destSess@(userId, destSrv, qId) = do
+getSMPProxyClient :: AgentClient -> Maybe SMPServerWithAuth -> SMPTransportSession -> AM (SMPConnectedClient, Either AgentErrorType ProxiedRelay)
+getSMPProxyClient c@AgentClient {active, smpClients, smpProxiedRelays, workerSeq} proxySrv_ destSess@(userId, destSrv, qId) = do
   unlessM (readTVarIO active) $ throwE INACTIVE
-  proxySrv <- getNextServer c userId [destSrv]
+  proxySrv <- maybe (getNextServer c userId [destSrv]) pure proxySrv_
   ts <- liftIO getCurrentTime
   atomically (getClientVar proxySrv ts) >>= \(tSess, auth, v) ->
     either (newProxyClient tSess auth ts) (waitForProxyClient tSess auth) v
@@ -993,9 +993,9 @@ withClient_ c tSess@(_, srv, _) action = do
       logServer "<--" c srv "" $ bshow e
       throwE e
 
-withProxySession :: AgentClient -> SMPTransportSession -> SMP.SenderId -> ByteString -> ((SMPConnectedClient, ProxiedRelay) -> AM a) -> AM a
-withProxySession c destSess@(_, destSrv, _) entId cmdStr action = do
-  (cl, sess_) <- getSMPProxyClient c destSess
+withProxySession :: AgentClient -> Maybe SMPServerWithAuth -> SMPTransportSession -> SMP.SenderId -> ByteString -> ((SMPConnectedClient, ProxiedRelay) -> AM a) -> AM a
+withProxySession c proxySrv_ destSess@(_, destSrv, _) entId cmdStr action = do
+  (cl, sess_) <- getSMPProxyClient c proxySrv_ destSess
   logServer ("--> " <> proxySrv cl <> " >") c destSrv entId cmdStr
   case sess_ of
     Right sess -> do
@@ -1053,7 +1053,7 @@ sendOrProxySMPCommand ::
   AM (Maybe SMPServer)
 sendOrProxySMPCommand c userId destSrv cmdStr senderId sendCmdViaProxy sendCmdDirectly = do
   sess <- liftIO $ mkTransportSession c userId destSrv senderId
-  ifM (atomically shouldUseProxy) (sendViaProxy sess) (sendDirectly sess $> Nothing)
+  ifM (atomically shouldUseProxy) (sendViaProxy Nothing sess) (sendDirectly sess $> Nothing)
   where
     shouldUseProxy = do
       cfg <- getNetworkConfig c
@@ -1071,22 +1071,31 @@ sendOrProxySMPCommand c userId destSrv cmdStr senderId sendCmdViaProxy sendCmdDi
         SPFAllowProtected -> ipAddressProtected cfg destSrv
         SPFProhibit -> False
     unknownServer = maybe True (notElem destSrv . knownSrvs) <$> TM.lookup userId (smpServers c)
-    sendViaProxy destSess@(_, _, qId) = do
-      r <- tryAgentError . withProxySession c destSess senderId ("PFWD " <> cmdStr) $ \(SMPConnectedClient smp _, proxySess) -> do
+    sendViaProxy :: Maybe SMPServerWithAuth -> SMPTransportSession -> AM (Maybe SMPServer)
+    sendViaProxy proxySrv_ destSess@(_, _, qId) = do
+      r <- tryAgentError . withProxySession c proxySrv_ destSess senderId ("PFWD " <> cmdStr) $ \(SMPConnectedClient smp _, proxySess@ProxiedRelay {prBasicAuth}) -> do
         r' <- liftClient SMP (clientServer smp) $ sendCmdViaProxy smp proxySess
+        let proxySrv = protocolClientServer' smp
         case r' of
-          Right () -> pure . Just $ protocolClientServer' smp
+          Right () -> pure $ Just proxySrv
           Left proxyErr -> do
             case proxyErr of
-              (ProxyProtocolError (SMP.PROXY SMP.NO_SESSION)) -> atomically deleteRelaySession
-              _ -> pure ()
-            throwE
-              PROXY
-                { proxyServer = protocolClientServer smp,
-                  relayServer = B.unpack $ strEncode destSrv,
-                  proxyErr
-                }
+              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)
diff --git a/src/Simplex/Messaging/Client.hs b/src/Simplex/Messaging/Client.hs
index 1837a256b..4b1e673b0 100644
--- a/src/Simplex/Messaging/Client.hs
+++ b/src/Simplex/Messaging/Client.hs
@@ -823,7 +823,7 @@ connectSMPProxiedRelay c@ProtocolClient {client_ = PClient {tcpConnectTimeout, t
         PKEY sId vr (chain, key) ->
           case supportedClientSMPRelayVRange `compatibleVersion` vr of
             Nothing -> throwE $ transportErr TEVersion
-            Just (Compatible v) -> liftEitherWith (const $ transportErr $ TEHandshake IDENTITY) $ ProxiedRelay sId v <$> validateRelay chain key
+            Just (Compatible v) -> liftEitherWith (const $ transportErr $ TEHandshake IDENTITY) $ ProxiedRelay sId v proxyAuth <$> validateRelay chain key
         r -> throwE $ unexpectedResponse r
   | otherwise = throwE $ PCETransportError TEVersion
   where
@@ -842,6 +842,7 @@ connectSMPProxiedRelay c@ProtocolClient {client_ = PClient {tcpConnectTimeout, t
 data ProxiedRelay = ProxiedRelay
   { prSessionId :: SessionId,
     prVersion :: VersionSMP,
+    prBasicAuth :: Maybe BasicAuth, -- auth is included here to allow reconnecting via the same proxy after NO_SESSION error
     prServerKey :: C.PublicKeyX25519
   }
 
@@ -902,7 +903,7 @@ proxySMPCommand ::
   SenderId ->
   Command 'Sender ->
   ExceptT SMPClientError IO (Either ProxyClientError ())
-proxySMPCommand c@ProtocolClient {thParams = proxyThParams, client_ = PClient {clientCorrId = g, tcpTimeout}} (ProxiedRelay sessionId v serverKey) spKey sId command = do
+proxySMPCommand c@ProtocolClient {thParams = proxyThParams, client_ = PClient {clientCorrId = g, tcpTimeout}} (ProxiedRelay sessionId v _ serverKey) spKey sId command = do
   -- prepare params
   let serverThAuth = (\ta -> ta {serverPeerPubKey = serverKey}) <$> thAuth proxyThParams
       serverThParams = smpTHParamsSetVersion v proxyThParams {sessionId, thAuth = serverThAuth}
diff --git a/tests/SMPProxyTests.hs b/tests/SMPProxyTests.hs
index 625dfbda7..7505ef977 100644
--- a/tests/SMPProxyTests.hs
+++ b/tests/SMPProxyTests.hs
@@ -34,7 +34,8 @@ import Simplex.Messaging.Client
 import qualified Simplex.Messaging.Crypto as C
 import Simplex.Messaging.Crypto.Ratchet (pattern PQSupportOn)
 import qualified Simplex.Messaging.Crypto.Ratchet as CR
-import Simplex.Messaging.Protocol as SMP
+import Simplex.Messaging.Protocol (EncRcvMsgBody (..), MsgBody, RcvMessage (..), SubscriptionMode (..), maxMessageLength, noMsgFlags)
+import qualified Simplex.Messaging.Protocol as SMP
 import Simplex.Messaging.Server.Env.STM (ServerConfig (..))
 import Simplex.Messaging.Transport
 import Simplex.Messaging.Util (bshow, tshow)
@@ -122,6 +123,8 @@ smpProxyTests = do
           agentViaProxyVersionError
         it "retries sending when destination or proxy relay is offline" $
           agentViaProxyRetryOffline
+        it "retries sending when destination relay session disconnects in proxy" $
+          agentViaProxyRetryNoSession
       describe "stress test 1k" $ do
         let deliver nAgents nMsgs = agentDeliverMessagesViaProxyConc (replicate nAgents [srv1]) (map bshow [1 :: Int .. nMsgs])
         it "2 agents, 250 messages" . oneServer $ deliver 2 250
@@ -157,7 +160,7 @@ deliverMessagesViaProxy proxyServ relayServ alg unsecuredMsgs securedMsgs = do
   -- prepare receiving queue
   (rPub, rPriv) <- atomically $ C.generateAuthKeyPair alg g
   (rdhPub, rdhPriv :: C.PrivateKeyX25519) <- atomically $ C.generateKeyPair g
-  QIK {rcvId, sndId, rcvPublicDhKey = srvDh} <- runExceptT' $ createSMPQueue rc (rPub, rPriv) rdhPub (Just "correct") SMSubscribe False
+  SMP.QIK {rcvId, sndId, rcvPublicDhKey = srvDh} <- runExceptT' $ createSMPQueue rc (rPub, rPriv) rdhPub (Just "correct") SMSubscribe False
   let dec = decryptMsgV3 $ C.dh' srvDh rdhPriv
   -- get proxy session
   sess0 <- runExceptT' $ connectSMPProxiedRelay pc relayServ (Just "correct")
@@ -374,18 +377,38 @@ agentViaProxyRetryOffline = do
     msgId = subtract baseId . fst
     servers srv = (initAgentServersProxy SPMAlways SPFProhibit) {smp = userServers [srv]}
 
+agentViaProxyRetryNoSession :: IO ()
+agentViaProxyRetryNoSession = do
+  let srv1 = SMPServer testHost testPort testKeyHash
+      srv2 = SMPServer testHost testPort2 testKeyHash
+  withAgent 1 agentCfg (servers srv1) testDB $ \a ->
+    withAgent 2 agentCfg (servers srv2) testDB2 $ \b -> do
+      withSmpServerConfigOn (transport @TLS) proxyCfg testPort $ \_ -> do
+        (aId, _) <- withServer2 $ \_ -> runRight $ makeConnection a b
+        nGet b =##> \case ("", "", DOWN _ [c]) -> c == aId; _ -> False
+        withServer2 $ \_ -> do
+          nGet b =##> \case ("", "", UP _ [c]) -> c == aId; _ -> False
+          -- to test retry in case of NO_SESSION error,
+          -- the client using server 1 as proxy and server 2 as destination
+          -- should be joining the connection, so the order is swapped here.
+          _ <- runRight $ makeConnection b a
+          pure ()
+  where
+    withServer2 = withSmpServerConfigOn (transport @TLS) proxyCfg {storeLogFile = Just testStoreLogFile2, storeMsgsFile = Just testStoreMsgsFile2} testPort2
+    servers srv = (initAgentServersProxy SPMAlways SPFProhibit) {smp = userServers [srv]}
+
 testNoProxy :: IO ()
 testNoProxy = do
   withSmpServerConfigOn (transport @TLS) cfg testPort2 $ \_ -> do
     testSMPClient_ "127.0.0.1" testPort2 proxyVRangeV8 $ \(th :: THandleSMP TLS 'TClient) -> do
-      (_, _, (_corrId, _entityId, reply)) <- sendRecv th (Nothing, "0", "", PRXY testSMPServer Nothing)
+      (_, _, (_corrId, _entityId, reply)) <- sendRecv th (Nothing, "0", "", SMP.PRXY testSMPServer Nothing)
       reply `shouldBe` Right (SMP.ERR $ SMP.PROXY SMP.BASIC_AUTH)
 
 testProxyAuth :: IO ()
 testProxyAuth = do
   withSmpServerConfigOn (transport @TLS) proxyCfgAuth testPort $ \_ -> do
     testSMPClient_ "127.0.0.1" testPort proxyVRangeV8 $ \(th :: THandleSMP TLS 'TClient) -> do
-      (_, _s, (_corrId, _entityId, reply)) <- sendRecv th (Nothing, "0", "", PRXY testSMPServer2 $ Just "wrong")
+      (_, _s, (_corrId, _entityId, reply)) <- sendRecv th (Nothing, "0", "", SMP.PRXY testSMPServer2 $ Just "wrong")
       reply `shouldBe` Right (SMP.ERR $ SMP.PROXY SMP.BASIC_AUTH)
   where
     proxyCfgAuth = proxyCfg {newQueueBasicAuth = Just "correct"}

From 492d2f86bc5fa23f33f0b5e743c9b44e1d19049b Mon Sep 17 00:00:00 2001
From: Evgeny Poberezkin <evgeny@poberezkin.com>
Date: Sat, 13 Jul 2024 22:34:10 +0100
Subject: [PATCH 04/10] smp server: additional control port commands to monitor
 server state (#1228)

* smp server: additional control port commands to monitor server state

* fix

* space
---
 src/Simplex/Messaging/Server.hs         | 46 ++++++++++++++++++++++---
 src/Simplex/Messaging/Server/Control.hs |  4 +++
 2 files changed, 46 insertions(+), 4 deletions(-)

diff --git a/src/Simplex/Messaging/Server.hs b/src/Simplex/Messaging/Server.hs
index 014bcd366..d88b2349a 100644
--- a/src/Simplex/Messaging/Server.hs
+++ b/src/Simplex/Messaging/Server.hs
@@ -6,7 +6,6 @@
 {-# LANGUAGE GADTs #-}
 {-# LANGUAGE KindSignatures #-}
 {-# LANGUAGE LambdaCase #-}
-{-# LANGUAGE MultiWayIf #-}
 {-# LANGUAGE NamedFieldPuns #-}
 {-# LANGUAGE NumericUnderscores #-}
 {-# LANGUAGE OverloadedLists #-}
@@ -60,6 +59,7 @@ import Data.List.NonEmpty (NonEmpty (..))
 import qualified Data.List.NonEmpty as L
 import qualified Data.Map.Strict as M
 import Data.Maybe (catMaybes, fromMaybe, isJust, isNothing)
+import qualified Data.Set as S
 import qualified Data.Text as T
 import Data.Text.Encoding (decodeLatin1)
 import Data.Time.Clock (UTCTime (..), diffTimeToPicoseconds, getCurrentTime)
@@ -377,7 +377,7 @@ smpServer started cfg@ServerConfig {transports, transportConfig = tCfg} = do
                   let age = systemSeconds now - systemSeconds createdAt
                   subscriptions' <- bshow . M.size <$> readTVarIO subscriptions
                   hPutStrLn h . B.unpack $ B.intercalate "," [bshow cid, encode sessionId, connected', strEncode createdAt, rcvActiveAt', sndActiveAt', bshow age, subscriptions']
-              CPStats -> withAdminRole $ do
+              CPStats -> withUserRole $ do
                 ss <- unliftIO u $ asks serverStats
                 let putStat :: Show a => ByteString -> (ServerStats -> TVar a) -> IO ()
                     putStat label var = readTVarIO (var ss) >>= \v -> B.hPutStr h $ label <> ": " <> bshow v <> "\n"
@@ -391,6 +391,7 @@ smpServer started cfg@ServerConfig {transports, transportConfig = tCfg} = do
                 putStat "qDeletedAll" qDeletedAll
                 putStat "qDeletedNew" qDeletedNew
                 putStat "qDeletedSecured" qDeletedSecured
+                readTVarIO (day $ activeQueues ss) >>= \v -> B.hPutStr h $ "dayMsgQueues" <> ": " <> bshow (S.size v) <> "\n"
                 putStat "msgSent" msgSent
                 putStat "msgRecv" msgRecv
                 putStat "msgSentNtf" msgSentNtf
@@ -414,7 +415,7 @@ smpServer started cfg@ServerConfig {transports, transportConfig = tCfg} = do
 #else
                 hPutStrLn h "Not available on GHC 8.10"
 #endif
-              CPSockets -> withAdminRole $ do
+              CPSockets -> withUserRole $ do
                 (accepted', closed', active') <- unliftIO u $ asks sockets
                 (accepted, closed, active) <- atomically $ (,,) <$> readTVar accepted' <*> readTVar closed' <*> readTVar active'
                 hPutStrLn h "Sockets: "
@@ -436,6 +437,43 @@ smpServer started cfg@ServerConfig {transports, transportConfig = tCfg} = do
 #else
                 hPutStrLn h "Not available on GHC 8.10"
 #endif
+              CPServerInfo -> readTVarIO role >>= \case
+                CPRNone -> do
+                  logError "Unauthorized control port command"
+                  hPutStrLn h "AUTH"
+                r -> do
+#if MIN_VERSION_base(4,18,0)
+                  threads <- liftIO listThreads
+                  hPutStrLn h $ "Threads: " <> show (length threads)
+#else
+                  hPutStrLn h "Threads: not available on GHC 8.10"
+#endif
+                  Env {clients, server = Server {subscribers, notifiers}} <- unliftIO u ask
+                  activeClients <- readTVarIO clients
+                  hPutStrLn h $ "Clients: " <> show (IM.size activeClients)
+                  when (r == CPRAdmin) $ do
+                    (smpSubCnt, smpClCnt) <- countClientSubs subscriptions activeClients
+                    (ntfSubCnt, ntfClCnt) <- countClientSubs ntfSubscriptions activeClients
+                    hPutStrLn h $ "SMP subscriptions (via clients, slow): " <> show smpSubCnt
+                    hPutStrLn h $ "SMP subscribed clients (via clients, slow): " <> show smpClCnt
+                    hPutStrLn h $ "Ntf subscriptions (via clients, slow): " <> show ntfSubCnt
+                    hPutStrLn h $ "Ntf subscribed clients (via clients, slow): " <> show ntfClCnt
+                  activeSubs <- readTVarIO subscribers
+                  activeNtfSubs <- readTVarIO notifiers
+                  hPutStrLn h $ "SMP subscriptions: " <> show (M.size activeSubs)
+                  hPutStrLn h $ "SMP subscribed clients: " <> show (countSubClients activeSubs)
+                  hPutStrLn h $ "Ntf subscriptions: " <> show (M.size activeNtfSubs)
+                  hPutStrLn h $ "Ntf subscribed clients: " <> show (countSubClients activeNtfSubs)
+                  where
+                    countClientSubs :: (Client -> TMap QueueId a) -> IM.IntMap Client -> IO (Int, Int)
+                    countClientSubs subSel = foldM addSubs (0, 0)
+                      where
+                        addSubs :: (Int, Int) -> Client -> IO (Int, Int)
+                        addSubs (subCnt, clCnt) cl = do
+                          subs <- readTVarIO $ subSel cl
+                          let cnt = M.size subs
+                          pure (subCnt + cnt, clCnt + if cnt == 0 then 0 else 1)
+                    countSubClients = S.size . M.foldr' (S.insert . clientId) S.empty
               CPDelete queueId' -> withUserRole $ unliftIO u $ do
                 st <- asks queueStore
                 ms <- asks msgStore
@@ -455,7 +493,7 @@ smpServer started cfg@ServerConfig {transports, transportConfig = tCfg} = do
                 hPutStrLn h "saving server state..."
                 unliftIO u $ saveServer True
                 hPutStrLn h "server state saved!"
-              CPHelp -> hPutStrLn h "commands: stats, stats-rts, clients, sockets, socket-threads, threads, delete, save, help, quit"
+              CPHelp -> hPutStrLn h "commands: stats, stats-rts, clients, sockets, socket-threads, threads, server-info, delete, save, help, quit"
               CPQuit -> pure ()
               CPSkip -> pure ()
               where
diff --git a/src/Simplex/Messaging/Server/Control.hs b/src/Simplex/Messaging/Server/Control.hs
index 9463fa777..b4c74e4ac 100644
--- a/src/Simplex/Messaging/Server/Control.hs
+++ b/src/Simplex/Messaging/Server/Control.hs
@@ -9,6 +9,7 @@ import Simplex.Messaging.Encoding.String
 import Simplex.Messaging.Protocol (BasicAuth)
 
 data CPClientRole = CPRNone | CPRUser | CPRAdmin
+  deriving (Eq)
 
 data ControlProtocol
   = CPAuth BasicAuth
@@ -20,6 +21,7 @@ data ControlProtocol
   | CPThreads
   | CPSockets
   | CPSocketThreads
+  | CPServerInfo
   | CPDelete ByteString
   | CPSave
   | CPHelp
@@ -37,6 +39,7 @@ instance StrEncoding ControlProtocol where
     CPThreads -> "threads"
     CPSockets -> "sockets"
     CPSocketThreads -> "socket-threads"
+    CPServerInfo -> "server-info"
     CPDelete bs -> "delete " <> strEncode bs
     CPSave -> "save"
     CPHelp -> "help"
@@ -53,6 +56,7 @@ instance StrEncoding ControlProtocol where
       "threads" -> pure CPThreads
       "sockets" -> pure CPSockets
       "socket-threads" -> pure CPSocketThreads
+      "server-info" -> pure CPServerInfo
       "delete" -> CPDelete <$> (A.space *> strP)
       "save" -> pure CPSave
       "help" -> pure CPHelp

From d4fa0af350dc44edbeaeff1d4881be1c830cb94b Mon Sep 17 00:00:00 2001
From: Evgeny Poberezkin <evgeny@poberezkin.com>
Date: Sun, 14 Jul 2024 17:57:34 +0100
Subject: [PATCH 05/10] ntf: additional tests for token registration when
 server and device are restarted (#1230)

* ntf: additional tests for token registration when server and device are restarted

* test response timeouts
---
 tests/AgentTests/NotificationTests.hs | 152 ++++++++++++++++++++++++--
 1 file changed, 141 insertions(+), 11 deletions(-)

diff --git a/tests/AgentTests/NotificationTests.hs b/tests/AgentTests/NotificationTests.hs
index a104c6cf5..92d97d641 100644
--- a/tests/AgentTests/NotificationTests.hs
+++ b/tests/AgentTests/NotificationTests.hs
@@ -5,6 +5,7 @@
 {-# LANGUAGE NamedFieldPuns #-}
 {-# LANGUAGE OverloadedStrings #-}
 {-# LANGUAGE PatternSynonyms #-}
+{-# LANGUAGE QuasiQuotes #-}
 {-# LANGUAGE TypeApplications #-}
 {-# OPTIONS_GHC -fno-warn-ambiguous-fields #-}
 
@@ -49,6 +50,7 @@ import qualified Data.ByteString.Base64.URL as U
 import Data.ByteString.Char8 (ByteString)
 import qualified Data.ByteString.Char8 as B
 import Data.Text.Encoding (encodeUtf8)
+import Database.SQLite.Simple.QQ (sql)
 import NtfClient
 import SMPAgentClient (agentCfg, initAgentServers, initAgentServers2, testDB, testDB2, testNtfServer, testNtfServer2)
 import SMPClient (cfg, cfgVPrev, testPort, testPort2, testStoreLogFile2, withSmpServer, withSmpServerConfigOn, withSmpServerStoreLogOn)
@@ -56,13 +58,14 @@ import Simplex.Messaging.Agent hiding (createConnection, joinConnection, sendMes
 import Simplex.Messaging.Agent.Client (ProtocolTestFailure (..), ProtocolTestStep (..), withStore')
 import Simplex.Messaging.Agent.Env.SQLite (AgentConfig, Env (..), InitialAgentServers)
 import Simplex.Messaging.Agent.Protocol hiding (CON, CONF, INFO, SENT)
-import Simplex.Messaging.Agent.Store.SQLite (closeSQLiteStore, getSavedNtfToken, reopenSQLiteStore)
+import Simplex.Messaging.Agent.Store.SQLite (closeSQLiteStore, getSavedNtfToken, reopenSQLiteStore, withTransaction)
+import qualified Simplex.Messaging.Agent.Store.SQLite.DB as DB
 import qualified Simplex.Messaging.Crypto as C
 import Simplex.Messaging.Encoding.String
 import Simplex.Messaging.Notifications.Protocol
 import Simplex.Messaging.Notifications.Server.Env (NtfServerConfig (..))
 import Simplex.Messaging.Notifications.Server.Push.APNS
-import Simplex.Messaging.Notifications.Types (NtfToken (..))
+import Simplex.Messaging.Notifications.Types (NtfTknAction (..), NtfToken (..))
 import Simplex.Messaging.Protocol (ErrorType (AUTH), MsgFlags (MsgFlags), NtfServer, ProtocolServer (..), SMPMsgMeta (..), SubscriptionMode (..))
 import qualified Simplex.Messaging.Protocol as SMP
 import Simplex.Messaging.Server.Env.STM (ServerConfig (..))
@@ -88,9 +91,21 @@ notificationTests t = do
     it "should allow the second registration with different credentials and delete the first after verification" $
       withAPNSMockServer $ \apns ->
         withNtfServer t $ testNtfTokenSecondRegistration apns
-    it "should re-register token when notification server is restarted" $
+    it "should verify token after notification server is restarted" $
       withAPNSMockServer $ \apns ->
         testNtfTokenServerRestart t apns
+    it "should re-verify token after notification server is restarted" $
+      withAPNSMockServer $ \apns ->
+        testNtfTokenServerRestartReverify t apns
+    it "should re-verify token after notification server is restarted when first request timed-out" $
+      withAPNSMockServer $ \apns ->
+        testNtfTokenServerRestartReverifyTimeout t apns
+    it "should re-register token when notification server is restarted" $
+      withAPNSMockServer $ \apns ->
+        testNtfTokenServerRestartReregister t apns
+    it "should re-register token when notification server is restarted when first request timed-out" $
+      withAPNSMockServer $ \apns ->
+        testNtfTokenServerRestartReregisterTimeout t apns
     it "should work with multiple configured servers" $
       withAPNSMockServer $ \apns ->
         testNtfTokenMultipleServers t apns
@@ -251,7 +266,7 @@ testNtfTokenServerRestart :: ATransport -> APNSMockServer -> IO ()
 testNtfTokenServerRestart t APNSMockServer {apnsQ} = do
   let tkn = DeviceToken PPApnsTest "abcd"
   ntfData <- withAgent 1 agentCfg initAgentServers testDB $ \a ->
-    withNtfServer t . runRight $ do
+    withNtfServerStoreLog t $ \_ -> runRight $ do
       NTRegistered <- registerNtfToken a tkn NMPeriodic
       APNSMockRequest {notification = APNSNotification {aps = APNSBackground _, notificationData = Just ntfData}, sendApnsResponse} <-
         atomically $ readTBQueue apnsQ
@@ -262,16 +277,131 @@ testNtfTokenServerRestart t APNSMockServer {apnsQ} = do
   withAgent 2 agentCfg initAgentServers testDB $ \a' ->
     -- server stopped before token is verified, so now the attempt to verify it will return AUTH error but re-register token,
     -- so that repeat verification happens without restarting the clients, when notification arrives
-    withNtfServer t . runRight_ $ do
+    withNtfServerStoreLog t $ \_ -> runRight_ $ do
       verification <- ntfData .-> "verification"
       nonce <- C.cbNonce <$> ntfData .-> "nonce"
-      Left (NTF _ AUTH) <- tryE $ verifyNtfToken a' tkn nonce verification
-      APNSMockRequest {notification = APNSNotification {aps = APNSBackground _, notificationData = Just ntfData'}, sendApnsResponse = sendApnsResponse'} <-
+      verifyNtfToken a' tkn nonce verification
+      NTActive <- checkNtfToken a' tkn
+      pure ()
+
+testNtfTokenServerRestartReverify :: ATransport -> APNSMockServer -> IO ()
+testNtfTokenServerRestartReverify t APNSMockServer {apnsQ} = do
+  let tkn = DeviceToken PPApnsTest "abcd"
+  withAgent 1 agentCfg initAgentServers testDB $ \a -> do
+    ntfData <- withNtfServerStoreLog t $ \_ -> runRight $ do
+      NTRegistered <- registerNtfToken a tkn NMPeriodic
+      APNSMockRequest {notification = APNSNotification {aps = APNSBackground _, notificationData = Just ntfData}, sendApnsResponse} <-
         atomically $ readTBQueue apnsQ
-      verification' <- ntfData' .-> "verification"
-      nonce' <- C.cbNonce <$> ntfData' .-> "nonce"
-      liftIO $ sendApnsResponse' APNSRespOk
-      verifyNtfToken a' tkn nonce' verification'
+      liftIO $ sendApnsResponse APNSRespOk
+      pure ntfData
+    runRight_ $ do
+      verification <- ntfData .-> "verification"
+      nonce <- C.cbNonce <$> ntfData .-> "nonce"
+      Left (BROKER _ NETWORK) <- tryE $ verifyNtfToken a tkn nonce verification
+      pure ()
+  threadDelay 1000000
+  withAgent 2 agentCfg initAgentServers testDB $ \a' ->
+    -- server stopped before token is verified, so now the attempt to verify it will return AUTH error but re-register token,
+    -- so that repeat verification happens without restarting the clients, when notification arrives
+    withNtfServerStoreLog t $ \_ -> runRight_ $ do
+      NTActive <- registerNtfToken a' tkn NMPeriodic
+      NTActive <- checkNtfToken a' tkn
+      pure ()
+
+testNtfTokenServerRestartReverifyTimeout :: ATransport -> APNSMockServer -> IO ()
+testNtfTokenServerRestartReverifyTimeout t APNSMockServer {apnsQ} = do
+  let tkn = DeviceToken PPApnsTest "abcd"
+  withAgent 1 agentCfg initAgentServers testDB $ \a@AgentClient {agentEnv = Env {store}} -> do
+    (nonce, verification) <- withNtfServerStoreLog t $ \_ -> runRight $ do
+      NTRegistered <- registerNtfToken a tkn NMPeriodic
+      APNSMockRequest {notification = APNSNotification {aps = APNSBackground _, notificationData = Just ntfData}, sendApnsResponse} <-
+        atomically $ readTBQueue apnsQ
+      liftIO $ sendApnsResponse APNSRespOk
+      verification <- ntfData .-> "verification"
+      nonce <- C.cbNonce <$> ntfData .-> "nonce"
+      verifyNtfToken a tkn nonce verification
+      pure (nonce, verification)
+    -- this emulates the situation when server verified token but the client did not receive the response
+    Just NtfToken {ntfTknStatus = NTActive, ntfTknAction = Just NTACheck, ntfDhSecret = Just dhSecret} <- withTransaction store getSavedNtfToken
+    Right code <- pure $ NtfRegCode <$> C.cbDecrypt dhSecret nonce verification
+    withTransaction store $ \db ->
+      DB.execute
+        db
+        [sql|
+          UPDATE ntf_tokens
+          SET tkn_status = ?, tkn_action = ?
+          WHERE provider = ? AND device_token = ?
+        |]
+        (NTConfirmed, Just (NTAVerify code), PPApnsTest, "abcd" :: ByteString)
+    Just NtfToken {ntfTknStatus = NTConfirmed, ntfTknAction = Just (NTAVerify _)} <- withTransaction store getSavedNtfToken
+    pure ()
+  threadDelay 1000000
+  withAgent 2 agentCfg initAgentServers testDB $ \a' ->
+    -- server stopped before token is verified, so now the attempt to verify it will return AUTH error but re-register token,
+    -- so that repeat verification happens without restarting the clients, when notification arrives
+    withNtfServerStoreLog t $ \_ -> runRight_ $ do
+      NTActive <- registerNtfToken a' tkn NMPeriodic
+      NTActive <- checkNtfToken a' tkn
+      pure ()
+
+testNtfTokenServerRestartReregister :: ATransport -> APNSMockServer -> IO ()
+testNtfTokenServerRestartReregister t APNSMockServer {apnsQ} = do
+  let tkn = DeviceToken PPApnsTest "abcd"
+  withAgent 1 agentCfg initAgentServers testDB $ \a ->
+    withNtfServerStoreLog t $ \_ -> runRight $ do
+      NTRegistered <- registerNtfToken a tkn NMPeriodic
+      APNSMockRequest {notification = APNSNotification {aps = APNSBackground _, notificationData = Just _}, sendApnsResponse} <-
+        atomically $ readTBQueue apnsQ
+      liftIO $ sendApnsResponse APNSRespOk
+  -- the new agent is created as otherwise when running the tests in CI the old agent was keeping the connection to the server
+  threadDelay 1000000
+  withAgent 2 agentCfg initAgentServers testDB $ \a' ->
+    -- server stopped before token is verified, and client might have lost verification notification.
+    -- so that repeat registration happens when client is restarted.
+    withNtfServerStoreLog t $ \_ -> runRight_ $ do
+      NTRegistered <- registerNtfToken a' tkn NMPeriodic
+      APNSMockRequest {notification = APNSNotification {aps = APNSBackground _, notificationData = Just ntfData}, sendApnsResponse} <-
+        atomically $ readTBQueue apnsQ
+      liftIO $ sendApnsResponse APNSRespOk
+      verification <- ntfData .-> "verification"
+      nonce <- C.cbNonce <$> ntfData .-> "nonce"
+      verifyNtfToken a' tkn nonce verification
+      NTActive <- checkNtfToken a' tkn
+      pure ()
+
+testNtfTokenServerRestartReregisterTimeout :: ATransport -> APNSMockServer -> IO ()
+testNtfTokenServerRestartReregisterTimeout t APNSMockServer {apnsQ} = do
+  let tkn = DeviceToken PPApnsTest "abcd"
+  withAgent 1 agentCfg initAgentServers testDB $ \a@AgentClient {agentEnv = Env {store}} -> do
+    withNtfServerStoreLog t $ \_ -> runRight $ do
+      NTRegistered <- registerNtfToken a tkn NMPeriodic
+      APNSMockRequest {notification = APNSNotification {aps = APNSBackground _, notificationData = Just _}, sendApnsResponse} <-
+        atomically $ readTBQueue apnsQ
+      liftIO $ sendApnsResponse APNSRespOk
+    -- this emulates the situation when server registered token but the client did not receive the response
+    withTransaction store $ \db ->
+      DB.execute
+        db
+        [sql|
+          UPDATE ntf_tokens
+          SET tkn_id = NULL, tkn_dh_secret = NULL, tkn_status = ?, tkn_action = ?
+          WHERE provider = ? AND device_token = ?
+        |]
+        (NTNew, Just NTARegister, PPApnsTest, "abcd" :: ByteString)
+    Just NtfToken {ntfTokenId = Nothing, ntfTknStatus = NTNew, ntfTknAction = Just NTARegister} <- withTransaction store getSavedNtfToken
+    pure ()
+  threadDelay 1000000
+  withAgent 2 agentCfg initAgentServers testDB $ \a' ->
+    -- server stopped before token is verified, and client might have lost verification notification.
+    -- so that repeat registration happens when client is restarted.
+    withNtfServerStoreLog t $ \_ -> runRight_ $ do
+      NTRegistered <- registerNtfToken a' tkn NMPeriodic
+      APNSMockRequest {notification = APNSNotification {aps = APNSBackground _, notificationData = Just ntfData}, sendApnsResponse} <-
+        atomically $ readTBQueue apnsQ
+      liftIO $ sendApnsResponse APNSRespOk
+      verification <- ntfData .-> "verification"
+      nonce <- C.cbNonce <$> ntfData .-> "nonce"
+      verifyNtfToken a' tkn nonce verification
       NTActive <- checkNtfToken a' tkn
       pure ()
 

From 291039159fbba8e3aa0ba244890a4f1085756cfc Mon Sep 17 00:00:00 2001
From: Evgeny Poberezkin <evgeny@poberezkin.com>
Date: Sun, 14 Jul 2024 23:19:02 +0100
Subject: [PATCH 06/10] ntf server: use SOCKS proxy to connect to onion-only
 SMP servers (#1229)

* ntf server: use SOCKS proxy to connect to onion-only SMP servers

* fix test
---
 src/Simplex/Messaging/Client.hs               |  2 +-
 .../Messaging/Notifications/Server/Main.hs    | 23 ++++++++++++++++++-
 src/Simplex/Messaging/Server/CLI.hs           |  7 ++++++
 src/Simplex/Messaging/Server/Main.hs          |  5 ----
 tests/AgentTests/FunctionalAPITests.hs        |  3 +--
 5 files changed, 31 insertions(+), 9 deletions(-)

diff --git a/src/Simplex/Messaging/Client.hs b/src/Simplex/Messaging/Client.hs
index 4b1e673b0..4cfef92dd 100644
--- a/src/Simplex/Messaging/Client.hs
+++ b/src/Simplex/Messaging/Client.hs
@@ -240,7 +240,7 @@ data SocksMode
   = -- | always use SOCKS proxy when enabled
     SMAlways
   | -- | use SOCKS proxy only for .onion hosts when no public host is available
-    -- This mode is used in SMP proxy to minimize SOCKS proxy usage.
+    -- This mode is used in SMP proxy and in notifications server to minimize SOCKS proxy usage.
     SMOnion
   deriving (Eq, Show)
 
diff --git a/src/Simplex/Messaging/Notifications/Server/Main.hs b/src/Simplex/Messaging/Notifications/Server/Main.hs
index 351fe6d72..477b12bfc 100644
--- a/src/Simplex/Messaging/Notifications/Server/Main.hs
+++ b/src/Simplex/Messaging/Notifications/Server/Main.hs
@@ -7,6 +7,7 @@
 
 module Simplex.Messaging.Notifications.Server.Main where
 
+import Control.Monad ((<$!>))
 import Data.Functor (($>))
 import Data.Ini (lookupValue, readIniFile)
 import Data.Maybe (fromMaybe)
@@ -14,6 +15,7 @@ import qualified Data.Text as T
 import qualified Data.Text.IO as T
 import Network.Socket (HostName)
 import Options.Applicative
+import Simplex.Messaging.Client (NetworkConfig (..), ProtocolClientConfig (..), SocksMode (..), defaultNetworkConfig)
 import Simplex.Messaging.Client.Agent (SMPClientAgentConfig (..), defaultSMPClientAgentConfig)
 import qualified Simplex.Messaging.Crypto as C
 import Simplex.Messaging.Notifications.Server (runNtfServer)
@@ -87,6 +89,14 @@ ntfServerCLI cfgPath logPath =
             <> ("port: " <> T.pack defaultServerPort <> "\n")
             <> "log_tls_errors: off\n"
             <> "websockets: off\n\n\
+               \[SUBSCRIBER]\n\
+               \# Network configuration for notification server client.\n\
+               \# SOCKS proxy port for subscribing to SMP servers.\n\
+               \# You may need a separate instance of SOCKS proxy for incoming single-hop requests.\n\
+               \# socks_proxy: localhost:9050\n\n\
+               \# `socks_mode` can be 'onion' for SOCKS proxy to be used for .onion destination hosts only (default)\n\
+               \# or 'always' to be used for all destination hosts (can be used if it is an .onion server).\n\
+               \# socks_mode: onion\n\n\
                \[INACTIVE_CLIENTS]\n\
                \# TTL and interval to check inactive clients\n\
                \disconnect: off\n"
@@ -115,7 +125,18 @@ ntfServerCLI cfgPath logPath =
               clientQSize = 64,
               subQSize = 512,
               pushQSize = 1048,
-              smpAgentCfg = defaultSMPClientAgentConfig {persistErrorInterval = 0},
+              smpAgentCfg =
+                defaultSMPClientAgentConfig
+                  { smpCfg =
+                      (smpCfg defaultSMPClientAgentConfig)
+                        { networkConfig =
+                            defaultNetworkConfig
+                              { socksProxy = either error id <$!> strDecodeIni "SUBSCRIBER" "socks_proxy" ini,
+                                socksMode = either (const SMOnion) textToSocksMode $ lookupValue "SUBSCRIBER" "socks_mode" ini
+                              }
+                        },
+                    persistErrorInterval = 0 -- seconds
+                  },
               apnsConfig = defaultAPNSPushClientConfig,
               subsBatchSize = 900,
               inactiveClientExpiration =
diff --git a/src/Simplex/Messaging/Server/CLI.hs b/src/Simplex/Messaging/Server/CLI.hs
index 956b30816..c22c1a161 100644
--- a/src/Simplex/Messaging/Server/CLI.hs
+++ b/src/Simplex/Messaging/Server/CLI.hs
@@ -24,6 +24,7 @@ import qualified Data.X509.File as XF
 import Data.X509.Validation (Fingerprint (..))
 import Network.Socket (HostName, ServiceName)
 import Options.Applicative
+import Simplex.Messaging.Client (SocksMode (..))
 import Simplex.Messaging.Encoding.String
 import Simplex.Messaging.Protocol (ProtoServerWithAuth (..), ProtocolServer (..), ProtocolTypeI)
 import Simplex.Messaging.Transport (ATransport (..), TLS, Transport (..))
@@ -301,3 +302,9 @@ clearDirIfExists path = whenM (doesDirectoryExist path) $ listDirectory path >>=
 
 getEnvPath :: String -> FilePath -> IO FilePath
 getEnvPath name def = maybe def (\case "" -> def; f -> f) <$> lookupEnv name
+
+textToSocksMode :: Text -> SocksMode
+textToSocksMode = \case
+  "always" -> SMAlways
+  "onion" -> SMOnion
+  s -> error . T.unpack $ "Invalid socks_mode: " <> s
diff --git a/src/Simplex/Messaging/Server/Main.hs b/src/Simplex/Messaging/Server/Main.hs
index 04e14544c..a135565bf 100644
--- a/src/Simplex/Messaging/Server/Main.hs
+++ b/src/Simplex/Messaging/Server/Main.hs
@@ -317,11 +317,6 @@ smpServerCLI_ generateSite serveStaticFiles cfgPath logPath =
               serverClientConcurrency = readIniDefault defaultProxyClientConcurrency "PROXY" "client_concurrency" ini,
               information = serverPublicInfo ini
             }
-        textToSocksMode :: Text -> SocksMode
-        textToSocksMode = \case
-          "always" -> SMAlways
-          "onion" -> SMOnion
-          s -> error . T.unpack $ "Invalid socks_mode: " <> s
         textToHostMode :: Text -> HostMode
         textToHostMode = \case
           "public" -> HMPublic
diff --git a/tests/AgentTests/FunctionalAPITests.hs b/tests/AgentTests/FunctionalAPITests.hs
index 1893ae105..3f70ad6ab 100644
--- a/tests/AgentTests/FunctionalAPITests.hs
+++ b/tests/AgentTests/FunctionalAPITests.hs
@@ -1347,12 +1347,11 @@ testExpireManyMessages t =
         liftIO $ threadDelay 2000000
         5 <- sendMessage a bId SMP.noMsgFlags "4" -- this won't expire
         get a =##> \case ("", c, MERR 2 (BROKER _ e)) -> bId == c && (e == TIMEOUT || e == NETWORK); _ -> False
-        -- get a =##> \case ("", c, MERRS [5, 6] (BROKER _ e)) -> bId == c && (e == TIMEOUT || e == NETWORK); _ -> False
         let expected c e = bId == c && (e == TIMEOUT || e == NETWORK)
         get a >>= \case
           ("", c, MERR 3 (BROKER _ e)) -> do
             liftIO $ expected c e `shouldBe` True
-            get a =##> \case ("", c', MERR 4 (BROKER _ e')) -> expected c' e'; ("", c', MERRS [6] (BROKER _ e')) -> expected c' e'; _ -> False
+            get a =##> \case ("", c', MERR 4 (BROKER _ e')) -> expected c' e'; ("", c', MERRS [4] (BROKER _ e')) -> expected c' e'; _ -> False
           ("", c, MERRS [3] (BROKER _ e)) -> do
             liftIO $ expected c e `shouldBe` True
             get a =##> \case ("", c', MERR 4 (BROKER _ e')) -> expected c' e'; _ -> False

From 1bdfc8ae00c201eaad97d13c7ab55bd7539984f4 Mon Sep 17 00:00:00 2001
From: Evgeny Poberezkin <evgeny@poberezkin.com>
Date: Sun, 14 Jul 2024 23:21:14 +0100
Subject: [PATCH 07/10] 6.0.0.1

---
 package.yaml    | 2 +-
 simplexmq.cabal | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/package.yaml b/package.yaml
index fbf95a7e5..303e84ef7 100644
--- a/package.yaml
+++ b/package.yaml
@@ -1,5 +1,5 @@
 name: simplexmq
-version: 6.0.0.0
+version: 6.0.0.1
 synopsis: SimpleXMQ message broker
 description: |
   This package includes <./docs/Simplex-Messaging-Server.html server>,
diff --git a/simplexmq.cabal b/simplexmq.cabal
index b7b1ca76c..0bfd62ce3 100644
--- a/simplexmq.cabal
+++ b/simplexmq.cabal
@@ -5,7 +5,7 @@ cabal-version: 1.12
 -- see: https://github.com/sol/hpack
 
 name:           simplexmq
-version:        6.0.0.0
+version:        6.0.0.1
 synopsis:       SimpleXMQ message broker
 description:    This package includes <./docs/Simplex-Messaging-Server.html server>,
                 <./docs/Simplex-Messaging-Client.html client> and

From 784d36d691381290e72884cc30c59ce646166134 Mon Sep 17 00:00:00 2001
From: Evgeny Poberezkin <evgeny@poberezkin.com>
Date: Mon, 15 Jul 2024 20:28:24 +0100
Subject: [PATCH 08/10] socks mode encoding (#1231)

* socks mode encoding

* test delays

* test delay

* unused import
---
 src/Simplex/Messaging/Client.hs                    | 10 ++++++++++
 src/Simplex/Messaging/Notifications/Server/Main.hs |  2 +-
 src/Simplex/Messaging/Server/CLI.hs                |  7 -------
 src/Simplex/Messaging/Server/Main.hs               |  2 +-
 tests/AgentTests/NotificationTests.hs              |  6 +++---
 5 files changed, 15 insertions(+), 12 deletions(-)

diff --git a/src/Simplex/Messaging/Client.hs b/src/Simplex/Messaging/Client.hs
index 4cfef92dd..80fd65ffc 100644
--- a/src/Simplex/Messaging/Client.hs
+++ b/src/Simplex/Messaging/Client.hs
@@ -244,6 +244,16 @@ data SocksMode
     SMOnion
   deriving (Eq, Show)
 
+instance StrEncoding SocksMode where
+  strEncode = \case
+    SMAlways -> "always"
+    SMOnion -> "onion"
+  strP =
+    A.takeTill (== ' ') >>= \case
+      "always" -> pure SMAlways
+      "onion" -> pure SMOnion
+      _ -> fail "Invalid Socks mode"
+
 -- | network configuration for the client
 data NetworkConfig = NetworkConfig
   { -- | use SOCKS5 proxy
diff --git a/src/Simplex/Messaging/Notifications/Server/Main.hs b/src/Simplex/Messaging/Notifications/Server/Main.hs
index 477b12bfc..dadeb82fc 100644
--- a/src/Simplex/Messaging/Notifications/Server/Main.hs
+++ b/src/Simplex/Messaging/Notifications/Server/Main.hs
@@ -132,7 +132,7 @@ ntfServerCLI cfgPath logPath =
                         { networkConfig =
                             defaultNetworkConfig
                               { socksProxy = either error id <$!> strDecodeIni "SUBSCRIBER" "socks_proxy" ini,
-                                socksMode = either (const SMOnion) textToSocksMode $ lookupValue "SUBSCRIBER" "socks_mode" ini
+                                socksMode = maybe SMOnion (either error id) $! strDecodeIni "SUBSCRIBER" "socks_mode" ini
                               }
                         },
                     persistErrorInterval = 0 -- seconds
diff --git a/src/Simplex/Messaging/Server/CLI.hs b/src/Simplex/Messaging/Server/CLI.hs
index c22c1a161..956b30816 100644
--- a/src/Simplex/Messaging/Server/CLI.hs
+++ b/src/Simplex/Messaging/Server/CLI.hs
@@ -24,7 +24,6 @@ import qualified Data.X509.File as XF
 import Data.X509.Validation (Fingerprint (..))
 import Network.Socket (HostName, ServiceName)
 import Options.Applicative
-import Simplex.Messaging.Client (SocksMode (..))
 import Simplex.Messaging.Encoding.String
 import Simplex.Messaging.Protocol (ProtoServerWithAuth (..), ProtocolServer (..), ProtocolTypeI)
 import Simplex.Messaging.Transport (ATransport (..), TLS, Transport (..))
@@ -302,9 +301,3 @@ clearDirIfExists path = whenM (doesDirectoryExist path) $ listDirectory path >>=
 
 getEnvPath :: String -> FilePath -> IO FilePath
 getEnvPath name def = maybe def (\case "" -> def; f -> f) <$> lookupEnv name
-
-textToSocksMode :: Text -> SocksMode
-textToSocksMode = \case
-  "always" -> SMAlways
-  "onion" -> SMOnion
-  s -> error . T.unpack $ "Invalid socks_mode: " <> s
diff --git a/src/Simplex/Messaging/Server/Main.hs b/src/Simplex/Messaging/Server/Main.hs
index a135565bf..784d0504a 100644
--- a/src/Simplex/Messaging/Server/Main.hs
+++ b/src/Simplex/Messaging/Server/Main.hs
@@ -305,7 +305,7 @@ smpServerCLI_ generateSite serveStaticFiles cfgPath logPath =
                           networkConfig =
                             defaultNetworkConfig
                               { socksProxy = either error id <$!> strDecodeIni "PROXY" "socks_proxy" ini,
-                                socksMode = either (const SMOnion) textToSocksMode $ lookupValue "PROXY" "socks_mode" ini,
+                                socksMode = maybe SMOnion (either error id) $! strDecodeIni "PROXY" "socks_mode" ini,
                                 hostMode = either (const HMPublic) textToHostMode $ lookupValue "PROXY" "host_mode" ini,
                                 requiredHostMode = fromMaybe False $ iniOnOff "PROXY" "required_host_mode" ini
                               }
diff --git a/tests/AgentTests/NotificationTests.hs b/tests/AgentTests/NotificationTests.hs
index 92d97d641..fd737e913 100644
--- a/tests/AgentTests/NotificationTests.hs
+++ b/tests/AgentTests/NotificationTests.hs
@@ -507,14 +507,14 @@ testNotificationSubscriptionExistingConnection APNSMockServer {apnsQ} baseId ali
   threadDelay 500000
   suspendAgent alice 0
   closeSQLiteStore store
-  threadDelay 500000
+  threadDelay 1000000
 
   -- aliceNtf client doesn't have subscription and is allowed to get notification message
   withAgent 3 aliceCfg initAgentServers testDB $ \aliceNtf -> runRight_ $ do
     (_, [SMPMsgMeta {msgFlags = MsgFlags True}]) <- getNotificationMessage aliceNtf nonce message
     pure ()
 
-  threadDelay 500000
+  threadDelay 1000000
   reopenSQLiteStore store
   foregroundAgent alice
   threadDelay 500000
@@ -524,7 +524,7 @@ testNotificationSubscriptionExistingConnection APNSMockServer {apnsQ} baseId ali
     ackMessage alice bobId (baseId + 1) Nothing
     -- delete notification subscription
     toggleConnectionNtfs alice bobId False
-    liftIO $ threadDelay 250000
+    liftIO $ threadDelay 500000
     -- send message
     2 <- msgId <$> sendMessage bob aliceId (SMP.MsgFlags True) "hello again"
     get bob ##> ("", aliceId, SENT $ baseId + 2)

From b40d55c358ebabeb6cf1eca8e64016e2bae09a15 Mon Sep 17 00:00:00 2001
From: Evgeny Poberezkin <evgeny@poberezkin.com>
Date: Mon, 15 Jul 2024 20:30:39 +0100
Subject: [PATCH 09/10] 6.0.0.2

---
 package.yaml    | 2 +-
 simplexmq.cabal | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/package.yaml b/package.yaml
index 303e84ef7..2604ad56e 100644
--- a/package.yaml
+++ b/package.yaml
@@ -1,5 +1,5 @@
 name: simplexmq
-version: 6.0.0.1
+version: 6.0.0.2
 synopsis: SimpleXMQ message broker
 description: |
   This package includes <./docs/Simplex-Messaging-Server.html server>,
diff --git a/simplexmq.cabal b/simplexmq.cabal
index 0bfd62ce3..6c8c924d5 100644
--- a/simplexmq.cabal
+++ b/simplexmq.cabal
@@ -5,7 +5,7 @@ cabal-version: 1.12
 -- see: https://github.com/sol/hpack
 
 name:           simplexmq
-version:        6.0.0.1
+version:        6.0.0.2
 synopsis:       SimpleXMQ message broker
 description:    This package includes <./docs/Simplex-Messaging-Server.html server>,
                 <./docs/Simplex-Messaging-Client.html client> and

From 8d56b0ba8519bdc68a3777d8d0672fbe01957479 Mon Sep 17 00:00:00 2001
From: Evgeny Poberezkin <evgeny@poberezkin.com>
Date: Wed, 17 Jul 2024 13:01:17 +0100
Subject: [PATCH 10/10] agent: allow starting only XFTP sending workers (#1232)

* agent: allow starting only XFTP sending workers

* use Bool param

* flip bool

* update
---
 src/Simplex/FileTransfer/Agent.hs | 15 ++++++++++++---
 src/Simplex/Messaging/Agent.hs    |  7 ++++++-
 2 files changed, 18 insertions(+), 4 deletions(-)

diff --git a/src/Simplex/FileTransfer/Agent.hs b/src/Simplex/FileTransfer/Agent.hs
index 4683143c5..725444086 100644
--- a/src/Simplex/FileTransfer/Agent.hs
+++ b/src/Simplex/FileTransfer/Agent.hs
@@ -12,6 +12,7 @@
 
 module Simplex.FileTransfer.Agent
   ( startXFTPWorkers,
+    startXFTPSndWorkers,
     closeXFTPAgent,
     toFSFilePath,
     -- Receiving files
@@ -82,13 +83,21 @@ import UnliftIO.Directory
 import qualified UnliftIO.Exception as E
 
 startXFTPWorkers :: AgentClient -> Maybe FilePath -> AM ()
-startXFTPWorkers c workDir = do
+startXFTPWorkers = startXFTPWorkers_ True
+{-# INLINE startXFTPWorkers #-}
+
+startXFTPSndWorkers :: AgentClient -> Maybe FilePath -> AM ()
+startXFTPSndWorkers = startXFTPWorkers_ False
+{-# INLINE startXFTPSndWorkers #-}
+
+startXFTPWorkers_ :: Bool -> AgentClient -> Maybe FilePath -> AM ()
+startXFTPWorkers_ allWorkers c workDir = do
   wd <- asks $ xftpWorkDir . xftpAgent
   atomically $ writeTVar wd workDir
   cfg <- asks config
-  startRcvFiles cfg
+  when allWorkers $ startRcvFiles cfg
   startSndFiles cfg
-  startDelFiles cfg
+  when allWorkers $ startDelFiles cfg
   where
     startRcvFiles :: AgentConfig -> AM ()
     startRcvFiles AgentConfig {rcvFilesTTL} = do
diff --git a/src/Simplex/Messaging/Agent.hs b/src/Simplex/Messaging/Agent.hs
index c08e04298..94ed7f44f 100644
--- a/src/Simplex/Messaging/Agent.hs
+++ b/src/Simplex/Messaging/Agent.hs
@@ -92,6 +92,7 @@ module Simplex.Messaging.Agent
     getNtfTokenData,
     toggleConnectionNtfs,
     xftpStartWorkers,
+    xftpStartSndWorkers,
     xftpReceiveFile,
     xftpDeleteRcvFile,
     xftpDeleteRcvFiles,
@@ -146,7 +147,7 @@ import Data.Time.Clock
 import Data.Time.Clock.System (systemToUTCTime)
 import Data.Traversable (mapAccumL)
 import Data.Word (Word16)
-import Simplex.FileTransfer.Agent (closeXFTPAgent, deleteSndFileInternal, deleteSndFileRemote, deleteSndFilesInternal, deleteSndFilesRemote, startXFTPWorkers, toFSFilePath, xftpDeleteRcvFile', xftpDeleteRcvFiles', xftpReceiveFile', xftpSendDescription', xftpSendFile')
+import Simplex.FileTransfer.Agent (closeXFTPAgent, deleteSndFileInternal, deleteSndFileRemote, deleteSndFilesInternal, deleteSndFilesRemote, startXFTPSndWorkers, startXFTPWorkers, toFSFilePath, xftpDeleteRcvFile', xftpDeleteRcvFiles', xftpReceiveFile', xftpSendDescription', xftpSendFile')
 import Simplex.FileTransfer.Description (ValidFileDescription)
 import Simplex.FileTransfer.Protocol (FileParty (..))
 import Simplex.FileTransfer.Types (RcvFileId, SndFileId)
@@ -522,6 +523,10 @@ xftpStartWorkers :: AgentClient -> Maybe FilePath -> AE ()
 xftpStartWorkers c = withAgentEnv c . startXFTPWorkers c
 {-# INLINE xftpStartWorkers #-}
 
+xftpStartSndWorkers :: AgentClient -> Maybe FilePath -> AE ()
+xftpStartSndWorkers c = withAgentEnv c . startXFTPSndWorkers c
+{-# INLINE xftpStartSndWorkers #-}
+
 -- | Receive XFTP file
 xftpReceiveFile :: AgentClient -> UserId -> ValidFileDescription 'FRecipient -> Maybe CryptoFileArgs -> Bool -> AE RcvFileId
 xftpReceiveFile c = withAgentEnv c .:: xftpReceiveFile' c