KeyPackage -> Conversation Internal API (#2375)

* New internal brig endpoints for galley use

* Tests

* chore(changelog)

changelog.d/5-internal/fs-532-brig

@@ -0,0 +1 @@
+New internal brig endpoints for MLS KeyPackage -> Conversation association query/update

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

@@ -34,6 +34,7 @@ where
 import Control.Lens ((.~))
 import qualified Data.Code as Code
 import Data.Id as Id
+import Data.Qualified (Qualified)
 import Data.Swagger (HasInfo (info), HasTitle (title), Swagger)
 import Imports hiding (head)
 import Servant hiding (Handler, JSON, addHeader, respond)
@@ -144,24 +145,61 @@ type AccountAPI =
         :> MultiVerb 'POST '[Servant.JSON] RegisterInternalResponses (Either RegisterError SelfProfile)
     )
 
-type MLSAPI = GetClientByKeyPackageRef :<|> GetMLSClients :<|> MapKeyPackageRefs
+type MLSAPI =
+  "mls"
+    :> ( ( "key-packages" :> Capture "ref" KeyPackageRef
+             :> ( Named
+                    "get-client-by-key-package-ref"
+                    ( Summary "Resolve an MLS key package ref to a qualified client ID"
+                        :> MultiVerb
+                             'GET
+                             '[Servant.JSON]
+                             '[ RespondEmpty 404 "Key package ref not found",
+                                Respond 200 "Key package ref found" ClientIdentity
+                              ]
+                             (Maybe ClientIdentity)
+                    )
+                    :<|> ( "conversation"
+                             :> ( PutConversationByKeyPackageRef
+                                    :<|> GetConversationByKeyPackageRef
+                                )
+                         )
+                )
+         )
+           :<|> GetMLSClients
+           :<|> MapKeyPackageRefs
+       )
+
+type PutConversationByKeyPackageRef =
+  Named
+    "put-conversation-by-key-package-ref"
+    ( Summary "Associate a conversation with a key package"
+        :> ReqBody '[Servant.JSON] (Qualified ConvId)
+        :> MultiVerb
+             'PUT
+             '[Servant.JSON]
+             [ RespondEmpty 404 "No key package found by reference",
+               RespondEmpty 204 "Converstaion associated"
+             ]
+             Bool
+    )
 
-type GetClientByKeyPackageRef =
-  Summary "Resolve an MLS key package ref to a qualified client ID"
-    :> "mls"
-    :> "key-packages"
-    :> Capture "ref" KeyPackageRef
-    :> MultiVerb
-         'GET
-         '[Servant.JSON]
-         '[ RespondEmpty 404 "Key package ref not found",
-            Respond 200 "Key package ref found" ClientIdentity
-          ]
-         (Maybe ClientIdentity)
+type GetConversationByKeyPackageRef =
+  Named
+    "get-conversation-by-key-package-ref"
+    ( Summary
+        "Retrieve the conversation associated with a key package"
+        :> MultiVerb
+             'GET
+             '[Servant.JSON]
+             [ RespondEmpty 404 "No associated conversation or bad key package",
+               Respond 200 "Conversation found" (Qualified ConvId)
+             ]
+             (Maybe (Qualified ConvId))
+    )
 
 type GetMLSClients =
   Summary "Return all MLS-enabled clients of a user"
-    :> "mls"
     :> "clients"
     :> CanThrow 'UserNotFound
     :> QualifiedCapture "user" UserId
@@ -173,7 +211,6 @@ type GetMLSClients =
 
 type MapKeyPackageRefs =
   Summary "Insert bundle into the KeyPackage ref mapping. Only for tests."
-    :> "mls"
     :> "key-package-refs"
     :> ReqBody '[Servant.JSON] KeyPackageBundle
     :> MultiVerb 'PUT '[Servant.JSON] '[RespondEmpty 204 "Mapping was updated"] ()

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

@@ -108,7 +108,14 @@ ejpdAPI =
     :<|> getConnectionsStatus
 
 mlsAPI :: ServerT BrigIRoutes.MLSAPI (Handler r)
-mlsAPI = getClientByKeyPackageRef :<|> getMLSClients :<|> mapKeyPackageRefsInternal
+mlsAPI =
+  ( \ref ->
+      Named @"get-client-by-key-package-ref" (getClientByKeyPackageRef ref)
+        :<|> Named @"put-conversation-by-key-package-ref" (putConvIdByKeyPackageRef ref)
+        :<|> Named @"get-conversation-by-key-package-ref" (getConvIdByKeyPackageRef ref)
+  )
+    :<|> getMLSClients
+    :<|> mapKeyPackageRefsInternal
 
 accountAPI :: ServerT BrigIRoutes.AccountAPI (Handler r)
 accountAPI = Named @"createUserNoVerify" createUserNoVerify
@@ -133,6 +140,14 @@ deleteAccountFeatureConfig uid =
 getClientByKeyPackageRef :: KeyPackageRef -> Handler r (Maybe ClientIdentity)
 getClientByKeyPackageRef = runMaybeT . mapMaybeT wrapClientE . Data.derefKeyPackage
 
+-- Used by galley to update conversation id in mls_key_package_ref
+putConvIdByKeyPackageRef :: KeyPackageRef -> Qualified ConvId -> Handler r Bool
+putConvIdByKeyPackageRef ref = lift . wrapClient . Data.keyPackageRefSetConvId ref
+
+-- Used by galley to retrieve conversation id from mls_key_package_ref
+getConvIdByKeyPackageRef :: KeyPackageRef -> Handler r (Maybe (Qualified ConvId))
+getConvIdByKeyPackageRef = runMaybeT . mapMaybeT wrapClientE . Data.keyPackageRefConvId
+
 getMLSClients :: Qualified UserId -> SignatureSchemeTag -> Handler r (Set ClientId)
 getMLSClients qusr ss = do
   usr <- lift $ tUnqualified <$> ensureLocal qusr

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

@@ -21,13 +21,18 @@ module Brig.Data.MLS.KeyPackage
     mapKeyPackageRef,
     countKeyPackages,
     derefKeyPackage,
+    keyPackageRefConvId,
+    keyPackageRefSetConvId,
   )
 where
 
 import Brig.App
 import Cassandra
+import Cassandra.Settings
 import Control.Error
+import Control.Exception
 import Control.Lens
+import Control.Monad.Catch
 import Control.Monad.Random (randomRIO)
 import Data.Domain
 import Data.Functor
@@ -97,6 +102,34 @@ derefKeyPackage ref = do
     q :: PrepQuery R (Identity KeyPackageRef) (Domain, UserId, ClientId)
     q = "SELECT domain, user, client from mls_key_package_refs WHERE ref = ?"
 
+keyPackageRefConvId :: MonadClient m => KeyPackageRef -> MaybeT m (Qualified ConvId)
+keyPackageRefConvId ref = MaybeT $ do
+  qr <- retry x1 $ query1 q (params LocalSerial (Identity ref))
+  pure $ do
+    (domain, cid) <- qr
+    Qualified <$> cid <*> domain
+  where
+    q :: PrepQuery R (Identity KeyPackageRef) (Maybe Domain, Maybe ConvId)
+    q = "SELECT conv_domain, conv FROM mls_key_package_refs WHERE ref = ?"
+
+-- We want to proper update, not an upsert, to avoid "ghost" refs without user+client
+keyPackageRefSetConvId :: MonadClient m => KeyPackageRef -> Qualified ConvId -> m Bool
+keyPackageRefSetConvId ref convId = do
+  updated <-
+    retry x5 $
+      trans
+        q
+        (params LocalQuorum (qDomain convId, qUnqualified convId, ref))
+          { serialConsistency = Just LocalSerialConsistency
+          }
+  case updated of
+    [] -> return False
+    [_] -> return True
+    _ -> throwM $ ErrorCall "Primary key violation detected mls_key_package_refs.ref"
+  where
+    q :: PrepQuery W (Domain, ConvId, KeyPackageRef) x
+    q = "UPDATE mls_key_package_refs SET conv_domain = ?, conv = ? WHERE ref = ? IF EXISTS"
+
 --------------------------------------------------------------------------------
 -- Utilities
 

services/brig/test/integration/API/Internal.hs

@@ -21,27 +21,35 @@ module API.Internal
 where
 
 import API.Internal.Util
+import API.MLS (createClient)
+import API.MLS.Util (SetKey (SetKey), uploadKeyPackages)
 import Bilge
 import Bilge.Assert
 import Brig.Data.User (lookupFeatureConferenceCalling, lookupStatus, userExists)
 import qualified Brig.Options as Opt
 import Brig.Types.Intra
-import Brig.Types.User (userId)
+import Brig.Types.User (User (userQualifiedId), userId)
 import qualified Cassandra as Cass
 import Control.Exception (ErrorCall (ErrorCall), throwIO)
 import Control.Lens ((^.), (^?!))
 import Control.Monad.Catch
+import Data.Aeson (decode)
 import qualified Data.Aeson.Lens as Aeson
 import qualified Data.Aeson.Types as Aeson
 import Data.ByteString.Conversion (toByteString')
 import Data.Id
+import Data.Qualified (Qualified (qDomain, qUnqualified))
 import qualified Data.Set as Set
 import Imports
+import Servant.API (ToHttpApiData (toUrlPiece))
+import Test.QuickCheck (Arbitrary (arbitrary), generate)
 import Test.Tasty
 import Test.Tasty.HUnit
+import UnliftIO (withSystemTempFile)
 import Util
 import Util.Options (Endpoint)
 import qualified Wire.API.Connection as Conn
+import Wire.API.MLS.KeyPackage
 import Wire.API.Routes.Internal.Brig.EJPD as EJPD
 import qualified Wire.API.Team.Feature as ApiFt
 import qualified Wire.API.Team.Member as Team
@@ -51,9 +59,17 @@ tests opts mgr db brig brigep gundeck galley = do
   return $
     testGroup "api/internal" $
       [ test mgr "ejpd requests" $ testEJPDRequest mgr brig brigep gundeck,
-        test mgr "account features: conferenceCalling" $ testFeatureConferenceCallingByAccount opts mgr db brig brigep galley,
+        test mgr "account features: conferenceCalling" $
+          testFeatureConferenceCallingByAccount opts mgr db brig brigep galley,
         test mgr "suspend and unsuspend user" $ testSuspendUser db brig,
-        test mgr "suspend non existing user and verify no db entry" $ testSuspendNonExistingUser db brig
+        test mgr "suspend non existing user and verify no db entry" $
+          testSuspendNonExistingUser db brig,
+        testGroup "mls/key-packages" $
+          [ test mgr "fresh get" $ testKpcFreshGet brig,
+            test mgr "put,get" $ testKpcPutGet brig,
+            test mgr "get,get" $ testKpcGetGet brig,
+            test mgr "put,put" $ testKpcPutPut brig
+          ]
       ]
 
 testSuspendUser :: forall m. TestConstraints m => Cass.ClientState -> Brig -> m ()
@@ -201,6 +217,87 @@ testFeatureConferenceCallingByAccount (Opt.optSettings -> settings) mgr db brig
   check $ ApiFt.TeamFeatureStatusNoConfig ApiFt.TeamFeatureDisabled
   check'
 
+keyPackageCreate :: HasCallStack => Brig -> Http KeyPackageRef
+keyPackageCreate brig = do
+  uid <- userQualifiedId <$> randomUser brig
+  clid <- createClient brig uid 0
+  withSystemTempFile "api.internal.kpc" $ \store _ ->
+    uploadKeyPackages brig store SetKey uid clid 2
+
+  uid2 <- userQualifiedId <$> randomUser brig
+  claimResp <-
+    post
+      ( brig
+          . paths
+            [ "mls",
+              "key-packages",
+              "claim",
+              toByteString' (qDomain uid),
+              toByteString' (qUnqualified uid)
+            ]
+          . zUser (qUnqualified uid2)
+          . contentJson
+      )
+  liftIO $
+    assertEqual "POST mls/key-packages/claim/:domain/:user failed" 200 (statusCode claimResp)
+  case responseBody claimResp >>= decode of
+    Nothing -> liftIO $ assertFailure "Claim response empty"
+    Just bundle -> case toList $ kpbEntries bundle of
+      [] -> liftIO $ assertFailure "Claim response held no bundles"
+      (h : _) -> return $ kpbeRef h
+
+kpcPut :: HasCallStack => Brig -> KeyPackageRef -> Qualified ConvId -> Http ()
+kpcPut brig ref qConv = do
+  resp <-
+    put
+      ( brig
+          . paths ["i", "mls", "key-packages", toByteString' $ toUrlPiece ref, "conversation"]
+          . contentJson
+          . json qConv
+      )
+  liftIO $ assertEqual "PUT i/mls/key-packages/:ref/conversation failed" 204 (statusCode resp)
+
+kpcGet :: HasCallStack => Brig -> KeyPackageRef -> Http (Maybe (Qualified ConvId))
+kpcGet brig ref = do
+  resp <-
+    get (brig . paths ["i", "mls", "key-packages", toByteString' $ toUrlPiece ref, "conversation"])
+  liftIO $ case statusCode resp of
+    404 -> return Nothing
+    200 -> return $ responseBody resp >>= decode
+    _ -> assertFailure "GET i/mls/key-packages/:ref/conversation failed"
+
+testKpcFreshGet :: Brig -> Http ()
+testKpcFreshGet brig = do
+  ref <- keyPackageCreate brig
+  mqConv <- kpcGet brig ref
+  liftIO $ assertEqual "(fresh) Get ~= Nothing" Nothing mqConv
+
+testKpcPutGet :: Brig -> Http ()
+testKpcPutGet brig = do
+  ref <- keyPackageCreate brig
+  qConv <- liftIO $ generate arbitrary
+  kpcPut brig ref qConv
+  mqConv <- kpcGet brig ref
+  liftIO $ assertEqual "Put x; Get ~= x" (Just qConv) mqConv
+
+testKpcGetGet :: Brig -> Http ()
+testKpcGetGet brig = do
+  ref <- keyPackageCreate brig
+  liftIO (generate arbitrary) >>= kpcPut brig ref
+  mqConv1 <- kpcGet brig ref
+  mqConv2 <- kpcGet brig ref
+  liftIO $ assertEqual "Get; Get ~= Get" mqConv1 mqConv2
+
+testKpcPutPut :: Brig -> Http ()
+testKpcPutPut brig = do
+  ref <- keyPackageCreate brig
+  qConv <- liftIO $ generate arbitrary
+  qConv2 <- liftIO $ generate arbitrary
+  kpcPut brig ref qConv
+  kpcPut brig ref qConv2
+  mqConv <- kpcGet brig ref
+  liftIO $ assertEqual "Put x; Put y ~= Put y" (Just qConv2) mqConv
+
 getFeatureConfig :: (MonadIO m, MonadHttp m, HasCallStack) => ApiFt.TeamFeatureName -> (Request -> Request) -> UserId -> m ResponseLBS
 getFeatureConfig feature galley uid = do
   get $ galley . paths ["feature-configs", toByteString' feature] . zUser uid