wehub-resource-sync
130 lines
4.4 KiB
Python
130 lines
4.4 KiB
Python
import logging
|
|
from types import ModuleType
|
|
from typing import Dict, List, Optional
|
|
|
|
from ray._private.authentication import (
|
|
authentication_constants,
|
|
authentication_utils as auth_utils,
|
|
)
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def get_token_auth_middleware(
|
|
aiohttp_module: ModuleType,
|
|
whitelisted_exact_paths: Optional[List[str]] = None,
|
|
whitelisted_path_prefixes: Optional[List[str]] = None,
|
|
):
|
|
"""Internal helper to create token auth middleware with provided modules.
|
|
|
|
Args:
|
|
aiohttp_module: The aiohttp module to use
|
|
whitelisted_exact_paths: List of exact paths that don't require authentication
|
|
whitelisted_path_prefixes: List of path prefixes that don't require authentication
|
|
Returns:
|
|
An aiohttp middleware function
|
|
"""
|
|
|
|
@aiohttp_module.web.middleware
|
|
async def token_auth_middleware(request, handler):
|
|
"""Middleware to validate bearer tokens when token authentication is enabled.
|
|
|
|
In minimal Ray installations (without ray._raylet), this middleware is a no-op
|
|
and passes all requests through without authentication.
|
|
"""
|
|
# No-op if token auth is not enabled or raylet is not available
|
|
if not auth_utils.is_token_auth_enabled():
|
|
return await handler(request)
|
|
|
|
# skip authentication for whitelisted paths
|
|
if (whitelisted_exact_paths and request.path in whitelisted_exact_paths) or (
|
|
whitelisted_path_prefixes
|
|
and request.path.startswith(tuple(whitelisted_path_prefixes))
|
|
):
|
|
return await handler(request)
|
|
|
|
# Try to get authentication token from multiple sources (in priority order):
|
|
# 1. Standard "Authorization" header (for API clients, SDKs)
|
|
# 2. Fallback "X-Ray-Authorization" header (for proxies and KubeRay)
|
|
# 3. Cookie (for web dashboard sessions)
|
|
|
|
auth_header = request.headers.get(
|
|
authentication_constants.AUTHORIZATION_HEADER_NAME, ""
|
|
)
|
|
|
|
if not auth_header:
|
|
auth_header = request.headers.get(
|
|
authentication_constants.RAY_AUTHORIZATION_HEADER_NAME, ""
|
|
)
|
|
|
|
if not auth_header:
|
|
token = request.cookies.get(
|
|
authentication_constants.AUTHENTICATION_TOKEN_COOKIE_NAME
|
|
)
|
|
if token:
|
|
# Format as Bearer token for validation
|
|
auth_header = (
|
|
authentication_constants.AUTHORIZATION_BEARER_PREFIX + token
|
|
)
|
|
|
|
if not auth_header:
|
|
return aiohttp_module.web.Response(
|
|
status=401, text="Unauthorized: Missing authentication token"
|
|
)
|
|
|
|
if not auth_utils.validate_request_token(auth_header):
|
|
return aiohttp_module.web.Response(
|
|
status=403, text="Forbidden: Invalid authentication token"
|
|
)
|
|
|
|
return await handler(request)
|
|
|
|
return token_auth_middleware
|
|
|
|
|
|
def get_auth_headers_if_auth_enabled(user_headers: Dict[str, str]) -> Dict[str, str]:
|
|
|
|
if not auth_utils.is_token_auth_enabled():
|
|
return {}
|
|
|
|
from ray._raylet import AuthenticationTokenLoader
|
|
|
|
# Check if user provided their own Authorization header (case-insensitive)
|
|
has_user_auth = any(
|
|
key.lower() == authentication_constants.AUTHORIZATION_HEADER_NAME
|
|
for key in user_headers.keys()
|
|
)
|
|
if has_user_auth:
|
|
# User has provided their own auth header, don't override
|
|
return {}
|
|
|
|
token_loader = AuthenticationTokenLoader.instance()
|
|
auth_headers = token_loader.get_token_for_http_header()
|
|
|
|
if not auth_headers:
|
|
# Token auth enabled but no token found
|
|
logger.warning(
|
|
"Token authentication is enabled but no token was found. "
|
|
"Requests to authenticated clusters will fail."
|
|
)
|
|
|
|
return auth_headers
|
|
|
|
|
|
def format_authentication_http_error(status: int, body: str) -> Optional[str]:
|
|
"""Return a user-friendly authentication error message, if applicable."""
|
|
|
|
if status == 401:
|
|
return "Authentication required: {body}\n\n{details}".format(
|
|
body=body,
|
|
details=authentication_constants.TOKEN_AUTH_ENABLED_BUT_NO_TOKEN_FOUND_ERROR_MESSAGE,
|
|
)
|
|
|
|
if status == 403:
|
|
return "Authentication failed: {body}\n\n{details}".format(
|
|
body=body,
|
|
details=authentication_constants.TOKEN_INVALID_ERROR_MESSAGE,
|
|
)
|
|
|
|
return None
|