From af02a9244241efb588b2d262c8b258c3046ac948 Mon Sep 17 00:00:00 2001
From: Evgeny Poberezkin <2769109+epoberezkin@users.noreply.github.com>
Date: Sat, 2 Sep 2023 23:34:00 +0100
Subject: [PATCH] core: fix WebRTC encryption, test (#3005)

---
 src/Simplex/Chat/Mobile/WebRTC.hs |  2 +-
 tests/MobileTests.hs              | 50 ++++++++++++++++++++++++++++++-
 2 files changed, 50 insertions(+), 2 deletions(-)

diff --git a/src/Simplex/Chat/Mobile/WebRTC.hs b/src/Simplex/Chat/Mobile/WebRTC.hs
index 3fd5f018e..19ba2b751 100644
--- a/src/Simplex/Chat/Mobile/WebRTC.hs
+++ b/src/Simplex/Chat/Mobile/WebRTC.hs
@@ -34,7 +34,7 @@ cTransformMedia f cKey cFrame cFrameLen = do
   frame <- getByteString cFrame cFrameLen
   runExceptT (f key frame >>= liftIO . putFrame) >>= newCAString . fromLeft ""
   where
-    putFrame s = when (B.length s < fromIntegral cFrameLen) $ putByteString cFrame s
+    putFrame s = when (B.length s <= fromIntegral cFrameLen) $ putByteString cFrame s
 {-# INLINE cTransformMedia #-}
 
 chatEncryptMedia :: ByteString -> ByteString -> ExceptT String IO ByteString
diff --git a/tests/MobileTests.hs b/tests/MobileTests.hs
index 31c080354..e11496ef4 100644
--- a/tests/MobileTests.hs
+++ b/tests/MobileTests.hs
@@ -1,22 +1,37 @@
 {-# LANGUAGE CPP #-}
+{-# LANGUAGE ScopedTypeVariables #-}
 
 module MobileTests where
 
 import ChatTests.Utils
 import Control.Monad.Except
+import Crypto.Random (getRandomBytes)
+import Data.ByteString (ByteString)
+import qualified Data.ByteString as B
+import qualified Data.ByteString.Char8 as BS
+import Data.Word (Word8)
+import Foreign.C
+import Foreign.Marshal.Alloc (mallocBytes)
+import Foreign.Ptr
 import Simplex.Chat.Mobile
+import Simplex.Chat.Mobile.Shared
+import Simplex.Chat.Mobile.WebRTC
 import Simplex.Chat.Store
 import Simplex.Chat.Store.Profiles
 import Simplex.Chat.Types (AgentUserId (..), Profile (..))
 import Simplex.Messaging.Agent.Store.SQLite (MigrationConfirmation (..))
+import qualified Simplex.Messaging.Crypto as C
+import Simplex.Messaging.Encoding.String
 import System.FilePath ((</>))
 import Test.Hspec
 
-mobileTests :: SpecWith FilePath
+mobileTests :: HasCallStack => SpecWith FilePath
 mobileTests = do
   describe "mobile API" $ do
     it "start new chat without user" testChatApiNoUser
     it "start new chat with existing user" testChatApi
+    fit "should encrypt/decrypt WebRTC frames" testMediaApi
+    fit "should encrypt/decrypt WebRTC frames via C API" testMediaCApi
 
 noActiveUser :: String
 #if defined(darwin_HOST_OS) && defined(swiftJSON)
@@ -113,3 +128,36 @@ testChatApi tmp = do
   chatRecvMsgWait cc 10000 `shouldReturn` ""
   chatParseMarkdown "hello" `shouldBe` "{}"
   chatParseMarkdown "*hello*" `shouldBe` parsedMarkdown
+
+testMediaApi :: HasCallStack => FilePath -> IO ()
+testMediaApi _ = do
+  key :: ByteString <- getRandomBytes 32
+  frame <- getRandomBytes 100
+  let keyStr = strEncode key
+      reserved = B.replicate (C.authTagSize + C.gcmIVSize) 0
+      frame' = frame <> reserved
+  Right encrypted <- runExceptT $ chatEncryptMedia keyStr frame'
+  encrypted `shouldNotBe` frame'
+  B.length encrypted `shouldBe` B.length frame'
+  runExceptT (chatDecryptMedia keyStr encrypted) `shouldReturn` Right frame'
+
+testMediaCApi :: HasCallStack => FilePath -> IO ()
+testMediaCApi _ = do
+  key :: ByteString <- getRandomBytes 32
+  frame <- getRandomBytes 100
+  let keyStr = strEncode key
+      reserved = B.replicate (C.authTagSize + C.gcmIVSize) 0
+      frame' = frame <> reserved
+  encrypted <- test cChatEncryptMedia keyStr frame'
+  encrypted `shouldNotBe` frame'
+  test cChatDecryptMedia keyStr encrypted `shouldReturn` frame'
+  where
+    test :: HasCallStack => (CString -> Ptr Word8 -> CInt -> IO CString) -> ByteString -> ByteString -> IO ByteString
+    test f keyStr frame = do
+      let len = B.length frame
+          cLen = fromIntegral len
+      ptr <- mallocBytes len
+      putByteString ptr frame
+      cKeyStr <- newCString $ BS.unpack keyStr
+      (f cKeyStr ptr cLen >>= peekCString) `shouldReturn` ""
+      getByteString ptr cLen