Remove handle_exc_for_ws function from security utils

This commit is contained in:
Mix 2024-04-02 23:01:26 +08:00 committed by HexMix
parent f40bbafd8d
commit cb368d439b
5 changed files with 3 additions and 44 deletions

View File

@ -3,7 +3,6 @@ from typing import Optional
from annotated_doc import Doc
from fastapi.openapi.models import APIKey, APIKeyIn
from fastapi.security.base import SecurityBase
from fastapi.security.utils import handle_exc_for_ws
from starlette.exceptions import HTTPException
from starlette.requests import HTTPConnection
from starlette.status import HTTP_403_FORBIDDEN
@ -109,7 +108,6 @@ class APIKeyQuery(APIKeyBase):
self.scheme_name = scheme_name or self.__class__.__name__
self.auto_error = auto_error
@handle_exc_for_ws
async def __call__(self, request: HTTPConnection) -> Optional[str]:
api_key = request.query_params.get(self.model.name)
return self.check_api_key(api_key, self.auto_error)
@ -198,7 +196,6 @@ class APIKeyHeader(APIKeyBase):
self.scheme_name = scheme_name or self.__class__.__name__
self.auto_error = auto_error
@handle_exc_for_ws
async def __call__(self, request: HTTPConnection) -> Optional[str]:
api_key = request.headers.get(self.model.name)
return self.check_api_key(api_key, self.auto_error)
@ -287,7 +284,6 @@ class APIKeyCookie(APIKeyBase):
self.scheme_name = scheme_name or self.__class__.__name__
self.auto_error = auto_error
@handle_exc_for_ws
async def __call__(self, request: HTTPConnection) -> Optional[str]:
api_key = request.cookies.get(self.model.name)
return self.check_api_key(api_key, self.auto_error)

View File

@ -7,7 +7,7 @@ from fastapi.exceptions import HTTPException
from fastapi.openapi.models import HTTPBase as HTTPBaseModel
from fastapi.openapi.models import HTTPBearer as HTTPBearerModel
from fastapi.security.base import SecurityBase
from fastapi.security.utils import get_authorization_scheme_param, handle_exc_for_ws
from fastapi.security.utils import get_authorization_scheme_param
from pydantic import BaseModel
from starlette.requests import HTTPConnection
from starlette.status import HTTP_401_UNAUTHORIZED, HTTP_403_FORBIDDEN
@ -80,7 +80,6 @@ class HTTPBase(SecurityBase):
self.scheme_name = scheme_name or self.__class__.__name__
self.auto_error = auto_error
@handle_exc_for_ws
async def __call__(
self, request: HTTPConnection
) -> Optional[HTTPAuthorizationCredentials]:
@ -186,7 +185,6 @@ class HTTPBasic(HTTPBase):
self.realm = realm
self.auto_error = auto_error
@handle_exc_for_ws
async def __call__( # type: ignore
self, request: HTTPConnection
) -> Optional[HTTPBasicCredentials]:
@ -301,7 +299,6 @@ class HTTPBearer(HTTPBase):
self.scheme_name = scheme_name or self.__class__.__name__
self.auto_error = auto_error
@handle_exc_for_ws
async def __call__(
self, request: HTTPConnection
) -> Optional[HTTPAuthorizationCredentials]:
@ -404,7 +401,6 @@ class HTTPDigest(HTTPBase):
self.scheme_name = scheme_name or self.__class__.__name__
self.auto_error = auto_error
@handle_exc_for_ws
async def __call__(
self, request: HTTPConnection
) -> Optional[HTTPAuthorizationCredentials]:

View File

@ -6,7 +6,7 @@ from fastapi.openapi.models import OAuth2 as OAuth2Model
from fastapi.openapi.models import OAuthFlows as OAuthFlowsModel
from fastapi.param_functions import Form
from fastapi.security.base import SecurityBase
from fastapi.security.utils import get_authorization_scheme_param, handle_exc_for_ws
from fastapi.security.utils import get_authorization_scheme_param
from starlette.requests import HTTPConnection
from starlette.status import HTTP_401_UNAUTHORIZED, HTTP_403_FORBIDDEN
@ -377,7 +377,6 @@ class OAuth2(SecurityBase):
self.scheme_name = scheme_name or self.__class__.__name__
self.auto_error = auto_error
@handle_exc_for_ws
async def __call__(self, request: HTTPConnection) -> Optional[str]:
authorization = request.headers.get("Authorization")
if not authorization:
@ -487,7 +486,6 @@ class OAuth2PasswordBearer(OAuth2):
auto_error=auto_error,
)
@handle_exc_for_ws
async def __call__(self, request: HTTPConnection) -> Optional[str]:
authorization = request.headers.get("Authorization")
scheme, param = get_authorization_scheme_param(authorization)
@ -598,7 +596,6 @@ class OAuth2AuthorizationCodeBearer(OAuth2):
auto_error=auto_error,
)
@handle_exc_for_ws
async def __call__(self, request: HTTPConnection) -> Optional[str]:
authorization = request.headers.get("Authorization")
scheme, param = get_authorization_scheme_param(authorization)

View File

@ -3,7 +3,6 @@ from typing import Optional
from annotated_doc import Doc
from fastapi.openapi.models import OpenIdConnect as OpenIdConnectModel
from fastapi.security.base import SecurityBase
from fastapi.security.utils import handle_exc_for_ws
from starlette.exceptions import HTTPException
from starlette.requests import HTTPConnection
from starlette.status import HTTP_403_FORBIDDEN
@ -74,7 +73,6 @@ class OpenIdConnect(SecurityBase):
self.scheme_name = scheme_name or self.__class__.__name__
self.auto_error = auto_error
@handle_exc_for_ws
async def __call__(self, request: HTTPConnection) -> Optional[str]:
authorization = request.headers.get("Authorization")
if not authorization:

View File

@ -1,10 +1,4 @@
from functools import wraps
from typing import Any, Awaitable, Callable, Optional, Tuple, TypeVar
from fastapi.exceptions import HTTPException, WebSocketException
from starlette.requests import HTTPConnection
from starlette.status import WS_1008_POLICY_VIOLATION
from starlette.websockets import WebSocket
from typing import Optional, Tuple
def get_authorization_scheme_param(
@ -14,25 +8,3 @@ def get_authorization_scheme_param(
return "", ""
scheme, _, param = authorization_header_value.partition(" ")
return scheme, param
_SecurityDepFunc = TypeVar(
"_SecurityDepFunc", bound=Callable[[Any, HTTPConnection], Awaitable[Any]]
)
def handle_exc_for_ws(func: _SecurityDepFunc) -> _SecurityDepFunc:
@wraps(func)
async def wrapper(self: Any, request: HTTPConnection) -> Any:
try:
return await func(self, request)
except HTTPException as e:
if not isinstance(request, WebSocket):
raise e
# close before accepted with result a HTTP 403 so the exception argument is ignored
# ref: https://asgi.readthedocs.io/en/latest/specs/www.html#close-send-event
raise WebSocketException(
code=WS_1008_POLICY_VIOLATION, reason=e.detail
) from None
return wrapper # type: ignore