update clients' key package refs upon committing

changelog.d/5-internal/update-key-package-ref

@@ -0,0 +1,2 @@
+Add internal endpoint in Brig to update clients' key package refs in DB upon committing.
+Brig should be deployed before Galley.

libs/wire-api/src/Wire/API/MLS/Commit.hs

@@ -31,7 +31,7 @@ instance ParseMLS Commit where
   parseMLS = Commit <$> parseMLSVector @Word32 parseMLS <*> parseMLSOptional parseMLS
 
 data UpdatePath = UpdatePath
-  { upLeaf :: KeyPackage,
+  { upLeaf :: RawMLS KeyPackage,
     upNodes :: [UpdatePathNode]
   }
 

libs/wire-api/src/Wire/API/MLS/KeyPackage.hs

@@ -33,6 +33,7 @@ module Wire.API.MLS.KeyPackage
     kpRef',
     KeyPackageTBS (..),
     KeyPackageRef (..),
+    KeyPackageUpdate (..),
   )
 where
 
@@ -114,6 +115,7 @@ instance ToSchema KeyPackageCount where
 newtype KeyPackageRef = KeyPackageRef {unKeyPackageRef :: ByteString}
   deriving stock (Eq, Ord, Show)
   deriving (FromHttpApiData, ToHttpApiData, S.ToParamSchema) via Base64ByteString
+  deriving (ToJSON, FromJSON, S.ToSchema) via (Schema KeyPackageRef)
 
 instance ToSchema KeyPackageRef where
   schema = named "KeyPackageRef" $ unKeyPackageRef .= fmap KeyPackageRef base64Schema
@@ -191,3 +193,10 @@ instance ParseMLS KeyPackage where
     KeyPackage
       <$> parseRawMLS parseMLS
       <*> parseMLSBytes @Word16
+
+--------------------------------------------------------------------------------
+
+data KeyPackageUpdate = KeyPackageUpdate
+  { kpupPrevious :: KeyPackageRef,
+    kpupNext :: KeyPackageRef
+  }

libs/wire-api/src/Wire/API/Routes/Internal/Brig.hs

@@ -198,6 +198,16 @@ type MLSAPI =
                                     '[RespondEmpty 201 "Key package ref mapping created"]
                                     ()
                            )
+                    :<|> Named
+                           "post-key-package-ref"
+                           ( Summary "Update a KeyPackageRef in mapping"
+                               :> ReqBody '[Servant.JSON] KeyPackageRef
+                               :> MultiVerb
+                                    'POST
+                                    '[Servant.JSON]
+                                    '[RespondEmpty 201 "Key package ref mapping updated"]
+                                    ()
+                           )
                 )
          )
            :<|> GetMLSClients

services/brig/src/Brig/API/Internal.hs

@@ -120,6 +120,7 @@ mlsAPI =
                  :<|> Named @"get-conversation-by-key-package-ref" (getConvIdByKeyPackageRef ref)
              )
         :<|> Named @"put-key-package-ref" (putKeyPackageRef ref)
+        :<|> Named @"post-key-package-ref" (postKeyPackageRef ref)
   )
     :<|> getMLSClients
     :<|> mapKeyPackageRefsInternal
@@ -161,6 +162,10 @@ putKeyPackageRef ref = lift . wrapClient . Data.addKeyPackageRef ref
 getConvIdByKeyPackageRef :: KeyPackageRef -> Handler r (Maybe (Qualified ConvId))
 getConvIdByKeyPackageRef = runMaybeT . mapMaybeT wrapClientE . Data.keyPackageRefConvId
 
+-- Used by galley to update key packages in mls_key_package_ref on commits with update_path
+postKeyPackageRef :: KeyPackageRef -> KeyPackageRef -> Handler r ()
+postKeyPackageRef ref = lift . wrapClient . Data.updateKeyPackageRef ref
+
 getMLSClients :: UserId -> SignatureSchemeTag -> Handler r (Set ClientId)
 getMLSClients usr ss = do
   results <- lift (wrapClient (API.lookupUsersClientIds (pure usr))) >>= getResult

services/brig/src/Brig/Data/MLS/KeyPackage.hs

@@ -24,6 +24,7 @@ module Brig.Data.MLS.KeyPackage
     keyPackageRefConvId,
     keyPackageRefSetConvId,
     addKeyPackageRef,
+    updateKeyPackageRef,
   )
 where
 
@@ -194,9 +195,48 @@ addKeyPackageRef ref nkpr = do
     q :: PrepQuery W (ClientId, ConvId, Domain, Domain, UserId, KeyPackageRef) x
     q = "UPDATE mls_key_package_refs SET client = ?, conv = ?, conv_domain = ?, domain = ?, user = ? WHERE ref = ?"
 
+-- | Update key package ref, used in Galley when commit reveals key package ref update for the sender.
+-- Nothing is changed if the previous key package ref is not found in the table.
+-- Updating amounts to INSERT the new key package ref, followed by DELETE the
+-- previous one.
+--
+-- FUTUREWORK: this function has to be extended if a table mapping (client,
+-- conversation) to key package ref is added, for instance, when implementing
+-- external delete proposals.
+updateKeyPackageRef :: MonadClient m => KeyPackageRef -> KeyPackageRef -> m ()
+updateKeyPackageRef prevRef newRef =
+  void . runMaybeT $ do
+    backup <- backupKeyPackageMeta prevRef
+    lift $ restoreKeyPackageMeta newRef backup >> deleteKeyPackage prevRef
+
 --------------------------------------------------------------------------------
 -- Utilities
 
+backupKeyPackageMeta :: MonadClient m => KeyPackageRef -> MaybeT m (ClientId, Qualified ConvId, Qualified UserId)
+backupKeyPackageMeta ref = do
+  (clientId, convId, convDomain, userDomain, userId) <- MaybeT . retry x1 $ query1 q (params LocalQuorum (Identity ref))
+  pure (clientId, Qualified convId convDomain, Qualified userId userDomain)
+  where
+    q :: PrepQuery R (Identity KeyPackageRef) (ClientId, ConvId, Domain, Domain, UserId)
+    q = "SELECT client, conv, conv_domain, domain, user FROM mls_key_package_refs WHERE ref = ?"
+
+restoreKeyPackageMeta :: MonadClient m => KeyPackageRef -> (ClientId, Qualified ConvId, Qualified UserId) -> m ()
+restoreKeyPackageMeta ref (clientId, convId, userId) = do
+  write q (params LocalQuorum (ref, clientId, qUnqualified convId, qDomain convId, qDomain userId, qUnqualified userId))
+  where
+    q :: PrepQuery W (KeyPackageRef, ClientId, ConvId, Domain, Domain, UserId) ()
+    q = "INSERT INTO mls_key_package_refs (ref, client, conv, conv_domain, domain, user) VALUES (?, ?, ?, ?, ?, ?)"
+
+deleteKeyPackage :: MonadClient m => KeyPackageRef -> m ()
+deleteKeyPackage ref =
+  retry x5 $
+    write
+      q
+      (params LocalQuorum (Identity ref))
+  where
+    q :: PrepQuery W (Identity KeyPackageRef) x
+    q = "DELETE FROM mls_key_package_refs WHERE ref = ?"
+
 pick :: [a] -> IO (Maybe a)
 pick [] = pure Nothing
 pick xs = do

services/galley/src/Galley/API/MLS/Message.hs

@@ -328,6 +328,7 @@ processCommit ::
     Member (ErrorS 'MissingLegalholdConsent) r,
     Member (Input (Local ())) r,
     Member ProposalStore r,
+    Member BrigAccess r,
     Member Resource r
   ) =>
   Qualified UserId ->
@@ -353,24 +354,47 @@ processCommit qusr con lconv epoch sender commit = do
   let ttlSeconds :: Int = 600 -- 10 minutes
   withCommitLock groupId epoch (fromIntegral ttlSeconds) $ do
     checkEpoch epoch (tUnqualified lconv)
-    when (epoch == Epoch 0) $ do
-      -- this is a newly created conversation, and it should contain exactly one
-      -- client (the creator)
-      case (sender, first (toList . lmMLSClients) self) of
-        (MemberSender ref, Left [creatorClient]) -> do
-          -- register the creator client
-          addKeyPackageRef
-            ref
-            qusr
-            creatorClient
-            (qUntagged (fmap Data.convId lconv))
-        -- remote clients cannot send the first commit
-        (_, Right _) -> throwS @'MLSStaleMessage
-        -- uninitialised conversations should contain exactly one client
-        (MemberSender _, _) ->
-          throw (InternalErrorWithDescription "Unexpected creator client set")
-        -- the sender of the first commit must be a member
-        _ -> throw (mlsProtocolError "Unexpected sender")
+    postponedKeyPackageRefUpdate <-
+      if epoch == Epoch 0
+        then do
+          -- this is a newly created conversation, and it should contain exactly one
+          -- client (the creator)
+          case (sender, first (toList . lmMLSClients) self) of
+            (MemberSender currentRef, Left [creatorClient]) -> do
+              -- use update path as sender reference and if not existing fall back to sender
+              senderRef <-
+                maybe
+                  (pure currentRef)
+                  ( (& note (mlsProtocolError "Could not compute key package ref"))
+                      . kpRef'
+                      . upLeaf
+                  )
+                  $ cPath commit
+              -- register the creator client
+              addKeyPackageRef
+                senderRef
+                qusr
+                creatorClient
+                (qUntagged (fmap Data.convId lconv))
+            -- remote clients cannot send the first commit
+            (_, Right _) -> throwS @'MLSStaleMessage
+            -- uninitialised conversations should contain exactly one client
+            (MemberSender _, _) ->
+              throw (InternalErrorWithDescription "Unexpected creator client set")
+            -- the sender of the first commit must be a member
+            _ -> throw (mlsProtocolError "Unexpected sender")
+          pure $ pure () -- no key package ref update necessary
+        else case (sender, upLeaf <$> cPath commit) of
+          (MemberSender senderRef, Just updatedKeyPackage) -> do
+            updatedRef <- kpRef' updatedKeyPackage & note (mlsProtocolError "Could not compute key package ref")
+            -- postpone key package ref update until other checks/processing passed
+            pure . updateKeyPackageRef $
+              KeyPackageUpdate
+                { kpupPrevious = senderRef,
+                  kpupNext = updatedRef
+                }
+          (_, Nothing) -> pure $ pure () -- ignore commits without update path
+          _ -> throw (mlsProtocolError "Unexpected sender")
 
     -- check all pending proposals are referenced in the commit
     allPendingProposals <- getAllPendingProposals groupId epoch
@@ -382,6 +406,8 @@ processCommit qusr con lconv epoch sender commit = do
     action <- foldMap (applyProposalRef (tUnqualified lconv) groupId epoch) (cProposals commit)
     updates <- executeProposalAction qusr con lconv action
 
+    -- update key package ref if necessary
+    postponedKeyPackageRefUpdate
     -- increment epoch number
     setConversationEpoch (Data.convId (tUnqualified lconv)) (succ epoch)
 

services/galley/src/Galley/Effects/BrigAccess.hs

@@ -51,6 +51,7 @@ module Galley.Effects.BrigAccess
     getClientByKeyPackageRef,
     getLocalMLSClients,
     addKeyPackageRef,
+    updateKeyPackageRef,
 
     -- * Features
     getAccountConferenceCallingConfigClient,
@@ -128,6 +129,7 @@ data BrigAccess m a where
   GetClientByKeyPackageRef :: KeyPackageRef -> BrigAccess m (Maybe ClientIdentity)
   GetLocalMLSClients :: Local UserId -> SignatureSchemeTag -> BrigAccess m (Set ClientId)
   AddKeyPackageRef :: KeyPackageRef -> Qualified UserId -> ClientId -> Qualified ConvId -> BrigAccess m ()
+  UpdateKeyPackageRef :: KeyPackageUpdate -> BrigAccess m ()
   UpdateSearchVisibilityInbound ::
     Multi.TeamStatus SearchVisibilityInboundConfig ->
     BrigAccess m ()

services/galley/src/Galley/Intra/Client.hs

@@ -25,6 +25,7 @@ module Galley.Intra.Client
     getClientByKeyPackageRef,
     getLocalMLSClients,
     addKeyPackageRef,
+    updateKeyPackageRef,
   )
 where
 
@@ -210,3 +211,14 @@ addKeyPackageRef ref qusr cl qcnv =
           . json (NewKeyPackageRef qusr cl qcnv)
           . expect2xx
       )
+
+updateKeyPackageRef :: KeyPackageUpdate -> App ()
+updateKeyPackageRef keyPackageRef =
+  void $
+    call
+      Brig
+      ( method POST
+          . paths ["i", "mls", "key-packages", toHeader $ kpupPrevious keyPackageRef]
+          . json (kpupNext keyPackageRef)
+          . expect2xx
+      )

services/galley/src/Galley/Intra/Effects.hs

@@ -82,6 +82,9 @@ interpretBrigAccess = interpret $ \case
   AddKeyPackageRef ref qusr cl qcnv ->
     embedApp $
       addKeyPackageRef ref qusr cl qcnv
+  UpdateKeyPackageRef update ->
+    embedApp $
+      updateKeyPackageRef update
   UpdateSearchVisibilityInbound status ->
     embedApp $ updateSearchVisibilityInbound status