core (pq): don't compress if message fits without compression; check compressed message fits size limit (#3888)

* core (pq): don't compress if message fits without compression; check compressed message fits size limit

* refactor

* errors

* fix tests

* envelope sizes

* refactor

* comment

* more flexible test

* refactor, comment

---------

Co-authored-by: Evgeny Poberezkin <evgeny@poberezkin.com>

cabal.project

@@ -12,7 +12,7 @@ constraints: zip +disable-bzip2 +disable-zstd
 source-repository-package
     type: git
     location: https://github.com/simplex-chat/simplexmq.git
-    tag: 851ed2d02e2a78c15893ad8bc9c5a4d917eb6a35
+    tag: b4c90781bba8cca3a8f7bea9e0c2b6707ff923af
 
 source-repository-package
     type: git

scripts/nix/sha256map.nix

@@ -1,5 +1,5 @@
 {
-  "https://github.com/simplex-chat/simplexmq.git"."851ed2d02e2a78c15893ad8bc9c5a4d917eb6a35" = "0rm13iknnqhdb42nmyjc2wj85z23p337bp026ihnychax5s1216j";
+  "https://github.com/simplex-chat/simplexmq.git"."b4c90781bba8cca3a8f7bea9e0c2b6707ff923af" = "0f4h1akgpkrg68lmhrnvrq6srr2c3gj0fyx4ghnsp5hmbyhn2mk2";
   "https://github.com/simplex-chat/hs-socks.git"."a30cc7a79a08d8108316094f8f2f82a0c5e1ac51" = "0yasvnr7g91k76mjkamvzab2kvlb1g5pspjyjn2fr6v83swjhj38";
   "https://github.com/simplex-chat/direct-sqlcipher.git"."f814ee68b16a9447fbb467ccc8f29bdd3546bfd9" = "1ql13f4kfwkbaq7nygkxgw84213i0zm7c1a8hwvramayxl38dq5d";
   "https://github.com/simplex-chat/sqlcipher-simple.git"."a46bd361a19376c5211f1058908fc0ae6bf42446" = "1z0r78d8f0812kxbgsm735qf6xx8lvaz27k1a0b4a2m0sshpd5gl";

src/Simplex/Chat.hs

@@ -3350,7 +3350,8 @@ processAgentMsgSndFile _corrId aFileId msg =
                   [] -> case xftpRedirectFor of
                     Nothing -> xftpSndFileRedirect user fileId rfd >>= toView . CRSndFileRedirectStartXFTP user ft
                     Just _ -> sendFileError "Prohibit chaining redirects" fileId vr ft
-                  rfds' -> do -- we have 1 chunk - use it as URI whether it is redirect or not
+                  rfds' -> do
+                    -- we have 1 chunk - use it as URI whether it is redirect or not
                     ft' <- maybe (pure ft) (\fId -> withStore $ \db -> getFileTransferMeta db user fId) xftpRedirectFor
                     toView $ CRSndStandaloneFileComplete user ft' $ map (decodeLatin1 . strEncode . FD.fileDescriptionURI) rfds'
             Just (AChatItem _ d cInfo _ci@ChatItem {meta = CIMeta {itemSharedMsgId = msgId_, itemDeleted}}) ->
@@ -6057,8 +6058,7 @@ sendDirectContactMessage user ct chatMsgEvent = do
   conn@Connection {connId, pqSupport} <- liftEither $ contactSendConn_ ct
   r <- sendDirectMessage_ conn pqSupport chatMsgEvent (ConnectionId connId)
   let (sndMessage, msgDeliveryId, pqEnc') = r
-  -- TODO PQ use updated ct' and conn'? check downstream if it may affect something, maybe it's not necessary
-  void $ createContactPQSndItem user ct conn pqEnc' -- (_ct', _conn')
+  void $ createContactPQSndItem user ct conn pqEnc'
   pure (sndMessage, msgDeliveryId)
 
 contactSendConn_ :: Contact -> Either ChatError Connection
@@ -6127,12 +6127,12 @@ processSndMessageBatch conn@Connection {connId} (MsgBatch batchBody sndMsgs) = d
 
 -- TODO v5.7 update batching for groups
 batchSndMessagesJSON :: NonEmpty SndMessage -> [Either ChatError MsgBatch]
-batchSndMessagesJSON = batchMessages maxRawMsgLength . L.toList
+batchSndMessagesJSON = batchMessages maxEncodedMsgLength . L.toList
 
 -- batchSndMessagesBinary :: forall m. ChatMonad m => NonEmpty SndMessage -> m [Either ChatError MsgBatch]
 -- batchSndMessagesBinary msgs = do
 --   compressed <- liftIO $ withCompressCtx maxChatMsgSize $ \cctx -> mapM (compressForBatch cctx) msgs
---   pure . map toMsgBatch . SMP.batchTransmissions_ (maxEncodedMsgLength PQEncOff) $ L.zip compressed msgs
+--   pure . map toMsgBatch . SMP.batchTransmissions_ (maxEncodedMsgLength) $ L.zip compressed msgs
 --   where
 --     compressForBatch cctx SndMessage {msgBody} = bimap (const TELargeMsg) smpEncode <$> compress cctx msgBody
 --     toMsgBatch :: SMP.TransportBatch SndMessage -> Either ChatError MsgBatch
@@ -6146,19 +6146,20 @@ encodeConnInfo chatMsgEvent = do
   vr <- chatVersionRange
   encodeConnInfoPQ PQSupportOff (maxVersion $ vr PQSupportOff) chatMsgEvent
 
--- TODO PQ check size after compression (in compressedBatchMsgBody_ ?)
 encodeConnInfoPQ :: (MsgEncodingI e, ChatMonad m) => PQSupport -> VersionChat -> ChatMsgEvent e -> m ByteString
 encodeConnInfoPQ pqSup v chatMsgEvent = do
   vr <- chatVersionRange
-  let msg = ChatMessage {chatVRange = vr pqSup, msgId = Nothing, chatMsgEvent}
-  case encodeChatMessage maxConnInfoLength msg of
-    ECMEncoded encodedBody -> case pqSup of
-      PQSupportOn | v >= pqEncryptionCompressionVersion -> liftIO $ compressedBatchMsgBody encodedBody
-      _ -> pure encodedBody
-    ECMLarge -> throwChatError $ CEException "large message"
-  where
-    compressedBatchMsgBody msgBody =
-      withCompressCtx (toEnum $ B.length msgBody) (`compressedBatchMsgBody_` msgBody)
+  let info = ChatMessage {chatVRange = vr pqSup, msgId = Nothing, chatMsgEvent}
+  case encodeChatMessage maxEncodedInfoLength info of
+    ECMEncoded connInfo -> case pqSup of
+      PQSupportOn | v >= pqEncryptionCompressionVersion && B.length connInfo > maxCompressedInfoLength -> do
+        connInfo' <- liftIO compressedBatchMsgBody
+        when (B.length connInfo' > maxCompressedInfoLength) $ throwChatError $ CEException "large compressed info"
+        pure connInfo'
+      _ -> pure connInfo
+      where
+        compressedBatchMsgBody = withCompressCtx (toEnum $ B.length connInfo) (`compressedBatchMsgBody_` connInfo)
+    ECMLarge -> throwChatError $ CEException "large info"
 
 deliverMessage :: ChatMonad m => Connection -> CMEventTag e -> MsgBody -> MessageId -> m (Int64, PQEncryption)
 deliverMessage conn cmEventTag msgBody msgId = do
@@ -6183,11 +6184,18 @@ deliverMessagesB msgReqs = do
   void $ withStoreBatch' $ \db -> map (updatePQSndEnabled db) (rights . L.toList $ sent)
   withStoreBatch $ \db -> L.map (bindRight $ createDelivery db) sent
   where
-    compressBodies = liftIO $ withCompressCtx (toEnum maxRawMsgLength) $ \cctx -> do
-      forME msgReqs $ \mr@(conn@Connection {pqSupport, connChatVersion}, msgFlags, msgBody, msgId) -> Right <$> case pqSupport of
-        PQSupportOn | connChatVersion >= pqEncryptionCompressionVersion ->
-          (\cBody -> (conn, msgFlags, cBody, msgId)) <$> compressedBatchMsgBody_ cctx msgBody
-        _ -> pure mr
+    compressBodies = liftIO $ withCompressCtx (toEnum maxEncodedMsgLength) $ \cxt ->
+      forME msgReqs $ \mr@(conn@Connection {pqSupport, connChatVersion = v}, msgFlags, msgBody, msgId) ->
+        runExceptT $ case pqSupport of
+          -- we only compress messages when:
+          -- 1) PQ support is enabled
+          -- 2) version is compatible with compression
+          -- 3) message is longer than max compressed size (as this function is not used for batched messages anyway)
+          PQSupportOn | v >= pqEncryptionCompressionVersion && B.length msgBody > maxCompressedMsgLength -> do
+            msgBody' <- liftIO $ compressedBatchMsgBody_ cxt msgBody
+            when (B.length msgBody' > maxCompressedMsgLength) $ throwError $ ChatError $ CEException "large compressed message"
+            pure (conn, msgFlags, msgBody', msgId)
+          _ -> pure mr
     toAgent = \case
       Right (conn@Connection {pqEncryption}, msgFlags, msgBody, _msgId) -> Right (aConnId conn, pqEncryption, msgFlags, msgBody)
       Left _ce -> Left (AP.INTERNAL "ChatError, skip") -- as long as it is Left, the agent batchers should just step over it

src/Simplex/Chat/Protocol.hs

@@ -531,29 +531,29 @@ $(JQ.deriveJSON defaultJSON ''QuotedMsg)
 
 -- this limit reserves space for metadata in forwarded messages
 -- 15780 (limit used for fileChunkSize) - 161 (x.grp.msg.forward overhead) = 15619, round to 15610
-maxRawMsgLength :: Int
-maxRawMsgLength = 15610
+maxEncodedMsgLength :: Int
+maxEncodedMsgLength = 15610
 
-maxEncodedMsgLength :: PQSupport -> Int
-maxEncodedMsgLength = \case
-  PQSupportOn -> 13410 -- reduced by 2200 (original message should be compressed)
-  PQSupportOff -> maxRawMsgLength
-{-# INLINE maxEncodedMsgLength #-}
+-- maxEncodedMsgLength - 2222, see e2eEncUserMsgLength in agent
+maxCompressedMsgLength :: Int
+maxCompressedMsgLength = 13388
 
-maxConnInfoLength :: PQSupport -> Int
-maxConnInfoLength = \case
-  PQSupportOn -> 10902 -- reduced by 3700
-  PQSupportOff -> 14602 -- 15610 - delta in agent between MSG and INFO
-{-# INLINE maxConnInfoLength #-}
+-- maxEncodedMsgLength - delta between MSG and INFO + 100 (returned for forward overhead)
+-- delta between MSG and INFO = e2eEncUserMsgLength (no PQ) - e2eEncConnInfoLength (no PQ) = 1008
+maxEncodedInfoLength :: Int
+maxEncodedInfoLength = 14702
+
+maxCompressedInfoLength :: Int
+maxCompressedInfoLength = 10976 -- maxEncodedInfoLength - 3726, see e2eEncConnInfoLength in agent
 
 data EncodedChatMessage = ECMEncoded ByteString | ECMLarge
 
-encodeChatMessage :: MsgEncodingI e => (PQSupport -> Int) -> ChatMessage e -> EncodedChatMessage
-encodeChatMessage getMaxSize msg = do
+encodeChatMessage :: MsgEncodingI e => Int -> ChatMessage e -> EncodedChatMessage
+encodeChatMessage maxSize msg = do
   case chatToAppMessage msg of
     AMJson m -> do
       let body = LB.toStrict $ J.encode m
-      if B.length body > getMaxSize PQSupportOff
+      if B.length body > maxSize
         then ECMLarge
         else ECMEncoded body
     AMBinary m -> ECMEncoded $ strEncode m
@@ -573,7 +573,8 @@ parseChatMessages s = case B.head s of
     decodeCompressed :: ByteString -> [Either String AChatMessage]
     decodeCompressed s' = case smpDecode s' of
       Left e -> [Left e]
-      Right compressed -> concatMap (either (pure . Left) parseChatMessages) . L.toList $ decompressBatch maxRawMsgLength compressed
+      -- TODO v5.7 don't reserve multiple large buffers when decoding batches
+      Right compressed -> concatMap (either (pure . Left) parseChatMessages) . L.toList $ decompressBatch maxEncodedMsgLength compressed
 
 compressedBatchMsgBody_ :: CompressCtx -> MsgBody -> IO ByteString
 compressedBatchMsgBody_ ctx msgBody = markCompressedBatch . smpEncode . (L.:| []) <$> compress ctx msgBody

tests/ChatTests/Direct.hs

@@ -2825,7 +2825,7 @@ runTestPQConnectViaAddress (alice, aPQ) (bob, bPQ) = do
 
 runTestPQVersionsViaLink :: HasCallStack => TestCC -> TestCC -> Bool -> VersionChat -> IO ()
 runTestPQVersionsViaLink alice bob pqExpected vExpected = do
-  img <- genProfileImgForLink
+  img <- genProfileImg
   let profileImage = "data:image/png;base64," <> B.unpack img
   alice `send` ("/set profile image " <> profileImage)
   _trimmedCmd1 <- getTermLine alice
@@ -2857,7 +2857,7 @@ runTestPQVersionsViaLink alice bob pqExpected vExpected = do
 
 runTestPQVersionsViaAddress :: HasCallStack => TestCC -> TestCC -> Bool -> VersionChat -> IO ()
 runTestPQVersionsViaAddress alice bob pqExpected vExpected = do
-  img <- genProfileImgForAddress
+  img <- genProfileImg
   let profileImage = "data:image/png;base64," <> B.unpack img
   alice `send` ("/set profile image " <> profileImage)
   _trimmedCmd1 <- getTermLine alice

tests/ChatTests/Utils.hs

@@ -245,30 +245,21 @@ sndRcvImg pqEnc enabled (cc1, msg, v1) (cc2, v2) = do
   img <- atomically $ B64.encode <$> C.randomBytes lrgLen g
   cc1 `send` ("/_send @2 json {\"msgContent\":{\"type\":\"image\",\"text\":\"" <> msg <> "\",\"image\":\"" <> B.unpack img <> "\"}}")
   cc1 .<## "}}"
-  when enabled $ cc1 <## (name2 <> ": quantum resistant end-to-end encryption enabled")
-  cc1 <# ("@" <> name2 <> " " <> msg)
+  cc1 <### ([ConsoleString (name2 <> ": quantum resistant end-to-end encryption enabled") | enabled] <> [WithTime ("@" <> name2 <> " " <> msg)])
   cc1 `pqSndForContact` 2 `shouldReturn` pqEnc
   cc1 `pqVerForContact` 2 `shouldReturn` v1
-  when enabled $ cc2 <## (name1 <> ": quantum resistant end-to-end encryption enabled")
-  cc2 <# (name1 <> "> " <> msg)
+  cc2 <### ([ConsoleString (name1 <> ": quantum resistant end-to-end encryption enabled") | enabled] <> [WithTime (name1 <> "> " <> msg)])
   cc2 `pqRcvForContact` 2 `shouldReturn` pqEnc
   cc2 `pqVerForContact` 2 `shouldReturn` v2
   where
-    lrgLen = maxEncodedMsgLength PQSupportOff * 3 `div` 4 - 110 -- 98 is ~ max size for binary image preview given the rest of the message
+    lrgLen = maxEncodedMsgLength * 3 `div` 4 - 110 -- 98 is ~ max size for binary image preview given the rest of the message
 
-genProfileImgForLink :: IO ByteString
-genProfileImgForLink = do
+genProfileImg :: IO ByteString
+genProfileImg = do
   g <- C.newRandom
   atomically $ B64.encode <$> C.randomBytes lrgLen g
   where
-    lrgLen = maxConnInfoLength PQSupportOff * 3 `div` 4 - 240 -- 214 is the magic number to make tests pass (10737)
-
-genProfileImgForAddress :: IO ByteString
-genProfileImgForAddress = do
-  g <- C.newRandom
-  atomically $ B64.encode <$> C.randomBytes lrgLen g
-  where
-    lrgLen = maxConnInfoLength PQSupportOff * 3 `div` 4 - 260 -- 238 is the magic number to make tests pass (10713)
+    lrgLen = maxEncodedInfoLength * 3 `div` 4 - 420
 
 -- PQ combinators /
 

tests/MessageBatching.hs

@@ -17,7 +17,7 @@ import Data.Text.Encoding (encodeUtf8)
 import Simplex.Chat.Messages.Batch
 import Simplex.Chat.Controller (ChatError (..), ChatErrorType (..))
 import Simplex.Chat.Messages (SndMessage (..))
-import Simplex.Chat.Protocol (SharedMsgId (..), maxRawMsgLength)
+import Simplex.Chat.Protocol (SharedMsgId (..), maxEncodedMsgLength)
 import Test.Hspec
 
 batchingTests :: Spec
@@ -99,7 +99,7 @@ testImageFitsSingleBatch = do
       msg s = SndMessage {msgId = 0, sharedMsgId = SharedMsgId "", msgBody = s}
       batched = "[" <> xMsgNewStr <> "," <> descrStr <> "]"
 
-  runBatcherTest' maxRawMsgLength [msg xMsgNewStr, msg descrStr] [] [batched]
+  runBatcherTest' maxEncodedMsgLength [msg xMsgNewStr, msg descrStr] [] [batched]
 
 runBatcherTest :: Int -> [SndMessage] -> [ChatError] -> [ByteString] -> Spec
 runBatcherTest maxLen msgs expectedErrors expectedBatches =