zilencer: Have push/notify endpoint return registrations to delete.

- The server sends the list of registrations it believes to have with the bouncer. - The bouncer includes in the response the registrations that it doesn't actually have and therefore the server should delete.
2025-11-04 05:53:43 +00:00 · 2023-12-12 00:06:37 +01:00
parent dd8a33f03e
commit b82ea179ac
3 changed files with 174 additions and 23 deletions
--- a/zerver/lib/push_notifications.py
+++ b/zerver/lib/push_notifications.py
@@ -586,6 +586,8 @@ def send_notifications_to_bouncer(
    apns_payload: Dict[str, Any],
    gcm_payload: Dict[str, Any],
    gcm_options: Dict[str, Any],
+    android_devices: Sequence[DeviceToken],
+    apple_devices: Sequence[DeviceToken],
 ) -> Tuple[int, int]:
    post_data = {
        "user_uuid": str(user_profile.uuid),
@@ -596,12 +598,32 @@ def send_notifications_to_bouncer(
        "apns_payload": apns_payload,
        "gcm_payload": gcm_payload,
        "gcm_options": gcm_options,
+        "android_devices": [device.token for device in android_devices],
+        "apple_devices": [device.token for device in apple_devices],
    }
    # Calls zilencer.views.remote_server_notify_push
    response_data = send_json_to_push_bouncer("POST", "push/notify", post_data)
    assert isinstance(response_data["total_android_devices"], int)
    assert isinstance(response_data["total_apple_devices"], int)

+    assert isinstance(response_data["deleted_devices"], dict)
+    assert isinstance(response_data["deleted_devices"]["android_devices"], list)
+    assert isinstance(response_data["deleted_devices"]["apple_devices"], list)
+    android_deleted_devices = response_data["deleted_devices"]["android_devices"]
+    apple_deleted_devices = response_data["deleted_devices"]["apple_devices"]
+    if android_deleted_devices or apple_deleted_devices:
+        logger.info(
+            "Deleting push tokens based on response from bouncer: Android: %s, Apple: %s",
+            sorted(android_deleted_devices),
+            sorted(apple_deleted_devices),
+        )
+        PushDeviceToken.objects.filter(
+            kind=PushDeviceToken.GCM, token__in=android_deleted_devices
+        ).delete()
+        PushDeviceToken.objects.filter(
+            kind=PushDeviceToken.APNS, token__in=apple_deleted_devices
+        ).delete()
+
    total_android_devices, total_apple_devices = (
        response_data["total_android_devices"],
        response_data["total_apple_devices"],
@@ -1162,16 +1184,18 @@ def handle_remove_push_notification(user_profile_id: int, message_ids: List[int]
    gcm_payload, gcm_options = get_remove_payload_gcm(user_profile, truncated_message_ids)
    apns_payload = get_remove_payload_apns(user_profile, truncated_message_ids)

+    android_devices = list(
+        PushDeviceToken.objects.filter(user=user_profile, kind=PushDeviceToken.GCM)
+    )
+    apple_devices = list(
+        PushDeviceToken.objects.filter(user=user_profile, kind=PushDeviceToken.APNS)
+    )
    if uses_notification_bouncer():
-        send_notifications_to_bouncer(user_profile, apns_payload, gcm_payload, gcm_options)
+        send_notifications_to_bouncer(
+            user_profile, apns_payload, gcm_payload, gcm_options, android_devices, apple_devices
+        )
    else:
        user_identity = UserPushIdentityCompat(user_id=user_profile_id)
-        android_devices = list(
-            PushDeviceToken.objects.filter(user=user_profile, kind=PushDeviceToken.GCM)
-        )
-        apple_devices = list(
-            PushDeviceToken.objects.filter(user=user_profile, kind=PushDeviceToken.APNS)
-        )

        android_successfully_sent_count = send_android_push_notification(
            user_identity, android_devices, gcm_payload, gcm_options
@@ -1326,9 +1350,16 @@ def handle_push_notification(user_profile_id: int, missed_message: Dict[str, Any
    )
    logger.info("Sending push notifications to mobile clients for user %s", user_profile_id)

+    android_devices = list(
+        PushDeviceToken.objects.filter(user=user_profile, kind=PushDeviceToken.GCM)
+    )
+
+    apple_devices = list(
+        PushDeviceToken.objects.filter(user=user_profile, kind=PushDeviceToken.APNS)
+    )
    if uses_notification_bouncer():
        total_android_devices, total_apple_devices = send_notifications_to_bouncer(
-            user_profile, apns_payload, gcm_payload, gcm_options
+            user_profile, apns_payload, gcm_payload, gcm_options, android_devices, apple_devices
        )
        logger.info(
            "Sent mobile push notifications for user %s through bouncer: %s via FCM devices, %s via APNs devices",
@@ -1338,14 +1369,6 @@ def handle_push_notification(user_profile_id: int, missed_message: Dict[str, Any
        )
        return

-    android_devices = list(
-        PushDeviceToken.objects.filter(user=user_profile, kind=PushDeviceToken.GCM)
-    )
-
-    apple_devices = list(
-        PushDeviceToken.objects.filter(user=user_profile, kind=PushDeviceToken.APNS)
-    )
-
    logger.info(
        "Sending mobile push notifications for local user %s: %s via FCM devices, %s via APNs devices",
        user_profile_id,
--- a/zerver/tests/test_push_notifications.py
+++ b/zerver/tests/test_push_notifications.py
@@ -93,6 +93,7 @@ from zerver.models import (
    get_stream,
 )
 from zilencer.models import RemoteZulipServerAuditLog
+from zilencer.views import DevicesToCleanUpDict

 if settings.ZILENCER_ENABLED:
    from zilencer.models import (
@@ -612,7 +613,13 @@ class PushBouncerNotificationTest(BouncerTestCase):
            )
        data = self.assert_json_success(result)
        self.assertEqual(
-            {"result": "success", "msg": "", "total_android_devices": 2, "total_apple_devices": 1},
+            {
+                "result": "success",
+                "msg": "",
+                "total_android_devices": 2,
+                "total_apple_devices": 1,
+                "deleted_devices": {"android_devices": [], "apple_devices": []},
+            },
            data,
        )
        self.assertEqual(
@@ -2146,9 +2153,38 @@ class HandlePushNotificationTest(PushNotificationTest):
                (b64_to_hex(device.token), device.ios_app_id, device.token)
                for device in RemotePushDeviceToken.objects.filter(kind=PushDeviceToken.GCM)
            ]
+
+            # Reset the local registrations for the user to make them compatible
+            # with the RemotePushDeviceToken entries.
+            PushDeviceToken.objects.filter(kind=PushDeviceToken.APNS).delete()
+            [
+                PushDeviceToken.objects.create(
+                    kind=PushDeviceToken.APNS,
+                    token=device.token,
+                    user=self.user_profile,
+                    ios_app_id=device.ios_app_id,
+                )
+                for device in RemotePushDeviceToken.objects.filter(kind=PushDeviceToken.APNS)
+            ]
+            PushDeviceToken.objects.filter(kind=PushDeviceToken.GCM).delete()
+            [
+                PushDeviceToken.objects.create(
+                    kind=PushDeviceToken.GCM,
+                    token=device.token,
+                    user=self.user_profile,
+                    ios_app_id=device.ios_app_id,
+                )
+                for device in RemotePushDeviceToken.objects.filter(kind=PushDeviceToken.GCM)
+            ]
+
            mock_gcm.json_request.return_value = {"success": {gcm_devices[0][2]: message.id}}
            send_notification.return_value.is_successful = False
            send_notification.return_value.description = "Unregistered"
+
+            # Ensure the setup is as expected:
+            self.assertNotEqual(
+                PushDeviceToken.objects.filter(kind=PushDeviceToken.APNS).count(), 0
+            )
            handle_push_notification(self.user_profile.id, missed_message)
            self.assertEqual(
                views_logger.output,
@@ -2167,9 +2203,16 @@ class HandlePushNotificationTest(PushNotificationTest):
                    f"APNs: Removing invalid/expired token {token} (Unregistered)",
                    pn_logger.output,
                )
+            self.assertIn(
+                "INFO:zerver.lib.push_notifications:Deleting push tokens based on response from bouncer: "
+                f"Android: [], Apple: {sorted([token for _, _ , token in apns_devices])}",
+                pn_logger.output,
+            )
            self.assertEqual(
                RemotePushDeviceToken.objects.filter(kind=PushDeviceToken.APNS).count(), 0
            )
+            # Local registrations have also been deleted:
+            self.assertEqual(PushDeviceToken.objects.filter(kind=PushDeviceToken.APNS).count(), 0)

    @override_settings(PUSH_NOTIFICATION_BOUNCER_URL="https://push.zulip.org.example.com")
    @responses.activate
@@ -2316,7 +2359,10 @@ class HandlePushNotificationTest(PushNotificationTest):
            )

    def test_send_notifications_to_bouncer(self) -> None:
-        user_profile = self.example_user("hamlet")
+        self.setup_apns_tokens()
+        self.setup_gcm_tokens()
+
+        user_profile = self.user_profile
        message = self.get_message(
            Recipient.PERSONAL,
            type_id=self.personal_recipient_user.id,
@@ -2347,6 +2393,8 @@ class HandlePushNotificationTest(PushNotificationTest):
                {"apns": True},
                {"gcm": True},
                {},
+                list(PushDeviceToken.objects.filter(user=user_profile, kind=PushDeviceToken.GCM)),
+                list(PushDeviceToken.objects.filter(user=user_profile, kind=PushDeviceToken.APNS)),
            )
            self.assertEqual(
                mock_logging_info.output,
@@ -2414,7 +2462,10 @@ class HandlePushNotificationTest(PushNotificationTest):
        )

    def test_send_remove_notifications_to_bouncer(self) -> None:
-        user_profile = self.example_user("hamlet")
+        self.setup_apns_tokens()
+        self.setup_gcm_tokens()
+
+        user_profile = self.user_profile
        message = self.get_message(
            Recipient.PERSONAL,
            type_id=self.personal_recipient_user.id,
@@ -2457,6 +2508,8 @@ class HandlePushNotificationTest(PushNotificationTest):
                    "zulip_message_id": message.id,
                },
                {"priority": "normal"},
+                list(PushDeviceToken.objects.filter(user=user_profile, kind=PushDeviceToken.GCM)),
+                list(PushDeviceToken.objects.filter(user=user_profile, kind=PushDeviceToken.APNS)),
            )
            user_message = UserMessage.objects.get(user_profile=self.user_profile, message=message)
            self.assertEqual(user_message.flags.active_mobile_push_notification, False)
@@ -3507,14 +3560,35 @@ class TestGetGCMPayload(PushNotificationTest):
        )


-class TestSendNotificationsToBouncer(ZulipTestCase):
+class TestSendNotificationsToBouncer(PushNotificationTest):
    @mock.patch("zerver.lib.remote_server.send_to_push_bouncer")
    def test_send_notifications_to_bouncer(self, mock_send: mock.MagicMock) -> None:
        user = self.example_user("hamlet")

-        mock_send.return_value = {"total_android_devices": 1, "total_apple_devices": 3}
+        self.setup_apns_tokens()
+        self.setup_gcm_tokens()
+
+        android_devices = PushDeviceToken.objects.filter(kind=PushDeviceToken.GCM)
+        apple_devices = PushDeviceToken.objects.filter(kind=PushDeviceToken.APNS)
+
+        self.assertNotEqual(android_devices.count(), 0)
+        self.assertNotEqual(apple_devices.count(), 0)
+
+        mock_send.return_value = {
+            "total_android_devices": 1,
+            "total_apple_devices": 3,
+            # This response tests the logic of the server taking
+            # deleted_devices from the bouncer and deleting the
+            # corresponding PushDeviceTokens - because the bouncer is
+            # communicating that those devices have been deleted due
+            # to failures from Apple/Google and have no further user.
+            "deleted_devices": DevicesToCleanUpDict(
+                android_devices=[device.token for device in android_devices],
+                apple_devices=[device.token for device in apple_devices],
+            ),
+        }
        total_android_devices, total_apple_devices = send_notifications_to_bouncer(
-            user, {"apns": True}, {"gcm": True}, {}
+            user, {"apns": True}, {"gcm": True}, {}, list(android_devices), list(apple_devices)
        )
        post_data = {
            "user_uuid": user.uuid,
@@ -3523,6 +3597,8 @@ class TestSendNotificationsToBouncer(ZulipTestCase):
            "apns_payload": {"apns": True},
            "gcm_payload": {"gcm": True},
            "gcm_options": {},
+            "android_devices": [device.token for device in android_devices],
+            "apple_devices": [device.token for device in apple_devices],
        }
        mock_send.assert_called_with(
            "POST",
@@ -3543,6 +3619,9 @@ class TestSendNotificationsToBouncer(ZulipTestCase):
            ),
        )

+        self.assertEqual(PushDeviceToken.objects.filter(kind=PushDeviceToken.APNS).count(), 0)
+        self.assertEqual(PushDeviceToken.objects.filter(kind=PushDeviceToken.GCM).count(), 0)
+

@override_settings(PUSH_NOTIFICATION_BOUNCER_URL="https://push.zulip.org.example.com")
 class TestSendToPushBouncer(ZulipTestCase):
--- a/zilencer/views.py
+++ b/zilencer/views.py
@@ -1,7 +1,7 @@
 import logging
 from collections import Counter
 from datetime import datetime, timezone
-from typing import Any, Dict, List, Optional, Type, TypeVar, Union
+from typing import Any, Dict, List, Optional, Type, TypedDict, TypeVar, Union
 from uuid import UUID

 import orjson
@@ -523,15 +523,64 @@ def remote_server_notify_push(
            increment=android_successfully_delivered + apple_successfully_delivered,
        )

+    deleted_devices = get_deleted_devices(
+        user_identity,
+        server,
+        android_devices=payload.get("android_devices", []),
+        apple_devices=payload.get("apple_devices", []),
+    )
+
    return json_success(
        request,
        data={
            "total_android_devices": len(android_devices),
            "total_apple_devices": len(apple_devices),
+            "deleted_devices": deleted_devices,
        },
    )


+class DevicesToCleanUpDict(TypedDict):
+    android_devices: List[str]
+    apple_devices: List[str]
+
+
+def get_deleted_devices(
+    user_identity: UserPushIdentityCompat,
+    server: RemoteZulipServer,
+    android_devices: List[str],
+    apple_devices: List[str],
+) -> DevicesToCleanUpDict:
+    """The remote server sends us a list of (tokens of) devices that it
+    believes it has registered. However some of them may have been
+    deleted by us due to errors received in the low level code
+    responsible for directly sending push notifications.
+
+    Query the database for the RemotePushDeviceTokens from these lists
+    that we do indeed have and return a list of the ones that we don't
+    have and thus presumably have already deleted - the remote server
+    will want to delete them too.
+    """
+
+    android_devices_we_have = RemotePushDeviceToken.objects.filter(
+        user_identity.filter_q(),
+        token__in=android_devices,
+        kind=RemotePushDeviceToken.GCM,
+        server=server,
+    ).values_list("token", flat=True)
+    apple_devices_we_have = RemotePushDeviceToken.objects.filter(
+        user_identity.filter_q(),
+        token__in=apple_devices,
+        kind=RemotePushDeviceToken.APNS,
+        server=server,
+    ).values_list("token", flat=True)
+
+    return DevicesToCleanUpDict(
+        android_devices=list(set(android_devices) - set(android_devices_we_have)),
+        apple_devices=list(set(apple_devices) - set(apple_devices_we_have)),
+    )
+
+
 def validate_incoming_table_data(
    server: RemoteZulipServer,
    model: Any,