parameterize type for link data by connection type

epoberezkin · epoberezkin · commit e7e28601f4db · 2025-04-07T18:03:49.000+01:00
diff --git a/src/Simplex/Messaging/Agent.hs b/src/Simplex/Messaging/Agent.hs
@@ -379,7 +379,7 @@ deleteContactShortLink c = withAgentEnv c . deleteContactShortLink' c
 {-# INLINE deleteContactShortLink #-}
 
 -- | Get and verify data from short link. For 1-time invitations it preserves the key to allow retries
-getConnShortLink :: AgentClient -> UserId -> ConnShortLink c -> AE (Maybe (ConnectionRequestUri c), ConnInfo)
+getConnShortLink :: AgentClient -> UserId -> ConnShortLink c -> AE (ConnectionRequestUri c, ConnInfo)
 getConnShortLink c = withAgentEnv c .: getConnShortLink' c
 {-# INLINE getConnShortLink #-}
 
@@ -838,7 +838,7 @@ setContactShortLink' c connId userData =
       SomeConn _ (ContactConnection _ rq) -> do
         (lnkId, linkKey, d) <- prepareLinkData rq
         addQueueLink c rq lnkId d
-        pure $ CSLContact (qServer rq) CCTContact linkKey
+        pure $ CSLContact CCTContact (qServer rq) linkKey
       _ -> throwE $ CMD PROHIBITED "setContactShortLink: not contact address"
   where
     prepareLinkData :: RcvQueue -> AM (SMP.LinkId, LinkKey, QueueLinkData)
@@ -870,7 +870,7 @@ deleteContactShortLink' c connId =
       _ -> throwE $ CMD PROHIBITED "deleteContactShortLink: 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 (Maybe (ConnectionRequestUri c), ConnInfo)
+getConnShortLink' :: forall c. AgentClient -> UserId -> ConnShortLink c -> AM (ConnectionRequestUri c, ConnInfo)
 getConnShortLink' c userId = \case
   CSLInvitation srv linkId linkKey -> do
     g <- asks random
@@ -886,17 +886,16 @@ getConnShortLink' c userId = \case
     ld@(sndId, _) <- secureGetQueueLink c userId invLink
     withStore' c $ \db -> setInvShortLinkSndId db invLink sndId
     decryptData srv linkKey k ld
-  CSLContact srv _ linkKey -> do
+  CSLContact _ srv linkKey -> do
     let (linkId, k) = SL.contactShortLinkKdf linkKey
     ld <- getQueueLink c userId srv linkId
     decryptData srv linkKey k ld
   where
-    decryptData :: ConnectionModeI c => SMPServer -> LinkKey -> C.SbKey -> (SMP.SenderId, QueueLinkData) -> AM (Maybe (ConnectionRequestUri c), ConnInfo)
+    decryptData :: ConnectionModeI c => SMPServer -> LinkKey -> C.SbKey -> (SMP.SenderId, QueueLinkData) -> AM (ConnectionRequestUri c, ConnInfo)
     decryptData srv linkKey k (sndId, d) = do
-      r@(cReq_, _) <- liftEither $ SL.decryptLinkData @c linkKey k d
-      forM_ cReq_ $ \cReq ->
-        unless ((srv, sndId) `sameQAddress` qAddress (connReqQueue cReq)) $
-          throwE $ AGENT $ A_LINK "different address"
+      r@(cReq, _) <- liftEither $ SL.decryptLinkData @c linkKey k d
+      unless ((srv, sndId) `sameQAddress` qAddress (connReqQueue cReq)) $
+        throwE $ AGENT $ A_LINK "different address"
       pure r
 
 deleteLocalInvShortLink' :: AgentClient -> ConnShortLink 'CMInvitation -> AM ()
@@ -979,7 +978,7 @@ newRcvConnSrv c userId connId enableNtfs cMode userData_ clientData pqInitKeys s
       Just ShortLinkCreds {shortLinkId, shortLinkKey}
         | qUri == qUri'  ->
             let link = case cReq of
-                  CRContactUri _ -> CSLContact srv CCTContact shortLinkKey
+                  CRContactUri _ -> CSLContact CCTContact srv shortLinkKey
                   CRInvitationUri {} -> CSLInvitation srv shortLinkId shortLinkKey
              in pure $ CCLink cReq (Just link)
         | otherwise -> throwE $ INTERNAL "different rcv queue address"
diff --git a/src/Simplex/Messaging/Agent/Protocol.hs b/src/Simplex/Messaging/Agent/Protocol.hs
@@ -109,7 +109,7 @@ module Simplex.Messaging.Agent.Protocol
     CRClientData,
     ServiceScheme,
     FixedLinkData (..),
-    UserLinkData (..),
+    ConnLinkData (..),
     OwnerAuth (..),
     OwnerId,
     ConnectionLink (..),
@@ -156,6 +156,7 @@ module Simplex.Messaging.Agent.Protocol
     updateSMPServerHosts,
     shortenShortLink,
     restoreShortLink,
+    linkUserData,
   )
 where
 
@@ -167,7 +168,7 @@ import qualified Data.Attoparsec.ByteString.Char8 as A
 import qualified Data.ByteString.Base64.URL as B64
 import Data.ByteString.Char8 (ByteString)
 import qualified Data.ByteString.Char8 as B
-import Data.Char (isDigit)
+import Data.Char (isDigit, toLower, toUpper)
 import Data.Foldable (find)
 import Data.Functor (($>))
 import Data.Int (Int64)
@@ -1349,7 +1350,7 @@ deriving instance Show AConnectionRequestUri
 
 data ConnShortLink (m :: ConnectionMode) where
   CSLInvitation :: SMPServer -> SMP.LinkId -> LinkKey -> ConnShortLink 'CMInvitation
-  CSLContact :: SMPServer -> ContactConnType -> LinkKey -> ConnShortLink 'CMContact
+  CSLContact :: ContactConnType -> SMPServer -> LinkKey -> ConnShortLink 'CMContact
 
 deriving instance Eq (ConnShortLink m)
 
@@ -1369,7 +1370,7 @@ instance ConnectionModeI c => ToField (ConnShortLink c) where toField = toField
 
 instance (Typeable c, ConnectionModeI c) => FromField (ConnShortLink c) where fromField = blobFieldDecoder strDecode
 
-data ContactConnType = CCTContact | CCTGroup deriving (Eq, Show)
+data ContactConnType = CCTContact | CCTChannel | CCTGroup deriving (Eq, Show)
 
 data AConnShortLink = forall m. ConnectionModeI m => ACSL (SConnectionMode m) (ConnShortLink m)
 
@@ -1406,50 +1407,84 @@ instance ConnectionModeI m => FromJSON (ConnectionLink m) where
 
 instance ConnectionModeI m => StrEncoding (ConnShortLink m) where
   strEncode = \case
-    CSLInvitation srv (SMP.EntityId lnkId) (LinkKey k) -> encLink srv (lnkId <> k) "i"
-    CSLContact srv ct (LinkKey k) -> encLink srv k $ case ct of CCTContact -> "c"; CCTGroup -> "g"
+    CSLInvitation srv (SMP.EntityId lnkId) (LinkKey k) -> encLink srv (lnkId <> k) 'i'
+    CSLContact ct srv (LinkKey k) -> encLink srv k $ toLower $ ctTypeChar ct
     where
       encLink (SMPServer (h :| hs) port (C.KeyHash kh)) linkUri linkType =
-        "https://" <> strEncode h <> port' <> "/" <> linkType <> "#" <> khStr <> hosts <> B64.encodeUnpadded linkUri
+        B.concat ["https://", strEncode h, port', "/", B.singleton linkType, "#", khStr, hosts, B64.encodeUnpadded linkUri]
         where
           port' = if null port then "" else B.pack (':' : port)
           hosts = if null hs then "" else strEncode (TransportHosts_ hs) <> "/"
           khStr = if B.null kh then "" else B64.encodeUnpadded kh <> "@"
-  strP = do
-    ACSL m l <- strP
-    case testEquality m $ sConnectionMode @m of
-      Just Refl -> pure l
-      _ -> fail "bad short link mode"
+  strP = (\(ACSL _ l) -> checkConnMode l) <$?> strP
+  {-# INLINE strP #-}
 
 instance StrEncoding AConnShortLink where
   strEncode (ACSL _ l) = strEncode l
+  {-# INLINE strEncode #-}
   strP = do
     h <- "https://" *> strP
     port <- A.char ':' *> (B.unpack <$> A.takeWhile1 isDigit) <|> pure ""
-    linkType <- A.char '/' *> A.anyChar
+    contactType <- A.char '/' *> contactTypeP
     keyHash <- optional (A.char '/') *> A.char '#' *> (strP <* A.char '@' <|> pure (C.KeyHash ""))
     TransportHosts_ hs <- strP <* "/" <|> pure (TransportHosts_ [])
     linkUri <- strP
     let srv = SMPServer (h :| hs) port keyHash
-    case linkType of
-      'i'
+    case contactType of
+      Nothing
         | B.length linkUri == 56 ->
             let (lnkId, k) = B.splitAt 24 linkUri
              in pure $ ACSL SCMInvitation $ CSLInvitation srv (SMP.EntityId lnkId) (LinkKey k)
         | otherwise -> fail "bad ConnShortLink: incorrect linkID and key length"
-      'c' -> contactP srv CCTContact linkUri
-      'g' -> contactP srv CCTGroup linkUri
-      _ -> fail "bad ConnShortLink: unknown link type"
+      Just ct
+        | B.length linkUri == 32 -> pure $ ACSL SCMContact $ CSLContact ct srv (LinkKey linkUri)
+        | otherwise -> fail "bad ConnShortLink: incorrect key length"
     where
-      contactP srv ct k
-        | B.length k == 32 = pure $ ACSL SCMContact $ CSLContact srv ct (LinkKey k)
-        | otherwise = fail "bad ConnShortLink: incorrect key length"
+      contactTypeP = do
+        Just <$> (A.anyChar >>= ctTypeP . toUpper)        
+          <|> A.char 'i' $> Nothing
+          <|> fail "unknown link type"
+
+instance ConnectionModeI m => Encoding (ConnShortLink m) where
+  smpEncode = \case
+    CSLInvitation srv lnkId (LinkKey k) -> smpEncode (CMInvitation, srv, lnkId, k)
+    CSLContact ct srv (LinkKey k) -> smpEncode (CMContact, ctTypeChar ct, srv, k)
+  smpP = (\(ACSL _ l) -> checkConnMode l) <$?> smpP
+  {-# INLINE smpP #-}
+
+instance Encoding AConnShortLink where
+  smpEncode (ACSL _ l) = smpEncode l
+  {-# INLINE smpEncode #-}
+  smpP =
+    smpP >>= \case
+      CMInvitation -> do
+        (srv, lnkId, k) <- smpP
+        pure $ ACSL SCMInvitation $ CSLInvitation srv lnkId (LinkKey k)
+      CMContact -> do
+        ct <- ctTypeP =<< A.anyChar
+        (srv, k) <- smpP
+        pure $ ACSL SCMContact $ CSLContact ct srv (LinkKey k)
+
+ctTypeP :: Char -> Parser ContactConnType
+ctTypeP = \case
+  'A' -> pure CCTContact
+  'C' -> pure CCTChannel
+  'G' -> pure CCTGroup 
+  _ -> fail "unknown contact address type"
+{-# INLINE ctTypeP #-}
+
+ctTypeChar :: ContactConnType -> Char
+ctTypeChar = \case
+  CCTContact -> 'A'
+  CCTChannel -> 'C'
+  CCTGroup -> 'G'
+{-# INLINE ctTypeChar #-}
 
 -- the servers passed to this function should be all preset servers, not servers configured by the user.
 shortenShortLink :: NonEmpty SMPServer -> ConnShortLink m -> ConnShortLink m
 shortenShortLink presetSrvs = \case
   CSLInvitation srv lnkId linkKey -> CSLInvitation (shortServer srv) lnkId linkKey
-  CSLContact srv ct linkKey -> CSLContact (shortServer srv) ct linkKey
+  CSLContact ct srv linkKey -> CSLContact ct (shortServer srv) linkKey
   where
     shortServer srv@(SMPServer hs@(h :| _) p kh) =
       if isPresetServer then SMPServer [h] "" (C.KeyHash "") else srv
@@ -1465,7 +1500,7 @@ shortenShortLink presetSrvs = \case
 restoreShortLink :: NonEmpty SMPServer -> ConnShortLink m -> ConnShortLink m
 restoreShortLink presetSrvs = \case
   CSLInvitation srv lnkId linkKey -> CSLInvitation (fullServer srv) lnkId linkKey
-  CSLContact srv ct linkKey -> CSLContact (fullServer srv) ct linkKey
+  CSLContact ct srv linkKey -> CSLContact ct (fullServer srv) linkKey
   where
     fullServer = \case
       s@(SMPServer [_] "" (C.KeyHash "")) -> fromMaybe s $ findPresetServer s presetSrvs
@@ -1500,14 +1535,28 @@ type CRClientData = Text
 data FixedLinkData c = FixedLinkData
   { agentVRange :: VersionRangeSMPA,
     rootKey :: C.PublicKeyEd25519,
-    connReq :: Maybe (ConnectionRequestUri c)
+    connReq :: ConnectionRequestUri c
   }
 
-data UserLinkData = UserLinkData
-  { agentVRange :: VersionRangeSMPA,
-    owners :: [OwnerAuth],
-    userData :: ConnInfo
-  }
+data ConnLinkData c where
+  InvitationLinkData :: VersionRangeSMPA -> ConnInfo -> 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 :: ConnInfo
+    } -> ConnLinkData 'CMContact
+
+data AConnLinkData = forall m. ConnectionModeI m => ACLD (SConnectionMode m) (ConnLinkData m)
+
+linkUserData :: ConnLinkData c -> ConnInfo
+linkUserData = \case
+  InvitationLinkData _ d -> d
+  ContactLinkData {userData} -> userData
 
 type OwnerId = ByteString
 
@@ -1540,14 +1589,28 @@ instance ConnectionModeI c => Encoding (FixedLinkData c) where
     (agentVRange, rootKey, connReq) <- smpP
     pure FixedLinkData {agentVRange, rootKey, connReq}
 
-instance Encoding UserLinkData where
-  smpEncode UserLinkData {agentVRange, owners, userData} =
-    smpEncode agentVRange <> smpEncodeList owners <> smpEncode (Large userData)
-  smpP = do
-    agentVRange <- smpP
-    owners <- smpListP
-    Large userData <- smpP
-    pure UserLinkData {agentVRange, owners, userData}
+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]
+  smpP = (\(ACLD _ d) -> checkConnMode d) <$?> smpP
+  {-# INLINE smpP #-}
+
+instance Encoding AConnLinkData where
+  smpEncode (ACLD _ d) = smpEncode d
+  {-# INLINE smpEncode #-}
+  smpP =
+    smpP >>= \case
+      CMInvitation -> do
+        (vr, userData) <- smpP
+        pure $ ACLD SCMInvitation $ InvitationLinkData vr userData
+      CMContact -> do
+        (agentVRange, direct) <- smpP
+        owners <- smpListP
+        relays <- smpListP
+        userData <- smpP
+        pure $ ACLD SCMContact ContactLinkData {agentVRange, direct, owners, relays, userData}
 
 -- | SMP queue status.
 data QueueStatus
diff --git a/src/Simplex/Messaging/Crypto/ShortLink.hs b/src/Simplex/Messaging/Crypto/ShortLink.hs
@@ -1,10 +1,12 @@
 {-# LANGUAGE DataKinds #-}
 {-# LANGUAGE DuplicateRecordFields #-}
+{-# LANGUAGE GADTs #-}
 {-# LANGUAGE MultiWayIf #-}
 {-# LANGUAGE NamedFieldPuns #-}
 {-# LANGUAGE OverloadedStrings #-}
 {-# LANGUAGE ScopedTypeVariables #-}
 {-# LANGUAGE TupleSections #-}
+{-# LANGUAGE TypeApplications #-}
 
 module Simplex.Messaging.Crypto.ShortLink
   ( contactShortLinkKdf,
@@ -46,15 +48,20 @@ contactShortLinkKdf (LinkKey k) =
 invShortLinkKdf :: LinkKey -> C.SbKey
 invShortLinkKdf (LinkKey k) = C.unsafeSbKey $ C.hkdf "" k "SimpleXInvLink" 32
 
-encodeSignLinkData :: ConnectionModeI c => C.KeyPair 'C.Ed25519 -> VersionRangeSMPA -> ConnectionRequestUri c -> ConnInfo -> (LinkKey, (ByteString, ByteString))
+encodeSignLinkData :: forall c. ConnectionModeI c => C.KeyPair 'C.Ed25519 -> VersionRangeSMPA -> ConnectionRequestUri c -> ConnInfo -> (LinkKey, (ByteString, ByteString))
 encodeSignLinkData (rootKey, pk) agentVRange connReq userData =
-  let fd = smpEncode FixedLinkData {agentVRange, rootKey, connReq = Just connReq}
-      ud = smpEncode UserLinkData {agentVRange, owners = [], userData}
-   in (LinkKey (C.sha3_256 fd), (encodeSign pk fd, encodeSign pk ud))
+  let fd = smpEncode FixedLinkData {agentVRange, rootKey, connReq}
+      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 UserLinkData {agentVRange, owners = [], userData}
+  encodeSign pk $ smpEncode $ connLinkData @'CMContact agentVRange userData
+
+connLinkData :: forall c. ConnectionModeI c => VersionRangeSMPA -> ConnInfo -> ConnLinkData c
+connLinkData agentVRange userData = case sConnectionMode @c of
+  SCMInvitation -> InvitationLinkData agentVRange userData
+  SCMContact -> ContactLinkData {agentVRange, direct = True, owners = [], relays = [], userData}
 
 encodeSign :: C.PrivateKeyEd25519 -> ByteString -> ByteString
 encodeSign pk s = smpEncode (C.sign' pk s) <> s
@@ -73,17 +80,17 @@ encryptData g k len s = do
   ct <- liftEitherWith cryptoError $ C.sbEncrypt k nonce s len
   pure $ EncDataBytes $ smpEncode nonce <> ct
 
-decryptLinkData :: ConnectionModeI c => LinkKey -> C.SbKey -> QueueLinkData -> Either AgentErrorType (Maybe (ConnectionRequestUri c), ConnInfo)
-decryptLinkData linkKey k (encFD, encUD) = do
+decryptLinkData :: forall c. ConnectionModeI c => LinkKey -> C.SbKey -> QueueLinkData -> Either AgentErrorType (ConnectionRequestUri c, ConnInfo)
+decryptLinkData linkKey k (encFD, encMD) = do
   (sig1, fd) <- decrypt encFD
-  (sig2, ud) <- decrypt encUD
+  (sig2, md) <- decrypt encMD
   FixedLinkData {rootKey, connReq} <- decode fd
-  UserLinkData {userData} <- decode ud
+  md' <- decode @(ConnLinkData c) md
   if
     | LinkKey (C.sha3_256 fd) /= linkKey -> linkErr "link data hash"
     | not (C.verify' rootKey sig1 fd) -> linkErr "link data signature"
-    | not (C.verify' rootKey sig2 ud) -> linkErr "user data signature"
-    | otherwise -> Right (connReq, userData)
+    | not (C.verify' rootKey sig2 md) -> linkErr "user data signature"
+    | otherwise -> Right (connReq, linkUserData md')
   where
     decrypt (EncDataBytes d) = do
       (nonce, Tail ct) <- decode d
diff --git a/tests/AgentTests/ConnectionRequestTests.hs b/tests/AgentTests/ConnectionRequestTests.hs
@@ -279,22 +279,22 @@ connectionRequestTests =
       smpEncodingTest contactAddressV2
       smpEncodingTest contactAddressClientData
     it "should serialize / parse short links" $ do
-      CSLContact srv CCTContact (LinkKey "0123456789abcdef0123456789abcdef") #==# "https://smp.simplex.im:5223/c#1234-w@jjbyvoemxysm7qxap7m5d5m35jzv5qq6gnlv7s4rsn7tdwwmuqciwpid.onion/MDEyMzQ1Njc4OWFiY2RlZjAxMjM0NTY3ODlhYmNkZWY"
-      CSLContact srv CCTGroup (LinkKey "0123456789abcdef0123456789abcdef") #==# "https://smp.simplex.im:5223/g#1234-w@jjbyvoemxysm7qxap7m5d5m35jzv5qq6gnlv7s4rsn7tdwwmuqciwpid.onion/MDEyMzQ1Njc4OWFiY2RlZjAxMjM0NTY3ODlhYmNkZWY"
-      CSLContact shortSrv CCTContact (LinkKey "0123456789abcdef0123456789abcdef") #==# "https://smp.simplex.im/c#MDEyMzQ1Njc4OWFiY2RlZjAxMjM0NTY3ODlhYmNkZWY"
+      CSLContact CCTContact srv (LinkKey "0123456789abcdef0123456789abcdef") #==# "https://smp.simplex.im:5223/a#1234-w@jjbyvoemxysm7qxap7m5d5m35jzv5qq6gnlv7s4rsn7tdwwmuqciwpid.onion/MDEyMzQ1Njc4OWFiY2RlZjAxMjM0NTY3ODlhYmNkZWY"
+      CSLContact CCTGroup srv (LinkKey "0123456789abcdef0123456789abcdef") #==# "https://smp.simplex.im:5223/g#1234-w@jjbyvoemxysm7qxap7m5d5m35jzv5qq6gnlv7s4rsn7tdwwmuqciwpid.onion/MDEyMzQ1Njc4OWFiY2RlZjAxMjM0NTY3ODlhYmNkZWY"
+      CSLContact CCTContact shortSrv (LinkKey "0123456789abcdef0123456789abcdef") #==# "https://smp.simplex.im/a#MDEyMzQ1Njc4OWFiY2RlZjAxMjM0NTY3ODlhYmNkZWY"
       CSLInvitation srv (EntityId "0123456789abcdef01234567") (LinkKey "0123456789abcdef0123456789abcdef") #==# "https://smp.simplex.im:5223/i#1234-w@jjbyvoemxysm7qxap7m5d5m35jzv5qq6gnlv7s4rsn7tdwwmuqciwpid.onion/MDEyMzQ1Njc4OWFiY2RlZjAxMjM0NTY3MDEyMzQ1Njc4OWFiY2RlZjAxMjM0NTY3ODlhYmNkZWY"
       CSLInvitation shortSrv (EntityId "0123456789abcdef01234567") (LinkKey "0123456789abcdef0123456789abcdef") #==# "https://smp.simplex.im/i#MDEyMzQ1Njc4OWFiY2RlZjAxMjM0NTY3MDEyMzQ1Njc4OWFiY2RlZjAxMjM0NTY3ODlhYmNkZWY"
     it "should shorten / restore short links" $ do
-      shortenShortLink [srv] (CSLContact srv CCTContact (LinkKey "0123456789abcdef0123456789abcdef"))
-        `shouldBe` CSLContact shortSrv CCTContact (LinkKey "0123456789abcdef0123456789abcdef")
-      shortenShortLink [srv] (CSLContact srv2 CCTContact (LinkKey "0123456789abcdef0123456789abcdef"))
-        `shouldBe` CSLContact srv2 CCTContact (LinkKey "0123456789abcdef0123456789abcdef")
-      restoreShortLink [srv] (CSLContact shortSrv CCTContact (LinkKey "0123456789abcdef0123456789abcdef"))
-        `shouldBe` CSLContact srv CCTContact (LinkKey "0123456789abcdef0123456789abcdef")
-      restoreShortLink [srv2] (CSLContact shortSrv CCTContact (LinkKey "0123456789abcdef0123456789abcdef"))
-        `shouldBe` CSLContact shortSrv CCTContact (LinkKey "0123456789abcdef0123456789abcdef")
-      restoreShortLink [srv] (CSLContact srv2 CCTContact (LinkKey "0123456789abcdef0123456789abcdef"))
-        `shouldBe` CSLContact srv2 CCTContact (LinkKey "0123456789abcdef0123456789abcdef")
+      shortenShortLink [srv] (CSLContact CCTContact srv (LinkKey "0123456789abcdef0123456789abcdef"))
+        `shouldBe` CSLContact CCTContact shortSrv (LinkKey "0123456789abcdef0123456789abcdef")
+      shortenShortLink [srv] (CSLContact CCTContact srv2 (LinkKey "0123456789abcdef0123456789abcdef"))
+        `shouldBe` CSLContact CCTContact srv2 (LinkKey "0123456789abcdef0123456789abcdef")
+      restoreShortLink [srv] (CSLContact CCTContact shortSrv (LinkKey "0123456789abcdef0123456789abcdef"))
+        `shouldBe` CSLContact CCTContact srv (LinkKey "0123456789abcdef0123456789abcdef")
+      restoreShortLink [srv2] (CSLContact CCTContact shortSrv (LinkKey "0123456789abcdef0123456789abcdef"))
+        `shouldBe` CSLContact CCTContact shortSrv (LinkKey "0123456789abcdef0123456789abcdef")
+      restoreShortLink [srv] (CSLContact CCTContact srv2 (LinkKey "0123456789abcdef0123456789abcdef"))
+        `shouldBe` CSLContact CCTContact srv2 (LinkKey "0123456789abcdef0123456789abcdef")
   where
     smpEncodingTest a = smpDecode (smpEncode a) `shouldBe` Right a
 
diff --git a/tests/AgentTests/FunctionalAPITests.hs b/tests/AgentTests/FunctionalAPITests.hs
diff --git a/tests/AgentTests/ShortLinkTests.hs b/tests/AgentTests/ShortLinkTests.hs
