From 656f2906602a20b541d11e7d1fffc34fa7950c8d Mon Sep 17 00:00:00 2001
From: Evgeny <evgeny@poberezkin.com>
Date: Fri, 13 Jun 2025 21:27:46 +0100
Subject: [PATCH] agent: fix updating short invitation link data (#1566)

---
 src/Simplex/Messaging/Agent.hs            | 4 ++--
 src/Simplex/Messaging/Crypto/ShortLink.hs | 6 +++---
 tests/AgentTests/ShortLinkTests.hs        | 6 +++---
 3 files changed, 8 insertions(+), 8 deletions(-)

diff --git a/src/Simplex/Messaging/Agent.hs b/src/Simplex/Messaging/Agent.hs
index 53e24c3a8..d178134c1 100644
--- a/src/Simplex/Messaging/Agent.hs
+++ b/src/Simplex/Messaging/Agent.hs
@@ -859,7 +859,7 @@ setConnShortLink' c connId cMode userData clientData =
         Just ShortLinkCreds {shortLinkId, shortLinkKey, linkPrivSigKey, linkEncFixedData} -> do
           let (linkId, k) = SL.contactShortLinkKdf shortLinkKey
           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
+          d <- liftError id $ SL.encryptUserData g k $ SL.encodeSignUserData SCMContact linkPrivSigKey smpAgentVRange userData
           pure (rq, linkId, cslContact shortLinkKey, (linkEncFixedData, d))
         Nothing -> do
           sigKeys@(_, privSigKey) <- atomically $ C.generateKeyPair @'C.Ed25519 g
@@ -877,7 +877,7 @@ setConnShortLink' c connId cMode userData clientData =
         g <- asks random
         AgentConfig {smpAgentVRange} <- asks config
         let k = SL.invShortLinkKdf shortLinkKey
-        d <- liftError id $ SL.encryptUserData g k $ SL.encodeSignUserData linkPrivSigKey smpAgentVRange userData
+        d <- liftError id $ SL.encryptUserData g k $ SL.encodeSignUserData SCMInvitation 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"
diff --git a/src/Simplex/Messaging/Crypto/ShortLink.hs b/src/Simplex/Messaging/Crypto/ShortLink.hs
index 42815f9ac..c32c05bcf 100644
--- a/src/Simplex/Messaging/Crypto/ShortLink.hs
+++ b/src/Simplex/Messaging/Crypto/ShortLink.hs
@@ -54,9 +54,9 @@ encodeSignLinkData (rootKey, pk) agentVRange connReq userData =
       md = smpEncode $ connLinkData @c agentVRange userData
    in (LinkKey (C.sha3_256 fd), (encodeSign pk fd, encodeSign pk md))
 
-encodeSignUserData :: C.PrivateKeyEd25519 -> VersionRangeSMPA -> ConnInfo -> ByteString
-encodeSignUserData pk agentVRange userData =
-  encodeSign pk $ smpEncode $ connLinkData @'CMContact agentVRange userData
+encodeSignUserData :: forall c. ConnectionModeI c => SConnectionMode c -> C.PrivateKeyEd25519 -> VersionRangeSMPA -> ConnInfo -> ByteString
+encodeSignUserData _ pk agentVRange userData =
+  encodeSign pk $ smpEncode $ connLinkData @c agentVRange userData
 
 connLinkData :: forall c. ConnectionModeI c => VersionRangeSMPA -> ConnInfo -> ConnLinkData c
 connLinkData agentVRange userData = case sConnectionMode @c of
diff --git a/tests/AgentTests/ShortLinkTests.hs b/tests/AgentTests/ShortLinkTests.hs
index d3aeca0cc..216b7c658 100644
--- a/tests/AgentTests/ShortLinkTests.hs
+++ b/tests/AgentTests/ShortLinkTests.hs
@@ -8,7 +8,7 @@ import AgentTests.ConnectionRequestTests (contactConnRequest, invConnRequest)
 import AgentTests.EqInstances ()
 import Control.Concurrent.STM
 import Control.Monad.Except
-import Simplex.Messaging.Agent.Protocol (AgentErrorType (..), ConnectionMode (..), LinkKey (..), SMPAgentError (..), linkUserData, supportedSMPAgentVRange)
+import Simplex.Messaging.Agent.Protocol (AgentErrorType (..), ConnectionMode (..), LinkKey (..), SConnectionMode (..), SMPAgentError (..), linkUserData, supportedSMPAgentVRange)
 import qualified Simplex.Messaging.Crypto as C
 import qualified Simplex.Messaging.Crypto.ShortLink as SL
 import Test.Hspec hiding (fit, it)
@@ -79,7 +79,7 @@ testUpdateContactShortLink = do
   Right (fd, _ud) <- runExceptT $ SL.encryptLinkData g k linkData
   -- encrypt updated user data
   let updatedUserData = "updated user data"
-      signed = SL.encodeSignUserData (snd sigKeys) supportedSMPAgentVRange updatedUserData
+      signed = SL.encodeSignUserData SCMContact (snd sigKeys) supportedSMPAgentVRange updatedUserData
   Right ud' <- runExceptT $ SL.encryptUserData g k signed
   -- decrypt
   Right (connReq, connData') <- pure $ SL.decryptLinkData linkKey k (fd, ud')
@@ -114,7 +114,7 @@ testContactShortLinkBadSignature = do
   let updatedUserData = "updated user data"
   -- another signature key
   (_, pk) <- atomically $ C.generateKeyPair @'C.Ed25519 g
-  let signed = SL.encodeSignUserData pk supportedSMPAgentVRange updatedUserData
+  let signed = SL.encodeSignUserData SCMContact pk supportedSMPAgentVRange updatedUserData
   Right ud' <- runExceptT $ SL.encryptUserData g k signed
   -- decryption fails
   SL.decryptLinkData @'CMContact linkKey k (fd, ud')