snapshot of https://github.com/zulip/zulip.git @ 48c69911c7dd4e03b20552eb3bd77ad8151c1357 (2026-09-15)

zerver/lib/push_registration.py

@@ -0,0 +1,220 @@
+import base64
+import binascii
+import hashlib
+import logging
+from typing import TypedDict
+
+from django.conf import settings
+from django.utils.translation import gettext as _
+
+from zerver.lib.devices import b64decode_token_id_base64
+from zerver.lib.exceptions import (
+    InvalidBouncerPublicKeyError,
+    InvalidEncryptedPushRegistrationError,
+    JsonableError,
+    MissingRemoteRealmError,
+    RequestExpiredError,
+)
+from zerver.lib.push_notifications import remove_push_device_token
+from zerver.lib.remote_server import (
+    PushNotificationBouncerError,
+    PushNotificationBouncerRetryLaterError,
+    PushNotificationBouncerServerError,
+    send_to_push_bouncer,
+)
+from zerver.models.devices import Device
+from zerver.models.push_notifications import PushDeviceToken
+from zerver.models.users import UserProfile, get_user_profile_by_id
+from zerver.tornado.django_api import send_event_on_commit
+
+if settings.ZILENCER_ENABLED:
+    from zilencer.views import do_register_remote_push_device
+
+logger = logging.getLogger(__name__)
+
+
+class RegisterPushDeviceToBouncerQueueItem(TypedDict):
+    user_profile_id: int
+    device_id: int
+    bouncer_public_key: str
+    encrypted_push_registration: str
+    token_id_base64: str
+
+
+def handle_registration_to_bouncer_failure(
+    user_profile: UserProfile, device_id: int, error_code: str
+) -> None:
+    """Handles a failed registration request to the bouncer by
+    notifying or preparing to notify clients.
+
+    * Sends a `device` event to notify online clients immediately.
+
+    * Stores the `error_code` in the `Device` table. This is later
+      used, along with other metadata, to notify offline clients the
+      next time they call `/register`. See the `devices` field in
+      the `/register` response.
+    """
+    Device.objects.filter(id=device_id).update(push_registration_error_code=error_code)
+    event = dict(
+        type="device",
+        op="update",
+        device_id=device_id,
+        push_registration_error_code=error_code,
+    )
+    send_event_on_commit(user_profile.realm, event, [user_profile.id])
+
+    # Report the `REQUEST_EXPIRED_ERROR` to the server admins as it indicates
+    # a long-lasting outage somewhere between the server and the bouncer,
+    # most likely in either the server or its local network configuration.
+    if error_code == Device.PushRegistrationErrorCode.REQUEST_EXPIRED:
+        logging.error("Push registration request for device_id=%s expired.", device_id)
+
+
+def compute_token_id_int(token: str) -> int:
+    hash_bytes = hashlib.sha256(token.encode()).digest()
+    return b64decode_token_id_base64(base64.b64encode(hash_bytes[0:8]).decode())
+
+
+def cleanup_legacy_push_device_token(user_profile: UserProfile, token_id_int: int) -> None:
+    MAX_REQUEST_RETRIES = 3
+    legacy_tokens = PushDeviceToken.objects.filter(user=user_profile)
+
+    for legacy_token in legacy_tokens:
+        # APNs treats its tokens (hex strings) as case-insensitive, and
+        # different client versions have sent them in different cases; the
+        # rest of the legacy code path accordingly normalizes to lowercase
+        # (see `remove_push_device_token`).  Match on the stored casing
+        # as well as its lower- and upper-case forms so that E2EE
+        # registration reliably cleans up the matching legacy token
+        # regardless of which case the client used.
+        token = legacy_token.token
+        if legacy_token.kind == PushDeviceToken.APNS:
+            candidate_token_ids = {
+                compute_token_id_int(token),
+                compute_token_id_int(token.lower()),
+                compute_token_id_int(token.upper()),
+            }
+        else:
+            candidate_token_ids = {compute_token_id_int(token)}
+
+        if token_id_int not in candidate_token_ids:
+            continue
+
+        for attempt in range(MAX_REQUEST_RETRIES):
+            try:
+                remove_push_device_token(user_profile, legacy_token.token, legacy_token.kind)
+                return
+            except (
+                PushNotificationBouncerRetryLaterError,
+                PushNotificationBouncerServerError,
+            ) as e:
+                # Network error or 5xx error response from bouncer server.
+                # TODO: Retry with backoff. Implement it while working on #35853.
+                if attempt == MAX_REQUEST_RETRIES - 1:
+                    logger.warning(
+                        "Failed to remove legacy push token from bouncer "
+                        "for user %s after retries: %s",
+                        user_profile.id,
+                        e,
+                    )
+                    return
+            except PushNotificationBouncerError as e:
+                # Invalid credentials or unexpected status code
+                logger.error(
+                    "Bouncer error cleaning up legacy push token for user %s: %s",
+                    user_profile.id,
+                    e,
+                )
+                return
+            except JsonableError:
+                # Token doesn't exist i.e. already cleaned up.
+                return
+
+
+def handle_register_push_device_to_bouncer(
+    queue_item: RegisterPushDeviceToBouncerQueueItem,
+) -> None:
+    user_profile_id = queue_item["user_profile_id"]
+    user_profile = get_user_profile_by_id(user_profile_id)
+    device_id = queue_item["device_id"]
+    bouncer_public_key = queue_item["bouncer_public_key"]
+    encrypted_push_registration = queue_item["encrypted_push_registration"]
+    token_id_base64 = queue_item["token_id_base64"]
+
+    try:
+        if settings.ZILENCER_ENABLED:
+            do_register_remote_push_device(
+                bouncer_public_key,
+                encrypted_push_registration,
+                token_id_base64,
+                realm=user_profile.realm,
+            )
+        else:
+            post_data: dict[str, str | int] = {
+                "realm_uuid": str(user_profile.realm.uuid),
+                "token_id": token_id_base64,
+                "encrypted_push_registration": encrypted_push_registration,
+                "bouncer_public_key": bouncer_public_key,
+            }
+            send_to_push_bouncer("POST", "push/e2ee/register", post_data)
+    except (
+        PushNotificationBouncerRetryLaterError,
+        PushNotificationBouncerServerError,
+    ) as e:  # nocoverage
+        # Network error or 5xx error response from bouncer server.
+        # Keep retrying to register until `RequestExpiredError` is raised.
+        raise PushNotificationBouncerRetryLaterError(e.msg)
+    except (
+        # Need to resubmit realm info - `manage.py register_server`
+        MissingRemoteRealmError,
+        # Invalid credentials or unexpected status code
+        PushNotificationBouncerError,
+    ):
+        # Server admins need to fix these set of errors, report them.
+        # Server should keep retrying to register until `RequestExpiredError` is raised.
+        error_msg = f"Push device registration request for device_id={device_id} failed."
+        logging.error(error_msg)
+        raise PushNotificationBouncerRetryLaterError(error_msg)
+    except (
+        InvalidBouncerPublicKeyError,
+        InvalidEncryptedPushRegistrationError,
+        RequestExpiredError,
+        # Any future or unexpected exceptions that we add.
+        JsonableError,
+    ) as e:
+        handle_registration_to_bouncer_failure(
+            user_profile, device_id, error_code=e.__class__.code.name
+        )
+        return
+
+    # Registration successful.
+    token_id_int = b64decode_token_id_base64(token_id_base64)
+    Device.objects.filter(id=device_id).update(
+        push_token_id=token_id_int, pending_push_token_id=None
+    )
+    event = dict(
+        type="device",
+        op="update",
+        device_id=device_id,
+        push_token_id=token_id_base64,
+        pending_push_token_id=None,
+    )
+    send_event_on_commit(user_profile.realm, event, [user_profile.id])
+
+    # When a client successfully registers a token for E2EE notifications,
+    # clear out any existing legacy registration on that account for the same token.
+    cleanup_legacy_push_device_token(user_profile, token_id_int)
+
+
+def check_push_key(push_key_str: str) -> bytes:
+    error_message = _("Invalid `push_key`")
+
+    try:
+        push_key_bytes = base64.b64decode(push_key_str, validate=True)
+    except binascii.Error:
+        raise JsonableError(error_message)
+
+    if len(push_key_bytes) != 33 or push_key_bytes[0] != 0x31:
+        raise JsonableError(error_message)
+
+    return push_key_bytes