From 4660ee9cff25a8372b27191d740a1669ab21c991 Mon Sep 17 00:00:00 2001
From: Evgeny Poberezkin <2769109+epoberezkin@users.noreply.github.com>
Date: Wed, 1 Nov 2023 09:15:51 +0000
Subject: [PATCH] agent: revert to TVar (#883)

---
 src/Simplex/Messaging/Agent/Env/SQLite.hs   |  5 ++-
 src/Simplex/Messaging/Agent/Store/SQLite.hs | 22 ++++++------
 src/Simplex/Messaging/Crypto.hs             |  6 ----
 tests/AgentTests/SQLiteTests.hs             | 39 ++++++++++-----------
 4 files changed, 32 insertions(+), 40 deletions(-)

diff --git a/src/Simplex/Messaging/Agent/Env/SQLite.hs b/src/Simplex/Messaging/Agent/Env/SQLite.hs
index c4cb19a59..cbdae289c 100644
--- a/src/Simplex/Messaging/Agent/Env/SQLite.hs
+++ b/src/Simplex/Messaging/Agent/Env/SQLite.hs
@@ -35,7 +35,6 @@ import Control.Monad.IO.Unlift
 import Control.Monad.Reader
 import Crypto.Random
 import Data.Int (Int64)
-import Data.IORef (IORef, newIORef)
 import Data.List.NonEmpty (NonEmpty)
 import Data.Map (Map)
 import Data.Time.Clock (NominalDiffTime, nominalDay)
@@ -179,7 +178,7 @@ defaultAgentConfig =
 data Env = Env
   { config :: AgentConfig,
     store :: SQLiteStore,
-    random :: IORef ChaChaDRG,
+    random :: TVar ChaChaDRG,
     clientCounter :: TVar Int,
     randomServer :: TVar StdGen,
     ntfSupervisor :: NtfSupervisor,
@@ -188,7 +187,7 @@ data Env = Env
 
 newSMPAgentEnv :: AgentConfig -> SQLiteStore -> IO Env
 newSMPAgentEnv config@AgentConfig {initialClientId} store = do
-  random <- newIORef =<< drgNew
+  random <- newTVarIO =<< drgNew
   clientCounter <- newTVarIO initialClientId
   randomServer <- newTVarIO =<< liftIO newStdGen
   ntfSupervisor <- atomically . newNtfSubSupervisor $ tbqSize config
diff --git a/src/Simplex/Messaging/Agent/Store/SQLite.hs b/src/Simplex/Messaging/Agent/Store/SQLite.hs
index f06db8a3e..72d4c261a 100644
--- a/src/Simplex/Messaging/Agent/Store/SQLite.hs
+++ b/src/Simplex/Messaging/Agent/Store/SQLite.hs
@@ -510,7 +510,7 @@ deleteUsersWithoutConns db = do
   pure userIds
 
 createConn_ ::
-  IORef ChaChaDRG ->
+  TVar ChaChaDRG ->
   ConnData ->
   (ByteString -> IO ()) ->
   IO (Either StoreError ByteString)
@@ -518,7 +518,7 @@ createConn_ gVar cData create = checkConstraint SEConnDuplicate $ case cData of
   ConnData {connId = ""} -> createWithRandomId gVar create
   ConnData {connId} -> create connId $> Right connId
 
-createNewConn :: DB.Connection -> IORef ChaChaDRG -> ConnData -> SConnectionMode c -> IO (Either StoreError ConnId)
+createNewConn :: DB.Connection -> TVar ChaChaDRG -> ConnData -> SConnectionMode c -> IO (Either StoreError ConnId)
 createNewConn db gVar cData@ConnData {userId, connAgentVersion, enableNtfs, duplexHandshake} cMode =
   createConn_ gVar cData $ \connId -> do
     DB.execute db "INSERT INTO connections (user_id, conn_id, conn_mode, smp_agent_version, enable_ntfs, duplex_handshake) VALUES (?,?,?,?,?,?)" (userId, connId, cMode, connAgentVersion, enableNtfs, duplexHandshake)
@@ -543,14 +543,14 @@ updateNewConnSnd db connId sq =
     updateConn :: IO (Either StoreError Int64)
     updateConn = Right <$> addConnSndQueue_ db connId sq
 
-createRcvConn :: DB.Connection -> IORef ChaChaDRG -> ConnData -> RcvQueue -> SConnectionMode c -> IO (Either StoreError ConnId)
+createRcvConn :: DB.Connection -> TVar ChaChaDRG -> ConnData -> RcvQueue -> SConnectionMode c -> IO (Either StoreError ConnId)
 createRcvConn db gVar cData@ConnData {userId, connAgentVersion, enableNtfs, duplexHandshake} q@RcvQueue {server} cMode =
   createConn_ gVar cData $ \connId -> do
     serverKeyHash_ <- createServer_ db server
     DB.execute db "INSERT INTO connections (user_id, conn_id, conn_mode, smp_agent_version, enable_ntfs, duplex_handshake) VALUES (?,?,?,?,?,?)" (userId, connId, cMode, connAgentVersion, enableNtfs, duplexHandshake)
     void $ insertRcvQueue_ db connId q serverKeyHash_
 
-createSndConn :: DB.Connection -> IORef ChaChaDRG -> ConnData -> SndQueue -> IO (Either StoreError ConnId)
+createSndConn :: DB.Connection -> TVar ChaChaDRG -> ConnData -> SndQueue -> IO (Either StoreError ConnId)
 createSndConn db gVar cData@ConnData {userId, connAgentVersion, enableNtfs, duplexHandshake} q@SndQueue {server} =
   -- check confirmed snd queue doesn't already exist, to prevent it being deleted by REPLACE in insertSndQueue_
   ifM (liftIO $ checkConfirmedSndQueueExists_ db q) (pure $ Left SESndQueueExists) $
@@ -769,7 +769,7 @@ smpConfirmation (senderKey, e2ePubKey, connInfo, smpReplyQueues_, smpClientVersi
       smpClientVersion = fromMaybe 1 smpClientVersion_
     }
 
-createConfirmation :: DB.Connection -> IORef ChaChaDRG -> NewConfirmation -> IO (Either StoreError ConfirmationId)
+createConfirmation :: DB.Connection -> TVar ChaChaDRG -> NewConfirmation -> IO (Either StoreError ConfirmationId)
 createConfirmation db gVar NewConfirmation {connId, senderConf = SMPConfirmation {senderKey, e2ePubKey, connInfo, smpReplyQueues, smpClientVersion}, ratchetState} =
   createWithRandomId gVar $ \confirmationId ->
     DB.execute
@@ -847,7 +847,7 @@ setHandshakeVersion :: DB.Connection -> ConnId -> Version -> Bool -> IO ()
 setHandshakeVersion db connId aVersion duplexHS =
   DB.execute db "UPDATE connections SET smp_agent_version = ?, duplex_handshake = ? WHERE conn_id = ?" (aVersion, duplexHS, connId)
 
-createInvitation :: DB.Connection -> IORef ChaChaDRG -> NewInvitation -> IO (Either StoreError InvitationId)
+createInvitation :: DB.Connection -> TVar ChaChaDRG -> NewInvitation -> IO (Either StoreError InvitationId)
 createInvitation db gVar NewInvitation {contactConnId, connReq, recipientConnInfo} =
   createWithRandomId gVar $ \invitationId ->
     DB.execute
@@ -2074,7 +2074,7 @@ updateHashSnd_ dbConn connId SndMsgData {..} =
     ]
 
 -- create record with a random ID
-createWithRandomId :: IORef ChaChaDRG -> (ByteString -> IO ()) -> IO (Either StoreError ByteString)
+createWithRandomId :: TVar ChaChaDRG -> (ByteString -> IO ()) -> IO (Either StoreError ByteString)
 createWithRandomId gVar create = tryCreate 3
   where
     tryCreate :: Int -> IO (Either StoreError ByteString)
@@ -2087,8 +2087,8 @@ createWithRandomId gVar create = tryCreate 3
           | SQL.sqlError e == SQL.ErrorConstraint -> tryCreate (n - 1)
           | otherwise -> pure . Left . SEInternal $ bshow e
 
-randomId :: IORef ChaChaDRG -> Int -> IO ByteString
-randomId gVar n = U.encode <$> C.pseudoRandomBytes' n gVar
+randomId :: TVar ChaChaDRG -> Int -> IO ByteString
+randomId gVar n = atomically $ U.encode <$> C.pseudoRandomBytes n gVar
 
 ntfSubAndSMPAction :: NtfSubAction -> (Maybe NtfSubNTFAction, Maybe NtfSubSMPAction)
 ntfSubAndSMPAction (NtfSubNTFAction action) = (Just action, Nothing)
@@ -2109,7 +2109,7 @@ getXFTPServerId_ db ProtocolServer {host, port, keyHash} = do
   firstRow fromOnly SEXFTPServerNotFound $
     DB.query db "SELECT xftp_server_id FROM xftp_servers WHERE xftp_host = ? AND xftp_port = ? AND xftp_key_hash = ?" (host, port, keyHash)
 
-createRcvFile :: DB.Connection -> IORef ChaChaDRG -> UserId -> FileDescription 'FRecipient -> FilePath -> FilePath -> CryptoFile -> IO (Either StoreError RcvFileId)
+createRcvFile :: DB.Connection -> TVar ChaChaDRG -> UserId -> FileDescription 'FRecipient -> FilePath -> FilePath -> CryptoFile -> IO (Either StoreError RcvFileId)
 createRcvFile db gVar userId fd@FileDescription {chunks} prefixPath tmpPath (CryptoFile savePath cfArgs) = runExceptT $ do
   (rcvFileEntityId, rcvFileId) <- ExceptT $ insertRcvFile fd
   liftIO $
@@ -2364,7 +2364,7 @@ getRcvFilesExpired db ttl = do
     |]
     (Only cutoffTs)
 
-createSndFile :: DB.Connection -> IORef ChaChaDRG -> UserId -> CryptoFile -> Int -> FilePath -> C.SbKey -> C.CbNonce -> IO (Either StoreError SndFileId)
+createSndFile :: DB.Connection -> TVar ChaChaDRG -> UserId -> CryptoFile -> Int -> FilePath -> C.SbKey -> C.CbNonce -> IO (Either StoreError SndFileId)
 createSndFile db gVar userId (CryptoFile path cfArgs) numRecipients prefixPath key nonce =
   createWithRandomId gVar $ \sndFileEntityId ->
     DB.execute
diff --git a/src/Simplex/Messaging/Crypto.hs b/src/Simplex/Messaging/Crypto.hs
index da4ede7d9..cfc8156cf 100644
--- a/src/Simplex/Messaging/Crypto.hs
+++ b/src/Simplex/Messaging/Crypto.hs
@@ -130,7 +130,6 @@ module Simplex.Messaging.Crypto
 
     -- * pseudo-random bytes
     pseudoRandomBytes,
-    pseudoRandomBytes',
 
     -- * digests
     sha256Hash,
@@ -192,10 +191,8 @@ import Data.ByteString.Char8 (ByteString)
 import qualified Data.ByteString.Char8 as B
 import Data.ByteString.Lazy (fromStrict, toStrict)
 import Data.Constraint (Dict (..))
-import Data.IORef (IORef, atomicModifyIORef')
 import Data.Kind (Constraint, Type)
 import Data.String
-import Data.Tuple (swap)
 import Data.Type.Equality
 import Data.Typeable (Proxy (Proxy), Typeable)
 import Data.Word (Word32)
@@ -1147,9 +1144,6 @@ pseudoRandomCbNonce gVar = CryptoBoxNonce <$> pseudoRandomBytes 24 gVar
 pseudoRandomBytes :: Int -> TVar ChaChaDRG -> STM ByteString
 pseudoRandomBytes n gVar = stateTVar gVar $ randomBytesGenerate n
 
-pseudoRandomBytes' :: Int -> IORef ChaChaDRG -> IO ByteString
-pseudoRandomBytes' n gVar = atomicModifyIORef' gVar $ swap . randomBytesGenerate n
-
 instance Encoding CbNonce where
   smpEncode = unCbNonce
   smpP = CryptoBoxNonce <$> A.take 24
diff --git a/tests/AgentTests/SQLiteTests.hs b/tests/AgentTests/SQLiteTests.hs
index b96574587..cf6e8373b 100644
--- a/tests/AgentTests/SQLiteTests.hs
+++ b/tests/AgentTests/SQLiteTests.hs
@@ -17,7 +17,6 @@ import Control.Exception (SomeException)
 import Control.Monad (replicateM_)
 import Crypto.Random (drgNew)
 import Data.ByteString.Char8 (ByteString)
-import Data.IORef (newIORef)
 import Data.List (isInfixOf)
 import qualified Data.Text as T
 import Data.Text.Encoding (encodeUtf8)
@@ -117,7 +116,7 @@ storeTests = do
 testConcurrentWrites :: SpecWith (SQLiteStore, SQLiteStore)
 testConcurrentWrites =
   it "should complete multiple concurrent write transactions w/t sqlite busy errors" $ \(s1, s2) -> do
-    g <- newIORef =<< drgNew
+    g <- newTVarIO =<< drgNew
     _ <- withTransaction s1 $ \db ->
       createRcvConn db g cData1 rcvQueue1 SCMInvitation
     let ConnData {connId} = cData1
@@ -207,7 +206,7 @@ sndQueue1 =
 testCreateRcvConn :: SpecWith SQLiteStore
 testCreateRcvConn =
   it "should create RcvConnection and add SndQueue" . withStoreTransaction $ \db -> do
-    g <- newIORef =<< drgNew
+    g <- newTVarIO =<< drgNew
     createRcvConn db g cData1 rcvQueue1 SCMInvitation
       `shouldReturn` Right "conn1"
     getConn db "conn1"
@@ -220,7 +219,7 @@ testCreateRcvConn =
 testCreateRcvConnRandomId :: SpecWith SQLiteStore
 testCreateRcvConnRandomId =
   it "should create RcvConnection and add SndQueue with random ID" . withStoreTransaction $ \db -> do
-    g <- newIORef =<< drgNew
+    g <- newTVarIO =<< drgNew
     Right connId <- createRcvConn db g cData1 {connId = ""} rcvQueue1 SCMInvitation
     let rq' = (rcvQueue1 :: RcvQueue) {connId}
         sq' = (sndQueue1 :: SndQueue) {connId}
@@ -234,7 +233,7 @@ testCreateRcvConnRandomId =
 testCreateRcvConnDuplicate :: SpecWith SQLiteStore
 testCreateRcvConnDuplicate =
   it "should throw error on attempt to create duplicate RcvConnection" . withStoreTransaction $ \db -> do
-    g <- newIORef =<< drgNew
+    g <- newTVarIO =<< drgNew
     _ <- createRcvConn db g cData1 rcvQueue1 SCMInvitation
     createRcvConn db g cData1 rcvQueue1 SCMInvitation
       `shouldReturn` Left SEConnDuplicate
@@ -242,7 +241,7 @@ testCreateRcvConnDuplicate =
 testCreateSndConn :: SpecWith SQLiteStore
 testCreateSndConn =
   it "should create SndConnection and add RcvQueue" . withStoreTransaction $ \db -> do
-    g <- newIORef =<< drgNew
+    g <- newTVarIO =<< drgNew
     createSndConn db g cData1 sndQueue1
       `shouldReturn` Right "conn1"
     getConn db "conn1"
@@ -255,7 +254,7 @@ testCreateSndConn =
 testCreateSndConnRandomID :: SpecWith SQLiteStore
 testCreateSndConnRandomID =
   it "should create SndConnection and add RcvQueue with random ID" . withStoreTransaction $ \db -> do
-    g <- newIORef =<< drgNew
+    g <- newTVarIO =<< drgNew
     Right connId <- createSndConn db g cData1 {connId = ""} sndQueue1
     let rq' = (rcvQueue1 :: RcvQueue) {connId}
         sq' = (sndQueue1 :: SndQueue) {connId}
@@ -269,7 +268,7 @@ testCreateSndConnRandomID =
 testCreateSndConnDuplicate :: SpecWith SQLiteStore
 testCreateSndConnDuplicate =
   it "should throw error on attempt to create duplicate SndConnection" . withStoreTransaction $ \db -> do
-    g <- newIORef =<< drgNew
+    g <- newTVarIO =<< drgNew
     _ <- createSndConn db g cData1 sndQueue1
     createSndConn db g cData1 sndQueue1
       `shouldReturn` Left SEConnDuplicate
@@ -279,7 +278,7 @@ testGetRcvConn =
   it "should get connection using rcv queue id and server" . withStoreTransaction $ \db -> do
     let smpServer = SMPServer "smp.simplex.im" "5223" testKeyHash
     let recipientId = "1234"
-    g <- newIORef =<< drgNew
+    g <- newTVarIO =<< drgNew
     _ <- createRcvConn db g cData1 rcvQueue1 SCMInvitation
     getRcvConn db smpServer recipientId
       `shouldReturn` Right (rcvQueue1, SomeConn SCRcv (RcvConnection cData1 rcvQueue1))
@@ -287,7 +286,7 @@ testGetRcvConn =
 testDeleteRcvConn :: SpecWith SQLiteStore
 testDeleteRcvConn =
   it "should create RcvConnection and delete it" . withStoreTransaction $ \db -> do
-    g <- newIORef =<< drgNew
+    g <- newTVarIO =<< drgNew
     _ <- createRcvConn db g cData1 rcvQueue1 SCMInvitation
     getConn db "conn1"
       `shouldReturn` Right (SomeConn SCRcv (RcvConnection cData1 rcvQueue1))
@@ -299,7 +298,7 @@ testDeleteRcvConn =
 testDeleteSndConn :: SpecWith SQLiteStore
 testDeleteSndConn =
   it "should create SndConnection and delete it" . withStoreTransaction $ \db -> do
-    g <- newIORef =<< drgNew
+    g <- newTVarIO =<< drgNew
     _ <- createSndConn db g cData1 sndQueue1
     getConn db "conn1"
       `shouldReturn` Right (SomeConn SCSnd (SndConnection cData1 sndQueue1))
@@ -311,7 +310,7 @@ testDeleteSndConn =
 testDeleteDuplexConn :: SpecWith SQLiteStore
 testDeleteDuplexConn =
   it "should create DuplexConnection and delete it" . withStoreTransaction $ \db -> do
-    g <- newIORef =<< drgNew
+    g <- newTVarIO =<< drgNew
     _ <- createRcvConn db g cData1 rcvQueue1 SCMInvitation
     _ <- upgradeRcvConnToDuplex db "conn1" sndQueue1
     getConn db "conn1"
@@ -324,7 +323,7 @@ testDeleteDuplexConn =
 testUpgradeRcvConnToDuplex :: SpecWith SQLiteStore
 testUpgradeRcvConnToDuplex =
   it "should throw error on attempt to add SndQueue to SndConnection or DuplexConnection" . withStoreTransaction $ \db -> do
-    g <- newIORef =<< drgNew
+    g <- newTVarIO =<< drgNew
     _ <- createSndConn db g cData1 sndQueue1
     let anotherSndQueue =
           SndQueue
@@ -352,7 +351,7 @@ testUpgradeRcvConnToDuplex =
 testUpgradeSndConnToDuplex :: SpecWith SQLiteStore
 testUpgradeSndConnToDuplex =
   it "should throw error on attempt to add RcvQueue to RcvConnection or DuplexConnection" . withStoreTransaction $ \db -> do
-    g <- newIORef =<< drgNew
+    g <- newTVarIO =<< drgNew
     _ <- createRcvConn db g cData1 rcvQueue1 SCMInvitation
     let anotherRcvQueue =
           RcvQueue
@@ -383,7 +382,7 @@ testUpgradeSndConnToDuplex =
 testSetRcvQueueStatus :: SpecWith SQLiteStore
 testSetRcvQueueStatus =
   it "should update status of RcvQueue" . withStoreTransaction $ \db -> do
-    g <- newIORef =<< drgNew
+    g <- newTVarIO =<< drgNew
     _ <- createRcvConn db g cData1 rcvQueue1 SCMInvitation
     getConn db "conn1"
       `shouldReturn` Right (SomeConn SCRcv (RcvConnection cData1 rcvQueue1))
@@ -395,7 +394,7 @@ testSetRcvQueueStatus =
 testSetSndQueueStatus :: SpecWith SQLiteStore
 testSetSndQueueStatus =
   it "should update status of SndQueue" . withStoreTransaction $ \db -> do
-    g <- newIORef =<< drgNew
+    g <- newTVarIO =<< drgNew
     _ <- createSndConn db g cData1 sndQueue1
     getConn db "conn1"
       `shouldReturn` Right (SomeConn SCSnd (SndConnection cData1 sndQueue1))
@@ -407,7 +406,7 @@ testSetSndQueueStatus =
 testSetQueueStatusDuplex :: SpecWith SQLiteStore
 testSetQueueStatusDuplex =
   it "should update statuses of RcvQueue and SndQueue in DuplexConnection" . withStoreTransaction $ \db -> do
-    g <- newIORef =<< drgNew
+    g <- newTVarIO =<< drgNew
     _ <- createRcvConn db g cData1 rcvQueue1 SCMInvitation
     _ <- upgradeRcvConnToDuplex db "conn1" sndQueue1
     getConn db "conn1"
@@ -459,7 +458,7 @@ testCreateRcvMsg_ db expectedPrevSndId expectedPrevHash connId rq rcvMsgData@Rcv
 testCreateRcvMsg :: SpecWith SQLiteStore
 testCreateRcvMsg =
   it "should reserve internal ids and create a RcvMsg" $ \st -> do
-    g <- newIORef =<< drgNew
+    g <- newTVarIO =<< drgNew
     let ConnData {connId} = cData1
     _ <- withTransaction st $ \db -> do
       createRcvConn db g cData1 rcvQueue1 SCMInvitation
@@ -490,7 +489,7 @@ testCreateSndMsg_ db expectedPrevHash connId sndMsgData@SndMsgData {..} = do
 testCreateSndMsg :: SpecWith SQLiteStore
 testCreateSndMsg =
   it "should create a SndMsg and return InternalId and PrevSndMsgHash" $ \st -> do
-    g <- newIORef =<< drgNew
+    g <- newTVarIO =<< drgNew
     let ConnData {connId} = cData1
     _ <- withTransaction st $ \db -> do
       createSndConn db g cData1 sndQueue1
@@ -503,7 +502,7 @@ testCreateRcvAndSndMsgs =
   it "should create multiple RcvMsg and SndMsg, correctly ordering internal Ids and returning previous state" $ \st -> do
     let ConnData {connId} = cData1
     _ <- withTransaction st $ \db -> do
-      g <- newIORef =<< drgNew
+      g <- newTVarIO =<< drgNew
       createRcvConn db g cData1 rcvQueue1 SCMInvitation
     withTransaction st $ \db -> do
       _ <- upgradeRcvConnToDuplex db "conn1" sndQueue1