rate_limit: Stop wrapping rate limited functions.

This refactors `rate_limit` so that we no longer use it as a decorator.
This is a workaround to python/mypy#12909 as
`rate_limit` previous expects different parameters than its callers.

Signed-off-by: Zixuan James Li <[email protected]>

Author: PIG208

zerver/decorator.py

@@ -508,7 +508,8 @@ def _wrapped_view_func(
         request: HttpRequest, /, *args: ParamT.args, **kwargs: ParamT.kwargs
     ) -> HttpResponse:
         process_client(request, request.user, is_browser_view=True, query=view_func.__name__)
-        return rate_limit(view_func)(request, *args, **kwargs)
+        rate_limit(request)
+        return view_func(request, *args, **kwargs)
 
     return _wrapped_view_func
 
@@ -681,10 +682,8 @@ def _wrapped_func_arguments(
         ) -> HttpResponse:
             user_profile = validate_api_key(request, None, api_key, False)
             if not skip_rate_limiting:
-                limited_func = rate_limit(view_func)
-            else:
-                limited_func = view_func
-            return limited_func(request, user_profile, *args, **kwargs)
+                rate_limit(request)
+            return view_func(request, user_profile, *args, **kwargs)
 
         return _wrapped_func_arguments
 
@@ -745,10 +744,8 @@ def _wrapped_func_arguments(
             try:
                 if not skip_rate_limiting:
                     # Apply rate limiting
-                    target_view_func = rate_limit(view_func)
-                else:
-                    target_view_func = view_func
-                return target_view_func(request, profile, *args, **kwargs)
+                    rate_limit(request)
+                return view_func(request, profile, *args, **kwargs)
             except Exception as err:
                 if not webhook_client_name:
                     raise err
@@ -822,9 +819,7 @@ def authenticate_log_and_execute_json(
     **kwargs: object,
 ) -> HttpResponse:
     if not skip_rate_limiting:
-        limited_view_func = rate_limit(view_func)
-    else:
-        limited_view_func = view_func
+        rate_limit(request)
 
     if not request.user.is_authenticated:
         if not allow_unauthenticated:
@@ -835,7 +830,7 @@ def authenticate_log_and_execute_json(
             is_browser_view=True,
             query=view_func.__name__,
         )
-        return limited_view_func(request, request.user, *args, **kwargs)
+        return view_func(request, request.user, *args, **kwargs)
 
     user_profile = request.user
     validate_account_and_subdomain(request, user_profile)
@@ -844,7 +839,7 @@ def authenticate_log_and_execute_json(
         raise JsonableError(_("Webhook bots can only access webhooks"))
 
     process_client(request, user_profile, is_browser_view=True, query=view_func.__name__)
-    return limited_view_func(request, user_profile, *args, **kwargs)
+    return view_func(request, user_profile, *args, **kwargs)
 
 
 # Checks if the user is logged in.  If not, return an error (the
@@ -1027,36 +1022,22 @@ def rate_limit_remote_server(
         raise e
 
 
-def rate_limit(func: ViewFuncT) -> ViewFuncT:
-    """Rate-limits a view."""
-
-    @wraps(func)
-    def wrapped_func(request: HttpRequest, *args: object, **kwargs: object) -> HttpResponse:
-
-        # It is really tempting to not even wrap our original function
-        # when settings.RATE_LIMITING is False, but it would make
-        # for awkward unit testing in some situations.
-        if not settings.RATE_LIMITING:
-            return func(request, *args, **kwargs)
-
-        if client_is_exempt_from_rate_limiting(request):
-            return func(request, *args, **kwargs)
-
-        user = request.user
-        remote_server = RequestNotes.get_notes(request).remote_server
+def rate_limit(request: HttpRequest) -> None:
+    if not settings.RATE_LIMITING:
+        return
 
-        if settings.ZILENCER_ENABLED and remote_server is not None:
-            rate_limit_remote_server(request, remote_server, domain="api_by_remote_server")
-        elif not user.is_authenticated:
-            rate_limit_request_by_ip(request, domain="api_by_ip")
-            return func(request, *args, **kwargs)
-        else:
-            assert isinstance(user, UserProfile)
-            rate_limit_user(request, user, domain="api_by_user")
+    if client_is_exempt_from_rate_limiting(request):
+        return
 
-        return func(request, *args, **kwargs)
+    remote_server = RequestNotes.get_notes(request).remote_server
 
-    return cast(ViewFuncT, wrapped_func)  # https://github.com/python/mypy/issues/1927
+    if settings.ZILENCER_ENABLED and remote_server is not None:
+        rate_limit_remote_server(request, remote_server, domain="api_by_remote_server")
+    elif not request.user.is_authenticated:
+        rate_limit_request_by_ip(request, domain="api_by_ip")
+    else:
+        assert isinstance(request.user, UserProfile)
+        rate_limit_user(request, request.user, domain="api_by_user")
 
 
 def return_success_on_head_request(

zerver/tests/test_decorators.py

@@ -519,7 +519,7 @@ def my_view(request: HttpRequest, user_profile: UserProfile) -> HttpResponse:
         request.method = "POST"
         request.user = self.example_user("hamlet")
         with mock.patch("zerver.decorator.rate_limit") as rate_limit_mock:
-            result = my_unlimited_view(request)
+            result = my_unlimited_view(request, request.user)
 
         self.assert_json_success(result)
         self.assertFalse(rate_limit_mock.called)
@@ -528,7 +528,7 @@ def my_view(request: HttpRequest, user_profile: UserProfile) -> HttpResponse:
         request.method = "POST"
         request.user = self.example_user("hamlet")
         with mock.patch("zerver.decorator.rate_limit") as rate_limit_mock:
-            result = my_rate_limited_view(request)
+            result = my_rate_limited_view(request, request.user)
 
         # Don't assert json_success, since it'll be the rate_limit mock object
         self.assertTrue(rate_limit_mock.called)
@@ -630,10 +630,9 @@ def test_authenticated_rest_api_view_errors(self) -> None:
 class RateLimitTestCase(ZulipTestCase):
     def get_ratelimited_view(self) -> Callable[..., HttpResponse]:
         def f(req: Any) -> HttpResponse:
+            rate_limit(req)
             return json_response(msg="some value")
 
-        f = rate_limit(f)
-
         return f
 
     def errors_disallowed(self) -> Any: