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.

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)
 
-    if uses_notification_bouncer():
-        send_notifications_to_bouncer(user_profile, apns_payload, gcm_payload, gcm_options)
-    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)
     )
+    if uses_notification_bouncer():
+        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_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,

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):

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,