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

@@ -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

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

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)
 

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

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);
+|]

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);
 
 
 

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) =>

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'