agent: encrypt messages on delivery (#1446)

* agent: save message body once (plan, schema)

* split

* new type

* bs

* encrypt on delivery

* schema

* fix test

* check pad size

* rename

---------

Co-authored-by: Evgeny Poberezkin <[email protected]>

Author: spaced4ndy

simplexmq.cabal

@@ -198,6 +198,7 @@ library
           Simplex.Messaging.Agent.Store.SQLite.Migrations.M20240930_ntf_tokens_to_delete
           Simplex.Messaging.Agent.Store.SQLite.Migrations.M20241007_rcv_queues_last_broker_ts
           Simplex.Messaging.Agent.Store.SQLite.Migrations.M20241224_ratchet_e2e_snd_params
+          Simplex.Messaging.Agent.Store.SQLite.Migrations.M20250203_msg_bodies
   if !flag(client_library)
       exposed-modules:
           Simplex.FileTransfer.Client.Main

src/Simplex/Messaging/Agent.hs

@@ -1379,6 +1379,7 @@ enqueueMessage c cData sq msgFlags aMessage =
   ExceptT $ fmap fst . runIdentity <$> enqueueMessageB c (Identity (Right (cData, [sq], Nothing, msgFlags, aMessage)))
 {-# INLINE enqueueMessage #-}
 
+-- TODO [save once] IntMap of msg bodies.
 -- this function is used only for sending messages in batch, it returns the list of successes to enqueue additional deliveries
 enqueueMessageB :: forall t. Traversable t => AgentClient -> t (Either AgentErrorType (ConnData, NonEmpty SndQueue, Maybe PQEncryption, MsgFlags, AMessage)) -> AM' (t (Either AgentErrorType ((AgentMsgId, PQEncryption), Maybe (ConnData, [SndQueue], AgentMsgId))))
 enqueueMessageB c reqs = do
@@ -1391,19 +1392,21 @@ enqueueMessageB c reqs = do
   where
     storeSentMsg :: DB.Connection -> AgentConfig -> (ConnData, NonEmpty SndQueue, Maybe PQEncryption, MsgFlags, AMessage) -> IO (Either AgentErrorType ((ConnData, NonEmpty SndQueue, Maybe PQEncryption, MsgFlags, AMessage), InternalId, PQEncryption))
     storeSentMsg db cfg req@(cData@ConnData {connId}, sq :| _, pqEnc_, msgFlags, aMessage) = fmap (first storeError) $ runExceptT $ do
-      let AgentConfig {smpAgentVRange, e2eEncryptVRange} = cfg
+      let AgentConfig {e2eEncryptVRange} = cfg
       internalTs <- liftIO getCurrentTime
       (internalId, internalSndId, prevMsgHash) <- ExceptT $ updateSndIds db connId
       let privHeader = APrivHeader (unSndId internalSndId) prevMsgHash
           agentMsg = AgentMessage privHeader aMessage
           agentMsgStr = smpEncode agentMsg
           internalHash = C.sha256Hash agentMsgStr
           currentE2EVersion = maxVersion e2eEncryptVRange
-      (encAgentMessage, pqEnc) <- agentRatchetEncrypt db cData agentMsgStr e2eEncAgentMsgLength pqEnc_ currentE2EVersion
-      let agentVersion = maxVersion smpAgentVRange
-          msgBody = smpEncode $ AgentMsgEnvelope {agentVersion, encAgentMessage}
-          msgType = agentMessageType agentMsg
-          msgData = SndMsgData {internalId, internalSndId, internalTs, msgType, msgFlags, msgBody, pqEncryption = pqEnc, internalHash, prevMsgHash}
+      -- TODO [save once] Save single MsgBody / enveloped body agentMsgStr (outside of withStoreBatch ... storeSentMsg).
+      -- TODO             Link messages to it, save encryption data per message.
+      -- TODO             'msg_body' field is not nullable - use default empty strings?
+      (mek, paddedLen, pqEnc) <- agentRatchetEncryptHeader db cData e2eEncAgentMsgLength pqEnc_ currentE2EVersion
+      withExceptT (SEAgentError . cryptoError) $ CR.rcCheckCanPad paddedLen agentMsgStr
+      let msgType = agentMessageType agentMsg
+          msgData = SndMsgData {internalId, internalSndId, internalTs, msgType, msgFlags, msgBody = agentMsgStr, pqEncryption = pqEnc, internalHash, prevMsgHash, encryptKey_ = Just mek, paddedLen_ = Just paddedLen}
       liftIO $ createSndMsg db connId msgData
       liftIO $ createSndMsgDelivery db connId sq internalId
       pure (req, internalId, pqEnc)
@@ -1451,7 +1454,7 @@ runSmpQueueMsgDelivery c@AgentClient {subQ} ConnData {connId} sq@SndQueue {userI
     liftIO $ throwWhenNoDelivery c sq
     atomically $ beginAgentOperation c AOSndNetwork
     withWork c doWork (\db -> getPendingQueueMsg db connId sq) $
-      \(rq_, PendingMsgData {msgId, msgType, msgBody, pqEncryption, msgFlags, msgRetryState, internalTs}) -> do
+      \(rq_, PendingMsgData {msgId, msgType, msgBody, pqEncryption, msgFlags, msgRetryState, internalTs, encryptKey_, paddedLen_}) -> do
         atomically $ endAgentOperation c AOMsgDelivery -- this operation begins in submitPendingMsg
         let mId = unId msgId
             ri' = maybe id updateRetryInterval2 msgRetryState ri
@@ -1461,7 +1464,15 @@ runSmpQueueMsgDelivery c@AgentClient {subQ} ConnData {connId} sq@SndQueue {userI
           resp <- tryError $ case msgType of
             AM_CONN_INFO -> sendConfirmation c sq msgBody
             AM_CONN_INFO_REPLY -> sendConfirmation c sq msgBody
-            _ -> sendAgentMessage c sq msgFlags msgBody
+            _ -> case (encryptKey_, paddedLen_) of
+              (Nothing, Nothing) -> sendAgentMessage c sq msgFlags msgBody
+              (Just mek, Just paddedLen) -> do
+                AgentConfig {smpAgentVRange} <- asks config
+                encAgentMessage <- liftError cryptoError $ CR.rcEncryptMsg mek paddedLen msgBody
+                let agentVersion = maxVersion smpAgentVRange
+                    msgBody' = smpEncode $ AgentMsgEnvelope {agentVersion, encAgentMessage}
+                sendAgentMessage c sq msgFlags msgBody'
+              _ -> throwE $ INTERNAL "runSmpQueueMsgDelivery: missing encryption data"
           case resp of
             Left e -> do
               let err = if msgType == AM_A_MSG_ then MERR mId e else ERR e
@@ -1833,7 +1844,7 @@ deleteConnQueues c waitDelivery ntf rqs = do
     deleteQueueRecs rs = do
       maxErrs <- asks $ deleteErrorCount . config
       rs' <- rights <$> withStoreBatch' c (\db -> map (deleteQueueRec db maxErrs) rs)
-      let delQ ((rq, _), err_) =  (qConnId rq,qServer rq,queueId rq,) <$> err_
+      let delQ ((rq, _), err_) = (qConnId rq,qServer rq,queueId rq,) <$> err_
           delQs_ = L.nonEmpty $ mapMaybe delQ rs'
       forM_ delQs_ $ \delQs -> notify ("", "", AEvt SAEConn $ DEL_RCVQS delQs)
       pure $ map fst rs'
@@ -2952,7 +2963,7 @@ storeConfirmation c cData@ConnData {connId, pqSupport, connAgentVersion = v} sq
     (encConnInfo, pqEncryption) <- agentRatchetEncrypt db cData agentMsgStr e2eEncConnInfoLength (Just pqEnc) currentE2EVersion
     let msgBody = smpEncode $ AgentConfirmation {agentVersion = v, e2eEncryption_, encConnInfo}
         msgType = agentMessageType agentMsg
-        msgData = SndMsgData {internalId, internalSndId, internalTs, msgType, msgBody, pqEncryption, msgFlags = SMP.MsgFlags {notification = True}, internalHash, prevMsgHash}
+        msgData = SndMsgData {internalId, internalSndId, internalTs, msgType, msgBody, pqEncryption, msgFlags = SMP.MsgFlags {notification = True}, internalHash, prevMsgHash, encryptKey_ = Nothing, paddedLen_ = Nothing}
     liftIO $ createSndMsg db connId msgData
     liftIO $ createSndMsgDelivery db connId sq internalId
 
@@ -2978,19 +2989,25 @@ enqueueRatchetKey c cData@ConnData {connId} sq e2eEncryption = do
       let msgBody = smpEncode $ AgentRatchetKey {agentVersion, e2eEncryption, info = agentMsgStr}
           msgType = agentMessageType agentMsg
           -- this message is e2e encrypted with queue key, not with double ratchet
-          msgData = SndMsgData {internalId, internalSndId, internalTs, msgType, msgBody, pqEncryption = PQEncOff, msgFlags = SMP.MsgFlags {notification = True}, internalHash, prevMsgHash}
+          msgData = SndMsgData {internalId, internalSndId, internalTs, msgType, msgBody, pqEncryption = PQEncOff, msgFlags = SMP.MsgFlags {notification = True}, internalHash, prevMsgHash, encryptKey_ = Nothing, paddedLen_ = Nothing}
       liftIO $ createSndMsg db connId msgData
       liftIO $ createSndMsgDelivery db connId sq internalId
       pure internalId
 
 -- encoded AgentMessage -> encoded EncAgentMessage
 agentRatchetEncrypt :: DB.Connection -> ConnData -> ByteString -> (VersionSMPA -> PQSupport -> Int) -> Maybe PQEncryption -> CR.VersionE2E -> ExceptT StoreError IO (ByteString, PQEncryption)
-agentRatchetEncrypt db ConnData {connId, connAgentVersion = v, pqSupport} msg getPaddedLen pqEnc_ currentE2EVersion = do
+agentRatchetEncrypt db cData msg getPaddedLen pqEnc_ currentE2EVersion = do
+  (mek, paddedLen, pqEnc) <- agentRatchetEncryptHeader db cData getPaddedLen pqEnc_ currentE2EVersion
+  encMsg <- withExceptT (SEAgentError . cryptoError) $ CR.rcEncryptMsg mek paddedLen msg
+  pure (encMsg, pqEnc)
+
+agentRatchetEncryptHeader :: DB.Connection -> ConnData -> (VersionSMPA -> PQSupport -> Int) -> Maybe PQEncryption -> CR.VersionE2E -> ExceptT StoreError IO (CR.MsgEncryptKeyX448, Int, PQEncryption)
+agentRatchetEncryptHeader db ConnData {connId, connAgentVersion = v, pqSupport} getPaddedLen pqEnc_ currentE2EVersion = do
   rc <- ExceptT $ getRatchet db connId
   let paddedLen = getPaddedLen v pqSupport
-  (encMsg, rc') <- withExceptT (SEAgentError . cryptoError) $ CR.rcEncrypt rc paddedLen msg pqEnc_ currentE2EVersion
+  (mek, rc') <- withExceptT (SEAgentError . cryptoError) $ CR.rcEncryptHeader rc pqEnc_ currentE2EVersion
   liftIO $ updateRatchet db connId rc' CR.SMDNoChange
-  pure (encMsg, CR.rcSndKEM rc')
+  pure (mek, paddedLen, CR.rcSndKEM rc')
 
 -- encoded EncAgentMessage -> encoded AgentMessage
 agentRatchetDecrypt :: TVar ChaChaDRG -> DB.Connection -> ConnId -> ByteString -> ExceptT StoreError IO (ByteString, PQEncryption)

src/Simplex/Messaging/Agent/Store.hs

@@ -33,7 +33,7 @@ import Simplex.Messaging.Agent.Store.Common
 import Simplex.Messaging.Agent.Store.Interface (DBOpts, appMigrations, createDBStore)
 import Simplex.Messaging.Agent.Store.Shared (MigrationConfirmation (..), MigrationError (..))
 import qualified Simplex.Messaging.Crypto as C
-import Simplex.Messaging.Crypto.Ratchet (PQEncryption, PQSupport, RatchetX448)
+import Simplex.Messaging.Crypto.Ratchet (MsgEncryptKeyX448, PQEncryption, PQSupport, RatchetX448)
 import Simplex.Messaging.Encoding.String
 import Simplex.Messaging.Protocol
   ( MsgBody,
@@ -542,7 +542,9 @@ data SndMsgData = SndMsgData
     msgBody :: MsgBody,
     pqEncryption :: PQEncryption,
     internalHash :: MsgHash,
-    prevMsgHash :: MsgHash
+    prevMsgHash :: MsgHash,
+    encryptKey_ :: Maybe MsgEncryptKeyX448,
+    paddedLen_ :: Maybe Int
   }
 
 data SndMsg = SndMsg
@@ -560,7 +562,9 @@ data PendingMsgData = PendingMsgData
     msgBody :: MsgBody,
     pqEncryption :: PQEncryption,
     msgRetryState :: Maybe RI2State,
-    internalTs :: InternalTs
+    internalTs :: InternalTs,
+    encryptKey_ :: Maybe MsgEncryptKeyX448,
+    paddedLen_ :: Maybe Int
   }
   deriving (Show)
 

src/Simplex/Messaging/Agent/Store/AgentStore.hs

@@ -844,18 +844,18 @@ getPendingQueueMsg db connId SndQueue {dbQueueId} =
           DB.query
             db
             [sql|
-              SELECT m.msg_type, m.msg_flags, m.msg_body, m.pq_encryption, m.internal_ts, s.retry_int_slow, s.retry_int_fast
+              SELECT m.msg_type, m.msg_flags, m.msg_body, m.pq_encryption, m.internal_ts, s.retry_int_slow, s.retry_int_fast, s.msg_encrypt_key, s.padded_msg_len
               FROM messages m
               JOIN snd_messages s ON s.conn_id = m.conn_id AND s.internal_id = m.internal_id
               WHERE m.conn_id = ? AND m.internal_id = ?
             |]
             (connId, msgId)
         err = SEInternal $ "msg delivery " <> bshow msgId <> " returned []"
-        pendingMsgData :: (AgentMessageType, Maybe MsgFlags, MsgBody, PQEncryption, InternalTs, Maybe Int64, Maybe Int64) -> PendingMsgData
-        pendingMsgData (msgType, msgFlags_, msgBody, pqEncryption, internalTs, riSlow_, riFast_) =
+        pendingMsgData :: (AgentMessageType, Maybe MsgFlags, MsgBody, PQEncryption, InternalTs, Maybe Int64, Maybe Int64, Maybe CR.MsgEncryptKeyX448, Maybe Int) -> PendingMsgData
+        pendingMsgData (msgType, msgFlags_, msgBody, pqEncryption, internalTs, riSlow_, riFast_, encryptKey_, paddedLen_) =
           let msgFlags = fromMaybe SMP.noMsgFlags msgFlags_
               msgRetryState = RI2State <$> riSlow_ <*> riFast_
-           in PendingMsgData {msgId, msgType, msgFlags, msgBody, pqEncryption, msgRetryState, internalTs}
+           in PendingMsgData {msgId, msgType, msgFlags, msgBody, pqEncryption, msgRetryState, internalTs, encryptKey_, paddedLen_}
     markMsgFailed msgId = DB.execute db "UPDATE snd_message_deliveries SET failed = 1 WHERE conn_id = ? AND internal_id = ?" (connId, msgId)
 
 getWorkItem :: Show i => ByteString -> IO (Maybe i) -> (i -> IO (Either StoreError a)) -> (i -> IO ()) -> IO (Either StoreError (Maybe a))
@@ -997,6 +997,7 @@ deleteDeliveredSndMsg db connId msgId = do
   cnt <- countPendingSndDeliveries_ db connId msgId
   when (cnt == 0) $ deleteMsg db connId msgId
 
+-- TODO [save once] Delete from shared message bodies if no deliveries reference it. (`when (cnt == 0)`)
 deleteSndMsgDelivery :: DB.Connection -> ConnId -> SndQueue -> InternalId -> Bool -> IO ()
 deleteSndMsgDelivery db connId SndQueue {dbQueueId} msgId keepForReceipt = do
   DB.execute
@@ -2206,11 +2207,11 @@ insertSndMsgDetails_ dbConn connId SndMsgData {..} =
     dbConn
     [sql|
       INSERT INTO snd_messages
-        ( conn_id, internal_snd_id, internal_id, internal_hash, previous_msg_hash)
+        ( conn_id, internal_snd_id, internal_id, internal_hash, previous_msg_hash, msg_encrypt_key, padded_msg_len)
       VALUES
-        (?,?,?,?,?)
+        (?,?,?,?,?,?,?)
     |]
-    (connId, internalSndId, internalId, Binary internalHash, Binary prevMsgHash)
+    (connId, internalSndId, internalId, Binary internalHash, Binary prevMsgHash, encryptKey_, paddedLen_)
 
 updateSndMsgHash :: DB.Connection -> ConnId -> InternalSndId -> MsgHash -> IO ()
 updateSndMsgHash db connId internalSndId internalHash =

src/Simplex/Messaging/Agent/Store/SQLite/Migrations.hs

@@ -66,6 +66,7 @@ import Simplex.Messaging.Agent.Store.SQLite.Migrations.M20240702_servers_stats
 import Simplex.Messaging.Agent.Store.SQLite.Migrations.M20240930_ntf_tokens_to_delete
 import Simplex.Messaging.Agent.Store.SQLite.Migrations.M20241007_rcv_queues_last_broker_ts
 import Simplex.Messaging.Agent.Store.SQLite.Migrations.M20241224_ratchet_e2e_snd_params
+import Simplex.Messaging.Agent.Store.SQLite.Migrations.M20250203_msg_bodies
 import Simplex.Messaging.Agent.Store.Shared
 import Simplex.Messaging.Encoding.String
 import Simplex.Messaging.Transport.Client (TransportHost)
@@ -108,7 +109,8 @@ schemaMigrations =
     ("m20240702_servers_stats", m20240702_servers_stats, Just down_m20240702_servers_stats),
     ("m20240930_ntf_tokens_to_delete", m20240930_ntf_tokens_to_delete, Just down_m20240930_ntf_tokens_to_delete),
     ("m20241007_rcv_queues_last_broker_ts", m20241007_rcv_queues_last_broker_ts, Just down_m20241007_rcv_queues_last_broker_ts),
-    ("m20241224_ratchet_e2e_snd_params", m20241224_ratchet_e2e_snd_params, Just down_m20241224_ratchet_e2e_snd_params)
+    ("m20241224_ratchet_e2e_snd_params", m20241224_ratchet_e2e_snd_params, Just down_m20241224_ratchet_e2e_snd_params),
+    ("m20250203_msg_bodies", m20250203_msg_bodies, Just down_m20250203_msg_bodies)
   ]
 
 -- | The list of migrations in ascending order by date

src/Simplex/Messaging/Agent/Store/SQLite/Migrations/M20250203_msg_bodies.hs

@@ -0,0 +1,36 @@
+{-# LANGUAGE QuasiQuotes #-}
+
+module Simplex.Messaging.Agent.Store.SQLite.Migrations.M20250203_msg_bodies where
+
+import Database.SQLite.Simple (Query)
+import Database.SQLite.Simple.QQ (sql)
+
+m20250203_msg_bodies :: Query
+m20250203_msg_bodies =
+    [sql|
+ALTER TABLE snd_messages ADD COLUMN msg_encrypt_key BLOB;
+ALTER TABLE snd_messages ADD COLUMN padded_msg_len INTEGER;
+
+
+-- CREATE TABLE msg_bodies (
+--   msg_body_id INTEGER PRIMARY KEY,
+--   msg_body BLOB NOT NULL DEFAULT x''
+-- )
+
+-- ALTER TABLE snd_messages ADD COLUMN msg_body_id INTEGER REFERENCES msg_bodies ON DELETE CASCADE;
+
+-- fkey to msg_bodies
+-- on each delivery check if other deliveries reference the same msg_body_id, if not delete it
+|]
+
+down_m20250203_msg_bodies :: Query
+down_m20250203_msg_bodies =
+    [sql|
+ALTER TABLE snd_messages DROP COLUMN msg_encrypt_key;
+ALTER TABLE snd_messages DROP COLUMN padded_msg_len;
+
+
+-- ALTER TABLE snd_messages DROP COLUMN msg_body_id;
+
+-- DROP TABLE msg_bodies;
+|]

src/Simplex/Messaging/Agent/Store/SQLite/Migrations/agent_schema.sql

@@ -127,6 +127,8 @@ CREATE TABLE snd_messages(
   retry_int_fast INTEGER,
   rcpt_internal_id INTEGER,
   rcpt_status TEXT,
+  msg_encrypt_key BLOB,
+  padded_msg_len INTEGER,
   PRIMARY KEY(conn_id, internal_snd_id),
   FOREIGN KEY(conn_id, internal_id) REFERENCES messages
   ON DELETE CASCADE

src/Simplex/Messaging/Crypto.hs

@@ -169,6 +169,7 @@ module Simplex.Messaging.Crypto
     sha512Hash,
 
     -- * Message padding / un-padding
+    canPad,
     pad,
     unPad,
 
@@ -1010,6 +1011,11 @@ decryptAEADNoPad aesKey iv ad msg (AuthTag tag) = do
 maxMsgLen :: Int
 maxMsgLen = 2 ^ (16 :: Int) - 3
 
+canPad :: Int -> Int -> Bool
+canPad msgLen paddedLen = msgLen <= maxMsgLen && padLen >= 0
+  where
+    padLen = paddedLen - msgLen - 2
+
 pad :: ByteString -> Int -> Either CryptoError ByteString
 pad msg paddedLen
   | len <= maxMsgLen && padLen >= 0 = Right $ encodeWord16 (fromIntegral len) <> msg <> B.replicate padLen '#'