zilencer: Tie RemotePushDeviceToken to RemoteRealm at registration.

This consists of the following pieces:
1. Makes servers using the bouncer send realm_uuid in requests for token
   registration. (Sidenote: realm_uuid is already sent in the "send
   notification" codepath as of
   48db4bf854)
2. This allows the bouncer to tie RemotePushDeviceToken to the
   RemoteRealm with matching realm_uuid at registration time.
3. Introduce handling of some potential weird edge cases around the
   realm_uuid and RemoteRealm objects in get_remote_realm_helper.

zerver/lib/push_notifications.py

@@ -641,6 +641,7 @@ def add_push_device_token(
         post_data = {
             "server_uuid": settings.ZULIP_ORG_ID,
             "user_uuid": str(user_profile.uuid),
+            "realm_uuid": str(user_profile.realm.uuid),
             # user_id is sent so that the bouncer can delete any pre-existing registrations
             # for this user+device to avoid duplication upon adding the uuid registration.
             "user_id": str(user_profile.id),

zerver/tests/test_push_notifications.py

@@ -98,6 +98,7 @@ if settings.ZILENCER_ENABLED:
         RemoteRealmCount,
         RemoteZulipServer,
     )
+    from zilencer.views import update_remote_realm_data_for_server
 
 
 class SendTestPushNotificationEndpointTest(BouncerTestCase):
@@ -735,6 +736,51 @@ class PushBouncerNotificationTest(BouncerTestCase):
             result = self.uuid_post(self.server_uuid, endpoint, payload)
             self.assert_json_error(result, "Invalid APNS token")
 
+    @override_settings(PUSH_NOTIFICATION_BOUNCER_URL="https://push.zulip.org.example.com")
+    @responses.activate
+    def test_register_token_realm_uuid_belongs_to_different_server(self) -> None:
+        self.add_mock_response()
+        user = self.example_user("cordelia")
+        self.login_user(user)
+
+        # Create a simulated second server. We will take user's RemoteRealm registration
+        # and change its server to this second server. This means that when the bouncer
+        # is processing the token registration request, it will find a RemoteRealm matching
+        # the realm_uuid in the request, but that RemoteRealm will be registered to a
+        # different server than the one making the request (self.server).
+        # This will make it log a warning and register the token, but of course without
+        # assigning the token registration to that RemoteRealm.
+        second_server = RemoteZulipServer.objects.create(
+            uuid=uuid.uuid4(),
+            api_key="magic_secret_api_key2",
+            hostname="demo2.example.com",
+            last_updated=now(),
+        )
+
+        remote_realm = RemoteRealm.objects.get(server=self.server, uuid=user.realm.uuid)
+        remote_realm.server = second_server
+        remote_realm.save()
+
+        endpoint = "/json/users/me/apns_device_token"
+        token = "apple-tokenaz"
+        with self.assertLogs("zilencer.views", level="WARN") as warn_log:
+            result = self.client_post(
+                endpoint, {"token": token, "appid": "org.zulip.Zulip"}, subdomain="zulip"
+            )
+            self.assert_json_success(result)
+        self.assertEqual(
+            warn_log.output,
+            [
+                "WARNING:zilencer.views:/api/v1/remotes/push/register: "
+                f"Realm {remote_realm.uuid!s} exists, but not registered to server {self.server.id}"
+            ],
+        )
+
+        remote_token = RemotePushDeviceToken.objects.get(token=token)
+        self.assertEqual(remote_token.server, self.server)
+        self.assertEqual(remote_token.user_uuid, user.uuid)
+        self.assertEqual(remote_token.remote_realm, None)
+
     @override_settings(PUSH_NOTIFICATION_BOUNCER_URL="https://push.zulip.org.example.com")
     @responses.activate
     def test_push_bouncer_api(self) -> None:
@@ -811,7 +857,38 @@ class PushBouncerNotificationTest(BouncerTestCase):
 
         # Add tokens
         for endpoint, token, kind, appid in endpoints:
-            # Test that we can push twice
+            # First register a token without having a RemoteRealm registration:
+            RemoteRealm.objects.all().delete()
+            with self.assertLogs("zilencer.views", level="INFO") as info_log:
+                result = self.client_post(endpoint, {"token": token, **appid}, subdomain="zulip")
+            self.assert_json_success(result)
+            self.assertIn(
+                "INFO:zilencer.views:/api/v1/remotes/push/register: Received request for "
+                f"unknown realm {user.realm.uuid!s}, server {server.id}, "
+                f"user {user.uuid!s}",
+                info_log.output,
+            )
+
+            # The registration succeeded, but RemotePushDeviceToken doesn't have remote_realm set:
+            tokens = list(
+                RemotePushDeviceToken.objects.filter(
+                    user_uuid=user.uuid, token=token, server=server
+                )
+            )
+            self.assert_length(tokens, 1)
+            self.assertEqual(tokens[0].kind, kind)
+            self.assertEqual(tokens[0].user_uuid, user.uuid)
+
+            # Delete it to clean up.
+            RemotePushDeviceToken.objects.filter(
+                user_uuid=user.uuid, token=token, server=server
+            ).delete()
+
+            # Create the expected RemoteRealm registration and proceed with testing with a
+            # normal setup.
+            update_remote_realm_data_for_server(self.server, get_realms_info_for_push_bouncer())
+
+            # Test that we can push more times
             result = self.client_post(endpoint, {"token": token, **appid}, subdomain="zulip")
             self.assert_json_success(result)
 
@@ -825,6 +902,9 @@ class PushBouncerNotificationTest(BouncerTestCase):
             )
             self.assert_length(tokens, 1)
             self.assertEqual(tokens[0].kind, kind)
+            # These new registrations have .remote_realm set properly.
+            assert tokens[0].remote_realm is not None
+            self.assertEqual(tokens[0].remote_realm.uuid, user.realm.uuid)
             self.assertEqual(tokens[0].ios_app_id, appid.get("appid"))
 
         # User should have tokens for both devices now.

zilencer/migrations/0043_remotepushdevicetoken_remote_realm.py

@@ -0,0 +1,20 @@
+# Generated by Django 4.2.7 on 2023-12-01 15:52
+
+import django.db.models.deletion
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+    dependencies = [
+        ("zilencer", "0042_alter_remoterealmauditlog_realm_id"),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name="remotepushdevicetoken",
+            name="remote_realm",
+            field=models.ForeignKey(
+                null=True, on_delete=django.db.models.deletion.SET_NULL, to="zilencer.remoterealm"
+            ),
+        ),
+    ]

zilencer/models.py

@@ -80,6 +80,8 @@ class RemotePushDeviceToken(AbstractPushDeviceToken):
     user_id = models.BigIntegerField(null=True)
     user_uuid = models.UUIDField(null=True)
 
+    remote_realm = models.ForeignKey("RemoteRealm", on_delete=models.SET_NULL, null=True)
+
     class Meta:
         unique_together = [
             # These indexes rely on the property that in Postgres,

zilencer/views.py

@@ -152,6 +152,7 @@ def register_remote_push_device(
     server: RemoteZulipServer,
     user_id: Optional[int] = REQ(json_validator=check_int, default=None),
     user_uuid: Optional[str] = REQ(default=None),
+    realm_uuid: Optional[str] = REQ(default=None),
     token: str = REQ(),
     token_kind: int = REQ(json_validator=check_int),
     ios_app_id: Optional[str] = REQ(str_validator=check_app_id, default=None),
@@ -173,6 +174,17 @@ def register_remote_push_device(
         # One of these is None, so these kwargs will lead to a proper registration
         # of either user_id or user_uuid type
         kwargs = {"user_id": user_id, "user_uuid": user_uuid}
+
+    if realm_uuid is not None:
+        # Servers 8.0+ also send the realm.uuid of the user.
+        assert isinstance(
+            user_uuid, str
+        ), "Servers new enough to send realm_uuid, should also have user_uuid"
+        remote_realm = get_remote_realm_helper(request, server, realm_uuid, user_uuid)
+        if remote_realm is not None:
+            # We want to associate the RemotePushDeviceToken with the RemoteRealm.
+            kwargs["remote_realm_id"] = remote_realm.id
+
     try:
         with transaction.atomic():
             RemotePushDeviceToken.objects.create(
@@ -328,6 +340,39 @@ def remote_server_send_test_notification(
     return json_success(request)
 
 
+def get_remote_realm_helper(
+    request: HttpRequest, server: RemoteZulipServer, realm_uuid: str, user_uuid: str
+) -> Optional[RemoteRealm]:
+    """
+    Tries to fetch RemoteRealm for the given realm_uuid and server. Otherwise,
+    returns None and logs what happened using request and user_uuid args to make
+    the output more informative.
+    """
+
+    try:
+        remote_realm = RemoteRealm.objects.get(uuid=realm_uuid)
+    except RemoteRealm.DoesNotExist:
+        logger.info(
+            "%s: Received request for unknown realm %s, server %s, user %s",
+            request.path,
+            realm_uuid,
+            server.id,
+            user_uuid,
+        )
+        return None
+
+    if remote_realm.server_id != server.id:
+        logger.warning(
+            "%s: Realm %s exists, but not registered to server %s",
+            request.path,
+            realm_uuid,
+            server.id,
+        )
+        return None
+
+    return remote_realm
+
+
 @has_request_variables
 def remote_server_notify_push(
     request: HttpRequest,
@@ -345,12 +390,10 @@ def remote_server_notify_push(
     realm_uuid = payload.get("realm_uuid")
     remote_realm = None
     if realm_uuid is not None:
-        try:
-            remote_realm = RemoteRealm.objects.get(uuid=realm_uuid, server=server)
-        except RemoteRealm.DoesNotExist:
-            # We don't yet have a RemoteRealm for this realm. E.g. the server hasn't yet
-            # submitted analytics data since the realm's creation.
-            remote_realm = None
+        assert isinstance(
+            user_uuid, str
+        ), "Servers new enough to send realm_uuid, should also have user_uuid"
+        remote_realm = get_remote_realm_helper(request, server, realm_uuid, user_uuid)
 
     android_devices = list(
         RemotePushDeviceToken.objects.filter(