agent: allow changing invitation link data before it is secured (#1552)

* agent: setInvitationShortLink api

* Eq instance

* allow changing link data on server, refactor

* fix

* encodings

* remove link data after connection

* Revert "encodings"

This reverts commit f8e254c.

---------

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

Author: spaced4ndy

src/Simplex/Messaging/Agent.hs

@@ -56,8 +56,8 @@ module Simplex.Messaging.Agent
     deleteConnectionAsync,
     deleteConnectionsAsync,
     createConnection,
-    setContactShortLink,
-    deleteContactShortLink,
+    setConnShortLink,
+    deleteConnShortLink,
     getConnShortLink,
     deleteLocalInvShortLink,
     changeConnectionUser,
@@ -372,13 +372,13 @@ createConnection c userId enableNtfs = withAgentEnv c .::. newConn c userId enab
 {-# INLINE createConnection #-}
 
 -- | Create or update user's contact connection short link
-setContactShortLink :: AgentClient -> ConnId -> ConnInfo -> Maybe CRClientData -> AE (ConnShortLink 'CMContact)
-setContactShortLink c = withAgentEnv c .:. setContactShortLink' c
-{-# INLINE setContactShortLink #-}
+setConnShortLink :: AgentClient -> ConnId -> SConnectionMode c -> ConnInfo -> Maybe CRClientData -> AE (ConnShortLink c)
+setConnShortLink c = withAgentEnv c .:: setConnShortLink' c
+{-# INLINE setConnShortLink #-}
 
-deleteContactShortLink :: AgentClient -> ConnId -> AE ()
-deleteContactShortLink c = withAgentEnv c . deleteContactShortLink' c
-{-# INLINE deleteContactShortLink #-}
+deleteConnShortLink :: AgentClient -> ConnId -> SConnectionMode c -> AE ()
+deleteConnShortLink c = withAgentEnv c .: deleteConnShortLink' c
+{-# INLINE deleteConnShortLink #-}
 
 -- | Get and verify data from short link. For 1-time invitations it preserves the key to allow retries
 getConnShortLink :: AgentClient -> UserId -> ConnShortLink c -> AE (ConnectionRequestUri c, ConnLinkData c)
@@ -833,26 +833,28 @@ newConn c userId enableNtfs cMode userData_ clientData pqInitKeys subMode = do
   (connId,) <$> newRcvConnSrv c userId connId enableNtfs cMode userData_ clientData pqInitKeys subMode srv
     `catchE` \e -> withStore' c (`deleteConnRecord` connId) >> throwE e
 
-setContactShortLink' :: AgentClient -> ConnId -> ConnInfo -> Maybe CRClientData -> AM (ConnShortLink 'CMContact)
-setContactShortLink' c connId userData clientData =
-  withConnLock c connId "setContactShortLink" $
-    withStore c (`getConn` connId) >>= \case
-      SomeConn _ (ContactConnection _ rq) -> do
-        (lnkId, linkKey, d) <- prepareLinkData rq
-        addQueueLink c rq lnkId d
-        pure $ CSLContact SLSServer CCTContact (qServer rq) linkKey
-      _ -> throwE $ CMD PROHIBITED "setContactShortLink: not contact address"
+setConnShortLink' :: AgentClient -> ConnId -> SConnectionMode c -> ConnInfo -> Maybe CRClientData -> AM (ConnShortLink c)
+setConnShortLink' c connId cMode userData clientData =
+  withConnLock c connId "setConnShortLink" $ do
+    SomeConn _ conn <- withStore c (`getConn` connId)
+    (rq, lnkId, sl, d) <- case (conn, cMode) of
+      (ContactConnection _ rq, SCMContact) -> prepareContactLinkData rq
+      (RcvConnection _ rq, SCMInvitation) -> prepareInvLinkData rq
+      _ -> throwE $ CMD PROHIBITED "setConnShortLink: invalid connection or mode"
+    addQueueLink c rq lnkId d
+    pure sl
   where
-    prepareLinkData :: RcvQueue -> AM (SMP.LinkId, LinkKey, QueueLinkData)
-    prepareLinkData rq@RcvQueue {server, sndId, e2ePrivKey, shortLink} = do
+    prepareContactLinkData :: RcvQueue -> AM (RcvQueue, SMP.LinkId, ConnShortLink 'CMContact, QueueLinkData)
+    prepareContactLinkData rq@RcvQueue {server, sndId, e2ePrivKey, shortLink} = do
       g <- asks random
       AgentConfig {smpClientVRange = vr, smpAgentVRange} <- asks config
+      let cslContact = CSLContact SLSServer CCTContact (qServer rq)
       case shortLink of
         Just ShortLinkCreds {shortLinkId, shortLinkKey, linkPrivSigKey, linkEncFixedData} -> do
           let (linkId, k) = SL.contactShortLinkKdf shortLinkKey
-          unless (shortLinkId == linkId) $ throwE $ INTERNAL "setContactShortLink: link ID is not derived from link"
+          unless (shortLinkId == linkId) $ throwE $ INTERNAL "setConnShortLink: link ID is not derived from link"
           d <- liftError id $ SL.encryptUserData g k $ SL.encodeSignUserData linkPrivSigKey smpAgentVRange userData
-          pure (linkId, shortLinkKey, (linkEncFixedData, d))
+          pure (rq, linkId, cslContact shortLinkKey, (linkEncFixedData, d))
         Nothing -> do
           sigKeys@(_, privSigKey) <- atomically $ C.generateKeyPair @'C.Ed25519 g
           let qUri = SMPQueueUri vr $ SMPQueueAddress server sndId (C.publicKey e2ePrivKey) (Just QMContact)
@@ -862,14 +864,26 @@ setContactShortLink' c connId userData clientData =
           srvData <- liftError id $ SL.encryptLinkData g k linkData
           let slCreds = ShortLinkCreds linkId linkKey privSigKey (fst srvData)
           withStore' c $ \db -> updateShortLinkCreds db rq slCreds
-          pure (linkId, linkKey, srvData)
-
-deleteContactShortLink' :: AgentClient -> ConnId -> AM ()
-deleteContactShortLink' c connId =
-  withConnLock c connId "deleteContactShortLink" $
-    withStore c (`getConn` connId) >>= \case
-      SomeConn _ (ContactConnection _ rq) -> deleteQueueLink c rq
-      _ -> throwE $ CMD PROHIBITED "deleteContactShortLink: not contact address"
+          pure (rq, linkId, cslContact linkKey, srvData)
+    prepareInvLinkData :: RcvQueue -> AM (RcvQueue, SMP.LinkId, ConnShortLink 'CMInvitation, QueueLinkData)
+    prepareInvLinkData rq@RcvQueue {shortLink} = case shortLink of
+      Just ShortLinkCreds {shortLinkId, shortLinkKey, linkPrivSigKey, linkEncFixedData} -> do
+        g <- asks random
+        AgentConfig {smpAgentVRange} <- asks config
+        let k = SL.invShortLinkKdf shortLinkKey
+        d <- liftError id $ SL.encryptUserData g k $ SL.encodeSignUserData linkPrivSigKey smpAgentVRange userData
+        let sl = CSLInvitation SLSServer (qServer rq) shortLinkId shortLinkKey
+        pure (rq, shortLinkId, sl, (linkEncFixedData, d))
+      Nothing -> throwE $ CMD PROHIBITED "setConnShortLink: no ShortLinkCreds in invitation"
+
+deleteConnShortLink' :: AgentClient -> ConnId -> SConnectionMode c -> AM ()
+deleteConnShortLink' c connId cMode =
+  withConnLock c connId "deleteConnShortLink" $ do
+    SomeConn _ conn <- withStore c (`getConn` connId)
+    case (conn, cMode) of
+      (ContactConnection _ rq, SCMContact) -> deleteQueueLink c rq
+      (RcvConnection _ rq, SCMInvitation) -> deleteQueueLink c rq
+      _ -> throwE $ CMD PROHIBITED "deleteConnShortLink: not contact address"
 
 -- TODO [short links] remove 1-time invitation data and link ID from the server after the message is sent.
 getConnShortLink' :: forall c. AgentClient -> UserId -> ConnShortLink c -> AM (ConnectionRequestUri c, ConnLinkData c)

src/Simplex/Messaging/Agent/Protocol.hs

@@ -1421,6 +1421,11 @@ data CreatedConnLink m = CCLink {connFullLink :: ConnectionRequestUri m, connSho
 
 data ACreatedConnLink = forall m. ConnectionModeI m => ACCL (SConnectionMode m) (CreatedConnLink m)
 
+instance Eq ACreatedConnLink where
+  ACCL m l == ACCL m' l' = case testEquality m m' of
+    Just Refl -> l == l'
+    _ -> False
+
 deriving instance Show ACreatedConnLink
 
 data AConnectionLink = forall m. ConnectionModeI m => ACL (SConnectionMode m) (ConnectionLink m)

src/Simplex/Messaging/Server.hs

@@ -1075,6 +1075,11 @@ isContactQueue QueueRec {queueMode, senderKey} = case queueMode of
   Just QMContact -> True
   Nothing -> isNothing senderKey -- for backward compatibility with pre-SKEY contact addresses
 
+isSecuredMsgQueue :: QueueRec -> Bool
+isSecuredMsgQueue QueueRec {queueMode, senderKey} = case queueMode of
+  Just QMContact -> False
+  _ -> isJust senderKey
+
 verifyCmdAuthorization :: Maybe (THandleAuth 'TServer, C.CbNonce) -> Maybe TransmissionAuth -> ByteString -> C.APublicAuthKey -> Bool
 verifyCmdAuthorization auth_ tAuth authorized key = maybe False (verify key) tAuth
   where
@@ -1249,13 +1254,14 @@ client
           KEY sKey -> withQueue $ \q _ -> either err (corrId,entId,) <$> secureQueue_ q sKey
           RKEY rKeys -> withQueue $ \q qr -> checkMode QMContact qr $ OK <$$ liftIO (updateKeys (queueStore ms) q rKeys)
           LSET lnkId d ->
-            withQueue $ \q qr -> checkContact qr $ liftIO $ case queueData qr of
-              Just (lnkId', _) | lnkId' /= lnkId -> pure $ Left AUTH
-              _ -> OK <$$ addQueueLinkData (queueStore ms) q lnkId d
+            withQueue $ \q qr -> case queueData qr of
+              _ | isSecuredMsgQueue qr -> pure $ err AUTH
+              Just (lnkId', _) | lnkId' /= lnkId -> pure $ err AUTH -- can't change link ID
+              _ -> liftIO $ either err (const ok) <$> addQueueLinkData (queueStore ms) q lnkId d
           LDEL ->
-            withQueue $ \q qr -> checkContact qr $ liftIO $ case queueData qr of
-              Just _ -> OK <$$ deleteQueueLinkData (queueStore ms) q
-              Nothing -> pure $ Right OK
+            withQueue $ \q qr -> case queueData qr of
+              Just _ -> liftIO $ either err (const ok) <$> deleteQueueLinkData (queueStore ms) q
+              Nothing -> pure ok
           NKEY nKey dhKey -> withQueue $ \q _ -> addQueueNotifier_ q nKey dhKey
           NDEL -> withQueue $ \q _ -> deleteQueueNotifier_ q
           OFF -> maybe (pure $ err INTERNAL) suspendQueue_ q_
@@ -1540,6 +1546,8 @@ client
                   case C.maxLenBS msgBody of
                     Left _ -> pure $ err LARGE_MSG
                     Right body -> do
+                      when (isJust (queueData qr) && isSecuredMsgQueue qr) $ void $ liftIO $
+                        deleteQueueLinkData (queueStore ms) q
                       ServerConfig {messageExpiration, msgIdBytes} <- asks config
                       msgId <- randomId' msgIdBytes
                       msg_ <- liftIO $ time "SEND" $ runExceptT $ do

tests/AgentTests/FunctionalAPITests.hs

@@ -1142,6 +1142,10 @@ testInviationShortLink viaProxy a b =
       Left (SMP _ AUTH) -> pure ()
       r -> liftIO $ expectationFailure ("unexpected result " <> show r)
     runRight $ testJoinConn_ viaProxy True a bId b connReq
+    -- invitation link data is removed after the connection is established
+    runExceptT (getConnShortLink b 1 shortLink) >>= \case
+      Left (SMP _ AUTH) -> pure ()
+      r -> liftIO $ expectationFailure ("unexpected result " <> show r)
 
 testJoinConn_ :: Bool -> Bool -> AgentClient -> ConnId -> AgentClient -> ConnectionRequestUri c -> ExceptT AgentErrorType IO ()
 testJoinConn_ viaProxy sndSecure a bId b connReq = do
@@ -1213,16 +1217,16 @@ testContactShortLink viaProxy a b =
       exchangeGreetingsViaProxy viaProxy a bId b aId
     -- update user data
     let updatedData = "updated user data"
-    shortLink' <- runRight $ setContactShortLink a contactId updatedData Nothing
+    shortLink' <- runRight $ setConnShortLink a contactId SCMContact updatedData Nothing
     shortLink' `shouldBe` shortLink
     (connReq4, updatedConnData') <- runRight $ getConnShortLink c 1 shortLink
     connReq4 `shouldBe` connReq
     linkUserData updatedConnData' `shouldBe` updatedData
     -- one more time
-    shortLink2 <- runRight $ setContactShortLink a contactId updatedData Nothing
+    shortLink2 <- runRight $ setConnShortLink a contactId SCMContact updatedData Nothing
     shortLink2 `shouldBe` shortLink
     -- delete short link
-    runRight_ $ deleteContactShortLink a contactId
+    runRight_ $ deleteConnShortLink a contactId SCMContact
     Left (SMP _ AUTH) <- runExceptT $ getConnShortLink c 1 shortLink
     pure ()
 
@@ -1232,7 +1236,7 @@ testAddContactShortLink viaProxy a b =
     (contactId, CCLink connReq0 Nothing) <- runRight $ A.createConnection a 1 True SCMContact Nothing Nothing CR.IKPQOn SMSubscribe
     Right connReq <- pure $ smpDecode (smpEncode connReq0) --
     let userData = "some user data"
-    shortLink <- runRight $ setContactShortLink a contactId userData Nothing
+    shortLink <- runRight $ setConnShortLink a contactId SCMContact userData Nothing
     (connReq', connData') <- runRight $ getConnShortLink b 1 shortLink
     strDecode (strEncode shortLink) `shouldBe` Right shortLink
     connReq' `shouldBe` connReq
@@ -1260,7 +1264,7 @@ testAddContactShortLink viaProxy a b =
       exchangeGreetingsViaProxy viaProxy a bId b aId
     -- update user data
     let updatedData = "updated user data"
-    shortLink' <- runRight $ setContactShortLink a contactId updatedData Nothing
+    shortLink' <- runRight $ setConnShortLink a contactId SCMContact updatedData Nothing
     shortLink' `shouldBe` shortLink
     (connReq4, updatedConnData') <- runRight $ getConnShortLink c 1 shortLink
     connReq4 `shouldBe` connReq
@@ -1291,7 +1295,7 @@ testContactShortLinkRestart ps = withAgentClients2 $ \a b -> do
     connReq' `shouldBe` connReq
     linkUserData connData' `shouldBe` userData
     -- update user data
-    shortLink' <- runRight $ setContactShortLink a contactId updatedData Nothing
+    shortLink' <- runRight $ setConnShortLink a contactId SCMContact updatedData Nothing
     shortLink' `shouldBe` shortLink
   withSmpServer ps $ do
     (connReq4, updatedConnData') <- runRight $ getConnShortLink b 1 shortLink
@@ -1303,7 +1307,7 @@ testAddContactShortLinkRestart ps = withAgentClients2 $ \a b -> do
   let userData = "some user data"
   ((contactId, CCLink connReq0 Nothing), shortLink) <- withSmpServer ps $ runRight $ do
     r@(contactId, _) <- A.createConnection a 1 True SCMContact Nothing Nothing CR.IKPQOn SMOnlyCreate
-    (r,) <$> setContactShortLink a contactId userData Nothing
+    (r,) <$> setConnShortLink a contactId SCMContact userData Nothing
   Right connReq <- pure $ smpDecode (smpEncode connReq0)
   let updatedData = "updated user data"
   withSmpServer ps $ do
@@ -1312,7 +1316,7 @@ testAddContactShortLinkRestart ps = withAgentClients2 $ \a b -> do
     connReq' `shouldBe` connReq
     linkUserData connData' `shouldBe` userData
     -- update user data
-    shortLink' <- runRight $ setContactShortLink a contactId updatedData Nothing
+    shortLink' <- runRight $ setConnShortLink a contactId SCMContact updatedData Nothing
     shortLink' `shouldBe` shortLink
   withSmpServer ps $ do
     (connReq4, updatedConnData') <- runRight $ getConnShortLink b 1 shortLink
@@ -1342,14 +1346,14 @@ testOldContactQueueShortLink ps@(_, msType) = withAgentClients2 $ \a b -> do
 
   withSmpServer ps $ do
     let userData = "some user data"
-    shortLink <- runRight $ setContactShortLink a contactId userData Nothing
+    shortLink <- runRight $ setConnShortLink a contactId SCMContact userData Nothing
     (connReq', connData') <- runRight $ getConnShortLink b 1 shortLink
     strDecode (strEncode shortLink) `shouldBe` Right shortLink
     connReq' `shouldBe` connReq
     linkUserData connData' `shouldBe` userData
     -- update user data
     let updatedData = "updated user data"
-    shortLink' <- runRight $ setContactShortLink a contactId updatedData Nothing
+    shortLink' <- runRight $ setConnShortLink a contactId SCMContact updatedData Nothing
     shortLink' `shouldBe` shortLink
     -- check updated
     (connReq'', updatedConnData') <- runRight $ getConnShortLink b 1 shortLink

tests/ServerTests.hs

@@ -1113,6 +1113,10 @@ testInvQueueLinkData =
       -- can't read link data with LGET
       Resp "2" lnkId' (ERR AUTH) <- sendRecv s ("", "2", lnkId, LGET)
       lnkId' `shouldBe` lnkId
+      -- can update link data before it is secured
+      let newLD = (EncDataBytes "fixed data", EncDataBytes "updated user data")
+      Resp "2a" rId' OK <- signSendRecv r rKey ("2a", rId, LSET lnkId newLD)
+      rId' `shouldBe` rId
 
       (sPub, sKey) <- atomically $ C.generateAuthKeyPair C.SEd25519 g
 
@@ -1128,23 +1132,22 @@ testInvQueueLinkData =
       Resp "5" lnkId2 (LNK sId2 ld') <- signSendRecv s sKey ("5", lnkId, LKEY sPub)
       (lnkId2, lnkId) #== "secures queue and returns link data, same link ID in response"
       (sId2, sId) #== "same sender ID in response"
-      (ld', ld) #== "returns stored data"
+      (ld', newLD) #== "returns updated stored data"
 
       (sPub', sKey') <- atomically $ C.generateAuthKeyPair C.SEd25519 g
       Resp "6" _ err4 <- signSendRecv s sKey' ("6", lnkId, LKEY sPub')
       (err4, ERR AUTH) #== "rejects if secured with different key"
 
       Resp "7" _ (LNK sId3 ld2) <- signSendRecv s sKey ("7", lnkId, LKEY sPub)
       sId3 `shouldBe` sId
-      ld2 `shouldBe` ld
+      ld2 `shouldBe` newLD
 
-      let newLD = (EncDataBytes "fixed data", EncDataBytes "updated user data")
-      Resp "8" rId' (ERR AUTH) <- signSendRecv r rKey ("8", rId, LSET lnkId newLD)
-      rId' `shouldBe` rId
-
-      Resp "9" rId2 (ERR AUTH) <- signSendRecv r rKey ("9", rId, LDEL)
+      Resp "8" rId2 (ERR AUTH) <- signSendRecv r rKey ("8", rId, LSET lnkId newLD)
       rId2 `shouldBe` rId
 
+      Resp "9" rId3 OK <- signSendRecv r rKey ("9", rId, LDEL)
+      rId3 `shouldBe` rId
+
 testContactQueueLinkData :: SpecWith (ASrvTransport, AStoreType)
 testContactQueueLinkData =
   it "create and access queue short link data for contact address" $ \(ATransport t, msType) ->