agent: make createConnection and setConnShortLink apis support setting all link data fields (#1665)

* agent: make createConnection and setConnShortLink apis support setting all link data fields

* add functions

* refactor

* refactor

* fix tests

---------

Co-authored-by: Evgeny Poberezkin <evgeny@poberezkin.com>

src/Simplex/Messaging/Agent.hs

@@ -381,12 +381,12 @@ deleteConnectionsAsync c waitDelivery = withAgentEnv c . deleteConnectionsAsync'
 {-# INLINE deleteConnectionsAsync #-}
 
 -- | Create SMP agent connection (NEW command)
-createConnection :: ConnectionModeI c => AgentClient -> NetworkRequestMode -> UserId -> Bool -> Bool -> SConnectionMode c -> Maybe UserLinkData -> Maybe CRClientData -> CR.InitialKeys -> SubscriptionMode -> AE (ConnId, (CreatedConnLink c, Maybe ClientServiceId))
+createConnection :: ConnectionModeI c => AgentClient -> NetworkRequestMode -> UserId -> Bool -> Bool -> SConnectionMode c -> Maybe (UserConnLinkData c) -> Maybe CRClientData -> CR.InitialKeys -> SubscriptionMode -> AE (ConnId, (CreatedConnLink c, Maybe ClientServiceId))
 createConnection c nm userId enableNtfs checkNotices = withAgentEnv c .::. newConn c nm userId enableNtfs checkNotices
 {-# INLINE createConnection #-}
 
 -- | Create or update user's contact connection short link
-setConnShortLink :: AgentClient -> NetworkRequestMode -> ConnId -> SConnectionMode c -> UserLinkData -> Maybe CRClientData -> AE (ConnShortLink c)
+setConnShortLink :: AgentClient -> NetworkRequestMode -> ConnId -> SConnectionMode c -> UserConnLinkData c -> Maybe CRClientData -> AE (ConnShortLink c)
 setConnShortLink c = withAgentEnv c .::. setConnShortLink' c
 {-# INLINE setConnShortLink #-}
 
@@ -865,12 +865,12 @@ switchConnectionAsync' c corrId connId =
             connectionStats c $ DuplexConnection cData rqs' sqs
       _ -> throwE $ CMD PROHIBITED "switchConnectionAsync: not duplex"
 
-newConn :: ConnectionModeI c => AgentClient -> NetworkRequestMode -> UserId -> Bool -> Bool -> SConnectionMode c -> Maybe UserLinkData -> Maybe CRClientData -> CR.InitialKeys -> SubscriptionMode -> AM (ConnId, (CreatedConnLink c, Maybe ClientServiceId))
-newConn c nm userId enableNtfs checkNotices cMode userData_ clientData pqInitKeys subMode = do
+newConn :: ConnectionModeI c => AgentClient -> NetworkRequestMode -> UserId -> Bool -> Bool -> SConnectionMode c -> Maybe (UserConnLinkData c) -> Maybe CRClientData -> CR.InitialKeys -> SubscriptionMode -> AM (ConnId, (CreatedConnLink c, Maybe ClientServiceId))
+newConn c nm userId enableNtfs checkNotices cMode linkData_ clientData pqInitKeys subMode = do
   srv <- getSMPServer c userId
   when (checkNotices && connMode cMode == CMContact) $ checkClientNotices c srv
   connId <- newConnNoQueues c userId enableNtfs cMode (CR.connPQEncryption pqInitKeys)
-  (connId,) <$> newRcvConnSrv c nm userId connId enableNtfs cMode userData_ clientData pqInitKeys subMode srv
+  (connId,) <$> newRcvConnSrv c nm userId connId enableNtfs cMode linkData_ clientData pqInitKeys subMode srv
     `catchE` \e -> withStore' c (`deleteConnRecord` connId) >> throwE e
 
 checkClientNotices :: AgentClient -> SMPServerWithAuth -> AM ()
@@ -886,19 +886,19 @@ checkClientNotices AgentClient {clientNotices, presetServers} (ProtoServerWithAu
         when (maybe True (ts <) expires_) $
           throwError NOTICE {server = safeDecodeUtf8 $ strEncode $ L.head host, preset = isNothing srvKey, expiresAt = roundedToUTCTime <$> expires_}
 
-setConnShortLink' :: AgentClient -> NetworkRequestMode -> ConnId -> SConnectionMode c -> UserLinkData -> Maybe CRClientData -> AM (ConnShortLink c)
-setConnShortLink' c nm connId cMode userData clientData =
+setConnShortLink' :: AgentClient -> NetworkRequestMode -> ConnId -> SConnectionMode c -> UserConnLinkData c -> Maybe CRClientData -> AM (ConnShortLink c)
+setConnShortLink' c nm connId cMode userLinkData 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
+    (rq, lnkId, sl, d) <- case (conn, cMode, userLinkData) of
+      (ContactConnection _ rq, SCMContact, d@UserContactLinkData {}) -> prepareContactLinkData rq d
+      (RcvConnection _ rq, SCMInvitation, d@UserInvLinkData {}) -> prepareInvLinkData rq d
       _ -> throwE $ CMD PROHIBITED "setConnShortLink: invalid connection or mode"
     addQueueLink c nm rq lnkId d
     pure sl
   where
-    prepareContactLinkData :: RcvQueue -> AM (RcvQueue, SMP.LinkId, ConnShortLink 'CMContact, QueueLinkData)
-    prepareContactLinkData rq@RcvQueue {shortLink} = do
+    prepareContactLinkData :: RcvQueue -> UserConnLinkData 'CMContact -> AM (RcvQueue, SMP.LinkId, ConnShortLink 'CMContact, QueueLinkData)
+    prepareContactLinkData rq@RcvQueue {shortLink} ud = do
       g <- asks random
       AgentConfig {smpClientVRange = vr, smpAgentVRange} <- asks config
       let cslContact = CSLContact SLSServer CCTContact (qServer rq)
@@ -906,25 +906,25 @@ setConnShortLink' c nm 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 SCMContact linkPrivSigKey smpAgentVRange userData
+          d <- liftError id $ SL.encryptUserData g k $ SL.encodeSignUserData SCMContact linkPrivSigKey smpAgentVRange ud
           pure (rq, linkId, cslContact shortLinkKey, (linkEncFixedData, d))
         Nothing -> do
           sigKeys@(_, privSigKey) <- atomically $ C.generateKeyPair @'C.Ed25519 g
           let qUri = SMPQueueUri vr $ (rcvSMPQueueAddress rq) {queueMode = Just QMContact}
               connReq = CRContactUri $ ConnReqUriData SSSimplex smpAgentVRange [qUri] clientData
-              (linkKey, linkData) = SL.encodeSignLinkData sigKeys smpAgentVRange connReq userData
+              (linkKey, linkData) = SL.encodeSignLinkData sigKeys smpAgentVRange connReq ud
               (linkId, k) = SL.contactShortLinkKdf linkKey
           srvData <- liftError id $ SL.encryptLinkData g k linkData
           let slCreds = ShortLinkCreds linkId linkKey privSigKey (fst srvData)
           withStore' c $ \db -> updateShortLinkCreds db rq slCreds
           pure (rq, linkId, cslContact linkKey, srvData)
-    prepareInvLinkData :: RcvQueue -> AM (RcvQueue, SMP.LinkId, ConnShortLink 'CMInvitation, QueueLinkData)
-    prepareInvLinkData rq@RcvQueue {shortLink} = case shortLink of
+    prepareInvLinkData :: RcvQueue -> UserConnLinkData 'CMInvitation -> AM (RcvQueue, SMP.LinkId, ConnShortLink 'CMInvitation, QueueLinkData)
+    prepareInvLinkData rq@RcvQueue {shortLink} ud = 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 SCMInvitation linkPrivSigKey smpAgentVRange userData
+        d <- liftError id $ SL.encryptUserData g k $ SL.encodeSignUserData SCMInvitation linkPrivSigKey smpAgentVRange ud
         let sl = CSLInvitation SLSServer (qServer rq) shortLinkId shortLinkKey
         pure (rq, shortLinkId, sl, (linkEncFixedData, d))
       Nothing -> throwE $ CMD PROHIBITED "setConnShortLink: no ShortLinkCreds in invitation"
@@ -989,13 +989,13 @@ changeConnectionUser' c oldUserId connId newUserId = do
   where
     updateConn = withStore' c $ \db -> setConnUserId db oldUserId connId newUserId
 
-newRcvConnSrv :: forall c. ConnectionModeI c => AgentClient -> NetworkRequestMode -> UserId -> ConnId -> Bool -> SConnectionMode c -> Maybe UserLinkData -> Maybe CRClientData -> CR.InitialKeys -> SubscriptionMode -> SMPServerWithAuth -> AM (CreatedConnLink c, Maybe ClientServiceId)
-newRcvConnSrv c nm userId connId enableNtfs cMode userData_ clientData pqInitKeys subMode srvWithAuth@(ProtoServerWithAuth srv _) = do
+newRcvConnSrv :: forall c. ConnectionModeI c => AgentClient -> NetworkRequestMode -> UserId -> ConnId -> Bool -> SConnectionMode c -> Maybe (UserConnLinkData c) -> Maybe CRClientData -> CR.InitialKeys -> SubscriptionMode -> SMPServerWithAuth -> AM (CreatedConnLink c, Maybe ClientServiceId)
+newRcvConnSrv c nm userId connId enableNtfs cMode userLinkData_ clientData pqInitKeys subMode srvWithAuth@(ProtoServerWithAuth srv _) = do
   case (cMode, pqInitKeys) of
     (SCMContact, CR.IKUsePQ) -> throwE $ CMD PROHIBITED "newRcvConnSrv"
     _ -> pure ()
   e2eKeys <- atomically . C.generateKeyPair =<< asks random
-  case userData_ of
+  case userLinkData_ of
     Just d -> do
       (nonce, qUri, cReq, qd) <- prepareLinkData d $ fst e2eKeys
       (rq, qUri') <- createRcvQueue (Just nonce) qd e2eKeys
@@ -1025,12 +1025,12 @@ newRcvConnSrv c nm userId connId enableNtfs cMode userData_ clientData pqInitKey
         SCMContact -> pure $ CRContactUri crData
         SCMInvitation -> do
           g <- asks random
-          let pqEnc = CR.initialPQEncryption (isJust userData_) pqInitKeys
+          let pqEnc = CR.initialPQEncryption (isJust userLinkData_) pqInitKeys
           (pk1, pk2, pKem, e2eRcvParams) <- liftIO $ CR.generateRcvE2EParams g (maxVersion e2eEncryptVRange) pqEnc
           withStore' c $ \db -> createRatchetX3dhKeys db connId pk1 pk2 pKem
           pure $ CRInvitationUri crData $ toVersionRangeT e2eRcvParams e2eEncryptVRange
-    prepareLinkData :: UserLinkData -> C.PublicKeyX25519 -> AM (C.CbNonce, SMPQueueUri, ConnectionRequestUri c, ClntQueueReqData)
-    prepareLinkData userData e2eDhKey = do
+    prepareLinkData :: UserConnLinkData c -> C.PublicKeyX25519 -> AM (C.CbNonce, SMPQueueUri, ConnectionRequestUri c, ClntQueueReqData)
+    prepareLinkData userLinkData e2eDhKey = do
       g <- asks random
       nonce@(C.CbNonce corrId) <- atomically $ C.randomCbNonce g
       sigKeys@(_, privSigKey) <- atomically $ C.generateKeyPair @'C.Ed25519 g
@@ -1040,7 +1040,7 @@ newRcvConnSrv c nm userId connId enableNtfs cMode userData_ clientData pqInitKey
           qm = case cMode of SCMContact -> QMContact; SCMInvitation -> QMMessaging
           qUri = SMPQueueUri vr $ SMPQueueAddress srv sndId e2eDhKey (Just qm)
       connReq <- createConnReq qUri
-      let (linkKey, linkData) = SL.encodeSignLinkData sigKeys smpAgentVRange connReq userData
+      let (linkKey, linkData) = SL.encodeSignLinkData sigKeys smpAgentVRange connReq userLinkData
       qd <- case cMode of
         SCMContact -> do
           let (linkId, k) = SL.contactShortLinkKdf linkKey

src/Simplex/Messaging/Agent/Protocol.hs

@@ -112,6 +112,8 @@ module Simplex.Messaging.Agent.Protocol
     ServiceScheme,
     FixedLinkData (..),
     ConnLinkData (..),
+    UserConnLinkData (..),
+    UserContactData (..),
     UserLinkData (..),
     OwnerAuth (..),
     OwnerId,
@@ -1690,25 +1692,30 @@ data FixedLinkData c = FixedLinkData
 
 data ConnLinkData c where
   InvitationLinkData :: VersionRangeSMPA -> UserLinkData -> ConnLinkData 'CMInvitation
-  ContactLinkData ::
-    { agentVRange :: VersionRangeSMPA,
-      -- direct connection via connReq in fixed data is allowed.
-      direct :: Bool,
-      -- additional owner keys to sign changes of mutable data.
-      owners :: [OwnerAuth],
-      -- alternative addresses of chat relays that receive requests for this contact address.
-      relays :: [ConnShortLink 'CMContact],
-      userData :: UserLinkData
-    } -> ConnLinkData 'CMContact
+  ContactLinkData :: VersionRangeSMPA -> UserContactData -> ConnLinkData 'CMContact
+
+data UserContactData = UserContactData
+  { -- direct connection via connReq in fixed data is allowed.
+    direct :: Bool,
+    -- additional owner keys to sign changes of mutable data.
+    owners :: [OwnerAuth],
+    -- alternative addresses of chat relays that receive requests for this contact address.
+    relays :: [ConnShortLink 'CMContact],
+    userData :: UserLinkData
+  }
 
 newtype UserLinkData = UserLinkData ByteString
 
 data AConnLinkData = forall m. ConnectionModeI m => ACLD (SConnectionMode m) (ConnLinkData m)
 
+data UserConnLinkData c where
+  UserInvLinkData :: UserLinkData -> UserConnLinkData 'CMInvitation
+  UserContactLinkData :: UserContactData -> UserConnLinkData 'CMContact
+
 linkUserData :: ConnLinkData c -> UserLinkData
 linkUserData = \case
   InvitationLinkData _ d -> d
-  ContactLinkData {userData} -> userData
+  ContactLinkData _ UserContactData {userData} -> userData
 {-# INLINE linkUserData #-}
 
 linkUserData' :: ConnLinkData c -> ByteString
@@ -1749,8 +1756,8 @@ instance ConnectionModeI c => Encoding (FixedLinkData c) where
 instance ConnectionModeI c => Encoding (ConnLinkData c) where
   smpEncode = \case
     InvitationLinkData vr userData -> smpEncode (CMInvitation, vr, userData)
-    ContactLinkData {agentVRange, direct, owners, relays, userData} ->
-      B.concat [smpEncode (CMContact, agentVRange, direct), smpEncodeList owners, smpEncodeList relays, smpEncode userData]
+    ContactLinkData vr UserContactData {direct, owners, relays, userData} ->
+      B.concat [smpEncode (CMContact, vr, direct), smpEncodeList owners, smpEncodeList relays, smpEncode userData]
   smpP = (\(ACLD _ d) -> checkConnMode d) <$?> smpP
   {-# INLINE smpP #-}
 
@@ -1763,11 +1770,12 @@ instance Encoding AConnLinkData where
         (vr, userData) <- smpP <* A.takeByteString -- ignoring tail for forward compatibility with the future link data encoding
         pure $ ACLD SCMInvitation $ InvitationLinkData vr userData
       CMContact -> do
-        (agentVRange, direct) <- smpP
+        (vr, direct) <- smpP
         owners <- smpListP
         relays <- smpListP
         userData <- smpP <* A.takeByteString -- ignoring tail for forward compatibility with the future link data encoding
-        pure $ ACLD SCMContact ContactLinkData {agentVRange, direct, owners, relays, userData}
+        let cd = UserContactData {direct, owners, relays, userData}
+        pure $ ACLD SCMContact $ ContactLinkData vr cd
 
 instance Encoding UserLinkData where
   smpEncode (UserLinkData s) = if B.length s <= 254 then smpEncode s else smpEncode ('\255', Large s)

src/Simplex/Messaging/Crypto/ShortLink.hs

@@ -1,6 +1,7 @@
 {-# LANGUAGE DataKinds #-}
 {-# LANGUAGE DuplicateRecordFields #-}
 {-# LANGUAGE GADTs #-}
+{-# LANGUAGE LambdaCase #-}
 {-# LANGUAGE MultiWayIf #-}
 {-# LANGUAGE NamedFieldPuns #-}
 {-# LANGUAGE OverloadedStrings #-}
@@ -48,20 +49,20 @@ contactShortLinkKdf (LinkKey k) =
 invShortLinkKdf :: LinkKey -> C.SbKey
 invShortLinkKdf (LinkKey k) = C.unsafeSbKey $ C.hkdf "" k "SimpleXInvLink" 32
 
-encodeSignLinkData :: forall c. ConnectionModeI c => C.KeyPairEd25519 -> VersionRangeSMPA -> ConnectionRequestUri c -> UserLinkData -> (LinkKey, (ByteString, ByteString))
+encodeSignLinkData :: ConnectionModeI c => C.KeyPairEd25519 -> VersionRangeSMPA -> ConnectionRequestUri c -> UserConnLinkData c -> (LinkKey, (ByteString, ByteString))
 encodeSignLinkData (rootKey, pk) agentVRange connReq userData =
   let fd = smpEncode FixedLinkData {agentVRange, rootKey, connReq}
-      md = smpEncode $ connLinkData @c agentVRange userData
+      md = smpEncode $ connLinkData agentVRange userData
    in (LinkKey (C.sha3_256 fd), (encodeSign pk fd, encodeSign pk md))
 
-encodeSignUserData :: forall c. ConnectionModeI c => SConnectionMode c -> C.PrivateKeyEd25519 -> VersionRangeSMPA -> UserLinkData -> ByteString
-encodeSignUserData _ pk agentVRange userData =
-  encodeSign pk $ smpEncode $ connLinkData @c agentVRange userData
+encodeSignUserData :: ConnectionModeI c => SConnectionMode c -> C.PrivateKeyEd25519 -> VersionRangeSMPA -> UserConnLinkData c -> ByteString
+encodeSignUserData _ pk agentVRange userLinkData =
+  encodeSign pk $ smpEncode $ connLinkData agentVRange userLinkData
 
-connLinkData :: forall c. ConnectionModeI c => VersionRangeSMPA -> UserLinkData -> ConnLinkData c
-connLinkData agentVRange userData = case sConnectionMode @c of
-  SCMInvitation -> InvitationLinkData agentVRange userData
-  SCMContact -> ContactLinkData {agentVRange, direct = True, owners = [], relays = [], userData}
+connLinkData :: VersionRangeSMPA -> UserConnLinkData c -> ConnLinkData c
+connLinkData vr = \case
+  UserInvLinkData d -> InvitationLinkData vr d
+  UserContactLinkData d -> ContactLinkData vr d
 
 encodeSign :: C.PrivateKeyEd25519 -> ByteString -> ByteString
 encodeSign pk s = smpEncode (C.sign' pk s) <> s