push_notifications: Add support for setting counts in iOS.

This adds a new function `get_apns_badge_count()` to
fetch count value for a user push notification and
then sends that value with the APNs payload.

Once a message is read from the web app, the count is
decremented accordingly and a push notification with
`event: remove` is sent to the iOS clients.

Fixes #10271.

zerver/lib/push_notifications.py

@@ -3,7 +3,7 @@ import binascii
 import logging
 import re
 import time
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Sequence, Tuple, Union
 
 import gcm
 import lxml.html
@@ -610,6 +610,19 @@ def get_apns_alert_subtitle(message: Message) -> str:
     # For group PMs, or regular messages to a stream, just use a colon to indicate this is the sender.
     return message.sender.full_name + ":"
 
+def get_apns_badge_count(user_profile: UserProfile, read_messages_ids: Optional[Sequence[int]]=[]) -> int:
+    return UserMessage.objects.filter(
+        user_profile=user_profile
+    ).extra(
+        where=[UserMessage.where_active_push_notification()]
+    ).exclude(
+        # If we've just marked some messages as read, they're still
+        # marked as having active notifications; we'll clear that flag
+        # only after we've sent that update to the devices.  So we need
+        # to exclude them explicitly from the count.
+        message_id__in=read_messages_ids
+    ).count()
+
 def get_message_payload_apns(user_profile: UserProfile, message: Message) -> Dict[str, Any]:
     '''A `message` payload for iOS, via APNs.'''
     zulip_data = get_message_payload(user_profile, message)
@@ -625,7 +638,7 @@ def get_message_payload_apns(user_profile: UserProfile, message: Message) -> Dic
             'body': content,
         },
         'sound': 'default',
-        'badge': 0,  # TODO: set badge count in a better way
+        'badge': get_apns_badge_count(user_profile),
         'custom': {'zulip': zulip_data},
     }
     return apns_data
@@ -664,6 +677,18 @@ def get_remove_payload_gcm(
     gcm_options = {'priority': 'normal'}
     return gcm_payload, gcm_options
 
+def get_remove_payload_apns(user_profile: UserProfile, message_ids: List[int]) -> Dict[str, Any]:
+    zulip_data = get_base_payload(user_profile)
+    zulip_data.update({
+        'event': 'remove',
+        'zulip_message_ids': ','.join(str(id) for id in message_ids),
+    })
+    apns_data = {
+        'badge': get_apns_badge_count(user_profile, message_ids),
+        'custom': {'zulip': zulip_data},
+    }
+    return apns_data
+
 def handle_remove_push_notification(user_profile_id: int, message_ids: List[int]) -> None:
     """This should be called when a message that had previously had a
     mobile push executed is read.  This triggers a mobile push notifica
@@ -674,17 +699,22 @@ def handle_remove_push_notification(user_profile_id: int, message_ids: List[int]
     user_profile = get_user_profile_by_id(user_profile_id)
     message_ids = bulk_access_messages_expect_usermessage(user_profile_id, message_ids)
     gcm_payload, gcm_options = get_remove_payload_gcm(user_profile, message_ids)
+    apns_payload = get_remove_payload_apns(user_profile, message_ids)
 
     if uses_notification_bouncer():
         send_notifications_to_bouncer(user_profile_id,
-                                      {},
+                                      apns_payload,
                                       gcm_payload,
                                       gcm_options)
     else:
         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 android_devices:
             send_android_push_notification(android_devices, gcm_payload, gcm_options)
+        if apple_devices:
+            send_apple_push_notification(user_profile_id, apple_devices, apns_payload)
 
     UserMessage.objects.filter(
         user_profile_id=user_profile_id,

zerver/tests/test_push_notifications.py

@@ -24,12 +24,14 @@ from zerver.lib.actions import (
     do_delete_messages,
     do_mark_stream_messages_as_read,
     do_regenerate_api_key,
+    do_update_message_flags,
 )
 from zerver.lib.push_notifications import (
     DeviceToken,
     absolute_avatar_url,
     b64_to_hex,
     datetime_to_timestamp,
+    get_apns_badge_count,
     get_apns_client,
     get_display_recipient,
     get_message_payload_apns,
@@ -923,11 +925,23 @@ class HandlePushNotificationTest(PushNotificationTest):
 
         with self.settings(PUSH_NOTIFICATION_BOUNCER_URL=True), \
                 mock.patch('zerver.lib.push_notifications'
-                           '.send_notifications_to_bouncer') as mock_send_android:
+                           '.send_notifications_to_bouncer') as mock_send:
             handle_remove_push_notification(user_profile.id, [message.id])
-            mock_send_android.assert_called_with(
+            mock_send.assert_called_with(
                 user_profile.id,
-                {},
+                {
+                    'badge': 0,
+                    'custom': {
+                        'zulip': {
+                            'server': 'testserver',
+                            'realm_id': self.sender.realm.id,
+                            'realm_uri': 'http://zulip.testserver',
+                            'user_id': self.user_profile.id,
+                            'event': 'remove',
+                            'zulip_message_ids': str(message.id),
+                        },
+                    },
+                },
                 {
                     'server': 'testserver',
                     'realm_id': self.sender.realm.id,
@@ -956,8 +970,14 @@ class HandlePushNotificationTest(PushNotificationTest):
             PushDeviceToken.objects.filter(user=self.user_profile,
                                            kind=PushDeviceToken.GCM))
 
+        apple_devices = list(
+            PushDeviceToken.objects.filter(user=self.user_profile,
+                                           kind=PushDeviceToken.APNS))
+
         with mock.patch('zerver.lib.push_notifications'
-                        '.send_android_push_notification') as mock_send_android:
+                        '.send_android_push_notification') as mock_send_android, \
+                mock.patch('zerver.lib.push_notifications'
+                           '.send_apple_push_notification') as mock_send_apple:
             handle_remove_push_notification(self.user_profile.id, [message.id])
             mock_send_android.assert_called_with(
                 android_devices,
@@ -971,6 +991,20 @@ class HandlePushNotificationTest(PushNotificationTest):
                     'zulip_message_id': message.id,
                 },
                 {'priority': 'normal'})
+            mock_send_apple.assert_called_with(
+                self.user_profile.id,
+                apple_devices,
+                {'badge': 0,
+                 'custom': {
+                     'zulip': {
+                         'server': 'testserver',
+                         'realm_id': self.sender.realm.id,
+                         'realm_uri': 'http://zulip.testserver',
+                         'user_id': self.user_profile.id,
+                         'event': 'remove',
+                         'zulip_message_ids': str(message.id),
+                     }
+                 }})
             user_message = UserMessage.objects.get(user_profile=self.user_profile,
                                                    message=message)
             self.assertEqual(user_message.flags.active_mobile_push_notification, False)
@@ -1152,12 +1186,39 @@ class TestAPNs(PushNotificationTest):
         self.assertEqual(
             modernize_apns_payload(
                 {'alert': 'Message from Hamlet',
-                 'message_ids': [3]}),
+                 'message_ids': [3],
+                 'badge': 0}),
             payload)
         self.assertEqual(
             modernize_apns_payload(payload),
             payload)
 
+    @mock.patch('zerver.lib.push_notifications.push_notifications_enabled', return_value = True)
+    def test_apns_badge_count(self, mock_push_notifications: mock.MagicMock) -> None:
+        user_profile = self.example_user('othello')
+        # Test APNs badge count for personal messages.
+        message_ids = [self.send_personal_message(self.sender,
+                                                  user_profile,
+                                                  'Content of message')
+                       for i in range(3)]
+        self.assertEqual(get_apns_badge_count(user_profile), 3)
+        # Similarly, test APNs badge count for stream mention.
+        stream = self.subscribe(user_profile, "Denmark")
+        message_ids += [self.send_stream_message(self.sender,
+                                                 stream.name,
+                                                 'Hi, @**Othello, the Moor of Venice**')
+                        for i in range(2)]
+        self.assertEqual(get_apns_badge_count(user_profile), 5)
+
+        num_messages = len(message_ids)
+        # Mark the messages as read and test whether
+        # the count decreases correctly.
+        for i, message_id in enumerate(message_ids):
+            do_update_message_flags(user_profile, get_client("website"), 'add', 'read', [message_id])
+            self.assertEqual(get_apns_badge_count(user_profile), num_messages - i - 1)
+
+        mock_push_notifications.assert_called()
+
 class TestGetAPNsPayload(PushNotificationTest):
     def test_get_message_payload_apns_personal_message(self) -> None:
         user_profile = self.example_user("othello")
@@ -1208,7 +1269,7 @@ class TestGetAPNsPayload(PushNotificationTest):
                 'body': message.content,
             },
             'sound': 'default',
-            'badge': 0,
+            'badge': 1,
             'custom': {
                 'zulip': {
                     'message_ids': [message.id],