From 18e73b8aa7996562f0f92db3967da02d960cb805 Mon Sep 17 00:00:00 2001
From: spaced4ndy <8711996+spaced4ndy@users.noreply.github.com>
Date: Fri, 23 May 2025 17:21:36 +0000
Subject: [PATCH] agent: pass CRClientData to setContactShortLink (#1546)

* agent: pass CRClientData to setContactShortLink

* fix

* fix
---
 src/Simplex/Messaging/Agent.hs         | 10 +++++-----
 src/Simplex/Messaging/TMap.hs          |  5 +++++
 tests/AgentTests/FunctionalAPITests.hs | 14 +++++++-------
 3 files changed, 17 insertions(+), 12 deletions(-)

diff --git a/src/Simplex/Messaging/Agent.hs b/src/Simplex/Messaging/Agent.hs
index 799fed250..792e1f607 100644
--- a/src/Simplex/Messaging/Agent.hs
+++ b/src/Simplex/Messaging/Agent.hs
@@ -371,8 +371,8 @@ 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 -> AE (ConnShortLink 'CMContact)
-setContactShortLink c = withAgentEnv c .: setContactShortLink' c
+setContactShortLink :: AgentClient -> ConnId -> ConnInfo -> Maybe CRClientData -> AE (ConnShortLink 'CMContact)
+setContactShortLink c = withAgentEnv c .:. setContactShortLink' c
 {-# INLINE setContactShortLink #-}
 
 deleteContactShortLink :: AgentClient -> ConnId -> AE ()
@@ -832,8 +832,8 @@ 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 -> AM (ConnShortLink 'CMContact)
-setContactShortLink' c connId userData = 
+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
@@ -855,7 +855,7 @@ setContactShortLink' c connId userData =
         Nothing -> do
           sigKeys@(_, privSigKey) <- atomically $ C.generateKeyPair @'C.Ed25519 g
           let qUri = SMPQueueUri vr $ SMPQueueAddress server sndId (C.publicKey e2ePrivKey) (Just QMContact)
-              connReq = CRContactUri $ ConnReqUriData SSSimplex smpAgentVRange [qUri] Nothing
+              connReq = CRContactUri $ ConnReqUriData SSSimplex smpAgentVRange [qUri] clientData
               (linkKey, linkData) = SL.encodeSignLinkData sigKeys smpAgentVRange connReq userData
               (linkId, k) = SL.contactShortLinkKdf linkKey
           srvData <- liftError id $ SL.encryptLinkData g k linkData
diff --git a/src/Simplex/Messaging/TMap.hs b/src/Simplex/Messaging/TMap.hs
index b743ce7bc..b09158d5f 100644
--- a/src/Simplex/Messaging/TMap.hs
+++ b/src/Simplex/Messaging/TMap.hs
@@ -13,6 +13,7 @@ module Simplex.Messaging.TMap
     insert,
     insertM,
     delete,
+    lookupInsert,
     lookupDelete,
     adjust,
     update,
@@ -72,6 +73,10 @@ delete :: Ord k => k -> TMap k a -> STM ()
 delete k m = modifyTVar' m $ M.delete k
 {-# INLINE delete #-}
 
+lookupInsert :: Ord k => k -> a -> TMap k a -> STM (Maybe a)
+lookupInsert k v m = stateTVar m $ M.alterF (,Just v) k
+{-# INLINE lookupInsert #-}
+
 lookupDelete :: Ord k => k -> TMap k a -> STM (Maybe a)
 lookupDelete k m = stateTVar m $ M.alterF (,Nothing) k
 {-# INLINE lookupDelete #-}
diff --git a/tests/AgentTests/FunctionalAPITests.hs b/tests/AgentTests/FunctionalAPITests.hs
index 0f2774862..12a289a0c 100644
--- a/tests/AgentTests/FunctionalAPITests.hs
+++ b/tests/AgentTests/FunctionalAPITests.hs
@@ -1200,13 +1200,13 @@ testContactShortLink viaProxy a b =
       exchangeGreetingsViaProxy viaProxy a bId b aId
     -- update user data
     let updatedData = "updated user data"
-    shortLink' <- runRight $ setContactShortLink a contactId updatedData
+    shortLink' <- runRight $ setContactShortLink a contactId 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
+    shortLink2 <- runRight $ setContactShortLink a contactId updatedData Nothing
     shortLink2 `shouldBe` shortLink
     -- delete short link
     runRight_ $ deleteContactShortLink a contactId
@@ -1219,7 +1219,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
+    shortLink <- runRight $ setContactShortLink a contactId userData Nothing
     (connReq', connData') <- runRight $ getConnShortLink b 1 shortLink
     strDecode (strEncode shortLink) `shouldBe` Right shortLink
     connReq' `shouldBe` connReq
@@ -1247,7 +1247,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
+    shortLink' <- runRight $ setContactShortLink a contactId updatedData Nothing
     shortLink' `shouldBe` shortLink
     (connReq4, updatedConnData') <- runRight $ getConnShortLink c 1 shortLink
     connReq4 `shouldBe` connReq
@@ -1278,7 +1278,7 @@ testContactShortLinkRestart ps = withAgentClients2 $ \a b -> do
     connReq' `shouldBe` connReq
     linkUserData connData' `shouldBe` userData
     -- update user data
-    shortLink' <- runRight $ setContactShortLink a contactId updatedData
+    shortLink' <- runRight $ setContactShortLink a contactId updatedData Nothing
     shortLink' `shouldBe` shortLink
   withSmpServer ps $ do
     (connReq4, updatedConnData') <- runRight $ getConnShortLink b 1 shortLink
@@ -1290,7 +1290,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
+    (r,) <$> setContactShortLink a contactId userData Nothing
   Right connReq <- pure $ smpDecode (smpEncode connReq0)
   let updatedData = "updated user data"
   withSmpServer ps $ do
@@ -1299,7 +1299,7 @@ testAddContactShortLinkRestart ps = withAgentClients2 $ \a b -> do
     connReq' `shouldBe` connReq
     linkUserData connData' `shouldBe` userData
     -- update user data
-    shortLink' <- runRight $ setContactShortLink a contactId updatedData
+    shortLink' <- runRight $ setContactShortLink a contactId updatedData Nothing
     shortLink' `shouldBe` shortLink
   withSmpServer ps $ do
     (connReq4, updatedConnData') <- runRight $ getConnShortLink b 1 shortLink