From 5ea3f72595951f2893d6d33c670dcf3717699632 Mon Sep 17 00:00:00 2001
From: Alexander Bondarenko <486682+dpwiz@users.noreply.github.com>
Date: Thu, 2 May 2024 19:10:58 +0300
Subject: [PATCH] switch to activeClientSession

---
 src/Simplex/Messaging/Agent.hs        | 3 +--
 src/Simplex/Messaging/Agent/Client.hs | 8 ++------
 src/Simplex/Messaging/Session.hs      | 9 ---------
 3 files changed, 3 insertions(+), 17 deletions(-)

diff --git a/src/Simplex/Messaging/Agent.hs b/src/Simplex/Messaging/Agent.hs
index 487f38eb9..9813b500c 100644
--- a/src/Simplex/Messaging/Agent.hs
+++ b/src/Simplex/Messaging/Agent.hs
@@ -172,7 +172,6 @@ import Simplex.Messaging.Parsers (parse)
 import Simplex.Messaging.Protocol (BrokerMsg, EntityId, ErrorType (AUTH), MsgBody, MsgFlags (..), NtfServer, ProtoServerWithAuth, ProtocolTypeI (..), SMPMsgMeta, SProtocolType (..), SndPublicAuthKey, SubscriptionMode (..), UserProtocol, VersionSMPC, XFTPServerWithAuth)
 import qualified Simplex.Messaging.Protocol as SMP
 import Simplex.Messaging.ServiceScheme (ServiceScheme (..))
-import Simplex.Messaging.Session (checkSessVar)
 import qualified Simplex.Messaging.TMap as TM
 import Simplex.Messaging.Transport (SMPVersion, THandleParams (sessionId))
 import Simplex.Messaging.Util
@@ -2044,7 +2043,7 @@ data ACKd = ACKd | ACKPending
 -- it cannot be finally, unfortunately, as sometimes it needs to be ACK+DEL
 processSMPTransmission :: AgentClient -> ServerTransmission SMPVersion BrokerMsg -> AM ()
 processSMPTransmission c@AgentClient {smpClients, subQ} (tSess@(_, srv, _), _v, sessId, isResponse, rId, cmd) = do
-  unlessM (atomically $ checkSessVar smpClients tSess $ either (const False) ((== sessId) . sessionId . thParams)) $
+  unlessM (atomically $ activeClientSession c tSess sessId) $
     throwE INTERNAL {internalErr = "Transmission from a repalced SMP client, skipping processing"}
   (rq, SomeConn _ conn) <- withStore c (\db -> getRcvConn db srv rId)
   processSMP rq conn $ toConnData conn
diff --git a/src/Simplex/Messaging/Agent/Client.hs b/src/Simplex/Messaging/Agent/Client.hs
index f338115da..ec3b6ac9d 100644
--- a/src/Simplex/Messaging/Agent/Client.hs
+++ b/src/Simplex/Messaging/Agent/Client.hs
@@ -157,7 +157,6 @@ import Data.ByteString.Base64
 import Data.ByteString.Char8 (ByteString)
 import qualified Data.ByteString.Char8 as B
 import Data.Either (lefts, partitionEithers)
-import Data.Function (on)
 import Data.Functor (($>))
 import Data.Int (Int64)
 import Data.List (deleteFirstsBy, foldl', partition, (\\))
@@ -1161,7 +1160,7 @@ temporaryOrHostError = \case
 
 -- | Subscribe to queues. The list of results can have a different order.
 subscribeQueues :: AgentClient -> [RcvQueue] -> AM [(RcvQueue, Either AgentErrorType ())]
-subscribeQueues c@AgentClient {smpClients} qs = do
+subscribeQueues c qs = do
   (errs, qs') <- partitionEithers <$> mapM checkQueue qs
   atomically $ do
     modifyTVar' (subscrConns c) (`S.union` S.fromList (map qConnId qs'))
@@ -1178,7 +1177,7 @@ subscribeQueues c@AgentClient {smpClients} qs = do
     subscribeQueues_ :: Env -> TVar Bool -> SMPClient -> NonEmpty RcvQueue -> IO (BatchResponses SMPClientError ())
     subscribeQueues_ env processed smp qs' = do
       rs <- sendBatch subscribeSMPQueues smp qs'
-      alive <- atomically $ checkSessVar smpClients (transportSession' smp) (either (const False) $ sameClient smp)
+      alive <- atomically $ activeClientSession c (transportSession' smp) (sessionId $ thParams smp)
       if alive
         then do
           atomically $ writeTVar processed True
@@ -1195,9 +1194,6 @@ activeClientSession c tSess sessId = sameSess <$> tryReadSessVar tSess (smpClien
       Just (Right smp) -> sessId == sessionId (thParams smp)
       _ -> False
 
-sameClient :: ProtocolClient v err msg -> ProtocolClient v err msg -> Bool
-sameClient = (==) `on` (sessionId . thParams)
-
 type BatchResponses e r = NonEmpty (RcvQueue, Either e r)
 
 -- statBatchSize is not used to batch the commands, only for traffic statistics
diff --git a/src/Simplex/Messaging/Session.hs b/src/Simplex/Messaging/Session.hs
index d1d29d1fb..75543b481 100644
--- a/src/Simplex/Messaging/Session.hs
+++ b/src/Simplex/Messaging/Session.hs
@@ -40,12 +40,3 @@ removeSessVar' v sessKey vs =
 
 tryReadSessVar :: Ord k => k -> TMap k (SessionVar a) -> STM (Maybe a)
 tryReadSessVar sessKey vs = TM.lookup sessKey vs $>>= (tryReadTMVar . sessionVar)
-
-checkSessVar :: Ord k => TMap k (SessionVar a) -> k -> (a -> Bool) -> STM Bool
-checkSessVar vs sessKey p =
-  TM.lookup sessKey vs >>= \case
-    Nothing -> pure False
-    Just SessionVar {sessionVar} ->
-      tryReadTMVar sessionVar >>= \case
-        Nothing -> pure False
-        Just x -> pure $ p x