From 8286303d0ffdb7a23f110fd373a01dfc4e962d10 Mon Sep 17 00:00:00 2001 From: sh <37271604+shumvgolove@users.noreply.github.com> Date: Fri, 23 Oct 2026 12:55:12 +0400 Subject: [PATCH] smp-server: fix notification store key leak (#1907) * smp-server: remove expired notification store keys * smp-server: remove delivered notification keys * tests: generate MsgId for ntf store test * smp-server: keep ntf store comments in place * smp-server: strict count of expired notifications --- src/Simplex/Messaging/Server.hs | 3 +- src/Simplex/Messaging/Server/NtfStore.hs | 38 +++++++++------- tests/CoreTests/MsgStoreTests.hs | 58 ++++++++++++++++++++++-- 3 files changed, 79 insertions(+), 20 deletions(-) diff --git a/src/Simplex/Messaging/Server.hs b/src/Simplex/Messaging/Server.hs index 3c6aa7241..b35b8d866 100644 --- a/src/Simplex/Messaging/Server.hs +++ b/src/Simplex/Messaging/Server.hs @@ -366,7 +366,7 @@ smpServer started cfg@ServerConfig {transports, transportConfig = tCfg, startOpt runDeliverNtfs ms ns' stats where runDeliverNtfs :: s -> NtfStore -> ServerStats -> IO () - runDeliverNtfs ms (NtfStore ns) stats = do + runDeliverNtfs ms ns'@(NtfStore ns) stats = do ntfs <- M.assocs <$> readTVarIO ns unless (null ntfs) $ getQueueNtfServices @(StoreQueue s) (queueStore ms) ntfs >>= \case @@ -378,6 +378,7 @@ smpServer started cfg@ServerConfig {transports, transportConfig = tCfg, startOpt cIds <- IS.toList <$> readTVarIO subClients forM_ cIds $ \cId -> getServerClient cId srv >>= mapM_ (deliverQueueNtfs ntfs') atomically $ modifyTVar' ns (`M.withoutKeys` S.fromList (map fst deleted)) + deleteEmptyNtfs ns' ntfs where deliverQueueNtfs ntfs' c@Client {ntfSubscriptions} = whenM (currentClient readTVarIO c) $ do diff --git a/src/Simplex/Messaging/Server/NtfStore.hs b/src/Simplex/Messaging/Server/NtfStore.hs index b73fd4860..7673d5b9a 100644 --- a/src/Simplex/Messaging/Server/NtfStore.hs +++ b/src/Simplex/Messaging/Server/NtfStore.hs @@ -10,10 +10,12 @@ module Simplex.Messaging.Server.NtfStore storeNtf, deleteNtfs, deleteExpiredNtfs, + deleteEmptyNtfs, ) where import Control.Concurrent.STM import Control.Monad (foldM) +import Data.Functor (($>)) import Data.Int (Int64) import qualified Data.Map.Strict as M import Data.Time.Clock.System (SystemTime (..)) @@ -22,6 +24,7 @@ import Simplex.Messaging.Encoding.String import Simplex.Messaging.Protocol (EncNMsgMeta, MsgId, NotifierId) import Simplex.Messaging.TMap (TMap) import qualified Simplex.Messaging.TMap as TM +import Simplex.Messaging.Util (whenM) newtype NtfStore = NtfStore (TMap NotifierId (TVar [MsgNtf])) @@ -33,13 +36,11 @@ data MsgNtf = MsgNtf } storeNtf :: NtfStore -> NotifierId -> MsgNtf -> IO () -storeNtf (NtfStore ns) nId ntf = do - TM.lookupIO nId ns >>= atomically . maybe newNtfs (`modifyTVar'` (ntf :)) +storeNtf (NtfStore ns) nId ntf = + atomically $ TM.lookup nId ns >>= maybe (TM.insertM nId (newTVar [ntf]) ns) (`modifyTVar'` (ntf :)) -- TODO [ntfdb] coalesce messages here once the client is updated to process multiple messages -- for single notification. -- when (isJust prevNtf) $ incStat $ msgNtfReplaced stats - where - newNtfs = TM.lookup nId ns >>= maybe (TM.insertM nId (newTVar [ntf]) ns) (`modifyTVar'` (ntf :)) deleteNtfs :: NtfStore -> NotifierId -> IO Int deleteNtfs (NtfStore ns) nId = atomically (TM.lookupDelete nId ns) >>= maybe (pure 0) (fmap length . readTVarIO) @@ -48,18 +49,23 @@ deleteExpiredNtfs :: NtfStore -> Int64 -> IO Int deleteExpiredNtfs (NtfStore ns) old = foldM (\expired -> fmap (expired +) . expireQueue) 0 . M.keys =<< readTVarIO ns where - expireQueue nId = TM.lookupIO nId ns >>= maybe (pure 0) expire - expire v = readTVarIO v >>= \case - [] -> pure 0 - _ -> - atomically $ readTVar v >>= \case - [] -> pure 0 - -- check the last message first, it is the earliest - ntfs | systemSeconds (ntfTs $ last $ ntfs) < old -> do - let !ntfs' = filter (\MsgNtf {ntfTs = ts} -> systemSeconds ts >= old) ntfs - writeTVar v ntfs' - pure $! length ntfs - length ntfs' - _ -> pure 0 + expireQueue nId = atomically $ TM.lookup nId ns >>= maybe (pure 0) (expire nId) + expire nId v = readTVar v >>= \case + [] -> TM.delete nId ns $> 0 + -- check the last message first, it is the earliest + ntfs | systemSeconds (ntfTs $ last $ ntfs) < old -> do + let !ntfs' = filter (\MsgNtf {ntfTs = ts} -> systemSeconds ts >= old) ntfs + if null ntfs' then TM.delete nId ns else writeTVar v ntfs' + pure $! length ntfs - length ntfs' + _ -> pure 0 + +deleteEmptyNtfs :: NtfStore -> [(NotifierId, TVar [MsgNtf])] -> IO () +deleteEmptyNtfs (NtfStore ns) = mapM_ deleteEmpty + where + deleteEmpty (nId, v) = + whenM (null <$> readTVarIO v) $ + atomically $ + TM.lookup nId ns >>= mapM_ (\v' -> whenM (null <$> readTVar v') $ TM.delete nId ns) data NtfLogRecord = NLRv1 NotifierId MsgNtf diff --git a/tests/CoreTests/MsgStoreTests.hs b/tests/CoreTests/MsgStoreTests.hs index 20914e6e1..badf7f2f1 100644 --- a/tests/CoreTests/MsgStoreTests.hs +++ b/tests/CoreTests/MsgStoreTests.hs @@ -24,24 +24,25 @@ import Crypto.Random (ChaChaDRG) import Data.ByteString.Char8 (ByteString) import Data.Int (Int64) import qualified Data.Map.Strict as M -import Data.Time.Clock.System (getSystemTime) +import Data.Time.Clock.System (SystemTime (..), getSystemTime) import Simplex.Messaging.Crypto (pattern MaxLenBS) import qualified Simplex.Messaging.Crypto as C -import Simplex.Messaging.Protocol (EncDataBytes (..), EntityId (..), ErrorType (..), LinkId, Message (..), QueueLinkData, RecipientId, SParty (..), noMsgFlags, randomMsgId) +import Simplex.Messaging.Protocol (EncDataBytes (..), EntityId (..), ErrorType (..), LinkId, Message (..), NotifierId, QueueLinkData, RecipientId, SParty (..), noMsgFlags, randomMsgId) import Simplex.Messaging.Server.MsgStore.STM import Simplex.Messaging.Server.MsgStore.Types +import Simplex.Messaging.Server.NtfStore import Simplex.Messaging.Server.QueueStore import Simplex.Messaging.Server.QueueStore.QueueInfo import Simplex.Messaging.Server.QueueStore.STM (STMQueueStore (..)) import Simplex.Messaging.Server.QueueStore.Types import Simplex.Messaging.TMap (TMap) +import qualified Simplex.Messaging.TMap as TM import Test.Hspec hiding (fit, it) import Util #if defined(dbServerPostgres) import AgentTests.FunctionalAPITests (runRight) import Control.Concurrent (threadDelay) -import Data.Time.Clock.System (SystemTime (..)) import Database.PostgreSQL.Simple (Only (..)) import qualified Database.PostgreSQL.Simple as DB import Simplex.Messaging.Agent.Store.Postgres.Common @@ -63,6 +64,9 @@ msgStoreTests = do it "tryDelPeekMsg (ACK not from NSE) should reset message counts when queue is empty" $ withMsgStore testPostgresStoreConfig testResetMessageCounts it "should expire messages across commit batches" $ withMsgStore testPostgresStoreConfig testExpireMessagesInBatches #endif + describe "Notification store" $ do + it "should remove keys of expired notifications" testExpireNtfs + it "should remove keys of delivered notifications" testDeleteEmptyNtfs where someMsgStoreTests :: MsgStoreClass s => (HasTestEnv => MsgStoreConfig s) -> Spec someMsgStoreTests cfg = do @@ -281,6 +285,54 @@ testDeleteQueueMaps ms = do where expectAuth = either (`shouldBe` AUTH) (\_ -> expectationFailure "deleted queue is still found") +mkMsgNtf :: TVar ChaChaDRG -> Int64 -> IO MsgNtf +mkMsgNtf g ts = do + ntfMsgId <- atomically $ randomMsgId 24 g + ntfNonce <- atomically $ C.randomCbNonce g + pure MsgNtf {ntfMsgId, ntfTs = MkSystemTime ts 0, ntfNonce, ntfEncMeta = "meta"} + +ntfCounts :: NtfStore -> IO [(NotifierId, Int)] +ntfCounts (NtfStore ns) = M.assocs <$> (mapM (fmap length . readTVarIO) =<< readTVarIO ns) + +testExpireNtfs :: IO () +testExpireNtfs = do + g <- C.newRandom + st@(NtfStore ns) <- NtfStore <$> TM.emptyIO + now <- systemSeconds <$> getSystemTime + let old = now - 100 + nId1 = EntityId "notifier 1" + nId2 = EntityId "notifier 2" + nId3 = EntityId "notifier 3" + storeNtf st nId1 =<< mkMsgNtf g old + storeNtf st nId2 =<< mkMsgNtf g old + storeNtf st nId2 =<< mkMsgNtf g now + atomically $ TM.insertM nId3 (newTVar []) ns + ntfCounts st `shouldReturn` [(nId1, 1), (nId2, 2), (nId3, 0)] + deleteExpiredNtfs st (now - 50) `shouldReturn` 2 + ntfCounts st `shouldReturn` [(nId2, 1)] + storeNtf st nId1 =<< mkMsgNtf g now + deleteExpiredNtfs st (now - 50) `shouldReturn` 0 + ntfCounts st `shouldReturn` [(nId1, 1), (nId2, 1)] + +testDeleteEmptyNtfs :: IO () +testDeleteEmptyNtfs = do + g <- C.newRandom + st@(NtfStore ns) <- NtfStore <$> TM.emptyIO + now <- systemSeconds <$> getSystemTime + let nId1 = EntityId "notifier 1" + nId2 = EntityId "notifier 2" + nId3 = EntityId "notifier 3" + nId4 = EntityId "notifier 4" + forM_ ([nId1, nId2, nId3, nId4] :: [NotifierId]) $ \nId -> storeNtf st nId =<< mkMsgNtf g now + ntfs <- M.assocs <$> readTVarIO ns + let flush nId = forM_ (lookup nId ntfs) $ \v -> atomically $ writeTVar v [] + mapM_ flush ([nId1, nId2, nId4] :: [NotifierId]) + storeNtf st nId2 =<< mkMsgNtf g now + deleteNtfs st nId4 `shouldReturn` 0 + storeNtf st nId4 =<< mkMsgNtf g now + deleteEmptyNtfs st ntfs + ntfCounts st `shouldReturn` [(nId2, 1), (nId3, 1), (nId4, 1)] + testLinkData :: QueueLinkData testLinkData = (EncDataBytes "fixed data", EncDataBytes "user data")