smp protocol: short links and other changes from RFC (#1489)

* smp protocol: short links types and other changes from RFC * add fields for queue link ID and data * create queue and ntf credentials with NEW command * all tests * simplfiy types, update rfc * update rfc * include SenderId in NEW request in case queue data is sent * store queue data and generate link ID if needed * update rfc * agent API and types * SMP commands and persistence for short links * SMP client functions for short links * agent client functions for short links * create rcv queue with short link (TODO secret_box) * encryption and encoding for link data, postgres client migration * test creating short link * get link and data, tests * comments * type signature
2026-07-20 20:01:47 +00:00 · 2025-03-26 17:26:27 +00:00
parent 0c3b25706a
commit b83d897650
44 changed files with 1701 additions and 338 deletions
@@ -35,7 +35,6 @@ import Control.Monad.Trans.Except
 import Data.ByteString.Builder (Builder)
 import qualified Data.ByteString.Builder as BB
 import Data.ByteString.Char8 (ByteString)
-import qualified Data.ByteString.Char8 as B
 import qualified Data.ByteString.Lazy as LB
 import Data.Bitraversable (bimapM)
 import Data.Either (fromRight)
@@ -43,10 +42,10 @@ import Data.Functor (($>))
 import Data.Int (Int64)
 import Data.List (intersperse)
 import qualified Data.Map.Strict as M
-import Data.Maybe (catMaybes)
+import Data.Maybe (catMaybes, fromMaybe)
 import qualified Data.Text as T
 import Data.Time.Clock.System (SystemTime (..), getSystemTime)
-import Database.PostgreSQL.Simple (Binary (..), Only (..), Query, SqlError)
+import Database.PostgreSQL.Simple (Binary (..), Only (..), Query, SqlError, (:.) (..))
 import qualified Database.PostgreSQL.Simple as DB
 import qualified Database.PostgreSQL.Simple.Copy as DB
 import Database.PostgreSQL.Simple.FromField (FromField (..))
@@ -146,10 +145,11 @@ instance StoreQueueClass q => QueueStoreClass q (PostgresQueueStore q) where
          >>= bimapM handleDuplicate pure
      atomically $ TM.insert rId sq queues
      atomically $ TM.insert (senderId qr) rId senders
+      forM_ (notifier qr) $ \NtfCreds {notifierId = nId} -> atomically $ TM.insert nId rId notifiers
      withLog "addStoreQueue" st $ \s -> logCreateQueue s rId qr
      pure sq
    where
-      PostgresQueueStore {queues, senders} = st
+      PostgresQueueStore {queues, senders, notifiers} = st
      -- Not doing duplicate checks in maps as the probability of duplicates is very low.
      -- It needs to be reconsidered when IDs are supplied by the users.
      -- hasId = anyM [TM.memberIO rId queues, TM.memberIO senderId senders, hasNotifier]
@@ -160,12 +160,14 @@ instance StoreQueueClass q => QueueStoreClass q (PostgresQueueStore q) where
    SRecipient -> getRcvQueue qId
    SSender -> TM.lookupIO qId senders >>= maybe loadSndQueue getRcvQueue
    SNotifier -> TM.lookupIO qId notifiers >>= maybe loadNtfQueue getRcvQueue
+    SLinkClient -> loadLinkQueue
    where
      PostgresQueueStore {queues, senders, notifiers} = st
      getRcvQueue rId = TM.lookupIO rId queues >>= maybe loadRcvQueue (pure . Right)
      loadRcvQueue = loadQueue " WHERE recipient_id = ?" $ \_ -> pure ()
      loadSndQueue = loadQueue " WHERE sender_id = ?" $ \rId -> TM.insert qId rId senders
      loadNtfQueue = loadQueue " WHERE notifier_id = ?" $ \_ -> pure () -- do NOT cache ref - ntf subscriptions are rare
+      loadLinkQueue = loadQueue " WHERE link_id = ?" $ \_ -> pure ()
      loadQueue condition insertRef =
        E.uninterruptibleMask_ $ runExceptT $ do
          (rId, qRec) <-
@@ -187,6 +189,43 @@ instance StoreQueueClass q => QueueStoreClass q (PostgresQueueStore q) where
                  TM.insert rId sq queues
                  pure sq

+  getQueueLinkData :: PostgresQueueStore q -> q -> LinkId -> IO (Either ErrorType QueueLinkData)
+  getQueueLinkData st sq lnkId = runExceptT $ do
+    qr <- ExceptT $ readQueueRecIO $ queueRec sq
+    case queueData qr of
+      Just (lnkId', _) | lnkId' == lnkId ->
+        withDB "getQueueLinkData" st $ \db -> firstRow id AUTH $
+          DB.query db "SELECT fixed_data, user_data FROM msg_queues WHERE link_id = ? AND deleted_at IS NULL" (Only lnkId)
+      _ -> throwE AUTH
+
+  addQueueLinkData :: PostgresQueueStore q -> q -> LinkId -> QueueLinkData -> IO (Either ErrorType ())
+  addQueueLinkData st sq lnkId d = 
+    withQueueRec sq "addQueueLinkData" $ \q -> case queueData q of
+      Nothing ->
+        addLink q $ \db -> DB.execute db qry (d :. (lnkId, rId))
+      Just (lnkId', _) | lnkId' == lnkId ->
+        addLink q $ \db -> DB.execute db (qry <> " AND (fixed_data IS NULL OR fixed_data != ?)") (d :. (lnkId, rId, fst d))
+      _ -> throwE AUTH
+    where
+      rId = recipientId sq
+      addLink q update = do
+        assertUpdated $ withDB' "addQueueLinkData" st update
+        atomically $ writeTVar (queueRec sq) $ Just q {queueData = Just (lnkId, d)}
+        withLog "addQueueLinkData" st $ \s -> logCreateLink s rId lnkId d
+      qry = "UPDATE msg_queues SET fixed_data = ?, user_data = ?, link_id = ? WHERE recipient_id = ? AND deleted_at IS NULL"
+
+  deleteQueueLinkData :: PostgresQueueStore q -> q -> IO (Either ErrorType ())
+  deleteQueueLinkData st sq =
+    withQueueRec sq "deleteQueueLinkData" $ \q -> case queueData q of
+      Just _ -> do
+        assertUpdated $ withDB' "deleteQueueLinkData" st $ \db ->
+          DB.execute db "UPDATE msg_queues SET link_id = NULL, fixed_data = NULL, user_data = NULL WHERE recipient_id = ? AND deleted_at IS NULL" (Only rId)
+        atomically $ writeTVar (queueRec sq) $ Just q {queueData = Nothing}
+        withLog "deleteQueueLinkData" st (`logDeleteLink` rId)
+      _ -> throwE AUTH
+    where
+      rId = recipientId sq
+
  secureQueue :: PostgresQueueStore q -> q -> SndPublicAuthKey -> IO (Either ErrorType ())
  secureQueue st sq sKey =
    withQueueRec sq "secureQueue" $ \q -> do
@@ -320,15 +359,15 @@ insertQueueQuery :: Query
 insertQueueQuery =
  [sql|
    INSERT INTO msg_queues
-      (recipient_id, recipient_key, rcv_dh_secret, sender_id, sender_key, snd_secure, notifier_id, notifier_key, rcv_ntf_dh_secret, status, updated_at)
-    VALUES (?,?,?,?,?,?,?,?,?,?,?)
+      (recipient_id, recipient_key, rcv_dh_secret, sender_id, sender_key, queue_mode, notifier_id, notifier_key, rcv_ntf_dh_secret, status, updated_at, link_id, fixed_data, user_data)
+    VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?,?)
  |]

-foldQueueRecs :: Monoid a => Bool -> PostgresQueueStore q -> Maybe Int64 -> ((RecipientId, QueueRec) -> IO a) -> IO a
-foldQueueRecs tty st skipOld_ f = do
+foldQueueRecs :: forall a q. Monoid a => Bool -> Bool -> PostgresQueueStore q -> Maybe Int64 -> ((RecipientId, QueueRec) -> IO a) -> IO a
+foldQueueRecs tty withData st skipOld_ f = do
  (n, r) <- withConnection (dbStore st) $ \db ->
-    foldRecs db (0 :: Int, mempty) $ \(i, acc) row -> do
-      r <- f $ rowToQueueRec row
+    foldRecs db (0 :: Int, mempty) $ \(i, acc) qr -> do
+      r <- f qr
      let !i' = i + 1
          !acc' = acc <> r
      when (tty && i' `mod` 100000 == 0) $ putStr (progress i' <> "\r") >> hFlush stdout
@@ -336,29 +375,49 @@ foldQueueRecs tty st skipOld_ f = do
  when tty $ putStrLn $ progress n
  pure r
  where
-    foldRecs db = case skipOld_ of
-      Nothing -> DB.fold_ db (queueRecQuery <> " WHERE deleted_at IS NULL")
-      Just old -> DB.fold db (queueRecQuery <> " WHERE deleted_at IS NULL AND updated_at > ?") (Only old)
+    foldRecs db acc f' = case skipOld_ of
+      Nothing
+        | withData -> DB.fold_ db (query <> " WHERE deleted_at IS NULL") acc $ \acc' -> f' acc' . rowToQueueRecWithData
+        | otherwise -> DB.fold_ db (query <> " WHERE deleted_at IS NULL") acc $ \acc' -> f' acc' . rowToQueueRec
+      Just old
+        | withData -> DB.fold db (query <> " WHERE deleted_at IS NULL AND updated_at > ?") (Only old) acc $ \acc' -> f' acc' . rowToQueueRecWithData
+        | otherwise -> DB.fold db (query <> " WHERE deleted_at IS NULL AND updated_at > ?") (Only old) acc $ \acc' -> f' acc' . rowToQueueRec
+    query = if withData then queueRecQueryWithData else queueRecQuery
    progress i = "Processed: " <> show i <> " records"

 queueRecQuery :: Query
 queueRecQuery =
  [sql|
    SELECT recipient_id, recipient_key, rcv_dh_secret,
-      sender_id, sender_key, snd_secure,
+      sender_id, sender_key, queue_mode,
      notifier_id, notifier_key, rcv_ntf_dh_secret,
-      status, updated_at
+      status, updated_at,
+      link_id
    FROM msg_queues
  |]

-type QueueRecRow = (RecipientId, RcvPublicAuthKey, RcvDhSecret, SenderId, Maybe SndPublicAuthKey, SenderCanSecure, Maybe NotifierId, Maybe NtfPublicAuthKey, Maybe RcvNtfDhSecret, ServerEntityStatus, Maybe RoundedSystemTime)
+queueRecQueryWithData :: Query
+queueRecQueryWithData =
+  [sql|
+    SELECT recipient_id, recipient_key, rcv_dh_secret,
+      sender_id, sender_key, queue_mode,
+      notifier_id, notifier_key, rcv_ntf_dh_secret,
+      status, updated_at,
+      link_id, fixed_data, user_data
+    FROM msg_queues
+  |]

-queueRecToRow :: (RecipientId, QueueRec) -> QueueRecRow
-queueRecToRow (rId, QueueRec {recipientKey, rcvDhSecret, senderId, senderKey, sndSecure, notifier = n, status, updatedAt}) =
-  (rId, recipientKey, rcvDhSecret, senderId, senderKey, sndSecure, notifierId <$> n, notifierKey <$> n, rcvNtfDhSecret <$> n, status, updatedAt)
+type QueueRecRow = (RecipientId, RcvPublicAuthKey, RcvDhSecret, SenderId, Maybe SndPublicAuthKey, Maybe QueueMode, Maybe NotifierId, Maybe NtfPublicAuthKey, Maybe RcvNtfDhSecret, ServerEntityStatus, Maybe RoundedSystemTime, Maybe LinkId)
+
+queueRecToRow :: (RecipientId, QueueRec) -> QueueRecRow :. (Maybe EncDataBytes, Maybe EncDataBytes)
+queueRecToRow (rId, QueueRec {recipientKey, rcvDhSecret, senderId, senderKey, queueMode, queueData, notifier = n, status, updatedAt}) =
+  (rId, recipientKey, rcvDhSecret, senderId, senderKey, queueMode, notifierId <$> n, notifierKey <$> n, rcvNtfDhSecret <$> n, status, updatedAt, linkId_)
+    :. (fst <$> queueData_, snd <$> queueData_)
+  where
+    (linkId_, queueData_) = queueDataColumns queueData

 queueRecToText :: (RecipientId, QueueRec) -> ByteString
-queueRecToText (rId, QueueRec {recipientKey, rcvDhSecret, senderId, senderKey, sndSecure, notifier = n, status, updatedAt}) =
+queueRecToText (rId, QueueRec {recipientKey, rcvDhSecret, senderId, senderKey, queueMode, queueData, notifier = n, status, updatedAt}) =
  LB.toStrict $ BB.toLazyByteString $ mconcat tabFields <> BB.char7 '\n'
  where
    tabFields = BB.char7 ',' `intersperse` fields
@@ -368,13 +427,17 @@ queueRecToText (rId, QueueRec {recipientKey, rcvDhSecret, senderId, senderKey, s
        renderField (toField rcvDhSecret),
        renderField (toField senderId),
        nullable senderKey,
-        renderField (toField sndSecure),
+        nullable queueMode,
        nullable (notifierId <$> n),
        nullable (notifierKey <$> n),
        nullable (rcvNtfDhSecret <$> n),
        BB.char7 '"' <> renderField (toField status) <> BB.char7 '"',
-        nullable updatedAt
+        nullable updatedAt,
+        nullable linkId_,
+        nullable (fst <$> queueData_),
+        nullable (snd <$> queueData_)
      ]
+    (linkId_, queueData_) = queueDataColumns queueData
    nullable :: ToField a => Maybe a -> Builder
    nullable = maybe mempty (renderField . toField)
    renderField :: Action -> Builder
@@ -385,10 +448,23 @@ queueRecToText (rId, QueueRec {recipientKey, rcvDhSecret, senderId, senderKey, s
      EscapeIdentifier s -> BB.byteString s -- Not used in COPY data
      Many as -> mconcat (map renderField as)

+queueDataColumns :: Maybe (LinkId, QueueLinkData) -> (Maybe LinkId, Maybe QueueLinkData)
+queueDataColumns = \case
+  Just (linkId, linkData) -> (Just linkId, Just linkData)
+  Nothing -> (Nothing, Nothing)
+
 rowToQueueRec :: QueueRecRow -> (RecipientId, QueueRec)
-rowToQueueRec (rId, recipientKey, rcvDhSecret, senderId, senderKey, sndSecure, notifierId_, notifierKey_, rcvNtfDhSecret_, status, updatedAt) =
+rowToQueueRec (rId, recipientKey, rcvDhSecret, senderId, senderKey, queueMode, notifierId_, notifierKey_, rcvNtfDhSecret_, status, updatedAt, linkId_) =
  let notifier = NtfCreds <$> notifierId_ <*> notifierKey_ <*> rcvNtfDhSecret_
-   in (rId, QueueRec {recipientKey, rcvDhSecret, senderId, senderKey, sndSecure, notifier, status, updatedAt})
+      queueData = (,(EncDataBytes "", EncDataBytes "")) <$> linkId_
+   in (rId, QueueRec {recipientKey, rcvDhSecret, senderId, senderKey, queueMode, queueData, notifier, status, updatedAt})
+
+rowToQueueRecWithData :: QueueRecRow :. (Maybe EncDataBytes, Maybe EncDataBytes) -> (RecipientId, QueueRec)
+rowToQueueRecWithData ((rId, recipientKey, rcvDhSecret, senderId, senderKey, queueMode, notifierId_, notifierKey_, rcvNtfDhSecret_, status, updatedAt, linkId_) :. (immutableData_, userData_)) =
+  let notifier = NtfCreds <$> notifierId_ <*> notifierKey_ <*> rcvNtfDhSecret_
+      encData =  fromMaybe (EncDataBytes "")
+      queueData = (,(encData immutableData_, encData userData_)) <$> linkId_
+   in (rId, QueueRec {recipientKey, rcvDhSecret, senderId, senderKey, queueMode, queueData, notifier, status, updatedAt})

 setStatusDB :: StoreQueueClass q => String -> PostgresQueueStore q -> q -> ServerEntityStatus -> ExceptT ErrorType IO () -> IO (Either ErrorType ())
 setStatusDB op st sq status writeLog =
@@ -441,4 +517,8 @@ instance FromField (C.DhSecret 'C.X25519) where fromField = blobFieldDecoder str
 instance ToField C.APublicAuthKey where toField = toField . Binary . C.encodePubKey

 instance FromField C.APublicAuthKey where fromField = blobFieldDecoder C.decodePubKey
+
+instance ToField EncDataBytes where toField (EncDataBytes s) = toField (Binary s)
+
+deriving newtype instance FromField EncDataBytes
 #endif