response: Replace json_unauthorized with UnauthorizedError.

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

zerver/decorator.py

@@ -52,6 +52,7 @@ from zerver.lib.exceptions import (
     RateLimited,
     RealmDeactivatedError,
     RemoteServerDeactivatedError,
+    UnauthorizedError,
     UnsupportedWebhookEventType,
     UserDeactivatedError,
     WebhookError,
@@ -59,7 +60,7 @@ from zerver.lib.exceptions import (
 from zerver.lib.queue import queue_json_publish
 from zerver.lib.rate_limiter import RateLimitedIPAddr, RateLimitedUser
 from zerver.lib.request import REQ, RequestNotes, has_request_variables
-from zerver.lib.response import json_method_not_allowed, json_success, json_unauthorized
+from zerver.lib.response import json_method_not_allowed, json_success
 from zerver.lib.subdomains import get_subdomain, user_matches_subdomain
 from zerver.lib.timestamp import datetime_to_timestamp, timestamp_to_datetime
 from zerver.lib.types import ViewFuncT
@@ -679,9 +680,9 @@ def authenticated_rest_api_view(
                     # So we ask the user to replace them with %40
                     role = role.replace("%40", "@")
             except ValueError:
-                return json_unauthorized(_("Invalid authorization header for basic auth"))
+                raise UnauthorizedError(_("Invalid authorization header for basic auth"))
             except KeyError:
-                return json_unauthorized(_("Missing authorization header for basic auth"))
+                raise UnauthorizedError(_("Missing authorization header for basic auth"))
 
             # Now we try to do authentication or die
             try:
@@ -694,7 +695,7 @@ def authenticated_rest_api_view(
                     client_name=full_webhook_client_name(webhook_client_name),
                 )
             except JsonableError as e:
-                return json_unauthorized(e.msg)
+                raise UnauthorizedError(e.msg)
             try:
                 if not skip_rate_limiting:
                     # Apply rate limiting
@@ -777,7 +778,7 @@ def authenticate_log_and_execute_json(
 
     if not request.user.is_authenticated:
         if not allow_unauthenticated:
-            return json_unauthorized()
+            raise UnauthorizedError()
 
         process_client(
             request,

zerver/lib/exceptions.py

@@ -36,6 +36,7 @@ class ErrorCode(Enum):
     PASSWORD_AUTH_DISABLED = auto()
     PASSWORD_RESET_REQUIRED = auto()
     AUTHENTICATION_FAILED = auto()
+    UNAUTHORIZED = auto()
 
 
 class JsonableError(Exception):
@@ -124,6 +125,28 @@ class JsonableError(Exception):
         return self.msg
 
 
+class UnauthorizedError(JsonableError):
+    code: ErrorCode = ErrorCode.UNAUTHORIZED
+    http_status_code: int = 401
+
+    def __init__(self, msg: Optional[str] = None, www_authenticate: Optional[str] = None) -> None:
+        if msg is None:
+            msg = _("Not logged in: API authentication or user session required")
+        super().__init__(msg)
+        if www_authenticate is None:
+            self.www_authenticate = 'Basic realm="zulip"'
+        elif www_authenticate == "session":
+            self.www_authenticate = 'Session realm="zulip"'
+        else:
+            raise AssertionError("Invalid www_authenticate value!")
+
+    @property
+    def extra_headers(self) -> Dict[str, Any]:
+        extra_headers_dict = super().extra_headers
+        extra_headers_dict["WWW-Authenticate"] = self.www_authenticate
+        return extra_headers_dict
+
+
 class StreamDoesNotExistError(JsonableError):
     code = ErrorCode.STREAM_DOES_NOT_EXIST
     data_fields = ["stream"]

zerver/lib/response.py

@@ -2,35 +2,16 @@ from typing import Any, List, Mapping, Optional
 
 import orjson
 from django.http import HttpRequest, HttpResponse, HttpResponseNotAllowed
-from django.utils.translation import gettext as _
 
-from zerver.lib.exceptions import JsonableError
-
-
-class HttpResponseUnauthorized(HttpResponse):
-    status_code = 401
-
-    def __init__(self, realm: str, www_authenticate: Optional[str] = None) -> None:
-        HttpResponse.__init__(self)
-        if www_authenticate is None:
-            self["WWW-Authenticate"] = f'Basic realm="{realm}"'
-        elif www_authenticate == "session":
-            self["WWW-Authenticate"] = f'Session realm="{realm}"'
-        else:
-            raise AssertionError("Invalid www_authenticate value!")
+from zerver.lib.exceptions import JsonableError, UnauthorizedError
 
 
 def json_unauthorized(
     message: Optional[str] = None, www_authenticate: Optional[str] = None
 ) -> HttpResponse:
-    if message is None:
-        message = _("Not logged in: API authentication or user session required")
-    resp = HttpResponseUnauthorized("zulip", www_authenticate=www_authenticate)
-    resp.content = orjson.dumps(
-        {"result": "error", "msg": message},
-        option=orjson.OPT_APPEND_NEWLINE,
+    return json_response_from_error(
+        UnauthorizedError(msg=message, www_authenticate=www_authenticate)
     )
-    return resp
 
 
 def json_method_not_allowed(methods: List[str]) -> HttpResponseNotAllowed: