From 64c5d19fc41aa985eb674c42934cfcb93fa4a801 Mon Sep 17 00:00:00 2001 From: Evgeny Date: Sun, 4 Oct 2026 10:18:41 +0100 Subject: [PATCH] agent: aggregate cap on skipped message keys (#1903) * agent: aggregate cap on skipped message keys * migration * update migration --------- Co-authored-by: Evgeny @ SimpleX Chat <259188159+evgeny-simplex@users.noreply.github.com> --- simplexmq.cabal | 1 + src/Simplex/Messaging/Agent.hs | 2 +- .../Messaging/Agent/Store/AgentStore.hs | 12 ++++++--- .../Agent/Store/Postgres/Migrations/App.hs | 4 ++- .../M20260929_skipped_messages_index.hs | 21 +++++++++++++++ .../Migrations/agent_postgres_schema.sql | 2 +- src/Simplex/Messaging/Crypto/Ratchet.hs | 4 +++ tests/AgentTests/SQLiteTests.hs | 26 ++++++++++++++++++- 8 files changed, 64 insertions(+), 8 deletions(-) create mode 100644 src/Simplex/Messaging/Agent/Store/Postgres/Migrations/M20260929_skipped_messages_index.hs diff --git a/simplexmq.cabal b/simplexmq.cabal index 39b40dd33..78fd61a1f 100644 --- a/simplexmq.cabal +++ b/simplexmq.cabal @@ -195,6 +195,7 @@ library Simplex.Messaging.Agent.Store.Postgres.Migrations.M20260712_address_dr_rpc Simplex.Messaging.Agent.Store.Postgres.Migrations.M20260823_snd_files_entitlement Simplex.Messaging.Agent.Store.Postgres.Migrations.M20260919_ratchet_verify_codes + Simplex.Messaging.Agent.Store.Postgres.Migrations.M20260929_skipped_messages_index else exposed-modules: Simplex.Messaging.Agent.Store.SQLite diff --git a/src/Simplex/Messaging/Agent.hs b/src/Simplex/Messaging/Agent.hs index be34deafb..1cf455114 100644 --- a/src/Simplex/Messaging/Agent.hs +++ b/src/Simplex/Messaging/Agent.hs @@ -4386,7 +4386,7 @@ agentRatchetDecrypt g db connId encAgentMsg = do agentRatchetDecrypt' :: TVar ChaChaDRG -> DB.Connection -> ConnId -> CR.RatchetX448 -> ByteString -> ExceptT StoreError IO (ByteString, PQEncryption) agentRatchetDecrypt' g db connId rc encAgentMsg = do - skipped <- liftIO $ getSkippedMsgKeys db connId + skipped <- liftIO $ getSkippedMsgKeys db connId CR.maxSkippedMsgKeys (agentMsgBody_, rc', skippedDiff) <- withExceptT (SEAgentError . cryptoError) $ CR.rcDecrypt g rc skipped encAgentMsg agentMsgBody <- liftEither $ first (SEAgentError . cryptoError) agentMsgBody_ liftIO $ updateRatchet db connId rc' skippedDiff diff --git a/src/Simplex/Messaging/Agent/Store/AgentStore.hs b/src/Simplex/Messaging/Agent/Store/AgentStore.hs index 311cdc14e..38b454c2b 100644 --- a/src/Simplex/Messaging/Agent/Store/AgentStore.hs +++ b/src/Simplex/Messaging/Agent/Store/AgentStore.hs @@ -1594,12 +1594,16 @@ getRatchet_ q db connId = where ratchet = maybe (Left SERatchetNotFound) Right . fromOnly -getSkippedMsgKeys :: DB.Connection -> ConnId -> IO SkippedMsgKeys -getSkippedMsgKeys db connId = - skipped <$> DB.query db "SELECT header_key, msg_n, msg_key FROM skipped_messages WHERE conn_id = ?" (Only connId) +getSkippedMsgKeys :: DB.Connection -> ConnId -> Int -> IO SkippedMsgKeys +getSkippedMsgKeys db connId maxKeys = do + (keys, oldKeys) <- splitAt maxKeys <$> DB.query db "SELECT skipped_message_id, header_key, msg_n, msg_key FROM skipped_messages WHERE conn_id = ? ORDER BY skipped_message_id DESC LIMIT ?" (connId, maxKeys + 1) + case oldKeys of + (skippedMsgId :: Int64, _, _, _) : _ -> DB.execute db "DELETE FROM skipped_messages WHERE conn_id = ? AND skipped_message_id <= ?" (connId, skippedMsgId) + [] -> pure () + pure $ skipped keys where skipped = foldl' addSkippedKey M.empty - addSkippedKey smks (hk, msgN, mk) = M.alter (Just . addMsgKey) hk smks + addSkippedKey smks (_, hk, msgN, mk) = M.alter (Just . addMsgKey) hk smks where addMsgKey = maybe (M.singleton msgN mk) (M.insert msgN mk) diff --git a/src/Simplex/Messaging/Agent/Store/Postgres/Migrations/App.hs b/src/Simplex/Messaging/Agent/Store/Postgres/Migrations/App.hs index f34b120e2..0173c4fb8 100644 --- a/src/Simplex/Messaging/Agent/Store/Postgres/Migrations/App.hs +++ b/src/Simplex/Messaging/Agent/Store/Postgres/Migrations/App.hs @@ -16,6 +16,7 @@ import Simplex.Messaging.Agent.Store.Postgres.Migrations.M20260411_service_certs import Simplex.Messaging.Agent.Store.Postgres.Migrations.M20260712_address_dr_rpc import Simplex.Messaging.Agent.Store.Postgres.Migrations.M20260823_snd_files_entitlement import Simplex.Messaging.Agent.Store.Postgres.Migrations.M20260919_ratchet_verify_codes +import Simplex.Messaging.Agent.Store.Postgres.Migrations.M20260929_skipped_messages_index import Simplex.Messaging.Agent.Store.Shared (Migration (..)) schemaMigrations :: [(String, Text, Maybe Text)] @@ -31,7 +32,8 @@ schemaMigrations = ("20260411_service_certs", m20260411_service_certs, Just down_m20260411_service_certs), ("20260712_address_dr_rpc", m20260712_address_dr_rpc, Just down_m20260712_address_dr_rpc), ("20260823_snd_files_entitlement", m20260823_snd_files_entitlement, Just down_m20260823_snd_files_entitlement), - ("20260919_ratchet_verify_codes", m20260919_ratchet_verify_codes, Just down_m20260919_ratchet_verify_codes) + ("20260919_ratchet_verify_codes", m20260919_ratchet_verify_codes, Just down_m20260919_ratchet_verify_codes), + ("20260929_skipped_messages_index", m20260929_skipped_messages_index, Just down_m20260929_skipped_messages_index) ] -- | The list of migrations in ascending order by date diff --git a/src/Simplex/Messaging/Agent/Store/Postgres/Migrations/M20260929_skipped_messages_index.hs b/src/Simplex/Messaging/Agent/Store/Postgres/Migrations/M20260929_skipped_messages_index.hs new file mode 100644 index 000000000..7f465cec7 --- /dev/null +++ b/src/Simplex/Messaging/Agent/Store/Postgres/Migrations/M20260929_skipped_messages_index.hs @@ -0,0 +1,21 @@ +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE QuasiQuotes #-} + +module Simplex.Messaging.Agent.Store.Postgres.Migrations.M20260929_skipped_messages_index where + +import Data.Text (Text) +import Text.RawString.QQ (r) + +m20260929_skipped_messages_index :: Text +m20260929_skipped_messages_index = + [r| +DROP INDEX idx_skipped_messages_conn_id; +CREATE INDEX idx_skipped_messages_conn_id ON skipped_messages(conn_id, skipped_message_id); +|] + +down_m20260929_skipped_messages_index :: Text +down_m20260929_skipped_messages_index = + [r| +DROP INDEX idx_skipped_messages_conn_id; +CREATE INDEX idx_skipped_messages_conn_id ON skipped_messages(conn_id); +|] diff --git a/src/Simplex/Messaging/Agent/Store/Postgres/Migrations/agent_postgres_schema.sql b/src/Simplex/Messaging/Agent/Store/Postgres/Migrations/agent_postgres_schema.sql index a0d408c36..c5be00a6e 100644 --- a/src/Simplex/Messaging/Agent/Store/Postgres/Migrations/agent_postgres_schema.sql +++ b/src/Simplex/Messaging/Agent/Store/Postgres/Migrations/agent_postgres_schema.sql @@ -1241,7 +1241,7 @@ CREATE UNIQUE INDEX idx_server_certs_user_id_host_port ON smp_agent_test_protoco -CREATE INDEX idx_skipped_messages_conn_id ON smp_agent_test_protocol_schema.skipped_messages USING btree (conn_id); +CREATE INDEX idx_skipped_messages_conn_id ON smp_agent_test_protocol_schema.skipped_messages USING btree (conn_id, skipped_message_id); diff --git a/src/Simplex/Messaging/Crypto/Ratchet.hs b/src/Simplex/Messaging/Crypto/Ratchet.hs index fe74d772a..067fedd89 100644 --- a/src/Simplex/Messaging/Crypto/Ratchet.hs +++ b/src/Simplex/Messaging/Crypto/Ratchet.hs @@ -72,6 +72,7 @@ module Simplex.Messaging.Crypto.Ratchet rcEncryptHeader, rcEncryptMsg, rcDecrypt, + maxSkippedMsgKeys, -- used in tests MsgHeader (..), RatchetInitParams (..), @@ -962,6 +963,9 @@ type DecryptResult a = (Either CryptoError ByteString, Ratchet a, SkippedMsgDiff maxSkip :: Word32 maxSkip = 512 +maxSkippedMsgKeys :: Int +maxSkippedMsgKeys = 2000 + rcDecrypt :: forall a. (AlgorithmI a, DhAlgorithm a) => diff --git a/tests/AgentTests/SQLiteTests.hs b/tests/AgentTests/SQLiteTests.hs index c22eddd7a..0a418cabb 100644 --- a/tests/AgentTests/SQLiteTests.hs +++ b/tests/AgentTests/SQLiteTests.hs @@ -20,12 +20,13 @@ import Control.Concurrent.Async (concurrently_) import Control.Concurrent.MVar import Control.Concurrent.STM import Control.Exception (SomeException) -import Control.Monad (replicateM_) +import Control.Monad (forM_, replicateM_) import Control.Monad.Trans.Except import Crypto.Random (ChaChaDRG) import Data.ByteArray (ScrubbedBytes) import Data.ByteString.Char8 (ByteString) import Data.List (isInfixOf) +import qualified Data.Map.Strict as M import qualified Data.Text as T import Data.Text.Encoding (encodeUtf8) import Data.Time @@ -51,6 +52,7 @@ import qualified Simplex.Messaging.Crypto as C import Simplex.Messaging.Crypto.File (CryptoFile (..)) import Simplex.Messaging.Crypto.Ratchet (pattern IKPQOn) import qualified Simplex.Messaging.Crypto.Ratchet as CR +import Simplex.Messaging.Encoding (Encoding (..)) import Simplex.Messaging.Encoding.String (StrEncoding (..)) import Simplex.Messaging.Protocol (EntityId (..), QueueMode (..), SubscriptionMode (..), pattern VersionSMPC) import qualified Simplex.Messaging.Protocol as SMP @@ -135,6 +137,7 @@ storeTests = do testCreateRcvMsg testCreateSndMsg testCreateRcvAndSndMsgs + it "should keep only the newest skipped message keys" testGetSkippedMsgKeys describe "Work items" $ do it "should getPendingQueueMsg" testGetPendingQueueMsg it "should getPendingServerCommand" testGetPendingServerCommand @@ -599,6 +602,27 @@ testCreateRcvAndSndMsgs = testCreateSndMsg_ db "snd_hash_1" connId sq $ mkSndMsgData (InternalId 5) (InternalSndId 2) "snd_hash_2" testCreateSndMsg_ db "snd_hash_2" connId sq $ mkSndMsgData (InternalId 6) (InternalSndId 3) "snd_hash_3" +testGetSkippedMsgKeys :: DBStore -> Expectation +testGetSkippedMsgKeys st = do + g <- C.newRandom + withTransaction st $ \db -> do + Right connId <- createNewConn db g cData1 {connId = ""} SCMInvitation + Right connId' <- createNewConn db g cData1 {connId = ""} SCMInvitation + createSkippedKeys db connId' + createSkippedKeys db connId + M.map M.keys <$> getSkippedMsgKeys db connId 4 + `shouldReturn` M.singleton (C.Key "header_key") [1, 2, 3, 4] + getMsgNs db connId `shouldReturn` [1 .. 4] + getMsgNs db connId' `shouldReturn` [1 .. 10] + where + createSkippedKeys :: DB.Connection -> ConnId -> IO () + createSkippedKeys db connId = do + DB.execute db "INSERT INTO ratchets (conn_id) VALUES (?)" (Only connId) + forM_ ([10, 9 .. 1] :: [Int]) $ \msgN -> + DB.execute db "INSERT INTO skipped_messages (conn_id, header_key, msg_n, msg_key) VALUES (?, ?, ?, ?)" (connId, "header_key" :: ByteString, msgN, smpEncode ("key" :: ByteString, "iv" :: ByteString)) + getMsgNs :: DB.Connection -> ConnId -> IO [Int] + getMsgNs db connId = map fromOnly <$> DB.query db "SELECT msg_n FROM skipped_messages WHERE conn_id = ? ORDER BY msg_n" (Only connId) + testCloseReopenStore :: IO () testCloseReopenStore = do st <- createStore'