FS-1517 Partial success on fetch prekeys

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

@@ -0,0 +1 @@
+Updating the V4 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

@@ -667,18 +667,33 @@ type PrekeyAPI =
                :> Post '[JSON] UserClientPrekeyMap
            )
     :<|> Named
-           "get-multi-user-prekey-bundle-qualified"
+           "get-multi-user-prekey-bundle-qualified@v3"
            ( 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 'V4
                :> "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 'V4
+               :> "users"
+               :> "list-prekeys"
+               :> ReqBody '[JSON] QualifiedUserClients
+               :> Post '[JSON] QualifiedUserClientPrekeyMapV4
+           )
 
 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 (..),
+    QualifiedUserClientPrekeyMapV4 (..),
     mkQualifiedUserClientPrekeyMap,
     qualifiedUserClientPrekeyMapFromList,
     UserClientsFull (..),
@@ -278,8 +279,34 @@ qualifiedUserClientMapSchema sch =
               (schemaDoc innerSchema ^. Swagger.schema . Swagger.example)
           )
 
+data QualifiedUserClientPrekeyMapV4 = QualifiedUserClientPrekeyMapV4
+  { qualifiedUserClientPrekeys :: QualifiedUserClientMap (Maybe Prekey),
+    failedToList :: Maybe [Qualified UserId]
+  }
+  deriving stock (Eq, Show)
+  deriving (FromJSON, ToJSON, Swagger.ToSchema) via Schema QualifiedUserClientPrekeyMapV4
+
+instance Arbitrary QualifiedUserClientPrekeyMapV4 where
+  arbitrary =
+    QualifiedUserClientPrekeyMapV4
+      <$> arbitrary
+      <*> arbitrary
+
+instance ToSchema QualifiedUserClientPrekeyMapV4 where
+  schema =
+    object "QualifiedUserClientPrekeyMapV4" $
+      QualifiedUserClientPrekeyMapV4
+        <$> fmap to' (from' .= field "qualified_user_client_prekeys" (map_ schema))
+        <*> failedToList .= maybe_ (optField "failed_to_list" (array schema))
+    where
+      from' :: QualifiedUserClientPrekeyMapV4 -> 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

libs/wire-api/test/golden/Test/Wire/API/Golden/Generated.hs

@@ -154,6 +154,7 @@ import qualified Test.Wire.API.Golden.Generated.PubClient_user
 import qualified Test.Wire.API.Golden.Generated.PushTokenList_user
 import qualified Test.Wire.API.Golden.Generated.PushToken_user
 import qualified Test.Wire.API.Golden.Generated.Push_2eToken_2eTransport_user
+import qualified Test.Wire.API.Golden.Generated.QualifiedUserClientPrekeyMapV4_user
 import qualified Test.Wire.API.Golden.Generated.QueuedNotificationList_user
 import qualified Test.Wire.API.Golden.Generated.QueuedNotification_user
 import qualified Test.Wire.API.Golden.Generated.RTCConfiguration_user
@@ -1438,5 +1439,11 @@ tests =
             (Test.Wire.API.Golden.Generated.Event_conversation.testObject_Event_conversation_9, "testObject_Event_conversation_9.json"),
             (Test.Wire.API.Golden.Generated.Event_conversation.testObject_Event_conversation_11, "testObject_Event_conversation_11.json"),
             (Test.Wire.API.Golden.Generated.Event_conversation.testObject_Event_conversation_10, "testObject_Event_conversation_10.json")
+          ],
+      testGroup "Golden: QualifiedUserClientPrekeyMapV4" $
+        testObjects
+          [ (Test.Wire.API.Golden.Generated.QualifiedUserClientPrekeyMapV4_user.testObject_QualifiedUserClientPrekeyMapV4_user_1, "testObject_QualifiedUserClientPrekeyMapV4_1.json"),
+            (Test.Wire.API.Golden.Generated.QualifiedUserClientPrekeyMapV4_user.testObject_QualifiedUserClientPrekeyMapV4_user_2, "testObject_QualifiedUserClientPrekeyMapV4_2.json"),
+            (Test.Wire.API.Golden.Generated.QualifiedUserClientPrekeyMapV4_user.testObject_QualifiedUserClientPrekeyMapV4_user_3, "testObject_QualifiedUserClientPrekeyMapV4_3.json")
           ]
     ]

libs/wire-api/test/golden/Test/Wire/API/Golden/Generated/QualifiedUserClientPrekeyMapV4_user.hs

@@ -0,0 +1,58 @@
+-- This file is part of the Wire Server implementation.
+--
+-- Copyright (C) 2022 Wire Swiss GmbH <opensource@wire.com>
+--
+-- This program is free software: you can redistribute it and/or modify it under
+-- the terms of the GNU Affero General Public License as published by the Free
+-- Software Foundation, either version 3 of the License, or (at your option) any
+-- later version.
+--
+-- This program is distributed in the hope that it will be useful, but WITHOUT
+-- ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
+-- FOR A PARTICULAR PURPOSE. See the GNU Affero General Public License for more
+-- details.
+--
+-- You should have received a copy of the GNU Affero General Public License along
+-- with this program. If not, see <https://www.gnu.org/licenses/>.
+
+module Test.Wire.API.Golden.Generated.QualifiedUserClientPrekeyMapV4_user where
+
+import Data.Domain (Domain (..))
+import Data.Id (ClientId (..), Id (Id), UserId)
+import qualified Data.Map as Map
+import Data.Qualified (Qualified (..))
+import qualified Data.UUID as UUID (fromString)
+import Imports
+import Wire.API.User.Client (QualifiedUserClientMap (..), QualifiedUserClientPrekeyMapV4 (..))
+
+domain1, domain2 :: Domain
+domain1 = Domain "example.com"
+domain2 = Domain "test.net"
+
+user1, user2 :: UserId
+user1 = Id . fromJust $ UUID.fromString "44f9c51e-0dce-4e7f-85ba-b4e5a545ce68"
+user2 = Id . fromJust $ UUID.fromString "284c4e8f-78ef-43f4-a77a-015c22e37960"
+
+clientId :: ClientId
+clientId = ClientId "0123456789ABCEF"
+
+testObject_QualifiedUserClientPrekeyMapV4_user_1 :: QualifiedUserClientPrekeyMapV4
+testObject_QualifiedUserClientPrekeyMapV4_user_1 =
+  QualifiedUserClientPrekeyMapV4
+    { qualifiedUserClientPrekeys = QualifiedUserClientMap mempty,
+      failedToList = Nothing
+    }
+
+testObject_QualifiedUserClientPrekeyMapV4_user_2 :: QualifiedUserClientPrekeyMapV4
+testObject_QualifiedUserClientPrekeyMapV4_user_2 =
+  QualifiedUserClientPrekeyMapV4
+    { qualifiedUserClientPrekeys = QualifiedUserClientMap $ Map.singleton domain1 $ Map.singleton user1 $ Map.singleton clientId Nothing,
+      failedToList = Just []
+    }
+
+testObject_QualifiedUserClientPrekeyMapV4_user_3 :: QualifiedUserClientPrekeyMapV4
+testObject_QualifiedUserClientPrekeyMapV4_user_3 =
+  QualifiedUserClientPrekeyMapV4
+    { qualifiedUserClientPrekeys = QualifiedUserClientMap mempty,
+      failedToList = Just [Qualified user1 domain1, Qualified user2 domain2]
+    }

libs/wire-api/test/golden/testObject_QualifiedUserClientPrekeyMapV4_1.json

@@ -0,0 +1,2 @@
+{ "qualified_user_client_prekeys" : {}
+}

libs/wire-api/test/golden/testObject_QualifiedUserClientPrekeyMapV4_2.json

@@ -0,0 +1,9 @@
+{ "qualified_user_client_prekeys" : {
+    "example.com" : {
+      "44f9c51e-0dce-4e7f-85ba-b4e5a545ce68" : {
+        "0123456789ABCEF" : null
+      }
+    }
+  }
+, "failed_to_list" : []
+}

libs/wire-api/test/golden/testObject_QualifiedUserClientPrekeyMapV4_3.json

@@ -0,0 +1,10 @@
+{ "qualified_user_client_prekeys" : {}
+, "failed_to_list" :
+  [ { "domain" : "example.com"
+    , "id" : "44f9c51e-0dce-4e7f-85ba-b4e5a545ce68"
+    }
+  , { "domain" : "test.net"
+    , "id" : "284c4e8f-78ef-43f4-a77a-015c22e37960"
+    }
+  ]
+}

libs/wire-api/wire-api.cabal

@@ -454,6 +454,7 @@ test-suite wire-api-golden-tests
     Test.Wire.API.Golden.Generated.PushToken_user
     Test.Wire.API.Golden.Generated.PushTokenList_user
     Test.Wire.API.Golden.Generated.QualifiedNewOtrMessage_user
+    Test.Wire.API.Golden.Generated.QualifiedUserClientPrekeyMapV4_user
     Test.Wire.API.Golden.Generated.QueuedNotification_user
     Test.Wire.API.Golden.Generated.QueuedNotificationList_user
     Test.Wire.API.Golden.Generated.ReceiptMode_user

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

@@ -42,6 +42,7 @@ module Brig.API.Client
     claimLocalPrekey,
     claimPrekeyBundle,
     claimMultiPrekeyBundles,
+    claimMultiPrekeyBundlesV3,
     Data.lookupClientIds,
   )
 where
@@ -288,13 +289,15 @@ claimRemotePrekeyBundle :: Qualified UserId -> ExceptT ClientError (AppT r) Prek
 claimRemotePrekeyBundle quser = do
   Federation.claimPrekeyBundle quser !>> ClientFederationError
 
-claimMultiPrekeyBundles ::
-  forall r.
-  (Member (Concurrency 'Unsafe) r) =>
+claimMultiPrekeyBundlesInternal ::
+  Member (Concurrency 'Unsafe) r =>
   LegalholdProtectee ->
   QualifiedUserClients ->
-  ExceptT ClientError (AppT r) QualifiedUserClientPrekeyMap
-claimMultiPrekeyBundles protectee quc = do
+  ExceptT
+    ClientError
+    (AppT r)
+    ([Qualified UserClientPrekeyMap], [Remote UserClients])
+claimMultiPrekeyBundlesInternal protectee quc = do
   loc <- qualifyLocal ()
   let (locals, remotes) =
         partitionQualifiedAndTag
@@ -304,6 +307,23 @@ claimMultiPrekeyBundles protectee quc = do
               (Map.assocs (qualifiedUserClients quc))
           )
   localPrekeys <- traverse claimLocal locals
+  pure (localPrekeys, remotes)
+  where
+    claimLocal ::
+      Member (Concurrency 'Unsafe) r =>
+      Local UserClients ->
+      ExceptT ClientError (AppT r) (Qualified UserClientPrekeyMap)
+    claimLocal luc =
+      tUntagged . qualifyAs luc
+        <$> claimLocalMultiPrekeyBundles protectee (tUnqualified luc)
+
+claimMultiPrekeyBundlesV3 ::
+  Member (Concurrency 'Unsafe) r =>
+  LegalholdProtectee ->
+  QualifiedUserClients ->
+  ExceptT ClientError (AppT r) QualifiedUserClientPrekeyMap
+claimMultiPrekeyBundlesV3 protectee quc = do
+  (localPrekeys, remotes) <- claimMultiPrekeyBundlesInternal protectee quc
   remotePrekeys <-
     mapExceptT wrapHttpClient $
       traverseConcurrentlyWithErrors
@@ -323,10 +343,38 @@ claimMultiPrekeyBundles protectee quc = do
       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)
+-- 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.
+  Member (Concurrency 'Unsafe) r =>
+  LegalholdProtectee ->
+  QualifiedUserClients ->
+  ExceptT ClientError (AppT r) QualifiedUserClientPrekeyMapV4
+claimMultiPrekeyBundles protectee quc = do
+  (localPrekeys, remotes) <- claimMultiPrekeyBundlesInternal protectee quc
+  remotePrekeys <- mapExceptT wrapHttpClient $ lift $ traverseConcurrentlySem claimRemote remotes
+  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
+    claimRemote :: Remote UserClients -> ExceptT FederationError HttpClientIO (Qualified UserClientPrekeyMap)
+    claimRemote ruc =
+      tUntagged . qualifyAs ruc
+        <$> Federation.claimMultiPrekeyBundle (tDomain ruc) (tUnqualified ruc)
 
 claimLocalMultiPrekeyBundles ::
   forall r.

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

@@ -73,6 +73,7 @@ import qualified Cassandra as Data
 import Control.Error hiding (bool)
 import Control.Lens (view, (.~), (?~), (^.))
 import Control.Monad.Catch (throwM)
+import Control.Monad.Except
 import Data.Aeson hiding (json)
 import Data.Bifunctor
 import qualified Data.ByteString.Lazy as Lazy
@@ -297,6 +298,7 @@ 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@v3" (callsFed (exposeAnnotations getMultiUserPrekeyBundleHV3))
         :<|> Named @"get-multi-user-prekey-bundle-qualified" (callsFed (exposeAnnotations getMultiUserPrekeyBundleH))
 
     userClientAPI :: ServerT UserClientAPI (Handler r)
@@ -475,19 +477,35 @@ getMultiUserPrekeyBundleUnqualifiedH zusr userClients = do
     throwStd (errorToWai @'E.TooManyClients)
   API.claimLocalMultiPrekeyBundles (ProtectedUser zusr) userClients !>> clientError
 
-getMultiUserPrekeyBundleH ::
-  (Member (Concurrency 'Unsafe) r) =>
-  UserId ->
+getMultiUserPrekeyBundleHInternal ::
+  (MonadReader Env m, MonadError Brig.API.Error.Error m) =>
   Public.QualifiedUserClients ->
-  (Handler r) Public.QualifiedUserClientPrekeyMap
-getMultiUserPrekeyBundleH zusr qualUserClients = do
+  m ()
+getMultiUserPrekeyBundleHInternal 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)
+
+getMultiUserPrekeyBundleHV3 ::
+  Member (Concurrency 'Unsafe) r =>
+  UserId ->
+  Public.QualifiedUserClients ->
+  (Handler r) Public.QualifiedUserClientPrekeyMap
+getMultiUserPrekeyBundleHV3 zusr qualUserClients = do
+  getMultiUserPrekeyBundleHInternal qualUserClients
+  API.claimMultiPrekeyBundlesV3 (ProtectedUser zusr) qualUserClients !>> clientError
+
+getMultiUserPrekeyBundleH ::
+  Member (Concurrency 'Unsafe) r =>
+  UserId ->
+  Public.QualifiedUserClients ->
+  (Handler r) Public.QualifiedUserClientPrekeyMapV4
+getMultiUserPrekeyBundleH zusr qualUserClients = do
+  getMultiUserPrekeyBundleHInternal qualUserClients
   API.claimMultiPrekeyBundles (ProtectedUser zusr) qualUserClients !>> clientError
 
 addClient ::

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

@@ -1451,7 +1451,7 @@ lookupProfilesV3 ::
   AppT r ([(Qualified UserId, FederationError)], [UserProfile])
 lookupProfilesV3 self others = do
   t <-
-    traverseConcurrently
+    traverseConcurrentlyAppT
       (lookupProfilesFromDomain self)
       (bucketQualified others)
   let (l, r) = partitionEithers t

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

@@ -22,7 +22,8 @@ module Brig.API.Util
     logInvitationCode,
     validateHandle,
     logEmail,
-    traverseConcurrently,
+    traverseConcurrentlyAppT,
+    traverseConcurrentlySem,
     traverseConcurrentlyWithErrors,
     traverseConcurrentlyWithErrorsSem,
     traverseConcurrentlyWithErrorsAppT,
@@ -99,12 +100,12 @@ logInvitationCode :: InvitationCode -> (Msg -> Msg)
 logInvitationCode code = Log.field "invitation_code" (toText $ fromInvitationCode code)
 
 -- | Traverse concurrently and collect errors.
-traverseConcurrently ::
+traverseConcurrentlyAppT ::
   (Traversable t, Member (C.Concurrency 'C.Unsafe) r) =>
   (a -> ExceptT e (AppT r) b) ->
   t a ->
   AppT r [Either (a, e) b]
-traverseConcurrently f t = do
+traverseConcurrentlyAppT f t = do
   env <- temporaryGetEnv
   AppT $
     lift $
@@ -126,6 +127,14 @@ traverseConcurrentlyWithErrors f =
           <=< pooledMapConcurrentlyN 8 (runExceptT . f)
       )
 
+traverseConcurrentlySem ::
+  (Traversable t, MonadUnliftIO m) =>
+  (a -> ExceptT e m b) ->
+  t a ->
+  m (t (Either (a, e) b))
+traverseConcurrentlySem f =
+  pooledMapConcurrentlyN 8 $ \a -> first (a,) <$> runExceptT (f a)
+
 -- | Traverse concurrently and fail on first error.
 traverseConcurrentlyWithErrorsSem ::
   forall t e a r b.