agent: reuse ratchet on repeat join (#1426)

* agent: reuse ratchet on repeat join

* check status

* update

---------

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

Author: spaced4ndy

simplexmq.cabal

@@ -130,6 +130,7 @@ library
       Simplex.Messaging.Agent.Store.SQLite.Migrations.M20240702_servers_stats
       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.TRcvQueues
       Simplex.Messaging.Client
       Simplex.Messaging.Client.Agent

src/Simplex/Messaging/Agent.hs

@@ -834,26 +834,39 @@ joinConn c userId connId enableNtfs cReq cInfo pqSupport subMode = do
 startJoinInvitation :: AgentClient -> UserId -> ConnId -> Maybe SndQueue -> Bool -> ConnectionRequestUri 'CMInvitation -> PQSupport -> AM (ConnData, SndQueue, CR.SndE2ERatchetParams 'C.X448)
 startJoinInvitation c userId connId sq_ enableNtfs cReqUri pqSup =
   lift (compatibleInvitationUri cReqUri) >>= \case
-    Just (qInfo, Compatible e2eRcvParams@(CR.E2ERatchetParams v _ rcDHRr kem_), Compatible connAgentVersion) -> do
-      g <- asks random
+    Just (qInfo, Compatible e2eRcvParams@(CR.E2ERatchetParams v _ _ _), Compatible connAgentVersion) -> do
+      -- this case avoids re-generating queue keys and subsequent failure of SKEY that timed out
+      -- e2ePubKey is always present, it's Maybe historically
       let pqSupport = pqSup `CR.pqSupportAnd` versionPQSupport_ connAgentVersion (Just v)
+      (sq', e2eSndParams) <- case sq_ of
+        Just sq@SndQueue {e2ePubKey = Just _k} -> do
+          e2eSndParams <-
+            withStore' c (\db -> getSndRatchet db connId v) >>= \case
+              Right r -> pure $ snd r
+              Left e -> do
+                atomically $ writeTBQueue (subQ c) ("", connId, AEvt SAEConn (ERR $ INTERNAL $ "no snd ratchet " <> show e))
+                createRatchet_ pqSupport e2eRcvParams
+          pure (sq, e2eSndParams)
+        _ -> do
+          q <- lift $ fst <$> newSndQueue userId "" qInfo
+          e2eSndParams <- createRatchet_ pqSupport e2eRcvParams
+          withStore c $ \db -> runExceptT $ do
+            sq' <- maybe (ExceptT $ updateNewConnSnd db connId q) pure sq_
+            pure (sq', e2eSndParams)
+      let cData = ConnData {userId, connId, connAgentVersion, enableNtfs, lastExternalSndId = 0, deleted = False, ratchetSyncState = RSOk, pqSupport}
+      pure (cData, sq', e2eSndParams)
+    Nothing -> throwE $ AGENT A_VERSION
+  where
+    createRatchet_ pqSupport e2eRcvParams@(CR.E2ERatchetParams v _ rcDHRr kem_) = do
+      g <- asks random
       (pk1, pk2, pKem, e2eSndParams) <- liftIO $ CR.generateSndE2EParams g v (CR.replyKEM_ v kem_ pqSupport)
       (_, rcDHRs) <- atomically $ C.generateKeyPair g
       rcParams <- liftEitherWith cryptoError $ CR.pqX3dhSnd pk1 pk2 pKem e2eRcvParams
       maxSupported <- asks $ maxVersion . e2eEncryptVRange . config
       let rcVs = CR.RatchetVersions {current = v, maxSupported}
           rc = CR.initSndRatchet rcVs rcDHRr rcDHRs rcParams
-      -- this case avoids re-generating queue keys and subsequent failure of SKEY that timed out
-      -- e2ePubKey is always present, it's Maybe historically
-      q <- case sq_ of
-        Just sq@SndQueue {e2ePubKey = Just _k} -> pure (sq :: SndQueue) {dbQueueId = DBNewQueue}
-        _ -> lift $ fst <$> newSndQueue userId "" qInfo
-      let cData = ConnData {userId, connId, connAgentVersion, enableNtfs, lastExternalSndId = 0, deleted = False, ratchetSyncState = RSOk, pqSupport}
-      sq' <- withStore c $ \db -> runExceptT $ do
-        liftIO $ createRatchet db connId rc
-        maybe (ExceptT $ updateNewConnSnd db connId q) pure sq_
-      pure (cData, sq', e2eSndParams)
-    Nothing -> throwE $ AGENT A_VERSION
+      withStore' c $ \db -> createSndRatchet db connId rc e2eSndParams
+      pure e2eSndParams
 
 connRequestPQSupport :: AgentClient -> PQSupport -> ConnectionRequestUri c -> IO (Maybe (VersionSMPA, PQSupport))
 connRequestPQSupport c pqSup cReq = withAgentEnv' c $ case cReq of
@@ -892,6 +905,7 @@ joinConnSrv c userId connId enableNtfs inv@CRInvitationUri {} cInfo pqSup subMod
     case conn of
       NewConnection _ -> doJoin Nothing
       SndConnection _ sq -> doJoin $ Just sq
+      DuplexConnection _ (RcvQueue {status = New} :| _) (sq@SndQueue {status = New} :| _) -> doJoin $ Just sq
       _ -> throwE $ CMD PROHIBITED $ "joinConnSrv: bad connection " <> show cType
   where
     doJoin :: Maybe SndQueue -> AM SndQueueSecured

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

@@ -131,6 +131,8 @@ module Simplex.Messaging.Agent.Store.SQLite
     -- Double ratchet persistence
     createRatchetX3dhKeys,
     getRatchetX3dhKeys,
+    createSndRatchet,
+    getSndRatchet,
     setRatchetX3dhKeys,
     createRatchet,
     deleteRatchet,
@@ -245,7 +247,6 @@ where
 
 import Control.Logger.Simple
 import Control.Monad
-import Control.Monad.Except
 import Control.Monad.IO.Class
 import Control.Monad.Trans.Except
 import Crypto.Random (ChaChaDRG)
@@ -586,7 +587,6 @@ updateNewConnSnd :: DB.Connection -> ConnId -> NewSndQueue -> IO (Either StoreEr
 updateNewConnSnd db connId sq =
   getConn db connId $>>= \case
     (SomeConn _ NewConnection {}) -> updateConn
-    (SomeConn _ SndConnection {}) -> updateConn -- to allow retries
     (SomeConn c _) -> pure . Left . SEBadConnType $ connType c
   where
     updateConn :: IO (Either StoreError SndQueue)
@@ -1250,19 +1250,48 @@ getRatchetX3dhKeys db connId =
       (Just k1, Just k2, pKem) -> Right (k1, k2, pKem)
       _ -> Left SEX3dhKeysNotFound
 
+createSndRatchet :: DB.Connection -> ConnId -> RatchetX448 -> CR.AE2ERatchetParams 'C.X448 -> IO ()
+createSndRatchet db connId ratchetState (CR.AE2ERatchetParams s (CR.E2ERatchetParams _ x3dhPubKey1 x3dhPubKey2 pqPubKem)) =
+  DB.execute
+    db
+    [sql|
+      INSERT INTO ratchets
+        (conn_id, ratchet_state, x3dh_pub_key_1, x3dh_pub_key_2, pq_pub_kem) VALUES (?, ?, ?, ?, ?)
+      ON CONFLICT (conn_id) DO UPDATE SET
+        ratchet_state = EXCLUDED.ratchet_state,
+        x3dh_priv_key_1 = NULL,
+        x3dh_priv_key_2 = NULL,
+        x3dh_pub_key_1 = EXCLUDED.x3dh_pub_key_1,
+        x3dh_pub_key_2 = EXCLUDED.x3dh_pub_key_2,
+        pq_priv_kem = NULL,
+        pq_pub_kem = EXCLUDED.pq_pub_kem
+    |]
+    (connId, ratchetState, x3dhPubKey1, x3dhPubKey2, CR.ARKP s <$> pqPubKem)
+
+getSndRatchet :: DB.Connection -> ConnId -> CR.VersionE2E -> IO (Either StoreError (RatchetX448, CR.AE2ERatchetParams 'C.X448))
+getSndRatchet db connId v =
+  firstRow' result SEX3dhKeysNotFound $
+    DB.query db "SELECT ratchet_state, x3dh_pub_key_1, x3dh_pub_key_2, pq_pub_kem FROM ratchets WHERE conn_id = ?" (Only connId)
+  where
+    result = \case
+      (Just ratchetState, Just k1, Just k2, pKem_) -> 
+        let params = case pKem_ of
+              Nothing -> CR.AE2ERatchetParams CR.SRKSProposed (CR.E2ERatchetParams v k1 k2 Nothing)
+              Just (CR.ARKP s pKem) -> CR.AE2ERatchetParams s (CR.E2ERatchetParams v k1 k2 (Just pKem))
+         in Right (ratchetState, params)
+      _ -> Left SEX3dhKeysNotFound
+
 -- used to remember new keys when starting ratchet re-synchronization
--- TODO remove the columns for public keys in v5.7.
--- Currently, the keys are not used but still stored to support app downgrade to the previous version.
 setRatchetX3dhKeys :: DB.Connection -> ConnId -> C.PrivateKeyX448 -> C.PrivateKeyX448 -> Maybe CR.RcvPrivRKEMParams -> IO ()
 setRatchetX3dhKeys db connId x3dhPrivKey1 x3dhPrivKey2 pqPrivKem =
   DB.execute
     db
     [sql|
       UPDATE ratchets
-      SET x3dh_priv_key_1 = ?, x3dh_priv_key_2 = ?, x3dh_pub_key_1 = ?, x3dh_pub_key_2 = ?, pq_priv_kem = ?
+      SET x3dh_priv_key_1 = ?, x3dh_priv_key_2 = ?, pq_priv_kem = ?
       WHERE conn_id = ?
     |]
-    (x3dhPrivKey1, x3dhPrivKey2, C.publicKey x3dhPrivKey1, C.publicKey x3dhPrivKey2, pqPrivKem, connId)
+    (x3dhPrivKey1, x3dhPrivKey2, pqPrivKem, connId)
 
 -- TODO remove the columns for public keys in v5.7.
 createRatchet :: DB.Connection -> ConnId -> RatchetX448 -> IO ()
@@ -1278,7 +1307,8 @@ createRatchet db connId rc =
         x3dh_priv_key_2 = NULL,
         x3dh_pub_key_1 = NULL,
         x3dh_pub_key_2 = NULL,
-        pq_priv_kem = NULL
+        pq_priv_kem = NULL,
+        pq_pub_kem = NULL
     |]
     [":conn_id" := connId, ":ratchet_state" := rc]
 

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

@@ -76,6 +76,7 @@ import Simplex.Messaging.Agent.Store.SQLite.Migrations.M20240624_snd_secure
 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.Encoding.String
 import Simplex.Messaging.Parsers (dropPrefix, sumTypeJSON)
 import Simplex.Messaging.Transport.Client (TransportHost)
@@ -120,7 +121,8 @@ schemaMigrations =
     ("m20240624_snd_secure", m20240624_snd_secure, Just down_m20240624_snd_secure),
     ("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)
+    ("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)
   ]
 
 -- | The list of migrations in ascending order by date

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

@@ -0,0 +1,18 @@
+{-# LANGUAGE QuasiQuotes #-}
+
+module Simplex.Messaging.Agent.Store.SQLite.Migrations.M20241224_ratchet_e2e_snd_params where
+
+import Database.SQLite.Simple (Query)
+import Database.SQLite.Simple.QQ (sql)
+
+m20241224_ratchet_e2e_snd_params :: Query
+m20241224_ratchet_e2e_snd_params =
+    [sql|
+ALTER TABLE ratchets ADD COLUMN pq_pub_kem BLOB;
+|]
+
+down_m20241224_ratchet_e2e_snd_params :: Query
+down_m20241224_ratchet_e2e_snd_params =
+    [sql|
+ALTER TABLE ratchets DROP COLUMN pq_pub_kem;
+|]

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

@@ -166,7 +166,8 @@ CREATE TABLE ratchets(
   ,
   x3dh_pub_key_1 BLOB,
   x3dh_pub_key_2 BLOB,
-  pq_priv_kem BLOB
+  pq_priv_kem BLOB,
+  pq_pub_kem BLOB
 ) WITHOUT ROWID;
 CREATE TABLE skipped_messages(
   skipped_message_id INTEGER PRIMARY KEY,

src/Simplex/Messaging/Crypto/Ratchet.hs

@@ -206,6 +206,10 @@ instance Encoding ARKEMParams where
       'A' -> ARKP SRKSAccepted .: RKParamsAccepted <$> smpP <*> smpP
       _ -> fail "bad ratchet KEM params"
 
+instance ToField ARKEMParams where toField = toField . smpEncode
+
+instance FromField ARKEMParams where fromField = blobFieldDecoder smpDecode
+
 data E2ERatchetParams (s :: RatchetKEMState) (a :: Algorithm)
   = E2ERatchetParams VersionE2E (PublicKey a) (PublicKey a) (Maybe (RKEMParams s))
   deriving (Show)