Fix MLS message notification bug

changelog.d/3-bug-fixes/mls-notification-bug

@@ -0,0 +1 @@
+Fix bug where notifications for MLS messages were not showing up in all notification streams of clients

integration/test/Test/MLS/Message.hs

@@ -1,5 +1,8 @@
+{-# OPTIONS_GHC -Wno-ambiguous-fields #-}
+
 module Test.MLS.Message where
 
+import API.Gundeck
 import MLS.Util
 import Notifications
 import SetupHelpers
@@ -50,3 +53,33 @@ testAppMessageSomeReachable = do
     pure alice1
 
   void $ createApplicationMessage alice1 "hi, bob!" >>= sendAndConsumeMessage
+
+testMessageNotifications :: HasCallStack => Domain -> App ()
+testMessageNotifications bobDomain = do
+  [alice, bob] <- createAndConnectUsers [OwnDomain, bobDomain]
+
+  [alice1, alice2, bob1, bob2] <- traverse (createMLSClient def) [alice, alice, bob, bob]
+  bobClient <- bob1 %. "client_id" & asString
+
+  traverse_ uploadNewKeyPackage [alice1, alice2, bob1, bob2]
+
+  void $ createNewGroup alice1
+
+  void $ withWebSocket bob $ \ws -> do
+    void $ createAddCommit alice1 [alice, bob] >>= sendAndConsumeCommitBundle
+    awaitMatch 10 isMemberJoinNotif ws
+
+  let get (opts :: GetNotifications) = do
+        notifs <- getNotifications bob opts {size = Just 10000} >>= getJSON 200
+        notifs %. "has_more" `shouldMatch` False
+        length <$> (notifs %. "notifications" & asList)
+
+  numNotifs <- get def
+  numNotifsClient <- get def {client = Just bobClient}
+
+  void $ withWebSocket bob $ \ws -> do
+    void $ createApplicationMessage alice1 "hi bob" >>= sendAndConsumeMessage
+    awaitMatch 10 isNewMLSMessageNotif ws
+
+  get def `shouldMatchInt` (numNotifs + 1)
+  get def {client = Just bobClient} `shouldMatchInt` (numNotifsClient + 1)

libs/gundeck-types/src/Gundeck/Types/Push/V2.hs

@@ -125,6 +125,12 @@ data RecipientClients
     RecipientClientsSome (List1 ClientId)
   deriving (Eq, Show, Ord)
 
+instance Semigroup RecipientClients where
+  RecipientClientsAll <> _ = RecipientClientsAll
+  _ <> RecipientClientsAll = RecipientClientsAll
+  RecipientClientsSome cs1 <> RecipientClientsSome cs2 =
+    RecipientClientsSome (cs1 <> cs2)
+
 makeLenses ''Recipient
 
 recipient :: UserId -> Route -> Recipient

libs/wire-api-federation/src/Wire/API/Federation/API/Galley.hs

@@ -21,6 +21,7 @@ import Data.Aeson (FromJSON, ToJSON)
 import Data.Domain
 import Data.Id
 import Data.Json.Util
+import Data.List.NonEmpty (NonEmpty)
 import Data.Misc (Milliseconds)
 import Data.Qualified
 import Data.Range
@@ -348,7 +349,7 @@ data RemoteMLSMessage = RemoteMLSMessage
     sender :: Qualified UserId,
     conversation :: ConvId,
     subConversation :: Maybe SubConvId,
-    recipients :: [(UserId, ClientId)],
+    recipients :: Map UserId (NonEmpty ClientId),
     message :: Base64ByteString
   }
   deriving stock (Eq, Show, Generic)

services/galley/src/Galley/API/Federation.hs

@@ -27,6 +27,7 @@ import Data.ByteString.Conversion (toByteString')
 import Data.Domain (Domain)
 import Data.Id
 import Data.Json.Util
+import Data.List1 (List1 (..))
 import Data.Map qualified as Map
 import Data.Map.Lens (toMapOf)
 import Data.Qualified
@@ -57,10 +58,12 @@ import Galley.Effects
 import Galley.Effects.ConversationStore qualified as E
 import Galley.Effects.FireAndForget qualified as E
 import Galley.Effects.MemberStore qualified as E
+import Galley.Intra.Push.Internal hiding (push)
 import Galley.Options
 import Galley.Types.Conversations.Members
 import Galley.Types.Conversations.One2One
 import Galley.Types.UserList (UserList (UserList))
+import Gundeck.Types.Push.V2 (RecipientClients (..))
 import Imports
 import Polysemy
 import Polysemy.Error
@@ -781,7 +784,7 @@ onMLSMessageSent domain rmm =
       assertMLSEnabled
       loc <- qualifyLocal ()
       let rcnv = toRemoteUnsafe domain rmm.conversation
-      let users = Set.fromList (map fst rmm.recipients)
+      let users = Map.keys rmm.recipients
       (members, allMembers) <-
         first Set.fromList
           <$> E.selectRemoteMembers (toList users) rcnv
@@ -794,7 +797,11 @@ onMLSMessageSent domain rmm =
                 \ users not in the conversation" ::
                   ByteString
               )
-      let recipients = filter (\(u, _) -> Set.member u members) rmm.recipients
+      let recipients =
+            filter (\r -> Set.member (_recipientUserId r) members)
+              . map (\(u, clts) -> Recipient u (RecipientClientsSome (List1 clts)))
+              . Map.assocs
+              $ rmm.recipients
       -- FUTUREWORK: support local bots
       let e =
             Event (tUntagged rcnv) rmm.subConversation rmm.sender rmm.time $

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

@@ -20,6 +20,8 @@ module Galley.API.MLS.Propagate where
 import Control.Comonad
 import Data.Id
 import Data.Json.Util
+import Data.List.NonEmpty (NonEmpty, nonEmpty)
+import Data.List1
 import Data.Map qualified as Map
 import Data.Qualified
 import Data.Time
@@ -29,7 +31,9 @@ import Galley.API.Util
 import Galley.Data.Services
 import Galley.Effects
 import Galley.Effects.BackendNotificationQueueAccess
+import Galley.Intra.Push.Internal
 import Galley.Types.Conversations.Members
+import Gundeck.Types.Push.V2 (RecipientClients (..))
 import Imports
 import Network.AMQP qualified as Q
 import Polysemy
@@ -80,7 +84,7 @@ propagateMessage qusr mSenderClient lConvOrSub con msg cm = do
       e = Event qcnv sconv qusr now $ EdMLSMessage msg.raw
 
   runMessagePush lConvOrSub (Just qcnv) $
-    newMessagePush botMap con mm (lmems >>= localMemberMLSClients mlsConv) e
+    newMessagePush botMap con mm (lmems >>= toList . localMemberRecipient mlsConv) e
 
   -- send to remotes
   (either (logRemoteNotificationError @"on-mls-message-sent") (const (pure ())) <=< enqueueNotificationsConcurrently Q.Persistent (map remoteMemberQualify rmems)) $
@@ -92,23 +96,27 @@ propagateMessage qusr mSenderClient lConvOrSub con msg cm = do
             metadata = mm,
             conversation = qUnqualified qcnv,
             subConversation = sconv,
-            recipients = tUnqualified rs >>= remoteMemberMLSClients,
+            recipients =
+              Map.fromList $
+                tUnqualified rs
+                  >>= toList . remoteMemberMLSClients,
             message = Base64ByteString msg.raw
           }
   where
     cmWithoutSender = maybe cm (flip cmRemoveClient cm . mkClientIdentity qusr) mSenderClient
-    localMemberMLSClients :: Local x -> LocalMember -> [(UserId, ClientId)]
-    localMemberMLSClients loc lm =
+
+    localMemberRecipient :: Local x -> LocalMember -> Maybe Recipient
+    localMemberRecipient loc lm = do
       let localUserQId = tUntagged (qualifyAs loc localUserId)
           localUserId = lmId lm
-       in map
-            (\(c, _) -> (localUserId, c))
-            (Map.assocs (Map.findWithDefault mempty localUserQId cmWithoutSender))
+      clients <- nonEmpty $ Map.keys (Map.findWithDefault mempty localUserQId cmWithoutSender)
+      pure $ Recipient localUserId (RecipientClientsSome (List1 clients))
 
-    remoteMemberMLSClients :: RemoteMember -> [(UserId, ClientId)]
-    remoteMemberMLSClients rm =
+    remoteMemberMLSClients :: RemoteMember -> Maybe (UserId, NonEmpty ClientId)
+    remoteMemberMLSClients rm = do
       let remoteUserQId = tUntagged (rmId rm)
           remoteUserId = qUnqualified remoteUserQId
-       in map
-            (\(c, _) -> (remoteUserId, c))
-            (Map.assocs (Map.findWithDefault mempty remoteUserQId cmWithoutSender))
+      clients <-
+        nonEmpty . map fst $
+          Map.assocs (Map.findWithDefault mempty remoteUserQId cmWithoutSender)
+      pure (remoteUserId, clients)

services/galley/src/Galley/API/Push.hs

@@ -53,22 +53,28 @@ data MessagePush
 
 type BotMap = Map UserId BotMember
 
+class ToRecipient a where
+  toRecipient :: a -> Recipient
+
+instance ToRecipient (UserId, ClientId) where
+  toRecipient (u, c) = Recipient u (RecipientClientsSome (List1.singleton c))
+
+instance ToRecipient Recipient where
+  toRecipient = id
+
 newMessagePush ::
+  ToRecipient r =>
   BotMap ->
   Maybe ConnId ->
   MessageMetadata ->
-  [(UserId, ClientId)] ->
+  [r] ->
   Event ->
   MessagePush
 newMessagePush botMap mconn mm userOrBots event =
-  let (recipients, botMembers) =
-        foldMap
-          ( \(u, c) ->
-              case Map.lookup u botMap of
-                Just botMember -> ([], [botMember])
-                Nothing -> ([Recipient u (RecipientClientsSome (List1.singleton c))], [])
-          )
-          userOrBots
+  let toPair r = case Map.lookup (_recipientUserId r) botMap of
+        Just botMember -> ([], [botMember])
+        Nothing -> ([r], [])
+      (recipients, botMembers) = foldMap (toPair . toRecipient) userOrBots
    in MessagePush mconn mm recipients botMembers event
 
 runMessagePush ::
@@ -90,9 +96,9 @@ runMessagePush loc mqcnv mp@(MessagePush _ _ _ botMembers event) = do
       else deliverAndDeleteAsync (qUnqualified qcnv) (map (,event) botMembers)
 
 toPush :: MessagePush -> Maybe Push
-toPush (MessagePush mconn mm userRecipients _ event) =
+toPush (MessagePush mconn mm rs _ event) =
   let usr = qUnqualified (evtFrom event)
-   in newPush ListComplete (Just usr) (ConvEvent event) userRecipients
+   in newPush ListComplete (Just usr) (ConvEvent event) rs
         <&> set pushConn mconn
           . set pushNativePriority (mmNativePriority mm)
           . set pushRoute (bool RouteDirect RouteAny (mmNativePush mm))

services/galley/src/Galley/Intra/Push/Internal.hs

@@ -24,6 +24,7 @@ import Control.Lens (makeLenses, set, view, (.~))
 import Data.Aeson (Object)
 import Data.Id (ConnId, UserId)
 import Data.Json.Util
+import Data.List.Extra
 import Data.List.NonEmpty (NonEmpty, nonEmpty)
 import Data.List1
 import Data.Qualified

services/galley/test/integration/API/MLS.hs

@@ -32,6 +32,7 @@ import Data.Aeson qualified as Aeson
 import Data.Domain
 import Data.Id
 import Data.Json.Util hiding ((#))
+import Data.List.NonEmpty (NonEmpty (..))
 import Data.List1 hiding (head)
 import Data.Map qualified as Map
 import Data.Qualified
@@ -890,7 +891,7 @@ testRemoteToRemoteInSub = do
   void $ runFedClient @"on-conversation-updated" fedGalleyClient bdom cu
 
   let txt = "Hello from another backend"
-      rcpts = [(alice, aliceC1), (alice, aliceC2), (eve, eveC)]
+      rcpts = Map.fromList [(alice, aliceC1 :| [aliceC2]), (eve, eveC :| [])]
       rm =
         RemoteMLSMessage
           { time = now,