zilencer: Add endpoint to register push device to bouncer.

This commit adds a zilencer endpoint to let self-hosted
servers register push devices to whom mobile push notifications
will be sent.

POST "/api/v1/remotes/push/e2ee/register"
Payload: realm_uuid, push_account_id, encrypted_push_registration,
bouncer_public_key

The post request needs to be authenticated with the server’s API key.

Note: For Zulip Cloud, a background fact about the push bouncer is
that it runs on the same server and database as the main application;
it’s not a separate service.
So, as an optimization, we plan to directly call the
`do_register_remote_push_device` function and skip the HTTP request.

zerver/lib/exceptions.py

@@ -59,6 +59,8 @@ class ErrorCode(Enum):
     CANNOT_ADMINISTER_CHANNEL = auto()
     REMOTE_SERVER_VERIFICATION_SECRET_NOT_PREPARED = auto()
     HOSTNAME_ALREADY_IN_USE_BOUNCER_ERROR = auto()
+    INVALID_BOUNCER_PUBLIC_KEY = auto()
+    REQUEST_EXPIRED = auto()
 
 
 class JsonableError(Exception):
@@ -827,3 +829,37 @@ class SlackImportInvalidFileError(Exception):
     def __init__(self, message: str) -> None:
         super().__init__(message)
         self.message = message
+
+
+class InvalidBouncerPublicKeyError(JsonableError):
+    code = ErrorCode.INVALID_BOUNCER_PUBLIC_KEY
+
+    def __init__(self) -> None:
+        pass
+
+    @staticmethod
+    @override
+    def msg_format() -> str:
+        return _("Invalid bouncer_public_key")
+
+
+class RequestExpiredError(JsonableError):
+    code = ErrorCode.REQUEST_EXPIRED
+
+    def __init__(self) -> None:
+        pass
+
+    @staticmethod
+    @override
+    def msg_format() -> str:
+        return _("Request expired")
+
+
+class InvalidEncryptedPushRegistrationError(JsonableError):
+    def __init__(self) -> None:
+        pass
+
+    @staticmethod
+    @override
+    def msg_format() -> str:
+        return _("Invalid encrypted_push_registration")

zerver/lib/push_notifications.py

@@ -73,6 +73,12 @@ logger = logging.getLogger(__name__)
 if settings.ZILENCER_ENABLED:
     from zilencer.models import RemotePushDeviceToken, RemoteZulipServer
 
+# Time (in seconds) for which the server should retry registering
+# a push device to the bouncer. 24 hrs is a good time limit because
+# a day is longer than any minor outage.
+PUSH_REGISTRATION_LIVENESS_TIMEOUT = 24 * 60 * 60
+
+
 DeviceToken: TypeAlias = Union[PushDeviceToken, "RemotePushDeviceToken"]
 
 

zerver/tests/test_push_registration.py

@@ -0,0 +1,192 @@
+import uuid
+from datetime import timedelta
+
+import orjson
+from django.conf import settings
+from django.utils.timezone import now
+from nacl.encoding import Base64Encoder
+from nacl.public import PublicKey, SealedBox
+
+from zerver.lib.test_classes import BouncerTestCase
+from zerver.lib.timestamp import datetime_to_timestamp
+from zilencer.models import RemotePushDevice, RemoteRealm
+
+
+class RegisterPushDeviceToBouncer(BouncerTestCase):
+    DEFAULT_SUBDOMAIN = ""
+
+    def get_register_push_device_payload(
+        self,
+        token: str = "apple-tokenaz",
+        token_kind: str = RemotePushDevice.TokenKind.APNS,
+        ios_app_id: str | None = "example.app",
+        timestamp: int | None = None,
+    ) -> dict[str, str | int]:
+        hamlet = self.example_user("hamlet")
+        remote_realm = RemoteRealm.objects.get(uuid=hamlet.realm.uuid)
+
+        if timestamp is None:
+            timestamp = datetime_to_timestamp(now())
+
+        push_registration = {
+            "token": token,
+            "token_kind": token_kind,
+            "ios_app_id": ios_app_id,
+            "timestamp": timestamp,
+        }
+
+        assert settings.PUSH_REGISTRATION_ENCRYPTION_KEYS
+        public_key_str: str = next(iter(settings.PUSH_REGISTRATION_ENCRYPTION_KEYS.keys()))
+        public_key = PublicKey(public_key_str.encode("utf-8"), Base64Encoder)
+        sealed_box = SealedBox(public_key)
+        encrypted_push_registration_bytes = sealed_box.encrypt(
+            orjson.dumps(push_registration), Base64Encoder
+        )
+        encrypted_push_registration = encrypted_push_registration_bytes.decode("utf-8")
+
+        payload: dict[str, str | int] = {
+            "realm_uuid": str(remote_realm.uuid),
+            "push_account_id": 2408,
+            "encrypted_push_registration": encrypted_push_registration,
+            "bouncer_public_key": public_key_str,
+        }
+        return payload
+
+    def test_register_push_device_success(self) -> None:
+        remote_push_devices_count = RemotePushDevice.objects.count()
+        self.assertEqual(remote_push_devices_count, 0)
+
+        payload = self.get_register_push_device_payload()
+
+        result = self.uuid_post(
+            self.server_uuid,
+            "/api/v1/remotes/push/e2ee/register",
+            payload,
+        )
+        response_dict = self.assert_json_success(result)
+        device_id = response_dict["device_id"]
+
+        remote_push_devices = RemotePushDevice.objects.all()
+        self.assert_length(remote_push_devices, 1)
+        self.assertEqual(device_id, remote_push_devices[0].device_id)
+
+        # Idempotent
+        result = self.uuid_post(
+            self.server_uuid,
+            "/api/v1/remotes/push/e2ee/register",
+            payload,
+        )
+        response_dict = self.assert_json_success(result)
+        self.assertEqual(response_dict["device_id"], device_id)
+        remote_push_devices_count = RemotePushDevice.objects.count()
+        self.assertEqual(remote_push_devices_count, 1)
+
+        # Android
+        payload = self.get_register_push_device_payload(
+            token="android-tokenaz", token_kind=RemotePushDevice.TokenKind.FCM, ios_app_id=None
+        )
+        result = self.uuid_post(
+            self.server_uuid,
+            "/api/v1/remotes/push/e2ee/register",
+            payload,
+        )
+        response_dict = self.assert_json_success(result)
+        device_id = response_dict["device_id"]
+
+        remote_push_devices = RemotePushDevice.objects.order_by("pk")
+        self.assert_length(remote_push_devices, 2)
+        self.assertEqual(device_id, remote_push_devices[1].device_id)
+
+    def test_register_push_device_error(self) -> None:
+        payload = self.get_register_push_device_payload()
+
+        invalid_realm_uuid_payload = {**payload, "realm_uuid": str(uuid.uuid4())}
+        with self.assertLogs("zilencer.views", level="INFO"):
+            result = self.uuid_post(
+                self.server_uuid,
+                "/api/v1/remotes/push/e2ee/register",
+                invalid_realm_uuid_payload,
+            )
+        self.assert_json_error(result, "Organization not registered", status_code=403)
+
+        invalid_bouncer_public_key_payload = {**payload, "bouncer_public_key": "invalid public key"}
+        result = self.uuid_post(
+            self.server_uuid,
+            "/api/v1/remotes/push/e2ee/register",
+            invalid_bouncer_public_key_payload,
+        )
+        self.assert_json_error(result, "Invalid bouncer_public_key")
+
+        liveness_timedout_payload = self.get_register_push_device_payload(
+            timestamp=datetime_to_timestamp(now() - timedelta(days=2))
+        )
+        result = self.uuid_post(
+            self.server_uuid,
+            "/api/v1/remotes/push/e2ee/register",
+            liveness_timedout_payload,
+        )
+        self.assert_json_error(result, "Request expired")
+
+        # Test the various cases resulting in InvalidEncryptedPushRegistrationError
+        payload = self.get_register_push_device_payload()
+        payload["encrypted_push_registration"] = "random-string-no-encryption"
+        result = self.uuid_post(
+            self.server_uuid,
+            "/api/v1/remotes/push/e2ee/register",
+            payload,
+        )
+        self.assert_json_error(result, "Invalid encrypted_push_registration")
+
+        invalid_ios_app_id_format_payload = self.get_register_push_device_payload(
+            ios_app_id="* -- +"
+        )
+        result = self.uuid_post(
+            self.server_uuid,
+            "/api/v1/remotes/push/e2ee/register",
+            invalid_ios_app_id_format_payload,
+        )
+        self.assert_json_error(result, "Invalid encrypted_push_registration")
+
+        invalid_token_kind_payload = self.get_register_push_device_payload(token_kind="xyz")
+        result = self.uuid_post(
+            self.server_uuid,
+            "/api/v1/remotes/push/e2ee/register",
+            invalid_token_kind_payload,
+        )
+        self.assert_json_error(result, "Invalid encrypted_push_registration")
+
+        missing_ios_app_id_payload = self.get_register_push_device_payload(ios_app_id=None)
+        result = self.uuid_post(
+            self.server_uuid,
+            "/api/v1/remotes/push/e2ee/register",
+            missing_ios_app_id_payload,
+        )
+        self.assert_json_error(result, "Invalid encrypted_push_registration")
+
+        set_ios_app_id_for_android_payload = self.get_register_push_device_payload(
+            token_kind=RemotePushDevice.TokenKind.FCM, ios_app_id="not-null"
+        )
+        result = self.uuid_post(
+            self.server_uuid,
+            "/api/v1/remotes/push/e2ee/register",
+            set_ios_app_id_for_android_payload,
+        )
+        self.assert_json_error(result, "Invalid encrypted_push_registration")
+
+        invalid_token_payload = self.get_register_push_device_payload(token="")
+        result = self.uuid_post(
+            self.server_uuid,
+            "/api/v1/remotes/push/e2ee/register",
+            invalid_token_payload,
+        )
+        self.assert_json_error(result, "Invalid encrypted_push_registration")
+
+        invalid_token_payload = self.get_register_push_device_payload(
+            token="xyz non-hex characters"
+        )
+        result = self.uuid_post(
+            self.server_uuid,
+            "/api/v1/remotes/push/e2ee/register",
+            invalid_token_payload,
+        )
+        self.assert_json_error(result, "Invalid encrypted_push_registration")

zilencer/models.py

@@ -643,6 +643,9 @@ class RemotePushDevice(models.Model):
                 # Each app install (token) can have multiple accounts (push_account_id).
                 # The (push_account_id, token) pair needs to be unique to avoid sending
                 # redundant notifications to the same account on a device.
+                #
+                # Also, the unique index created is used by a query in
+                # 'do_register_remote_push_device'.
                 fields=["push_account_id", "token"],
                 name="unique_remote_push_device_push_account_id_token",
             ),

zilencer/urls.py

@@ -8,6 +8,7 @@ from zilencer.auth import remote_server_path
 from zilencer.views import (
     deactivate_remote_server,
     register_remote_push_device,
+    register_remote_push_device_for_e2ee_push_notification,
     register_remote_server,
     remote_server_check_analytics,
     remote_server_notify_push,
@@ -24,6 +25,9 @@ i18n_urlpatterns: Any = []
 # Zilencer views following the REST API style
 push_bouncer_patterns = [
     remote_server_path("remotes/push/register", POST=register_remote_push_device),
+    remote_server_path(
+        "remotes/push/e2ee/register", POST=register_remote_push_device_for_e2ee_push_notification
+    ),
     remote_server_path("remotes/push/unregister", POST=unregister_remote_push_device),
     remote_server_path("remotes/push/unregister/all", POST=unregister_all_remote_push_devices),
     remote_server_path("remotes/push/notify", POST=remote_server_notify_push),

zilencer/views.py

@@ -22,7 +22,11 @@ from django.utils.translation import gettext as err_
 from django.views.decorators.csrf import csrf_exempt
 from dns import resolver as dns_resolver
 from dns.exception import DNSException
-from pydantic import BaseModel, ConfigDict, Json, StringConstraints
+from nacl.encoding import Base64Encoder
+from nacl.exceptions import CryptoError
+from nacl.public import PrivateKey, SealedBox
+from pydantic import BaseModel, ConfigDict, Json, StringConstraints, model_validator
+from pydantic import ValidationError as PydanticValidationError
 from pydantic.functional_validators import AfterValidator
 
 from analytics.lib.counts import (
@@ -38,16 +42,22 @@ from zerver.decorator import require_post
 from zerver.lib.email_validation import validate_is_not_disposable
 from zerver.lib.exceptions import (
     ErrorCode,
+    InvalidBouncerPublicKeyError,
+    InvalidEncryptedPushRegistrationError,
     JsonableError,
+    MissingRemoteRealmError,
     RateLimitedError,
     RemoteRealmServerMismatchError,
     RemoteServerDeactivatedError,
+    RequestExpiredError,
 )
 from zerver.lib.outgoing_http import OutgoingSession
 from zerver.lib.push_notifications import (
+    PUSH_REGISTRATION_LIVENESS_TIMEOUT,
     HostnameAlreadyInUseBouncerError,
     InvalidRemotePushDeviceTokenError,
     UserPushIdentityCompat,
+    b64_to_hex,
     send_android_push_notification,
     send_apple_push_notification,
     send_test_push_notification_directly_to_devices,
@@ -64,7 +74,7 @@ from zerver.lib.remote_server import (
 from zerver.lib.request import RequestNotes
 from zerver.lib.response import json_success
 from zerver.lib.send_email import EMAIL_DATE_FORMAT, FromAddress
-from zerver.lib.timestamp import timestamp_to_datetime
+from zerver.lib.timestamp import datetime_to_timestamp, timestamp_to_datetime
 from zerver.lib.typed_endpoint import (
     ApnsAppId,
     JsonBodyPayload,
@@ -75,7 +85,7 @@ from zerver.lib.typed_endpoint import (
 from zerver.lib.typed_endpoint_validators import check_string_fixed_length
 from zerver.lib.types import RemoteRealmDictValue
 from zerver.models.realm_audit_logs import AuditLogEventType
-from zerver.models.realms import DisposableEmailError
+from zerver.models.realms import DisposableEmailError, Realm
 from zilencer.auth import (
     InvalidZulipServerKeyError,
     generate_registration_transfer_verification_secret,
@@ -84,6 +94,7 @@ from zilencer.auth import (
 from zilencer.lib.remote_counts import MissingDataError
 from zilencer.models import (
     RemoteInstallationCount,
+    RemotePushDevice,
     RemotePushDeviceToken,
     RemoteRealm,
     RemoteRealmAuditLog,
@@ -475,6 +486,128 @@ def register_remote_push_device(
     return json_success(request)
 
 
+class PushRegistration(BaseModel):
+    token: str
+    token_kind: str
+    ios_app_id: ApnsAppId | None = None
+    timestamp: int
+
+    def is_valid_token(self) -> bool:
+        if self.token == "" or len(self.token) > 4096:
+            # Invalid token length
+            return False
+
+        if self.token_kind == RemotePushDevice.TokenKind.APNS:
+            # Validate that we can actually decode the token.
+            try:
+                b64_to_hex(self.token)
+            except Exception:
+                return False
+        return True
+
+    @model_validator(mode="after")
+    def validate_terms(self) -> "PushRegistration":
+        if self.token_kind not in [RemotePushDevice.TokenKind.APNS, RemotePushDevice.TokenKind.FCM]:
+            raise ValueError("Invalid token_kind")
+
+        if self.token_kind == RemotePushDevice.TokenKind.APNS and self.ios_app_id is None:
+            raise ValueError("Missing ios_app_id")
+
+        if self.token_kind == RemotePushDevice.TokenKind.FCM and self.ios_app_id is not None:
+            raise ValueError(
+                f"For token_kind={RemotePushDevice.TokenKind.FCM}, ios_app_id should be null"
+            )
+
+        if not self.is_valid_token():
+            raise ValueError("Invalid token")
+
+        return self
+
+
+def do_register_remote_push_device(
+    bouncer_public_key: str,
+    encrypted_push_registration: str,
+    push_account_id: int,
+    *,
+    realm: Realm | None = None,
+    remote_realm: RemoteRealm | None = None,
+) -> int:
+    assert (realm is None) ^ (remote_realm is None)
+
+    assert settings.PUSH_REGISTRATION_ENCRYPTION_KEYS
+    if bouncer_public_key not in settings.PUSH_REGISTRATION_ENCRYPTION_KEYS:
+        raise InvalidBouncerPublicKeyError
+
+    # Decrypt push_registration
+    bouncer_private_key: str = settings.PUSH_REGISTRATION_ENCRYPTION_KEYS[bouncer_public_key]
+    private_key = PrivateKey(bouncer_private_key.encode("utf-8"), encoder=Base64Encoder)
+    unseal_box = SealedBox(private_key)
+
+    try:
+        push_registration_bytes = unseal_box.decrypt(
+            Base64Encoder.decode(encrypted_push_registration.encode("utf-8"))
+        )
+    except (TypeError, CryptoError):
+        raise InvalidEncryptedPushRegistrationError
+
+    try:
+        push_registration = PushRegistration.model_validate_json(push_registration_bytes)
+    except PydanticValidationError:
+        raise InvalidEncryptedPushRegistrationError
+
+    if (
+        datetime_to_timestamp(timezone_now()) - push_registration.timestamp
+        > PUSH_REGISTRATION_LIVENESS_TIMEOUT
+    ):
+        raise RequestExpiredError
+
+    # If already registered, return the device_id.
+    # The query uses the unique index created by the
+    # 'unique_remote_push_device_push_account_id_token' constraint.
+    remote_push_device = RemotePushDevice.objects.filter(
+        token=push_registration.token, push_account_id=push_account_id
+    ).first()
+    if remote_push_device:
+        return remote_push_device.device_id
+
+    remote_push_device = RemotePushDevice.objects.create(
+        realm=realm,
+        remote_realm=remote_realm,
+        token=push_registration.token,
+        token_kind=push_registration.token_kind,
+        push_account_id=push_account_id,
+        ios_app_id=push_registration.ios_app_id,
+    )
+    return remote_push_device.device_id
+
+
+@typed_endpoint
+def register_remote_push_device_for_e2ee_push_notification(
+    request: HttpRequest,
+    server: RemoteZulipServer,
+    *,
+    realm_uuid: str,
+    push_account_id: Json[int],
+    encrypted_push_registration: str,
+    bouncer_public_key: str,
+) -> HttpResponse:
+    remote_realm = get_remote_realm_helper(request, server, realm_uuid)
+    if remote_realm is None:
+        raise MissingRemoteRealmError
+    else:
+        remote_realm.last_request_datetime = timezone_now()
+        remote_realm.save(update_fields=["last_request_datetime"])
+
+    device_id = do_register_remote_push_device(
+        bouncer_public_key,
+        encrypted_push_registration,
+        push_account_id,
+        remote_realm=remote_realm,
+    )
+
+    return json_success(request, {"device_id": device_id})
+
+
 @typed_endpoint
 def unregister_remote_push_device(
     request: HttpRequest,