rest: Extract remote_server_path from rest_path.

This allows us to separate the zilencer paths from other JSON paths,
with explicit type annotation expecting `RemoteZulipServer` as the
second parameter of the handler using
authenticated_remote_server_view.

The test case is also updated to remove a test for a situation that no
longer occurs anymore, since we don't perform subdomain checks on
remote servers.

Signed-off-by: Zixuan James Li <p359101898@gmail.com>

zerver/decorator.py

@@ -62,9 +62,6 @@ from zerver.lib.users import is_2fa_verified
 from zerver.lib.utils import has_api_key_format, statsd
 from zerver.models import UserProfile, get_client, get_user_profile_by_api_key
 
-if settings.ZILENCER_ENABLED:
-    from zilencer.models import RemoteZulipServer
-
 if TYPE_CHECKING:
     from django.http.request import _ImmutableQueryDict
 
@@ -242,18 +239,12 @@ def validate_api_key(
     api_key: str,
     allow_webhook_access: bool = False,
     client_name: Optional[str] = None,
-) -> Union[UserProfile, "RemoteZulipServer"]:
+) -> UserProfile:
     # Remove whitespace to protect users from trivial errors.
     api_key = api_key.strip()
     if role is not None:
         role = role.strip()
 
-    # If `role` doesn't look like an email, it might be a uuid.
-    if settings.ZILENCER_ENABLED and role is not None and "@" not in role:
-        from zilencer.auth import validate_remote_server
-
-        return validate_remote_server(request, role, api_key)
-
     user_profile = access_user_by_api_key(request, api_key, email=role)
     if user_profile.is_incoming_webhook and not allow_webhook_access:
         raise JsonableError(_("This API is not available to incoming webhook bots."))
@@ -749,7 +740,6 @@ def authenticated_rest_api_view(
 
             # Now we try to do authentication or die
             try:
-                # profile is a Union[UserProfile, RemoteZulipServer]
                 profile = validate_api_key(
                     request,
                     role,

zerver/tests/test_decorators.py

@@ -2015,6 +2015,11 @@ class RestAPITest(ZulipTestCase):
         self.assertEqual(result.status_code, 405)
         self.assert_in_response("Method Not Allowed", result)
 
+        with self.settings(ZILENCER_ENABLED=True):
+            result = self.client_patch("/json/remotes/push/register")
+            self.assertEqual(result.status_code, 405)
+            self.assert_in_response("Method Not Allowed", result)
+
     def test_options_method(self) -> None:
         self.login("hamlet")
         result = self.client_options("/json/users")

zerver/tests/test_push_notifications.py

@@ -205,22 +205,6 @@ class PushBouncerNotificationTest(BouncerTestCase):
 
         hamlet = self.example_user("hamlet")
 
-        with self.assertLogs(level="WARNING") as warn_log:
-            result = self.api_post(
-                hamlet,
-                endpoint,
-                dict(user_id=user_id, token_kin=token_kind, token=token),
-            )
-        self.assertEqual(
-            warn_log.output,
-            [
-                "WARNING:root:User hamlet@zulip.com (zulip) attempted to access API on wrong subdomain ()"
-            ],
-        )
-        self.assert_json_error(
-            result, "Account is not associated with this subdomain", status_code=401
-        )
-
         # We need the root ('') subdomain to be in use for this next
         # test, since the push bouncer API is only available there:
         realm = get_realm("zulip")

zilencer/auth.py

@@ -1,14 +1,29 @@
-from django.http import HttpRequest
+from functools import wraps
+from typing import Any, Callable
+
+from django.http import HttpRequest, HttpResponse
+from django.urls import path
+from django.urls.resolvers import URLPattern
 from django.utils.crypto import constant_time_compare
 from django.utils.translation import gettext as _
+from typing_extensions import Concatenate, ParamSpec
 
-from zerver.decorator import process_client
-from zerver.lib.exceptions import ErrorCode, JsonableError, RemoteServerDeactivatedError
+from zerver.decorator import get_basic_credentials, process_client
+from zerver.lib.exceptions import (
+    ErrorCode,
+    JsonableError,
+    RemoteServerDeactivatedError,
+    UnauthorizedError,
+)
+from zerver.lib.rate_limiter import rate_limit
 from zerver.lib.request import RequestNotes
+from zerver.lib.rest import get_target_view_function_or_response
 from zerver.lib.subdomains import get_subdomain
 from zerver.models import Realm
 from zilencer.models import RemoteZulipServer, get_remote_server_by_uuid
 
+ParamT = ParamSpec("ParamT")
+
 
 class InvalidZulipServerError(JsonableError):
     code = ErrorCode.INVALID_ZULIP_SERVER
@@ -48,3 +63,39 @@ def validate_remote_server(
     RequestNotes.get_notes(request).remote_server = remote_server
     process_client(request)
     return remote_server
+
+
+def authenticated_remote_server_view(
+    view_func: Callable[Concatenate[HttpRequest, RemoteZulipServer, ParamT], HttpResponse]
+) -> Callable[Concatenate[HttpRequest, ParamT], HttpResponse]:
+    @wraps(view_func)
+    def _wrapped_view_func(
+        request: HttpRequest, /, *args: ParamT.args, **kwargs: ParamT.kwargs
+    ) -> HttpResponse:
+        role, api_key = get_basic_credentials(request)
+        if "@" in role:
+            raise JsonableError(_("Must validate with valid Zulip server API key"))
+        try:
+            remote_server = validate_remote_server(request, role, api_key)
+        except JsonableError as e:
+            raise UnauthorizedError(e.msg)
+
+        rate_limit(request)
+        return view_func(request, remote_server, *args, **kwargs)
+
+    return _wrapped_view_func
+
+
+def remote_server_dispatch(request: HttpRequest, **kwargs: Any) -> HttpResponse:
+    result = get_target_view_function_or_response(request, kwargs)
+    if isinstance(result, HttpResponse):
+        return result
+    target_function, view_flags = result
+    return authenticated_remote_server_view(target_function)(request, **kwargs)
+
+
+def remote_server_path(
+    route: str,
+    **handlers: Callable[Concatenate[HttpRequest, RemoteZulipServer, ParamT], HttpResponse],
+) -> URLPattern:
+    return path(route, remote_server_dispatch, handlers)

zilencer/urls.py

@@ -3,7 +3,7 @@ from typing import Any
 from django.conf.urls import include
 from django.urls import path
 
-from zerver.lib.rest import rest_path
+from zilencer.auth import remote_server_path
 from zilencer.views import (
     deactivate_remote_server,
     register_remote_push_device,
@@ -19,16 +19,16 @@ i18n_urlpatterns: Any = []
 
 # Zilencer views following the REST API style
 v1_api_and_json_patterns = [
-    rest_path("remotes/push/register", POST=register_remote_push_device),
-    rest_path("remotes/push/unregister", POST=unregister_remote_push_device),
-    rest_path("remotes/push/unregister/all", POST=unregister_all_remote_push_devices),
-    rest_path("remotes/push/notify", POST=remote_server_notify_push),
+    remote_server_path("remotes/push/register", POST=register_remote_push_device),
+    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),
     # Push signup doesn't use the REST API, since there's no auth.
     path("remotes/server/register", register_remote_server),
-    rest_path("remotes/server/deactivate", POST=deactivate_remote_server),
+    remote_server_path("remotes/server/deactivate", POST=deactivate_remote_server),
     # For receiving table data used in analytics and billing
-    rest_path("remotes/server/analytics", POST=remote_server_post_analytics),
-    rest_path("remotes/server/analytics/status", GET=remote_server_check_analytics),
+    remote_server_path("remotes/server/analytics", POST=remote_server_post_analytics),
+    remote_server_path("remotes/server/analytics/status", GET=remote_server_check_analytics),
 ]
 
 urlpatterns = [