double ratchet versioning for post-quantum encryption (#1025)

* correctly parse new Ratchet fields when omitted

* rfc: migrating connection versions to pqdr

* update rfc

* WIP (dont commit)

* rename versions

* update ratchet version based on PQ encryption feature flag

* remove duplicate function

* synchronize ratchet, fix tests, refactor

* comments

* test

* pattern

src/Simplex/Messaging/Agent.hs

@@ -558,14 +558,14 @@ newConnAsync c userId corrId enableNtfs cMode pqInitKeys subMode = do
 newConnNoQueues :: AgentMonad m => AgentClient -> UserId -> ConnId -> Bool -> SConnectionMode c -> CR.PQEncryption -> m ConnId
 newConnNoQueues c userId connId enableNtfs cMode pqEncryption = do
   g <- asks random
-  connAgentVersion <- asks $ maxVersion . smpAgentVRange . config
+  connAgentVersion <- asks $ maxVersion . ($ pqEncryption) . smpAgentVRange . config
   let cData = ConnData {userId, connId, connAgentVersion, enableNtfs, lastExternalSndId = 0, deleted = False, ratchetSyncState = RSOk, pqEncryption}
   withStore c $ \db -> createNewConn db g cData cMode
 
 joinConnAsync :: AgentMonad m => AgentClient -> UserId -> ACorrId -> Bool -> ConnectionRequestUri c -> ConnInfo -> CR.PQEncryption -> SubscriptionMode -> m ConnId
 joinConnAsync c userId corrId enableNtfs cReqUri@(CRInvitationUri ConnReqUriData {crAgentVRange} _) cInfo pqEncryption subMode = do
   withInvLock c (strEncode cReqUri) "joinConnAsync" $ do
-    aVRange <- asks $ smpAgentVRange . config
+    aVRange <- asks $ ($ pqEncryption) . smpAgentVRange . config
     case crAgentVRange `compatibleVersion` aVRange of
       Just (Compatible connAgentVersion) -> do
         g <- asks random
@@ -667,14 +667,16 @@ newRcvConnSrv c userId connId enableNtfs cMode clientData pqInitKeys subMode srv
   when enableNtfs $ do
     ns <- asks ntfSupervisor
     atomically $ sendNtfSubCommand ns (connId, NSCCreate)
-  let crData = ConnReqUriData SSSimplex smpAgentVRange [qUri] clientData
+  let pqEnc = CR.connPQEncryption pqInitKeys
+      crData = ConnReqUriData SSSimplex (smpAgentVRange pqEnc) [qUri] clientData
+      e2eVRange = e2eEncryptVRange pqEnc
   case cMode of
     SCMContact -> pure (connId, CRContactUri crData)
     SCMInvitation -> do
       g <- asks random
-      (pk1, pk2, pKem, e2eRcvParams) <- liftIO $ CR.generateRcvE2EParams g (maxVersion e2eEncryptVRange) (CR.initialPQEncryption pqInitKeys)
+      (pk1, pk2, pKem, e2eRcvParams) <- liftIO $ CR.generateRcvE2EParams g (maxVersion e2eVRange) (CR.initialPQEncryption pqInitKeys)
       withStore' c $ \db -> createRatchetX3dhKeys db connId pk1 pk2 pKem
-      pure (connId, CRInvitationUri crData $ toVersionRangeT e2eRcvParams e2eEncryptVRange)
+      pure (connId, CRInvitationUri crData $ toVersionRangeT e2eRcvParams e2eVRange)
 
 joinConn :: AgentMonad m => AgentClient -> UserId -> ConnId -> Bool -> ConnectionRequestUri c -> ConnInfo -> CR.PQEncryption -> SubscriptionMode -> m ConnId
 joinConn c userId connId enableNtfs cReq cInfo pqEnc subMode = do
@@ -687,17 +689,18 @@ joinConn c userId connId enableNtfs cReq cInfo pqEnc subMode = do
 startJoinInvitation :: AgentMonad m => UserId -> ConnId -> Bool -> ConnectionRequestUri 'CMInvitation -> CR.PQEncryption -> m (Compatible VersionSMPA, ConnData, NewSndQueue, CR.Ratchet 'C.X448, CR.SndE2ERatchetParams 'C.X448)
 startJoinInvitation userId connId enableNtfs (CRInvitationUri ConnReqUriData {crAgentVRange, crSmpQueues = (qUri :| _)} e2eRcvParamsUri) pqEncryption = do
   AgentConfig {smpClientVRange, smpAgentVRange, e2eEncryptVRange} <- asks config
+  let e2eVRange = e2eEncryptVRange pqEncryption
   case ( qUri `compatibleVersion` smpClientVRange,
-         e2eRcvParamsUri `compatibleVersion` e2eEncryptVRange,
-         crAgentVRange `compatibleVersion` smpAgentVRange
+         e2eRcvParamsUri `compatibleVersion` e2eVRange,
+         crAgentVRange `compatibleVersion` smpAgentVRange pqEncryption
        ) of
     (Just qInfo, Just (Compatible e2eRcvParams@(CR.E2ERatchetParams v _ rcDHRr kem_)), Just aVersion@(Compatible connAgentVersion)) -> do
       g <- asks random
       (pk1, pk2, pKem, e2eSndParams) <- liftIO $ CR.generateSndE2EParams g v (CR.replyKEM_ pqEncryption kem_)
       (_, rcDHRs) <- atomically $ C.generateKeyPair g
-      -- TODO PQ generate KEM keypair if needed - is it done?
       rcParams <- liftEitherWith cryptoError $ CR.pqX3dhSnd pk1 pk2 pKem e2eRcvParams
-      let rc = CR.initSndRatchet e2eEncryptVRange rcDHRr rcDHRs rcParams
+      let rcVs = CR.RVersions {current = v, maxSupported = maxVersion e2eVRange}
+          rc = CR.initSndRatchet rcVs rcDHRr rcDHRs rcParams
       q <- newSndQueue userId "" qInfo
       let cData = ConnData {userId, connId, connAgentVersion, enableNtfs, lastExternalSndId = 0, deleted = False, ratchetSyncState = RSOk, pqEncryption}
       pure (aVersion, cData, q, rc, e2eSndParams)
@@ -720,7 +723,7 @@ joinConnSrv c userId connId enableNtfs inv@CRInvitationUri {} cInfo pqEnc subMod
         void $ withStore' c $ \db -> deleteConn db Nothing connId'
         throwError e
 joinConnSrv c userId connId enableNtfs (CRContactUri ConnReqUriData {crAgentVRange, crSmpQueues = (qUri :| _)}) cInfo pqEnc subMode srv = do
-  aVRange <- asks $ smpAgentVRange . config
+  aVRange <- asks $ ($ pqEnc) . smpAgentVRange . config
   clientVRange <- asks $ smpClientVRange . config
   case ( qUri `compatibleVersion` clientVRange,
          crAgentVRange `compatibleVersion` aVRange
@@ -1107,23 +1110,24 @@ enqueueMessage c cData sq pqEnc_ msgFlags aMessage =
 -- this function is used only for sending messages in batch, it returns the list of successes to enqueue additional deliveries
 enqueueMessageB :: forall m t. (AgentMonad' m, Traversable t) => AgentClient -> t (Either AgentErrorType (ConnData, NonEmpty SndQueue, Maybe CR.PQEncryption, MsgFlags, AMessage)) -> m (t (Either AgentErrorType ((AgentMsgId, CR.PQEncryption), Maybe (ConnData, [SndQueue], AgentMsgId))))
 enqueueMessageB c reqs = do
-  aVRange <- asks $ maxVersion . smpAgentVRange . config
-  reqMids <- withStoreBatch c $ \db -> fmap (bindRight $ storeSentMsg db aVRange) reqs
+  getAVRange <- asks $ smpAgentVRange . config
+  reqMids <- withStoreBatch c $ \db -> fmap (bindRight $ storeSentMsg db getAVRange) reqs
   forME reqMids $ \((cData, sq :| sqs, _, _, _), InternalId msgId, pqSecr) -> do
     submitPendingMsg c cData sq
     let sqs' = filter isActiveSndQ sqs
     pure $ Right ((msgId, pqSecr), if null sqs' then Nothing else Just (cData, sqs', msgId))
   where
-    storeSentMsg :: DB.Connection -> VersionSMPA -> (ConnData, NonEmpty SndQueue, Maybe CR.PQEncryption, MsgFlags, AMessage) -> IO (Either AgentErrorType ((ConnData, NonEmpty SndQueue, Maybe CR.PQEncryption, MsgFlags, AMessage), InternalId, CR.PQEncryption))
-    storeSentMsg db agentVersion req@(ConnData {connId}, sq :| _, pqEnc_, msgFlags, aMessage) = fmap (first storeError) $ runExceptT $ do
+    storeSentMsg :: DB.Connection -> (CR.PQEncryption -> VersionRangeSMPA) -> (ConnData, NonEmpty SndQueue, Maybe CR.PQEncryption, MsgFlags, AMessage) -> IO (Either AgentErrorType ((ConnData, NonEmpty SndQueue, Maybe CR.PQEncryption, MsgFlags, AMessage), InternalId, CR.PQEncryption))
+    storeSentMsg db getAVRange req@(ConnData {connId, connAgentVersion = v}, sq :| _, pqEnc_, msgFlags, aMessage) = fmap (first storeError) $ runExceptT $ do
       internalTs <- liftIO getCurrentTime
       (internalId, internalSndId, prevMsgHash) <- liftIO $ updateSndIds db connId
       let privHeader = APrivHeader (unSndId internalSndId) prevMsgHash
           agentMsg = AgentMessage privHeader aMessage
           agentMsgStr = smpEncode agentMsg
           internalHash = C.sha256Hash agentMsgStr
-      (encAgentMessage, pqEncryption) <- agentRatchetEncrypt db connId agentMsgStr e2eEncUserMsgLength pqEnc_
-      let msgBody = smpEncode $ AgentMsgEnvelope {agentVersion, encAgentMessage}
+      (encAgentMessage, pqEncryption) <- agentRatchetEncrypt db connId agentMsgStr (e2eEncUserMsgLength v) pqEnc_
+      let agentVersion = maxVersion . getAVRange $ fromMaybe CR.PQEncOff pqEnc_
+          msgBody = smpEncode $ AgentMsgEnvelope {agentVersion, encAgentMessage}
           msgType = agentMessageType agentMsg
           msgData = SndMsgData {internalId, internalSndId, internalTs, msgType, msgFlags, msgBody, pqEncryption, internalHash, prevMsgHash}
       liftIO $ createSndMsg db connId msgData
@@ -1402,18 +1406,19 @@ abortConnectionSwitch' c connId =
 synchronizeRatchet' :: AgentMonad m => AgentClient -> ConnId -> CR.PQEncryption -> Bool -> m ConnectionStats
 synchronizeRatchet' c connId pqEnc force = withConnLock c connId "synchronizeRatchet" $ do
   withStore c (`getConn` connId) >>= \case
-    SomeConn _ (DuplexConnection cData rqs sqs)
+    SomeConn _ (DuplexConnection cData@ConnData {pqEncryption} rqs sqs)
       | ratchetSyncAllowed cData || force -> do
           -- check queues are not switching?
+          cData' <- if pqEncryption == pqEnc then pure cData else withStore' c $ \db -> setConnPQEncryption db cData pqEnc
           AgentConfig {e2eEncryptVRange} <- asks config
           g <- asks random
-          (pk1, pk2, pKem, e2eParams) <- liftIO $ CR.generateRcvE2EParams g (maxVersion e2eEncryptVRange) pqEnc
-          enqueueRatchetKeyMsgs c cData sqs e2eParams
+          (pk1, pk2, pKem, e2eParams) <- liftIO $ CR.generateRcvE2EParams g (maxVersion $ e2eEncryptVRange pqEnc) pqEnc
+          enqueueRatchetKeyMsgs c cData' sqs e2eParams
           withStore' c $ \db -> do
             setConnRatchetSync db connId RSStarted
             setRatchetX3dhKeys db connId pk1 pk2 pKem
-          let cData' = cData {ratchetSyncState = RSStarted} :: ConnData
-              conn' = DuplexConnection cData' rqs sqs
+          let cData'' = cData' {ratchetSyncState = RSStarted} :: ConnData
+              conn' = DuplexConnection cData'' rqs sqs
           pure $ connectionStats conn'
       | otherwise -> throwError $ CMD PROHIBITED
     _ -> throwError $ CMD PROHIBITED
@@ -2064,8 +2069,8 @@ processSMPTransmission c@AgentClient {smpClients, subQ} (tSess@(_, srv, _), _v,
                       _ -> prohibited >> ack
                   _ -> prohibited >> ack
               updateConnVersion :: Connection c -> ConnData -> VersionSMPA -> m (Connection c)
-              updateConnVersion conn' cData' msgAgentVersion = do
-                aVRange <- asks $ smpAgentVRange . config
+              updateConnVersion conn' cData'@ConnData {pqEncryption} msgAgentVersion = do
+                aVRange <- asks $ ($ pqEncryption) . smpAgentVRange . config
                 let msgAVRange = fromMaybe (versionToRange msgAgentVersion) $ safeVersionRange (minVersion aVRange) msgAgentVersion
                 case msgAVRange `compatibleVersion` aVRange of
                   Just (Compatible av)
@@ -2126,21 +2131,27 @@ processSMPTransmission c@AgentClient {smpClients, subQ} (tSess@(_, srv, _), _v,
           parseMessage :: Encoding a => ByteString -> m a
           parseMessage = liftEither . parse smpP (AGENT A_MESSAGE)
 
+          -- TODO PQ make sure pqEncryption in conn' is set correctly
           smpConfirmation :: SMP.MsgId -> Connection c -> C.APublicAuthKey -> C.PublicKeyX25519 -> Maybe (CR.SndE2ERatchetParams 'C.X448) -> ByteString -> VersionSMPC -> VersionSMPA -> m ()
           smpConfirmation srvMsgId conn' senderKey e2ePubKey e2eEncryption encConnInfo smpClientVersion agentVersion = do
             logServer "<--" c srv rId $ "MSG <CONF>:" <> logSecret srvMsgId
             AgentConfig {smpClientVRange, smpAgentVRange, e2eEncryptVRange} <- asks config
+            let ConnData {pqEncryption} = toConnData conn'
+                aVRange = smpAgentVRange pqEncryption
+                e2eVRange = e2eEncryptVRange pqEncryption
             unless
-              (agentVersion `isCompatible` smpAgentVRange && smpClientVersion `isCompatible` smpClientVRange)
+              (agentVersion `isCompatible` aVRange && smpClientVersion `isCompatible` smpClientVRange)
               (throwError $ AGENT A_VERSION)
             case status of
               New -> case (conn', e2eEncryption) of
                 -- party initiating connection
-                (RcvConnection ConnData {pqEncryption} _, Just (CR.AE2ERatchetParams _ e2eSndParams@(CR.E2ERatchetParams e2eVersion _ _ _))) -> do
-                  unless (e2eVersion `isCompatible` e2eEncryptVRange) (throwError $ AGENT A_VERSION)
+                (RcvConnection _ _, Just (CR.AE2ERatchetParams _ e2eSndParams@(CR.E2ERatchetParams e2eVersion _ _ _))) -> do
+                  unless (e2eVersion `isCompatible` e2eVRange) (throwError $ AGENT A_VERSION)
                   (pk1, rcDHRs, pKem) <- withStore c (`getRatchetX3dhKeys` connId)
                   rcParams <- liftError cryptoError $ CR.pqX3dhRcv pk1 rcDHRs pKem e2eSndParams
-                  let rc = CR.initRcvRatchet e2eEncryptVRange rcDHRs rcParams pqEncryption
+                  -- TODO PQ combine isCompatible check and construction in one call
+                  let rcVs = CR.RVersions {current = e2eVersion, maxSupported = maxVersion e2eVRange}
+                      rc = CR.initRcvRatchet rcVs rcDHRs rcParams pqEncryption
                   g <- asks random
                   (agentMsgBody_, rc', skipped) <- liftError cryptoError $ CR.rcDecrypt g rc M.empty encConnInfo
                   case (agentMsgBody_, skipped) of
@@ -2153,7 +2164,7 @@ processSMPTransmission c@AgentClient {smpClients, subQ} (tSess@(_, srv, _), _v,
                         processConf connInfo senderConf = do
                           let newConfirmation = NewConfirmation {connId, senderConf, ratchetState = rc'}
                           confId <- withStore c $ \db -> do
-                            setConnectionVersion db connId agentVersion
+                            setConnAgentVersion db connId agentVersion
                             createConfirmation db g newConfirmation
                           let srvs = map qServer $ smpReplyQueues senderConf
                           notify $ CONF confId srvs connInfo
@@ -2182,8 +2193,6 @@ processSMPTransmission c@AgentClient {smpClients, subQ} (tSess@(_, srv, _), _v,
                     -- `sndStatus == Active` when HELLO was previously sent, and this is the reply HELLO
                     -- this branch is executed by the accepting party in duplexHandshake mode (v2)
                     -- (was executed by initiating party in v1 that is no longer supported)
-                    --
-                    -- TODO PQ encryption mode
                     | sndStatus == Active -> notify $ CON pqEncryption
                     | otherwise -> enqueueDuplexHello sq
                   _ -> pure ()
@@ -2328,13 +2337,17 @@ processSMPTransmission c@AgentClient {smpClients, subQ} (tSess@(_, srv, _), _v,
             DuplexConnection {} -> action conn'
             _ -> qError $ name <> ": message must be sent to duplex connection"
 
+          -- TODO PQ make sure pqEncryption is set correctly here
           newRatchetKey :: CR.RcvE2ERatchetParams 'C.X448 -> Connection 'CDuplex -> m ()
-          newRatchetKey e2eOtherPartyParams@(CR.E2ERatchetParams e2eVersion k1Rcv k2Rcv kem_) conn'@(DuplexConnection cData'@ConnData {lastExternalSndId} _ sqs) =
+          newRatchetKey e2eOtherPartyParams@(CR.E2ERatchetParams e2eVersion k1Rcv k2Rcv _) conn'@(DuplexConnection cData'@ConnData {lastExternalSndId, pqEncryption} _ sqs) =
             unlessM ratchetExists $ do
               AgentConfig {e2eEncryptVRange} <- asks config
-              unless (e2eVersion `isCompatible` e2eEncryptVRange) (throwError $ AGENT A_VERSION)
+              let connE2EVRange = e2eEncryptVRange pqEncryption
+              unless (e2eVersion `isCompatible` connE2EVRange) (throwError $ AGENT A_VERSION)
               keys <- getSendRatchetKeys
-              initRatchet e2eEncryptVRange keys
+              -- TODO PQ combine with `isCompatible` check above
+              let rcVs = CR.RVersions {current = e2eVersion, maxSupported = maxVersion connE2EVRange}
+              initRatchet rcVs keys
               notifyAgreed
             where
               rkHashRcv = rkHash k1Rcv k2Rcv
@@ -2360,8 +2373,7 @@ processSMPTransmission c@AgentClient {smpClients, subQ} (tSess@(_, srv, _), _v,
                 where
                   sendReplyKey = do
                     g <- asks random
-                    -- TODO PQ the decision to use KEM should depend on connection
-                    (pk1, pk2, pKem, e2eParams) <- liftIO $ CR.generateRcvE2EParams g e2eVersion CR.PQEncOn
+                    (pk1, pk2, pKem, e2eParams) <- liftIO $ CR.generateRcvE2EParams g e2eVersion pqEncryption
                     enqueueRatchetKeyMsgs c cData' sqs e2eParams
                     pure (pk1, pk2, pKem)
                   notifyRatchetSyncError = do
@@ -2380,16 +2392,15 @@ processSMPTransmission c@AgentClient {smpClients, subQ} (tSess@(_, srv, _), _v,
                 createRatchet db connId rc
               -- compare public keys `k1` in AgentRatchetKey messages sent by self and other party
               -- to determine ratchet initilization ordering
-              initRatchet :: CR.VersionRangeE2E -> (C.PrivateKeyX448, C.PrivateKeyX448, Maybe CR.RcvPrivRKEMParams) -> m ()
-              initRatchet e2eEncryptVRange (pk1, pk2, pKem)
+              initRatchet :: CR.RatchetVersions -> (C.PrivateKeyX448, C.PrivateKeyX448, Maybe CR.RcvPrivRKEMParams) -> m ()
+              initRatchet rcVs (pk1, pk2, pKem)
                 | rkHash (C.publicKey pk1) (C.publicKey pk2) <= rkHashRcv = do
                     rcParams <- liftError cryptoError $ CR.pqX3dhRcv pk1 pk2 pKem e2eOtherPartyParams
-                    -- TODO PQ the decision to use KEM should either depend on the global setting or on whether it was enabled in connection before
-                    recreateRatchet $ CR.initRcvRatchet e2eEncryptVRange pk2 rcParams $ CR.PQEncryption (isJust kem_)
+                    recreateRatchet $ CR.initRcvRatchet rcVs pk2 rcParams pqEncryption
                 | otherwise = do
                     (_, rcDHRs) <- atomically . C.generateKeyPair =<< asks random
                     rcParams <- liftEitherWith cryptoError $ CR.pqX3dhSnd pk1 pk2 (CR.APRKP CR.SRKSProposed <$> pKem) e2eOtherPartyParams
-                    recreateRatchet $ CR.initSndRatchet e2eEncryptVRange k2Rcv rcDHRs rcParams
+                    recreateRatchet $ CR.initSndRatchet rcVs k2Rcv rcDHRs rcParams
                     void . enqueueMessages' c cData' sqs Nothing SMP.MsgFlags {notification = True} $ EREADY lastExternalSndId
 
           checkMsgIntegrity :: PrevExternalSndId -> ExternalSndId -> PrevRcvMsgHash -> ByteString -> MsgIntegrity
@@ -2432,7 +2443,7 @@ confirmQueueAsync c cData sq srv connInfo e2eEncryption_ pqEnc subMode = do
   submitPendingMsg c cData sq
 
 confirmQueue :: forall m. AgentMonad m => Compatible VersionSMPA -> AgentClient -> ConnData -> SndQueue -> SMPServerWithAuth -> ConnInfo -> Maybe (CR.SndE2ERatchetParams 'C.X448) -> Maybe CR.PQEncryption -> SubscriptionMode -> m ()
-confirmQueue (Compatible agentVersion) c cData@ConnData {connId} sq srv connInfo e2eEncryption_ pqEnc_ subMode = do
+confirmQueue (Compatible agentVersion) c cData@ConnData {connId, connAgentVersion = v} sq srv connInfo e2eEncryption_ pqEnc_ subMode = do
   msg <- mkConfirmation =<< mkAgentConfirmation c cData sq srv connInfo subMode
   sendConfirmation c sq msg
   withStore' c $ \db -> setSndQueueStatus db sq Confirmed
@@ -2440,7 +2451,7 @@ confirmQueue (Compatible agentVersion) c cData@ConnData {connId} sq srv connInfo
     mkConfirmation :: AgentMessage -> m MsgBody
     mkConfirmation aMessage = withStore c $ \db -> runExceptT $ do
       void . liftIO $ updateSndIds db connId
-      (encConnInfo, _) <- agentRatchetEncrypt db connId (smpEncode aMessage) e2eEncConnInfoLength pqEnc_
+      (encConnInfo, _) <- agentRatchetEncrypt db connId (smpEncode aMessage) (e2eEncConnInfoLength v) pqEnc_
       pure . smpEncode $ AgentConfirmation {agentVersion, e2eEncryption_, encConnInfo}
 
 mkAgentConfirmation :: AgentMonad m => AgentClient -> ConnData -> SndQueue -> SMPServerWithAuth -> ConnInfo -> SubscriptionMode -> m AgentMessage
@@ -2454,13 +2465,13 @@ enqueueConfirmation c cData sq connInfo e2eEncryption_ pqEnc_ = do
   submitPendingMsg c cData sq
 
 storeConfirmation :: AgentMonad m => AgentClient -> ConnData -> SndQueue -> Maybe (CR.SndE2ERatchetParams 'C.X448) -> Maybe CR.PQEncryption -> AgentMessage -> m ()
-storeConfirmation c ConnData {connId, connAgentVersion} sq e2eEncryption_ pqEnc_ agentMsg = withStore c $ \db -> runExceptT $ do
+storeConfirmation c ConnData {connId, connAgentVersion = v} sq e2eEncryption_ pqEnc_ agentMsg = withStore c $ \db -> runExceptT $ do
   internalTs <- liftIO getCurrentTime
   (internalId, internalSndId, prevMsgHash) <- liftIO $ updateSndIds db connId
   let agentMsgStr = smpEncode agentMsg
       internalHash = C.sha256Hash agentMsgStr
-  (encConnInfo, pqEncryption) <- agentRatchetEncrypt db connId agentMsgStr e2eEncConnInfoLength pqEnc_
-  let msgBody = smpEncode $ AgentConfirmation {agentVersion = connAgentVersion, e2eEncryption_, encConnInfo}
+  (encConnInfo, pqEncryption) <- agentRatchetEncrypt db connId agentMsgStr (e2eEncConnInfoLength v) pqEnc_
+  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}
   liftIO $ createSndMsg db connId msgData
@@ -2472,8 +2483,8 @@ enqueueRatchetKeyMsgs c cData (sq :| sqs) e2eEncryption = do
   mapM_ (enqueueSavedMessage c cData msgId) $ filter isActiveSndQ sqs
 
 enqueueRatchetKey :: forall m. AgentMonad m => AgentClient -> ConnData -> SndQueue -> CR.RcvE2ERatchetParams 'C.X448 -> m AgentMsgId
-enqueueRatchetKey c cData@ConnData {connId} sq e2eEncryption = do
-  aVRange <- asks $ smpAgentVRange . config
+enqueueRatchetKey c cData@ConnData {connId, pqEncryption} sq e2eEncryption = do
+  aVRange <- asks $ ($ pqEncryption) . smpAgentVRange . config
   msgId <- storeRatchetKey $ maxVersion aVRange
   submitPendingMsg c cData sq
   pure $ unId msgId
@@ -2487,8 +2498,7 @@ enqueueRatchetKey c cData@ConnData {connId} sq e2eEncryption = do
           internalHash = C.sha256Hash agentMsgStr
       let msgBody = smpEncode $ AgentRatchetKey {agentVersion, e2eEncryption, info = agentMsgStr}
           msgType = agentMessageType agentMsg
-          -- TODO PQ set pqEncryption based on connection mode
-          msgData = SndMsgData {internalId, internalSndId, internalTs, msgType, msgBody, pqEncryption = CR.PQEncOff, msgFlags = SMP.MsgFlags {notification = True}, internalHash, prevMsgHash}
+          msgData = SndMsgData {internalId, internalSndId, internalTs, msgType, msgBody, pqEncryption, msgFlags = SMP.MsgFlags {notification = True}, internalHash, prevMsgHash}
       liftIO $ createSndMsg db connId msgData
       liftIO $ createSndMsgDelivery db connId sq internalId
       pure internalId

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

@@ -56,15 +56,14 @@ import qualified Simplex.Messaging.Agent.Store.SQLite.Migrations as Migrations
 import Simplex.Messaging.Client
 import Simplex.Messaging.Client.Agent ()
 import qualified Simplex.Messaging.Crypto as C
-import Simplex.Messaging.Crypto.Ratchet (VersionRangeE2E, supportedE2EEncryptVRange)
+import Simplex.Messaging.Crypto.Ratchet (PQEncryption, VersionRangeE2E, supportedE2EEncryptVRange)
 import Simplex.Messaging.Notifications.Client (defaultNTFClientConfig)
 import Simplex.Messaging.Notifications.Transport (NTFVersion)
 import Simplex.Messaging.Notifications.Types
 import Simplex.Messaging.Protocol (NtfServer, VersionRangeSMPC, XFTPServer, XFTPServerWithAuth, supportedSMPClientVRange)
-import Simplex.Messaging.Transport (SMPVersion)
 import Simplex.Messaging.TMap (TMap)
 import qualified Simplex.Messaging.TMap as TM
-import Simplex.Messaging.Transport (TLS, Transport (..))
+import Simplex.Messaging.Transport (SMPVersion, TLS, Transport (..))
 import Simplex.Messaging.Transport.Client (defaultSMPPort)
 import Simplex.Messaging.Util (allFinally, catchAllErrors, tryAllErrors)
 import System.Random (StdGen, newStdGen)
@@ -117,8 +116,8 @@ data AgentConfig = AgentConfig
     caCertificateFile :: FilePath,
     privateKeyFile :: FilePath,
     certificateFile :: FilePath,
-    e2eEncryptVRange :: VersionRangeE2E,
-    smpAgentVRange :: VersionRangeSMPA,
+    e2eEncryptVRange :: PQEncryption -> VersionRangeE2E,
+    smpAgentVRange :: PQEncryption -> VersionRangeSMPA,
     smpClientVRange :: VersionRangeSMPC
   }
 

src/Simplex/Messaging/Agent/Protocol.hs

@@ -39,6 +39,7 @@ module Simplex.Messaging.Agent.Protocol
     pattern VersionSMPA,
     ratchetSyncSMPAgentVersion,
     deliveryRcptsSMPAgentVersion,
+    pqdrSMPAgentVersion,
     supportedSMPAgentVRange,
     e2eEncConnInfoLength,
     e2eEncUserMsgLength,
@@ -187,7 +188,15 @@ import Simplex.FileTransfer.Protocol (FileParty (..))
 import Simplex.FileTransfer.Transport (XFTPErrorType)
 import Simplex.Messaging.Agent.QueryString
 import qualified Simplex.Messaging.Crypto as C
-import Simplex.Messaging.Crypto.Ratchet (InitialKeys (..), PQEncryption (..), pattern PQEncOff, RcvE2ERatchetParams, RcvE2ERatchetParamsUri, SndE2ERatchetParams)
+import Simplex.Messaging.Crypto.Ratchet
+  ( InitialKeys (..),
+    PQEncryption (..),
+    pattern PQEncOff,
+    pattern PQEncOn,
+    RcvE2ERatchetParams,
+    RcvE2ERatchetParamsUri,
+    SndE2ERatchetParams
+  )
 import Simplex.Messaging.Encoding
 import Simplex.Messaging.Encoding.String
 import Simplex.Messaging.Parsers
@@ -234,6 +243,7 @@ import UnliftIO.Exception (Exception)
 -- 2 - "duplex" (more efficient) connection handshake (6/9/2022)
 -- 3 - support ratchet renegotiation (6/30/2023)
 -- 4 - delivery receipts (7/13/2023)
+-- 5 - post-quantum double ratchet (3/14/2024)
 
 data SMPAgentVersion
 
@@ -255,24 +265,34 @@ ratchetSyncSMPAgentVersion = VersionSMPA 3
 deliveryRcptsSMPAgentVersion :: VersionSMPA
 deliveryRcptsSMPAgentVersion = VersionSMPA 4
 
+pqdrSMPAgentVersion :: VersionSMPA
+pqdrSMPAgentVersion = VersionSMPA 5
+
+-- TODO v5.7 increase to 5
 currentSMPAgentVersion :: VersionSMPA
 currentSMPAgentVersion = VersionSMPA 4
 
-supportedSMPAgentVRange :: VersionRangeSMPA
-supportedSMPAgentVRange = mkVersionRange duplexHandshakeSMPAgentVersion currentSMPAgentVersion
+-- TODO v5.7 remove dependency of version range on whether PQ encryption is used
+supportedSMPAgentVRange :: PQEncryption -> VersionRangeSMPA
+supportedSMPAgentVRange pq =
+  mkVersionRange duplexHandshakeSMPAgentVersion $ case pq of
+    PQEncOn -> pqdrSMPAgentVersion
+    PQEncOff -> currentSMPAgentVersion
 
 -- it is shorter to allow all handshake headers,
 -- including E2E (double-ratchet) parameters and
 -- signing key of the sender for the server
--- TODO PQ this should be version-dependent
--- previously it was 14848, reduced by 3700 (roughly the increase of message ratchet header size + key and ciphertext in reply link)
-e2eEncConnInfoLength :: Int
-e2eEncConnInfoLength = 11148
+e2eEncConnInfoLength :: VersionSMPA -> Int
+e2eEncConnInfoLength v
+  -- reduced by 3700 (roughly the increase of message ratchet header size + key and ciphertext in reply link)
+  | v >= pqdrSMPAgentVersion = 11148
+  | otherwise = 14848
 
--- TODO PQ this should be version-dependent
--- previously it was 15856, reduced by 2200 (roughly the increase of message ratchet header size)
-e2eEncUserMsgLength :: Int
-e2eEncUserMsgLength = 13656
+e2eEncUserMsgLength :: VersionSMPA -> Int
+e2eEncUserMsgLength v
+  -- reduced by 2200 (roughly the increase of message ratchet header size)
+  | v >= pqdrSMPAgentVersion = 13656
+  | otherwise = 15856
 
 -- | Raw (unparsed) SMP agent protocol transmission.
 type ARawTransmission = (ByteString, ByteString, ByteString)

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

@@ -58,6 +58,7 @@ module Simplex.Messaging.Agent.Store.SQLite
     getConnData,
     setConnDeleted,
     setConnAgentVersion,
+    setConnPQEncryption,
     getDeletedConnIds,
     getDeletedWaitingDeliveryConnIds,
     setConnRatchetSync,
@@ -93,7 +94,6 @@ module Simplex.Messaging.Agent.Store.SQLite
     getAcceptedConfirmation,
     removeConfirmations,
     -- Invitations - sent via Contact connections
-    setConnectionVersion,
     createInvitation,
     getInvitation,
     acceptInvitation,
@@ -889,10 +889,6 @@ removeConfirmations db connId =
     |]
     [":conn_id" := connId]
 
-setConnectionVersion :: DB.Connection -> ConnId -> VersionSMPA -> IO ()
-setConnectionVersion db connId aVersion =
-  DB.execute db "UPDATE connections SET smp_agent_version = ? WHERE conn_id = ?" (aVersion, connId)
-
 createInvitation :: DB.Connection -> TVar ChaChaDRG -> NewInvitation -> IO (Either StoreError InvitationId)
 createInvitation db gVar NewInvitation {contactConnId, connReq, recipientConnInfo} =
   createWithRandomId gVar $ \invitationId ->
@@ -1956,6 +1952,11 @@ setConnAgentVersion :: DB.Connection -> ConnId -> VersionSMPA -> IO ()
 setConnAgentVersion db connId aVersion =
   DB.execute db "UPDATE connections SET smp_agent_version = ? WHERE conn_id = ?" (aVersion, connId)
 
+setConnPQEncryption :: DB.Connection -> ConnData -> CR.PQEncryption -> IO ConnData
+setConnPQEncryption db cData@ConnData {connId} pqEnc = do
+  DB.execute db "UPDATE connections SET pq_encryption = ? WHERE conn_id = ?" (pqEnc, connId)
+  pure (cData :: ConnData) {pqEncryption = pqEnc}
+
 getDeletedConnIds :: DB.Connection -> IO [ConnId]
 getDeletedConnIds db = map fromOnly <$> DB.query db "SELECT conn_id FROM connections WHERE deleted = ?" (Only True)
 

src/Simplex/Messaging/Crypto/Ratchet.hs

@@ -23,6 +23,8 @@ module Simplex.Messaging.Crypto.Ratchet
     SkippedMsgDiff (..),
     SkippedMsgKeys,
     InitialKeys (..),
+    pattern IKPQOn,
+    pattern IKPQOff,
     PQEncryption (..),
     pattern PQEncOn,
     pattern PQEncOff,
@@ -40,8 +42,9 @@ module Simplex.Messaging.Crypto.Ratchet
     VersionE2E,
     VersionRangeE2E,
     pattern VersionE2E,
+    RatchetVersions (..),
     kdfX3DHE2EEncryptVersion,
-    pqRatchetVersion,
+    pqRatchetE2EEncryptVersion,
     currentE2EEncryptVersion,
     supportedE2EEncryptVRange,
     generateRcvE2EParams,
@@ -58,7 +61,6 @@ module Simplex.Messaging.Crypto.Ratchet
     rcDecrypt,
     -- used in tests
     MsgHeader (..),
-    RatchetVersions (..),
     RatchetInitParams (..),
     UseKEM (..),
     RKEMParams (..),
@@ -71,6 +73,8 @@ module Simplex.Messaging.Crypto.Ratchet
     ratchetVersions,
     fullHeaderLen,
     applySMDiff,
+    encodeMsgHeader,
+    msgHeaderP,
   )
 where
 
@@ -85,7 +89,7 @@ import Crypto.Random (ChaChaDRG)
 import Data.Aeson (FromJSON (..), ToJSON (..))
 import qualified Data.Aeson as J
 import qualified Data.Aeson.TH as JQ
-import Data.Attoparsec.ByteString (Parser)
+import Data.Attoparsec.ByteString (Parser, peekWord8')
 import qualified Data.Attoparsec.ByteString.Char8 as A
 import qualified Data.ByteArray as BA
 import Data.ByteString.Char8 (ByteString)
@@ -131,14 +135,19 @@ pattern VersionE2E v = Version v
 kdfX3DHE2EEncryptVersion :: VersionE2E
 kdfX3DHE2EEncryptVersion = VersionE2E 2
 
-pqRatchetVersion :: VersionE2E
-pqRatchetVersion = VersionE2E 3
+pqRatchetE2EEncryptVersion :: VersionE2E
+pqRatchetE2EEncryptVersion = VersionE2E 3
 
+-- TODO v5.7 increase to 3
 currentE2EEncryptVersion :: VersionE2E
-currentE2EEncryptVersion = VersionE2E 3
+currentE2EEncryptVersion = VersionE2E 2
 
-supportedE2EEncryptVRange :: VersionRangeE2E
-supportedE2EEncryptVRange = mkVersionRange kdfX3DHE2EEncryptVersion currentE2EEncryptVersion
+-- TODO v5.7 remove dependency of version range on whether PQ encryption is used
+supportedE2EEncryptVRange :: PQEncryption -> VersionRangeE2E
+supportedE2EEncryptVRange pq =
+  mkVersionRange kdfX3DHE2EEncryptVersion $ case pq of 
+    PQEncOn -> pqRatchetE2EEncryptVersion
+    PQEncOff -> currentE2EEncryptVersion
 
 data RatchetKEMState
   = RKSProposed -- only KEM encapsulation key
@@ -215,7 +224,7 @@ deriving instance Show AnyE2ERatchetParams
 
 instance (RatchetKEMStateI s, AlgorithmI a) => Encoding (E2ERatchetParams s a) where
   smpEncode (E2ERatchetParams v k1 k2 kem_)
-    | v >= pqRatchetVersion = smpEncode (v, k1, k2, kem_)
+    | v >= pqRatchetE2EEncryptVersion = smpEncode (v, k1, k2, kem_)
     | otherwise = smpEncode (v, k1, k2)
   smpP = toParams <$?> smpP
     where
@@ -243,7 +252,7 @@ instance Encoding AnyE2ERatchetParams where
     where
       kemP :: VersionE2E -> Parser (Maybe ARKEMParams)
       kemP v
-        | v >= pqRatchetVersion = smpP
+        | v >= pqRatchetE2EEncryptVersion = smpP
         | otherwise = pure Nothing
 
 instance VersionI E2EVersion (E2ERatchetParams s a) where
@@ -283,7 +292,7 @@ instance (RatchetKEMStateI s, AlgorithmI a) => StrEncoding (E2ERatchetParamsUri
         <> maybe [] encodeKem kem_
     where
       encodeKem kem
-        | maxVersion vs < pqRatchetVersion = []
+        | maxVersion vs < pqRatchetE2EEncryptVersion = []
         | otherwise = case kem of
             RKParamsProposed k -> [("kem_key", strEncode k)]
             RKParamsAccepted ct k -> [("kem_ct", strEncode ct), ("kem_key", strEncode k)]
@@ -313,7 +322,7 @@ instance StrEncoding AnyE2ERatchetParamsUri where
       _ -> fail "bad e2e params"
     where
       kemP vr query
-        | maxVersion vr >= pqRatchetVersion =
+        | maxVersion vr >= pqRatchetE2EEncryptVersion =
             queryParam_ "kem_key" query
               $>>= \k -> Just . kemParams k <$> queryParam_ "kem_ct" query
         | otherwise = pure Nothing
@@ -366,7 +375,7 @@ generateE2EParams g v useKEM_ = do
   where
     kemParams :: IO (Maybe (RKEMParams s, PrivRKEMParams s))
     kemParams = case useKEM_ of
-      Just useKem | v >= pqRatchetVersion -> Just <$> do
+      Just useKem | v >= pqRatchetE2EEncryptVersion -> Just <$> do
         ks@(k, _) <- sntrup761Keypair g
         case useKem of
           ProposeKEM -> pure (RKParamsProposed k, PrivateRKParamsProposed ks)
@@ -414,7 +423,7 @@ pqX3dhSnd spk1 spk2 spKem_ (E2ERatchetParams v rk1 rk2 rKem_) = do
   where
     sndPq :: Either CryptoError (Maybe KEMKeyPair, Maybe RatchetKEMAccepted)
     sndPq = case spKem_ of
-      Just (APRKP _ ps) | v >= pqRatchetVersion -> case (ps, rKem_) of
+      Just (APRKP _ ps) | v >= pqRatchetE2EEncryptVersion -> case (ps, rKem_) of
         (PrivateRKParamsAccepted ct shared ks, Just (RKParamsProposed k)) -> Right (Just ks, Just $ RatchetKEMAccepted k shared ct)
         (PrivateRKParamsProposed ks, _) -> Right (Just ks, Nothing) -- both parties can send "proposal" in case of ratchet renegotiation
         _ -> Left CERatchetKEMState
@@ -430,7 +439,7 @@ pqX3dhRcv rpk1 rpk2 rpKem_ (E2ERatchetParams v sk1 sk2 sKem_) = do
   where
     rcvPq :: ExceptT CryptoError IO (Maybe (KEMKeyPair, RatchetKEMAccepted))
     rcvPq = case sKem_ of
-      Just (RKParamsAccepted ct k') | v >= pqRatchetVersion -> case rpKem_ of
+      Just (RKParamsAccepted ct k') | v >= pqRatchetE2EEncryptVersion -> case rpKem_ of
         Just (PrivateRKParamsProposed ks@(_, pk)) -> do
           shared <- liftIO $ sntrup761Dec ct pk
           pure $ Just (ks, RatchetKEMAccepted k' shared ct)
@@ -457,7 +466,6 @@ data Ratchet a = Ratchet
     rcAD :: Str,
     rcDHRs :: PrivateKey a,
     rcKEM :: Maybe RatchetKEM,
-    -- TODO PQ make them optional via JSON parser for PQEncryption
     rcEnableKEM :: PQEncryption, -- will enable KEM on the next ratchet step
     rcSndKEM :: PQEncryption, -- used KEM hybrid secret for sending ratchet
     rcRcvKEM :: PQEncryption, -- used KEM hybrid secret for receiving ratchet
@@ -584,12 +592,12 @@ instance FromField MessageKey where fromField = blobFieldDecoder smpDecode
 -- // above added for KEM
 -- @
 initSndRatchet ::
-  forall a. (AlgorithmI a, DhAlgorithm a) => VersionRangeE2E -> PublicKey a -> PrivateKey a -> (RatchetInitParams, Maybe KEMKeyPair) -> Ratchet a
-initSndRatchet v rcDHRr rcDHRs (RatchetInitParams {assocData, ratchetKey, sndHK, rcvNextHK, kemAccepted}, rcPQRs_) = do
+  forall a. (AlgorithmI a, DhAlgorithm a) => RatchetVersions -> PublicKey a -> PrivateKey a -> (RatchetInitParams, Maybe KEMKeyPair) -> Ratchet a
+initSndRatchet rcVersion rcDHRr rcDHRs (RatchetInitParams {assocData, ratchetKey, sndHK, rcvNextHK, kemAccepted}, rcPQRs_) = do
   -- state.RK, state.CKs, state.NHKs = KDF_RK_HE(SK, DH(state.DHRs, state.DHRr) || state.PQRss)
   let (rcRK, rcCKs, rcNHKs) = rootKdf ratchetKey rcDHRr rcDHRs (rcPQRss <$> kemAccepted)
    in Ratchet
-        { rcVersion = ratchetVersions v,
+        { rcVersion,
           rcAD = assocData,
           rcDHRs,
           rcKEM = (`RatchetKEM` kemAccepted) <$> rcPQRs_,
@@ -613,10 +621,10 @@ initSndRatchet v rcDHRr rcDHRs (RatchetInitParams {assocData, ratchetKey, sndHK,
 -- Please note that the public part of rcDHRs was sent to the sender
 -- as part of the connection request and random salt was received from the sender.
 initRcvRatchet ::
-  forall a. (AlgorithmI a, DhAlgorithm a) => VersionRangeE2E -> PrivateKey a -> (RatchetInitParams, Maybe KEMKeyPair) -> PQEncryption -> Ratchet a
-initRcvRatchet v rcDHRs (RatchetInitParams {assocData, ratchetKey, sndHK, rcvNextHK, kemAccepted}, rcPQRs_) rcEnableKEM =
+  forall a. (AlgorithmI a, DhAlgorithm a) => RatchetVersions -> PrivateKey a -> (RatchetInitParams, Maybe KEMKeyPair) -> PQEncryption -> Ratchet a
+initRcvRatchet rcVersion rcDHRs (RatchetInitParams {assocData, ratchetKey, sndHK, rcvNextHK, kemAccepted}, rcPQRs_) rcEnableKEM =
   Ratchet
-    { rcVersion = ratchetVersions v,
+    { rcVersion,
       rcAD = assocData,
       rcDHRs,
       -- rcKEM:
@@ -654,56 +662,69 @@ data MsgHeader a = MsgHeader
 -- 69 = 2 (original size) + 2 + 1+56 (Curve448) + 4 + 4
 -- TODO PQ this must be version-dependent
 -- TODO this is the exact size, some reserve should be added
-paddedHeaderLen :: Int
-paddedHeaderLen = 2284
+paddedHeaderLen :: VersionE2E -> Int
+paddedHeaderLen v
+  | v >= pqRatchetE2EEncryptVersion = 2284
+  | otherwise = 88
 
 -- only used in tests to validate correct padding
 -- (2 bytes - version size, 1 byte - header size, not to have it fixed or version-dependent)
-fullHeaderLen :: Int
-fullHeaderLen = 2 + 1 + paddedHeaderLen + authTagSize + ivSize @AES256
+fullHeaderLen :: VersionE2E -> Int
+fullHeaderLen v = 2 + 1 + paddedHeaderLen v + authTagSize + ivSize @AES256
 
-instance AlgorithmI a => Encoding (MsgHeader a) where
-  smpEncode MsgHeader {msgMaxVersion, msgDHRs, msgKEM, msgPN, msgNs}
-    | msgMaxVersion >= pqRatchetVersion = smpEncode (msgMaxVersion, msgDHRs, msgKEM, msgPN, msgNs)
-    | otherwise = smpEncode (msgMaxVersion, msgDHRs, msgPN, msgNs)
-  smpP = do
-    msgMaxVersion <- smpP
-    msgDHRs <- smpP
-    msgKEM <- if msgMaxVersion >= pqRatchetVersion then smpP else pure Nothing
-    msgPN <- smpP
-    msgNs <- smpP
-    pure MsgHeader {msgMaxVersion, msgDHRs, msgKEM, msgPN, msgNs}
+-- pass the current version, as MsgHeader only includes the max supported version that can be different from the current
+encodeMsgHeader :: AlgorithmI a => VersionE2E -> MsgHeader a -> ByteString
+encodeMsgHeader v MsgHeader {msgMaxVersion, msgDHRs, msgKEM, msgPN, msgNs}
+  | v >= pqRatchetE2EEncryptVersion = smpEncode (msgMaxVersion, msgDHRs, msgKEM, msgPN, msgNs)
+  | otherwise = smpEncode (msgMaxVersion, msgDHRs, msgPN, msgNs)
+
+-- pass the current version, as MsgHeader only includes the max supported version that can be different from the current
+msgHeaderP :: AlgorithmI a => VersionE2E -> Parser (MsgHeader a)
+msgHeaderP v = do
+  msgMaxVersion <- smpP
+  msgDHRs <- smpP
+  msgKEM <- if v >= pqRatchetE2EEncryptVersion then smpP else pure Nothing
+  msgPN <- smpP
+  msgNs <- smpP
+  pure MsgHeader {msgMaxVersion, msgDHRs, msgKEM, msgPN, msgNs}
 
 data EncMessageHeader = EncMessageHeader
-  { ehVersion :: VersionE2E,
+  { ehVersion :: VersionE2E, -- this is current ratchet version
     ehIV :: IV,
     ehAuthTag :: AuthTag,
     ehBody :: ByteString
   }
 
+-- this encoding depends on version in EncMessageHeader because it is "current" ratchet version
 instance Encoding EncMessageHeader where
   smpEncode EncMessageHeader {ehVersion, ehIV, ehAuthTag, ehBody}
-    | ehVersion >= pqRatchetVersion = smpEncode (ehVersion, ehIV, ehAuthTag, Large ehBody)
+    | ehVersion >= pqRatchetE2EEncryptVersion = smpEncode (ehVersion, ehIV, ehAuthTag, Large ehBody)
     | otherwise = smpEncode (ehVersion, ehIV, ehAuthTag, ehBody)
   smpP = do
     (ehVersion, ehIV, ehAuthTag) <- smpP
-    ehBody <- if ehVersion >= pqRatchetVersion then unLarge <$> smpP else smpP
+    ehBody <- if ehVersion >= pqRatchetE2EEncryptVersion then unLarge <$> smpP else smpP
     pure EncMessageHeader {ehVersion, ehIV, ehAuthTag, ehBody}
 
+-- the header is length-prefixed to parse it as string and use as part of associated data for authenticated encryption
 data EncRatchetMessage = EncRatchetMessage
   { emHeader :: ByteString,
     emAuthTag :: AuthTag,
     emBody :: ByteString
   }
 
+-- the encoder always uses 2-byte lengths for the new version, even for short headers without PQ keys.
 encodeEncRatchetMessage :: VersionE2E -> EncRatchetMessage -> ByteString
 encodeEncRatchetMessage v EncRatchetMessage {emHeader, emBody, emAuthTag}
-  | v >= pqRatchetVersion = smpEncode (Large emHeader, emAuthTag, Tail emBody)
+  | v >= pqRatchetE2EEncryptVersion = smpEncode (Large emHeader, emAuthTag, Tail emBody)
   | otherwise = smpEncode (emHeader, emAuthTag, Tail emBody)
 
-encRatchetMessageP :: VersionE2E -> Parser EncRatchetMessage
-encRatchetMessageP v = do
-  emHeader <- if v >= pqRatchetVersion then unLarge <$> smpP else smpP
+-- This parser relies on the fact that header cannot be shorter than 32 bytes (it is ~69 bytes without PQ KEM),
+-- therefore if the first byte is less or equal to 31 (x1F), then we have 2 byte-length limited to 8191.
+-- This allows upgrading the current version in one message.
+encRatchetMessageP :: Parser EncRatchetMessage
+encRatchetMessageP = do
+  len1 <- peekWord8'
+  emHeader <- if len1 < 32 then unLarge <$> smpP else smpP
   (emAuthTag, Tail emBody) <- smpP
   pure EncRatchetMessage {emHeader, emBody, emAuthTag}
 
@@ -724,6 +745,7 @@ instance ToJSON PQEncryption where
 
 instance FromJSON PQEncryption where
   parseJSON v = PQEncryption <$> parseJSON v
+  omittedField = Just PQEncOff
 
 replyKEM_ :: PQEncryption -> Maybe (RKEMParams 'RKSProposed) -> Maybe AUseKEM
 replyKEM_ pqEnc kem_ = case pqEnc of
@@ -747,6 +769,12 @@ instance StrEncoding PQEncryption where
 data InitialKeys = IKUsePQ | IKNoPQ PQEncryption
   deriving (Eq, Show)
 
+pattern IKPQOn :: InitialKeys
+pattern IKPQOn = IKNoPQ PQEncOn
+
+pattern IKPQOff :: InitialKeys
+pattern IKPQOff = IKNoPQ PQEncOff
+
 instance StrEncoding InitialKeys where
   strEncode = \case
     IKUsePQ -> "pq=invitation"
@@ -773,25 +801,30 @@ joinContactInitialKeys = \case
 
 rcEncrypt :: AlgorithmI a => Ratchet a -> Int -> ByteString -> Maybe PQEncryption -> ExceptT CryptoError IO (ByteString, Ratchet a)
 rcEncrypt Ratchet {rcSnd = Nothing} _ _ _ = throwE CERatchetState
-rcEncrypt rc@Ratchet {rcSnd = Just sr@SndRatchet {rcCKs, rcHKs}, rcDHRs, rcKEM, rcNs, rcPN, rcAD = Str rcAD, rcVersion} paddedMsgLen msg pqMode_ = do
+rcEncrypt rc@Ratchet {rcSnd = Just sr@SndRatchet {rcCKs, rcHKs}, rcDHRs, rcKEM, rcNs, rcPN, rcAD = Str rcAD, rcVersion} paddedMsgLen msg pqEnc_ = do
   -- state.CKs, mk = KDF_CK(state.CKs)
   let (ck', mk, iv, ehIV) = chainKdf rcCKs
   -- enc_header = HENCRYPT(state.HKs, header)
-  (ehAuthTag, ehBody) <- encryptAEAD rcHKs ehIV paddedHeaderLen rcAD msgHeader
+  let v = current rcVersion
+  (ehAuthTag, ehBody) <- encryptAEAD rcHKs ehIV (paddedHeaderLen v) rcAD (msgHeader v)
   -- return enc_header, ENCRYPT(mk, plaintext, CONCAT(AD, enc_header))
-  -- TODO PQ versioning in Ratchet should change: we should use "current" version here
-  let emHeader = smpEncode EncMessageHeader {ehVersion = maxSupported rcVersion, ehBody, ehAuthTag, ehIV}
+  let emHeader = smpEncode EncMessageHeader {ehVersion = v, ehBody, ehAuthTag, ehIV}
   (emAuthTag, emBody) <- encryptAEAD mk iv paddedMsgLen (rcAD <> emHeader) msg
-  let msg' = encodeEncRatchetMessage (maxSupported rcVersion) EncRatchetMessage {emHeader, emBody, emAuthTag}
+  let msg' = encodeEncRatchetMessage v EncRatchetMessage {emHeader, emBody, emAuthTag}
       -- state.Ns += 1
-      rc' = rc {rcSnd = Just sr {rcCKs = ck'}, rcNs = rcNs + 1}
-      rc'' = case pqMode_ of
+      -- TODO v5.8 remove comments below
+      -- Note that maxSupported will not downgrade here below current.
+      -- TODO v5.7 remove comments below
+      -- It will downgrade when decrypting the message when the current version downgrades to remove support for PQ encryption.
+      -- TODO v5.8 replace `max v currentE2EEncryptVersion` with `v` (to allow downgrade when app downgraded)
+      rc' = rc {rcSnd = Just sr {rcCKs = ck'}, rcNs = rcNs + 1, rcVersion = rcVersion {maxSupported = max v currentE2EEncryptVersion}}
+      rc'' = case pqEnc_ of
         Nothing -> rc'
-        Just rcEnableKEM
-          | enablePQ rcEnableKEM -> rc' {rcEnableKEM}
-          | otherwise ->
-              let rcKEM' = (\rck -> rck {rcKEMs = Nothing}) <$> rcKEM
-               in rc' {rcEnableKEM, rcKEM = rcKEM'}
+        -- This sets max version to support PQ encryption.
+        -- Current version upgrade happens when peer decrypts the message.
+        -- TODO v5.7 remove version upgrade here, as it's already upgraded above
+        Just PQEncOn -> rc' {rcEnableKEM = PQEncOn, rcVersion = rcVersion {maxSupported = pqRatchetE2EEncryptVersion}}
+        Just PQEncOff -> rc' {rcEnableKEM = PQEncOff, rcKEM = (\rck -> rck {rcKEMs = Nothing}) <$> rcKEM}
   pure (msg', rc'')
   where
     -- header = HEADER_PQ2(
@@ -801,8 +834,9 @@ rcEncrypt rc@Ratchet {rcSnd = Just sr@SndRatchet {rcCKs, rcHKs}, rcDHRs, rcKEM,
     --   pn = state.PN,
     --   n = state.Ns
     -- )
-    msgHeader =
-      smpEncode
+    msgHeader v =
+      encodeMsgHeader
+        v 
         MsgHeader
           { msgMaxVersion = maxSupported rcVersion,
             msgDHRs = publicKey rcDHRs,
@@ -837,7 +871,7 @@ rcDecrypt ::
   ExceptT CryptoError IO (DecryptResult a)
 rcDecrypt g rc@Ratchet {rcRcv, rcAD = Str rcAD, rcVersion} rcMKSkipped msg' = do
   -- TODO PQ versioning should change
-  encMsg@EncRatchetMessage {emHeader} <- parseE CryptoHeaderError (encRatchetMessageP $ maxSupported rcVersion) msg'
+  encMsg@EncRatchetMessage {emHeader} <- parseE CryptoHeaderError encRatchetMessageP msg'
   encHdr <- parseE CryptoHeaderError smpP emHeader
   -- plaintext = TrySkippedMessageKeysHE(state, enc_header, cipher-text, AD)
   decryptSkipped encHdr encMsg >>= \case
@@ -851,9 +885,16 @@ rcDecrypt g rc@Ratchet {rcRcv, rcAD = Str rcAD, rcVersion} rcMKSkipped msg' = do
     SMMessage r -> pure r
   where
     decryptRcMessage :: RatchetStep -> MsgHeader a -> EncRatchetMessage -> ExceptT CryptoError IO (DecryptResult a)
-    decryptRcMessage rcStep MsgHeader {msgDHRs, msgKEM, msgPN, msgNs} encMsg = do
+    decryptRcMessage rcStep hdr@MsgHeader {msgMaxVersion, msgPN, msgNs} encMsg = do
       -- if dh_ratchet:
-      (rc', smks1) <- ratchetStep rcStep
+      (rc', smks1) <- case rcStep of
+        SameRatchet -> pure (upgradedRatchet, M.empty)
+        AdvanceRatchet -> do
+          -- SkipMessageKeysHE(state, header.pn)
+          (rc', hmks) <- liftEither $ skipMessageKeys msgPN upgradedRatchet
+          -- DHRatchetPQ2HE(state, header)
+          (,hmks) <$> ratchetStep rc' hdr
+      -- SkipMessageKeysHE(state, header.n)
       case skipMessageKeys msgNs rc' of
         Left e -> pure (Left e, rc', smkDiff smks1)
         Right (rc''@Ratchet {rcRcv = Just rr@RcvRatchet {rcCKr}, rcNr}, smks2) -> do
@@ -863,47 +904,44 @@ rcDecrypt g rc@Ratchet {rcRcv, rcAD = Str rcAD, rcVersion} rcMKSkipped msg' = do
           msg <- decryptMessage (MessageKey mk iv) encMsg
           -- state . Nr += 1
           pure (msg, rc'' {rcRcv = Just rr {rcCKr = rcCKr'}, rcNr = rcNr + 1}, smkDiff $ smks1 <> smks2)
-        Right (rc'', smks2) -> do
+        Right (rc'', smks2) ->
           pure (Left CERatchetState, rc'', smkDiff $ smks1 <> smks2)
       where
+        upgradedRatchet :: Ratchet a
+        upgradedRatchet
+          | msgMaxVersion > current rcVersion = rc {rcVersion = rcVersion {current = min msgMaxVersion $ maxSupported rcVersion}}
+          | otherwise = rc
         smkDiff :: SkippedMsgKeys -> SkippedMsgDiff
         smkDiff smks = if M.null smks then SMDNoChange else SMDAdd smks
-        ratchetStep :: RatchetStep -> ExceptT CryptoError IO (Ratchet a, SkippedMsgKeys)
-        ratchetStep SameRatchet = pure (rc, M.empty)
-        ratchetStep AdvanceRatchet =
-          -- SkipMessageKeysHE(state, header.pn)
-          case skipMessageKeys msgPN rc of
-            Left e -> throwE e
-            Right (rc'@Ratchet {rcDHRs, rcRK, rcNHKs, rcNHKr}, hmks) -> do
-              -- DHRatchetPQ2HE(state, header)
-              (kemSS, kemSS', rcKEM') <- pqRatchetStep rc' msgKEM
-              -- state.DHRs = GENERATE_DH()
-              (_, rcDHRs') <- atomically $ generateKeyPair @a g
-              -- state.RK, state.CKr, state.NHKr = KDF_RK_HE(state.RK, DH(state.DHRs, state.DHRr) || ss)
-              let (rcRK', rcCKr', rcNHKr') = rootKdf rcRK msgDHRs rcDHRs kemSS
-                  -- state.RK, state.CKs, state.NHKs = KDF_RK_HE(state.RK, DH(state.DHRs, state.DHRr) || state.PQRss)
-                  (rcRK'', rcCKs', rcNHKs') = rootKdf rcRK' msgDHRs rcDHRs' kemSS'
-                  sndKEM = isJust kemSS'
-                  rcvKEM = isJust kemSS
-                  rc'' =
-                    rc'
-                      { rcDHRs = rcDHRs',
-                        rcKEM = rcKEM',
-                        rcEnableKEM = PQEncryption $ sndKEM || rcvKEM,
-                        rcSndKEM = PQEncryption sndKEM,
-                        rcRcvKEM = PQEncryption rcvKEM,
-                        rcRK = rcRK'',
-                        rcSnd = Just SndRatchet {rcDHRr = msgDHRs, rcCKs = rcCKs', rcHKs = rcNHKs},
-                        rcRcv = Just RcvRatchet {rcCKr = rcCKr', rcHKr = rcNHKr},
-                        rcPN = rcNs rc,
-                        rcNs = 0,
-                        rcNr = 0,
-                        rcNHKs = rcNHKs',
-                        rcNHKr = rcNHKr'
-                      }
-              pure (rc'', hmks)
+        ratchetStep :: Ratchet a -> MsgHeader a -> ExceptT CryptoError IO (Ratchet a)
+        ratchetStep rc'@Ratchet {rcDHRs, rcRK, rcNHKs, rcNHKr} MsgHeader {msgDHRs, msgKEM} = do
+          (kemSS, kemSS', rcKEM') <- pqRatchetStep rc' msgKEM
+          -- state.DHRs = GENERATE_DH()
+          (_, rcDHRs') <- atomically $ generateKeyPair @a g
+          -- state.RK, state.CKr, state.NHKr = KDF_RK_HE(state.RK, DH(state.DHRs, state.DHRr) || ss)
+          let (rcRK', rcCKr', rcNHKr') = rootKdf rcRK msgDHRs rcDHRs kemSS
+              -- state.RK, state.CKs, state.NHKs = KDF_RK_HE(state.RK, DH(state.DHRs, state.DHRr) || state.PQRss)
+              (rcRK'', rcCKs', rcNHKs') = rootKdf rcRK' msgDHRs rcDHRs' kemSS'
+              sndKEM = isJust kemSS'
+              rcvKEM = isJust kemSS
+          pure
+            rc'
+              { rcDHRs = rcDHRs',
+                rcKEM = rcKEM',
+                rcEnableKEM = PQEncryption $ sndKEM || rcvKEM,
+                rcSndKEM = PQEncryption sndKEM,
+                rcRcvKEM = PQEncryption rcvKEM,
+                rcRK = rcRK'',
+                rcSnd = Just SndRatchet {rcDHRr = msgDHRs, rcCKs = rcCKs', rcHKs = rcNHKs},
+                rcRcv = Just RcvRatchet {rcCKr = rcCKr', rcHKr = rcNHKr},
+                rcPN = rcNs rc,
+                rcNs = 0,
+                rcNr = 0,
+                rcNHKs = rcNHKs',
+                rcNHKr = rcNHKr'
+              }
         pqRatchetStep :: Ratchet a -> Maybe ARKEMParams -> ExceptT CryptoError IO (Maybe KEMSharedKey, Maybe KEMSharedKey, Maybe RatchetKEM)
-        pqRatchetStep Ratchet {rcKEM, rcEnableKEM = PQEncryption pqEnc} = \case
+        pqRatchetStep Ratchet {rcKEM, rcEnableKEM = PQEncryption pqEnc, rcVersion = rv} = \case
           -- received message does not have KEM in header,
           -- but the user enabled KEM when sending previous message
           Nothing -> case rcKEM of
@@ -913,7 +951,7 @@ rcDecrypt g rc@Ratchet {rcRcv, rcAD = Str rcAD, rcVersion} rcMKSkipped msg' = do
             _ -> pure (Nothing, Nothing, Nothing)
           -- received message has KEM in header.
           Just (ARKP _ ps)
-            | pqEnc -> do
+            | pqEnc && current rv >= pqRatchetE2EEncryptVersion -> do
                 -- state.PQRr = header.kem
                 (ss, rcPQRr) <- sharedSecret
                 -- state.PQRct = PQKEM-ENC(state.PQRr, state.PQRss) // encapsulated additional shared secret KEM #1
@@ -981,9 +1019,9 @@ rcDecrypt g rc@Ratchet {rcRcv, rcAD = Str rcAD, rcVersion} rcMKSkipped msg' = do
         e -> throwE e
     -- header = HDECRYPT(state.NHKr, enc_header)
     decryptNextHeader hdr = (AdvanceRatchet,) <$> decryptHeader (rcNHKr rc) hdr
-    decryptHeader k EncMessageHeader {ehBody, ehAuthTag, ehIV} = do
+    decryptHeader k EncMessageHeader {ehVersion, ehBody, ehAuthTag, ehIV} = do
       header <- decryptAEAD k ehIV rcAD ehBody ehAuthTag `catchE` \_ -> throwE CERatchetHeader
-      parseE' CryptoHeaderError smpP header
+      parseE' CryptoHeaderError (msgHeaderP ehVersion) header
     decryptMessage :: MessageKey -> EncRatchetMessage -> ExceptT CryptoError IO (Either CryptoError ByteString)
     decryptMessage (MessageKey mk iv) EncRatchetMessage {emHeader, emBody, emAuthTag} =
       -- DECRYPT(mk, cipher-text, CONCAT(AD, enc_header))