WPB-5989 polysemy create client subsystem (part 4 - final)

changelog.d/5-internal/WPB-5989

@@ -1 +1 @@
-Moved code from galley to ClientSubsystem (#5154, #5147, #5157, #5156, #5165)
+Moved code from galley to ClientSubsystem (#5154, #5147, #5157, #5156, #5165, #5168)

libs/wire-api/src/Wire/API/User/Client/Prekey.hs

@@ -141,7 +141,7 @@ data UncheckedPrekeyBundle = UncheckedPrekeyBundle
     -- | Prekey bundle
     prekeyKey :: Text
   }
-  deriving stock (Eq, Show, Generic)
+  deriving stock (Eq, Show, Generic, Ord)
   deriving (Arbitrary) via (GenericUniform UncheckedPrekeyBundle)
   deriving (FromJSON, ToJSON, S.ToSchema) via Schema UncheckedPrekeyBundle
 
@@ -247,7 +247,7 @@ decodePrekeyBundlePrekeyPayload = do
 
 newtype LastPrekey = LastPrekey
   {unpackLastPrekey :: UncheckedPrekeyBundle}
-  deriving stock (Eq, Show, Generic)
+  deriving stock (Eq, Show, Generic, Ord)
   deriving (FromJSON, ToJSON, S.ToSchema) via Schema LastPrekey
 
 instance ToSchema LastPrekey where

libs/wire-subsystems/src/Wire/ClientSubsystem.hs

@@ -54,5 +54,10 @@ data ClientSubsystem m a where
   -- Prekeys
   ClaimPrekey :: LegalholdProtectee -> UserId -> Domain -> ClientId -> ClientSubsystem m (Maybe ClientPrekey)
   ClaimLocalPrekey :: LegalholdProtectee -> UserId -> ClientId -> ClientSubsystem m (Maybe ClientPrekey)
+  ClaimPrekeyBundle :: LegalholdProtectee -> Domain -> UserId -> ClientSubsystem m PrekeyBundle
+  ClaimLocalPrekeyBundle :: LegalholdProtectee -> UserId -> ClientSubsystem m PrekeyBundle
+  ClaimMultiPrekeyBundlesV3 :: LegalholdProtectee -> QualifiedUserClients -> ClientSubsystem m QualifiedUserClientPrekeyMap
+  ClaimMultiPrekeyBundles :: LegalholdProtectee -> QualifiedUserClients -> ClientSubsystem m QualifiedUserClientPrekeyMapV4
+  ClaimLocalMultiPrekeyBundles :: LegalholdProtectee -> UserClients -> ClientSubsystem m UserClientPrekeyMap
 
 makeSem ''ClientSubsystem

libs/wire-subsystems/src/Wire/ClientSubsystem/Interpreter.hs

@@ -7,11 +7,13 @@ module Wire.ClientSubsystem.Interpreter
 where
 
 import Control.Monad
+import Data.Bifunctor
 import Data.ByteString.Conversion
 import Data.Default
 import Data.Domain
 import Data.Id
 import Data.Json.Util (ToJSONObject (..), toUTCTimeMillis)
+import Data.List.Extra (chunksOf)
 import Data.Map qualified as Map
 import Data.Misc
 import Data.Qualified
@@ -51,6 +53,7 @@ import Wire.Events as Events
 import Wire.FederationAPIAccess
 import Wire.GalleyAPIAccess as GalleyAPIAccess
 import Wire.NotificationSubsystem
+import Wire.Sem.Concurrency
 import Wire.Sem.Logger qualified as Log
 import Wire.Sem.Now qualified as Now
 import Wire.UserSubsystem (UserSubsystem)
@@ -75,7 +78,8 @@ runClientSubsystem ::
     Member EmailSubsystem r,
     Member DeleteQueue r,
     Member (Input ClientSubsystemConfig) r,
-    Member (Error FederationError) r
+    Member (Error FederationError) r,
+    Member (Concurrency 'Unsafe) r
   ) =>
   InterpreterFor AuthenticationSubsystem (UserSubsystem ': r) ->
   InterpreterFor UserSubsystem r ->
@@ -101,6 +105,11 @@ runClientSubsystem runAuth runUser =
       UpdateClient uid cid payload -> updateClient uid cid payload
       ClaimPrekey protectee uid domain cid -> claimPrekey protectee uid domain cid
       ClaimLocalPrekey protectee uid cid -> claimLocalPrekey protectee uid cid
+      ClaimPrekeyBundle protectee domain uid -> claimPrekeyBundle protectee domain uid
+      ClaimLocalPrekeyBundle protectee uid -> claimLocalPrekeyBundle protectee uid
+      ClaimMultiPrekeyBundlesV3 protectee qucs -> claimMultiPrekeyBundlesV3 protectee qucs
+      ClaimMultiPrekeyBundles protectee qucs -> claimMultiPrekeyBundles protectee qucs
+      ClaimLocalMultiPrekeyBundles protectee ucs -> claimLocalMultiPrekeyBundles protectee ucs
 
 -- nb. We must ensure that the set of clients known to brig is always
 -- a superset of the clients known to galley.
@@ -474,6 +483,174 @@ claimRemotePrekey (Qualified user domain) client = do
   Log.info $ msg @Text "Brig-federation: claiming remote prekey"
   runFederated (toRemoteUnsafe domain ()) $ fedClient @'Brig @"claim-prekey" (user, client)
 
+claimPrekeyBundle ::
+  ( Member ClientStore r,
+    Member TinyLog r,
+    HasBrigFederationAccess m r,
+    Member GalleyAPIAccess r,
+    Member (Input (Local ())) r,
+    Member (Error ClientError) r
+  ) =>
+  LegalholdProtectee -> Domain -> UserId -> Sem r PrekeyBundle
+claimPrekeyBundle protectee domain uid = do
+  isDomainLocal <- isLocalDomain domain
+  if isDomainLocal
+    then claimLocalPrekeyBundle protectee uid
+    else claimRemotePrekeyBundle (Qualified uid domain)
+
+claimLocalPrekeyBundle :: (Member ClientStore r, Member GalleyAPIAccess r) => LegalholdProtectee -> UserId -> Sem r PrekeyBundle
+claimLocalPrekeyBundle protectee u = do
+  clients <- map (.clientId) <$> ClientStore.lookupClients u
+  GalleyAPIAccess.guardLegalHold protectee (mkUserClients [(u, clients)])
+  PrekeyBundle u . catMaybes <$> mapM (ClientStore.claimPrekey u) clients
+
+claimRemotePrekeyBundle ::
+  ( Member TinyLog r,
+    HasBrigFederationAccess m r,
+    Member (Error ClientError) r
+  ) =>
+  Qualified UserId ->
+  Sem r PrekeyBundle
+claimRemotePrekeyBundle (Qualified user domain) = do
+  Log.info $ msg @Text "Brig-federation: claiming remote prekey bundle"
+  mapError ClientFederationError $ runFederated (toRemoteUnsafe domain ()) $ fedClient @'Brig @"claim-prekey-bundle" user
+
+claimMultiPrekeyBundlesInternal ::
+  forall r.
+  ( Member (Concurrency 'Unsafe) r,
+    Member ClientStore r,
+    Member (Input (Local ())) r,
+    Member GalleyAPIAccess r,
+    Member TinyLog r,
+    Member DeleteQueue r,
+    Member AuthenticationSubsystem r
+  ) =>
+  LegalholdProtectee ->
+  QualifiedUserClients ->
+  Sem r ([Qualified UserClientPrekeyMap], [Remote UserClients])
+claimMultiPrekeyBundlesInternal protectee quc = do
+  loc <- qualifyLocal ()
+  let (locals, remotes) =
+        partitionQualifiedAndTag
+          loc
+          ( map
+              (fmap UserClients . uncurry (flip Qualified))
+              (Map.assocs (qualifiedUserClients quc))
+          )
+  localPrekeys <- traverse claimLocal locals
+  pure (localPrekeys, remotes)
+  where
+    claimLocal ::
+      Local UserClients ->
+      Sem r (Qualified UserClientPrekeyMap)
+    claimLocal luc =
+      tUntagged . qualifyAs luc
+        <$> claimLocalMultiPrekeyBundles protectee (tUnqualified luc)
+
+claimMultiPrekeyBundlesV3 ::
+  forall r m.
+  ( Member (Concurrency 'Unsafe) r,
+    Member ClientStore r,
+    Member GalleyAPIAccess r,
+    Member TinyLog r,
+    HasBrigFederationAccess m r,
+    Member (Input (Local ())) r,
+    Member AuthenticationSubsystem r,
+    Member DeleteQueue r,
+    Member (Error FederationError) r
+  ) =>
+  LegalholdProtectee ->
+  QualifiedUserClients ->
+  Sem r QualifiedUserClientPrekeyMap
+claimMultiPrekeyBundlesV3 protectee quc = do
+  (localPrekeys, remotes) <- claimMultiPrekeyBundlesInternal protectee quc
+  Log.info $ msg @Text "Brig-federation: claiming remote multi-user prekey bundle"
+  remotePrekeys :: [Remote UserClientPrekeyMap] <- runFederatedConcurrently remotes $ \rucs -> fedClient @'Brig @"claim-multi-prekey-bundle" (mconcat $ tUnqualified rucs)
+  pure . qualifiedUserClientPrekeyMapFromList $ localPrekeys <> (fmap tUntagged remotePrekeys)
+
+-- Similar to claimMultiPrekeyBundles except for the following changes
+-- 1) A new return type that contains both the client map and a list of
+--    users that prekeys couldn't be fetched for.
+-- 2) A semantic change on federation errors when gathering remote clients.
+--    Remote federation errors at this step no-longer cause the entire call
+--    to fail, allowing partial results to be returned.
+claimMultiPrekeyBundles ::
+  forall r m.
+  ( Member (Concurrency 'Unsafe) r,
+    Member ClientStore r,
+    Member GalleyAPIAccess r,
+    Member TinyLog r,
+    HasBrigFederationAccess m r,
+    Member (Input (Local ())) r,
+    Member AuthenticationSubsystem r,
+    Member DeleteQueue r
+  ) =>
+  LegalholdProtectee ->
+  QualifiedUserClients ->
+  Sem r QualifiedUserClientPrekeyMapV4
+claimMultiPrekeyBundles protectee quc = do
+  (localPrekeys, remotes) <- claimMultiPrekeyBundlesInternal protectee quc
+  Log.info $ msg @Text "Brig-federation: claiming remote multi-user prekey bundle"
+  remotePrekeys <-
+    fmap (fmap (bimap (first collapseRemoteUsers) tUntagged)) $
+      runFederatedConcurrentlyEither remotes $ \rucs ->
+        fedClient @'Brig @"claim-multi-prekey-bundle" (mconcat $ tUnqualified rucs)
+  let prekeys =
+        getQualifiedUserClientPrekeyMap $
+          qualifiedUserClientPrekeyMapFromList $
+            localPrekeys <> rights remotePrekeys
+      failed = lefts remotePrekeys >>= toQualifiedUser . fst
+  pure $
+    QualifiedUserClientPrekeyMapV4 prekeys $
+      if null failed
+        then Nothing
+        else pure failed
+  where
+    toQualifiedUser :: Remote UserClients -> [Qualified UserId]
+    toQualifiedUser r = fmap (\u -> Qualified u $ tDomain r) . Map.keys . userClients . qUnqualified $ tUntagged r
+
+    collapseRemoteUsers :: Remote [UserClients] -> Remote UserClients
+    collapseRemoteUsers rucs = toRemoteUnsafe (tDomain rucs) (mconcat $ tUnqualified rucs)
+
+claimLocalMultiPrekeyBundles ::
+  forall r.
+  ( Member (Concurrency 'Unsafe) r,
+    Member ClientStore r,
+    Member GalleyAPIAccess r,
+    Member TinyLog r,
+    Member DeleteQueue r,
+    Member AuthenticationSubsystem r
+  ) =>
+  LegalholdProtectee ->
+  UserClients ->
+  Sem r UserClientPrekeyMap
+claimLocalMultiPrekeyBundles protectee ucs = do
+  GalleyAPIAccess.guardLegalHold protectee ucs
+  fmap mkUserClientPrekeyMap
+    . foldMap (getChunk . Map.fromList)
+    . chunksOf 16
+    . Map.toList
+    . userClients
+    $ ucs
+  where
+    getChunk :: Map UserId (Set ClientId) -> Sem r (Map UserId (Map ClientId (Maybe UncheckedPrekeyBundle)))
+    getChunk m =
+      Map.fromListWith (<>)
+        <$> unsafePooledMapConcurrentlyN
+          16
+          (\(u, cids) -> (u,) <$> getUserKeys u cids)
+          (Map.toList m)
+
+    getUserKeys :: UserId -> Set ClientId -> Sem r (Map ClientId (Maybe UncheckedPrekeyBundle))
+    getUserKeys u =
+      sequenceA . Map.fromSet (getClientKeys u)
+
+    getClientKeys :: UserId -> ClientId -> Sem r (Maybe UncheckedPrekeyBundle)
+    getClientKeys u c = do
+      key <- fmap prekeyData <$> ClientStore.claimPrekey u c
+      when (isNothing key) $ noPrekeys u c
+      pure key
+
 -- Utilities
 
 -- | Defensive measure when no prekey is found for a

libs/wire-subsystems/test/unit/Wire/ClientSubsystem/InterpreterSpec.hs

@@ -4,6 +4,7 @@ import Data.Aeson qualified as A
 import Data.Default
 import Data.Id
 import Data.Json.Util (toUTCTimeMillis)
+import Data.Map qualified as Map
 import Data.Qualified
 import Data.Set qualified as Set
 import Imports
@@ -366,6 +367,114 @@ spec = describe "ClientSubsystem.Interpreter" do
           Nothing -> counterexample "expected a client prekey, but got nothing" False
           Just pk -> pk.prekeyClient === clientId
 
+  prop "claim local prekey" $ \user (FakeLastPrekey lpk) ->
+    let uid = user.id
+        domain = testDomain
+        luid = toLocalUnsafe domain uid
+        new = newClient PermanentClientType lpk
+        clientId = clientIdFromPrekey (unpackLastPrekey lpk)
+        testResult =
+          runClientSubsystemTest [user] do
+            void $ addClient luid Nothing new
+            claimLocalPrekey (ProtectedUser uid) uid clientId
+     in expectRight testResult.result $ \case
+          Nothing -> counterexample "expected a client prekey, but got nothing" False
+          Just pk -> pk.prekeyClient === clientId
+
+  prop "claim prekey bundle" $ \user (FakeLastPrekey lpk1) (FakeLastPrekey lpk2) ->
+    (lpk1 /= lpk2)
+      ==> let uid = user.id
+              domain = testDomain
+              luid = toLocalUnsafe domain uid
+              new1 = newClient PermanentClientType lpk1
+              new2 = newClient PermanentClientType lpk2
+              clientId1 = clientIdFromPrekey (unpackLastPrekey lpk1)
+              clientId2 = clientIdFromPrekey (unpackLastPrekey lpk2)
+              expectedClientIds = Set.fromList [clientId1, clientId2]
+              testResult =
+                runClientSubsystemTest [user] do
+                  void $ addClient luid Nothing new1
+                  void $ addClient luid Nothing new2
+                  claimPrekeyBundle (ProtectedUser uid) domain uid
+           in expectRight testResult.result $ \bundle ->
+                (bundle.prekeyUser === uid)
+                  .&&. (Set.fromList (fmap (.prekeyClient) bundle.prekeyClients) === expectedClientIds)
+
+  prop "claim local prekey bundle" $ \user (FakeLastPrekey lpk1) (FakeLastPrekey lpk2) ->
+    (lpk1 /= lpk2)
+      ==> let uid = user.id
+              domain = testDomain
+              luid = toLocalUnsafe domain uid
+              new1 = newClient PermanentClientType lpk1
+              new2 = newClient PermanentClientType lpk2
+              clientId1 = clientIdFromPrekey (unpackLastPrekey lpk1)
+              clientId2 = clientIdFromPrekey (unpackLastPrekey lpk2)
+              expectedClientIds = Set.fromList [clientId1, clientId2]
+              testResult =
+                runClientSubsystemTest [user] do
+                  void $ addClient luid Nothing new1
+                  void $ addClient luid Nothing new2
+                  claimLocalPrekeyBundle (ProtectedUser uid) uid
+           in expectRight testResult.result $ \bundle ->
+                (bundle.prekeyUser === uid)
+                  .&&. (Set.fromList (fmap (.prekeyClient) bundle.prekeyClients) === expectedClientIds)
+
+  prop "claim multi prekey bundles v3" $ \protectee testData ->
+    (unique testData)
+      ==> let domain = testDomain
+              testResult =
+                runClientSubsystemTest (fmap fst testData) do
+                  for_ testData $ \(user, (FakeLastPrekey lpk)) -> do
+                    let uid = user.id
+                        luid = toLocalUnsafe domain uid
+                        new = newClient PermanentClientType lpk
+                    addClient luid Nothing new
+                  let qUserClients = QualifiedUserClients $ Map.fromList [(domain, Map.fromList (fmap toUserClients testData))]
+                  claimMultiPrekeyBundlesV3 (ProtectedUser protectee) qUserClients
+           in expectRight testResult.result $ \m ->
+                let qClientMap = m.getQualifiedUserClientPrekeyMap.qualifiedUserClientMap
+                    userMap = fromMaybe mempty $ Map.lookup domain qClientMap
+                 in Map.size qClientMap === 1 .&&. Map.size userMap === length testData
+
+  prop "claim multi prekey bundles" $ \protectee testData ->
+    (unique testData)
+      ==> let domain = testDomain
+              testResult =
+                runClientSubsystemTest (fmap fst testData) do
+                  for_ testData $ \(user, (FakeLastPrekey lpk)) -> do
+                    let uid = user.id
+                        luid = toLocalUnsafe domain uid
+                        new = newClient PermanentClientType lpk
+                    addClient luid Nothing new
+                  let qUserClients = QualifiedUserClients $ Map.fromList [(domain, Map.fromList (fmap toUserClients testData))]
+                  claimMultiPrekeyBundles (ProtectedUser protectee) qUserClients
+           in expectRight testResult.result $ \m ->
+                let qClientMap = m.qualifiedUserClientPrekeys.qualifiedUserClientMap
+                    userMap = fromMaybe mempty $ Map.lookup domain qClientMap
+                 in Map.size qClientMap === 1 .&&. Map.size userMap === length testData
+
+  prop "claim local multi prekey bundles" $ \protectee testData ->
+    (unique testData)
+      ==> let domain = testDomain
+              testResult =
+                runClientSubsystemTest (fmap fst testData) do
+                  for_ testData $ \(user, (FakeLastPrekey lpk)) -> do
+                    let uid = user.id
+                        luid = toLocalUnsafe domain uid
+                        new = newClient PermanentClientType lpk
+                    addClient luid Nothing new
+                  let userClients = UserClients $ Map.fromList (fmap toUserClients testData)
+                  claimLocalMultiPrekeyBundles (ProtectedUser protectee) userClients
+           in expectRight testResult.result $ \m ->
+                let clientMap = m.getUserClientPrekeyMap.userClientMap
+                 in Map.size clientMap === length testData
+  where
+    toUserClients (user, FakeLastPrekey lpk) = (user.id, Set.fromList [clientIdFromPrekey (unpackLastPrekey lpk)])
+
+    unique testData =
+      length testData == length (Set.fromList (fmap ((.id) . fst) testData))
+        && length testData == length (Set.fromList ((fmap snd) testData))
+
 newtype FakeUpdateClient = FakeUpdateClient {unFakeUpdateClient :: UpdateClient}
   deriving (Show, Eq, Generic)
 

libs/wire-subsystems/test/unit/Wire/Util.hs

@@ -43,7 +43,7 @@ anyElementOf :: NonEmptyList a -> Gen a
 anyElementOf = elements . toList . getNonEmpty
 
 newtype FakeLastPrekey = FakeLastPrekey {unFakeLastPrekey :: LastPrekey}
-  deriving (Show, Eq, Generic)
+  deriving (Show, Eq, Generic, Ord)
 
 instance Arbitrary FakeLastPrekey where
   arbitrary = FakeLastPrekey <$> QC.elements someLastPrekeys