FS-1517 Partial success on fetch prekeys

changelog.d/1-api-changes/fs-1517

@@ -0,0 +1 @@
+Updating the V3 version of /users/list-prekeys to return partial successes, listing users that could not be listed.

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

@@ -651,18 +651,33 @@ type PrekeyAPI =
                :> Post '[JSON] UserClientPrekeyMap
            )
     :<|> Named
-           "get-multi-user-prekey-bundle-qualified"
+           "get-multi-user-prekey-bundle-qualified@v2"
            ( Summary
                "Given a map of domain to (map of user IDs to client IDs) return a \
                \prekey for each one. You can't request information for more users than \
                \maximum conversation size."
                :> MakesFederatedCall 'Brig "claim-multi-prekey-bundle"
                :> ZUser
+               :> Until 'V3
                :> "users"
                :> "list-prekeys"
                :> ReqBody '[JSON] QualifiedUserClients
                :> Post '[JSON] QualifiedUserClientPrekeyMap
            )
+    :<|> Named
+           "get-multi-user-prekey-bundle-qualified"
+           ( Summary
+               "Given a map of domain to (map of user IDs to client IDs) return a \
+               \prekey for each one. You can't request information for more users than \
+               \maximum conversation size."
+               :> MakesFederatedCall 'Brig "claim-multi-prekey-bundle"
+               :> ZUser
+               :> From 'V3
+               :> "users"
+               :> "list-prekeys"
+               :> ReqBody '[JSON] QualifiedUserClients
+               :> Post '[JSON] QualifiedUserClientPrekeyMapV3
+           )
 
 type UserClientAPI =
   -- User Client API ----------------------------------------------------

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

@@ -32,6 +32,7 @@ module Wire.API.User.Client
     mkUserClientPrekeyMap,
     QualifiedUserClientMap (..),
     QualifiedUserClientPrekeyMap (..),
+    QualifiedUserClientPrekeyMapV3 (..),
     mkQualifiedUserClientPrekeyMap,
     qualifiedUserClientPrekeyMapFromList,
     UserClientsFull (..),
@@ -278,8 +279,34 @@ qualifiedUserClientMapSchema sch =
               (schemaDoc innerSchema ^. Swagger.schema . Swagger.example)
           )
 
+data QualifiedUserClientPrekeyMapV3 = QualifiedUserClientPrekeyMapV3
+  { qualifiedUserClientPrekeys :: QualifiedUserClientMap (Maybe Prekey),
+    failedToList :: Maybe [Qualified UserId]
+  }
+  deriving stock (Eq, Show)
+  deriving (FromJSON, ToJSON, Swagger.ToSchema) via Schema QualifiedUserClientPrekeyMapV3
+
+instance Arbitrary QualifiedUserClientPrekeyMapV3 where
+  arbitrary =
+    QualifiedUserClientPrekeyMapV3
+      <$> arbitrary
+      <*> arbitrary
+
+instance ToSchema QualifiedUserClientPrekeyMapV3 where
+  schema =
+    object "QualifiedUserClientPrekeyMapV3" $
+      QualifiedUserClientPrekeyMapV3
+        <$> fmap to' (from' .= field "qualified_user_client_prekeys" (map_ schema))
+        <*> failedToList .= maybe_ (optField "failed_to_list" (array schema))
+    where
+      from' :: QualifiedUserClientPrekeyMapV3 -> Map Domain UserClientPrekeyMap
+      from' = coerce . qualifiedUserClientPrekeys
+      to' :: Map Domain UserClientPrekeyMap -> QualifiedUserClientMap (Maybe Prekey)
+      to' = coerce
+
 newtype QualifiedUserClientPrekeyMap = QualifiedUserClientPrekeyMap
-  {getQualifiedUserClientPrekeyMap :: QualifiedUserClientMap (Maybe Prekey)}
+  { getQualifiedUserClientPrekeyMap :: QualifiedUserClientMap (Maybe Prekey)
+  }
   deriving stock (Eq, Show)
   deriving newtype (Arbitrary)
   deriving (FromJSON, ToJSON, Swagger.ToSchema) via Schema QualifiedUserClientPrekeyMap

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

@@ -42,6 +42,7 @@ module Brig.API.Client
     claimLocalPrekey,
     claimPrekeyBundle,
     claimMultiPrekeyBundles,
+    claimMultiPrekeyBundlesV3,
     Data.lookupClientIds,
   )
 where
@@ -328,6 +329,52 @@ claimMultiPrekeyBundles protectee quc = do
       tUntagged . qualifyAs luc
         <$> claimLocalMultiPrekeyBundles protectee (tUnqualified luc)
 
+-- 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 messages couldn't be sent to.
+-- 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.
+claimMultiPrekeyBundlesV3 ::
+  forall r.
+  Member (Concurrency 'Unsafe) r =>
+  LegalholdProtectee ->
+  QualifiedUserClients ->
+  ExceptT ClientError (AppT r) QualifiedUserClientPrekeyMapV3
+claimMultiPrekeyBundlesV3 protectee quc = do
+  loc <- qualifyLocal ()
+  let (locals, remotes) =
+        partitionQualifiedAndTag
+          loc
+          ( map
+              (fmap UserClients . uncurry (flip Qualified))
+              (Map.assocs (qualifiedUserClients quc))
+          )
+  localPrekeys <- traverse claimLocal locals
+  remotePrekeys <- mapExceptT wrapHttpClient $ lift $ traverseConcurrently claimRemote remotes
+  let prekeys =
+        getQualifiedUserClientPrekeyMap $
+          qualifiedUserClientPrekeyMapFromList $
+            localPrekeys <> rights remotePrekeys
+      failed = lefts remotePrekeys >>= toQualifiedUser . fst
+  pure $
+    QualifiedUserClientPrekeyMapV3 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
+    claimRemote :: Remote UserClients -> ExceptT FederationError HttpClientIO (Qualified UserClientPrekeyMap)
+    claimRemote ruc =
+      tUntagged . qualifyAs ruc
+        <$> Federation.claimMultiPrekeyBundle (tDomain ruc) (tUnqualified ruc)
+
+    claimLocal :: Local UserClients -> ExceptT ClientError (AppT r) (Qualified UserClientPrekeyMap)
+    claimLocal luc =
+      tUntagged . qualifyAs luc
+        <$> claimLocalMultiPrekeyBundles protectee (tUnqualified luc)
+
 claimLocalMultiPrekeyBundles ::
   forall r.
   Member (Concurrency 'Unsafe) r =>

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

@@ -286,7 +286,8 @@ servantSitemap =
         :<|> Named @"get-users-prekey-bundle-unqualified" (callsFed (exposeAnnotations getPrekeyBundleUnqualifiedH))
         :<|> Named @"get-users-prekey-bundle-qualified" (callsFed (exposeAnnotations getPrekeyBundleH))
         :<|> Named @"get-multi-user-prekey-bundle-unqualified" getMultiUserPrekeyBundleUnqualifiedH
-        :<|> Named @"get-multi-user-prekey-bundle-qualified" (callsFed (exposeAnnotations getMultiUserPrekeyBundleH))
+        :<|> Named @"get-multi-user-prekey-bundle-qualified@v2" (callsFed (exposeAnnotations getMultiUserPrekeyBundleH))
+        :<|> Named @"get-multi-user-prekey-bundle-qualified" (callsFed (exposeAnnotations getMultiUserPrekeyBundleHV3))
 
     userClientAPI :: ServerT UserClientAPI (Handler r)
     userClientAPI =
@@ -477,6 +478,21 @@ getMultiUserPrekeyBundleH zusr qualUserClients = do
     throwStd (errorToWai @'E.TooManyClients)
   API.claimMultiPrekeyBundles (ProtectedUser zusr) qualUserClients !>> clientError
 
+getMultiUserPrekeyBundleHV3 ::
+  Member (Concurrency 'Unsafe) r =>
+  UserId ->
+  Public.QualifiedUserClients ->
+  (Handler r) Public.QualifiedUserClientPrekeyMapV3
+getMultiUserPrekeyBundleHV3 zusr qualUserClients = do
+  maxSize <- fromIntegral . setMaxConvSize <$> view settings
+  let Sum (size :: Int) =
+        Map.foldMapWithKey
+          (\_ v -> Sum . Map.size $ v)
+          (Public.qualifiedUserClients qualUserClients)
+  when (size > maxSize) $
+    throwStd (errorToWai @'E.TooManyClients)
+  API.claimMultiPrekeyBundlesV3 (ProtectedUser zusr) qualUserClients !>> clientError
+
 addClient ::
   (Member GalleyProvider r) =>
   UserId ->

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

@@ -22,6 +22,7 @@ module Brig.API.Util
     logInvitationCode,
     validateHandle,
     logEmail,
+    traverseConcurrently,
     traverseConcurrentlyWithErrors,
     traverseConcurrentlyWithErrorsSem,
     traverseConcurrentlyWithErrorsAppT,
@@ -44,6 +45,7 @@ import Brig.Types.Intra (accountUser)
 import Control.Lens (view)
 import Control.Monad.Catch (throwM)
 import Control.Monad.Trans.Except
+import Data.Bifunctor (Bifunctor (first))
 import Data.Domain (Domain)
 import Data.Handle (Handle, parseHandle)
 import Data.Id
@@ -109,6 +111,14 @@ traverseConcurrentlyWithErrors f =
           <=< pooledMapConcurrentlyN 8 (runExceptT . f)
       )
 
+traverseConcurrently ::
+  (Traversable t, MonadUnliftIO m) =>
+  (a -> ExceptT e m b) ->
+  t a ->
+  m (t (Either (a, e) b))
+traverseConcurrently f =
+  pooledMapConcurrentlyN 8 $ \a -> first (a,) <$> runExceptT (f a)
+
 -- | Traverse concurrently and fail on first error.
 traverseConcurrentlyWithErrorsSem ::
   forall t e a r b.