chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,14 @@
|
||||
# Authentication error messages
|
||||
TOKEN_AUTH_ENABLED_BUT_NO_TOKEN_FOUND_ERROR_MESSAGE = (
|
||||
"Token authentication is enabled but no authentication token was found"
|
||||
)
|
||||
|
||||
TOKEN_INVALID_ERROR_MESSAGE = "Token authentication is enabled but the authentication token is invalid or incorrect." # noqa: E501
|
||||
|
||||
# HTTP header and cookie constants
|
||||
AUTHORIZATION_HEADER_NAME = "authorization"
|
||||
AUTHORIZATION_BEARER_PREFIX = "Bearer "
|
||||
RAY_AUTHORIZATION_HEADER_NAME = "x-ray-authorization"
|
||||
|
||||
AUTHENTICATION_TOKEN_COOKIE_NAME = "ray-authentication-token"
|
||||
AUTHENTICATION_TOKEN_COOKIE_MAX_AGE = 30 * 24 * 60 * 60 # 30 days
|
||||
@@ -0,0 +1,9 @@
|
||||
import secrets
|
||||
|
||||
|
||||
def generate_new_authentication_token() -> str:
|
||||
"""Generate an authentication token for the cluster.
|
||||
|
||||
256 bits of entropy is considered sufficient to be durable to brute force attacks.
|
||||
"""
|
||||
return secrets.token_hex(32)
|
||||
@@ -0,0 +1,103 @@
|
||||
"""Authentication token setup for Ray.
|
||||
|
||||
This module provides functions to generate and save authentication tokens
|
||||
for Ray's token-based authentication system. Token loading and caching is
|
||||
handled by the C++ AuthenticationTokenLoader.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
from ray._private.authentication.authentication_constants import (
|
||||
TOKEN_AUTH_ENABLED_BUT_NO_TOKEN_FOUND_ERROR_MESSAGE,
|
||||
)
|
||||
from ray._private.authentication.authentication_token_generator import (
|
||||
generate_new_authentication_token,
|
||||
)
|
||||
from ray._raylet import (
|
||||
AuthenticationMode,
|
||||
AuthenticationTokenLoader,
|
||||
get_authentication_mode,
|
||||
)
|
||||
from ray.exceptions import AuthenticationError
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def generate_and_save_token() -> None:
|
||||
"""Generate a new random token and save it in the default token path.
|
||||
|
||||
Returns:
|
||||
The newly generated authentication token.
|
||||
"""
|
||||
# Generate a UUID-based token
|
||||
token = generate_new_authentication_token()
|
||||
|
||||
token_path = _get_default_token_path()
|
||||
try:
|
||||
# Create directory if it doesn't exist
|
||||
token_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# Write token to file with explicit flush and fsync
|
||||
with open(token_path, "w") as f:
|
||||
f.write(token)
|
||||
|
||||
logger.info(f"Generated new authentication token and saved to {token_path}")
|
||||
except Exception:
|
||||
raise
|
||||
|
||||
|
||||
def _get_default_token_path() -> Path:
|
||||
"""Get the default token file path (~/.ray/auth_token).
|
||||
|
||||
Returns:
|
||||
Path object pointing to ~/.ray/auth_token
|
||||
"""
|
||||
return Path.home() / ".ray" / "auth_token"
|
||||
|
||||
|
||||
def ensure_token_if_auth_enabled(
|
||||
system_config: Optional[Dict[str, Any]] = None, create_token_if_missing: bool = True
|
||||
) -> None:
|
||||
"""Check authentication settings and set up token resources if authentication is enabled.
|
||||
|
||||
Ray calls this early during ray.init() to do the following for token-based authentication:
|
||||
1. Check whether you enabled token-based authentication.
|
||||
2. Make sure a token is available if authentication is enabled.
|
||||
3. Generate and save a default token for new local clusters if one doesn't already exist.
|
||||
|
||||
Args:
|
||||
system_config: Ray raises an error if you set AUTH_MODE in system_config instead of the environment.
|
||||
create_token_if_missing: Generate a new token if one doesn't already exist.
|
||||
|
||||
Raises:
|
||||
RuntimeError: Ray raises this error if authentication is enabled but no token is found when connecting
|
||||
to an existing cluster.
|
||||
"""
|
||||
|
||||
# Check if you enabled token authentication.
|
||||
if get_authentication_mode() != AuthenticationMode.TOKEN:
|
||||
if (
|
||||
system_config
|
||||
and "AUTH_MODE" in system_config
|
||||
and system_config["AUTH_MODE"] != "disabled"
|
||||
):
|
||||
raise RuntimeError(
|
||||
"Set authentication mode can only be set with the `RAY_AUTH_MODE` environment variable, not using the system_config."
|
||||
)
|
||||
return
|
||||
|
||||
token_loader = AuthenticationTokenLoader.instance()
|
||||
|
||||
if not token_loader.has_token(ignore_auth_mode=True):
|
||||
if create_token_if_missing:
|
||||
# Generate a new token.
|
||||
generate_and_save_token()
|
||||
|
||||
# Reload the cache so subsequent calls to token_loader read the new token.
|
||||
token_loader.reset_cache()
|
||||
else:
|
||||
raise AuthenticationError(
|
||||
TOKEN_AUTH_ENABLED_BUT_NO_TOKEN_FOUND_ERROR_MESSAGE
|
||||
)
|
||||
@@ -0,0 +1,58 @@
|
||||
try:
|
||||
from ray._raylet import (
|
||||
AuthenticationMode,
|
||||
get_authentication_mode,
|
||||
validate_authentication_token,
|
||||
)
|
||||
|
||||
_RAYLET_AVAILABLE = True
|
||||
except ImportError:
|
||||
# ray._raylet not available during doc builds
|
||||
_RAYLET_AVAILABLE = False
|
||||
|
||||
|
||||
def is_token_auth_enabled() -> bool:
|
||||
"""Check if token authentication is enabled.
|
||||
|
||||
Returns:
|
||||
bool: True if AUTH_MODE is set to "token".
|
||||
"""
|
||||
if not _RAYLET_AVAILABLE:
|
||||
return False
|
||||
|
||||
return get_authentication_mode() == AuthenticationMode.TOKEN
|
||||
|
||||
|
||||
def validate_request_token(auth_header: str) -> bool:
|
||||
"""Validate the Authorization header from an HTTP request.
|
||||
|
||||
Args:
|
||||
auth_header: The Authorization header value (e.g., "Bearer <token>")
|
||||
|
||||
Returns:
|
||||
bool: True if token is valid, False otherwise
|
||||
"""
|
||||
if not _RAYLET_AVAILABLE or not auth_header:
|
||||
return False
|
||||
|
||||
# validate_authentication_token expects full "Bearer <token>" format
|
||||
# and performs equality comparison via C++ layer
|
||||
return validate_authentication_token(auth_header)
|
||||
|
||||
|
||||
def get_authentication_mode_name(mode: AuthenticationMode) -> str:
|
||||
"""Convert AuthenticationMode enum value to string name.
|
||||
|
||||
Args:
|
||||
mode: AuthenticationMode enum value from ray._raylet
|
||||
|
||||
Returns:
|
||||
String name: "disabled", "token"
|
||||
"""
|
||||
from ray._raylet import AuthenticationMode
|
||||
|
||||
_MODE_NAMES = {
|
||||
AuthenticationMode.DISABLED: "disabled",
|
||||
AuthenticationMode.TOKEN: "token",
|
||||
}
|
||||
return _MODE_NAMES.get(mode, "unknown")
|
||||
@@ -0,0 +1,155 @@
|
||||
"""gRPC client interceptor for token-based authentication."""
|
||||
|
||||
import logging
|
||||
from collections import namedtuple
|
||||
from typing import Tuple
|
||||
|
||||
import grpc
|
||||
from grpc import aio as aiogrpc
|
||||
|
||||
from ray._raylet import AuthenticationTokenLoader
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
# Named tuple to hold client call details
|
||||
_ClientCallDetails = namedtuple(
|
||||
"_ClientCallDetails",
|
||||
("method", "timeout", "metadata", "credentials", "wait_for_ready", "compression"),
|
||||
)
|
||||
|
||||
|
||||
def _get_authentication_metadata_tuple() -> Tuple[Tuple[str, str], ...]:
|
||||
"""Get gRPC metadata tuple for authentication. Currently only supported for token authentication.
|
||||
|
||||
Returns:
|
||||
tuple: Empty tuple or ((AUTHORIZATION_HEADER_NAME, "Bearer <token>"),)
|
||||
"""
|
||||
token_loader = AuthenticationTokenLoader.instance()
|
||||
if not token_loader.has_token():
|
||||
return ()
|
||||
|
||||
headers = token_loader.get_token_for_http_header()
|
||||
|
||||
# Convert HTTP header dict to gRPC metadata tuple
|
||||
# gRPC expects: (("key", "value"), ...)
|
||||
return tuple((k, v) for k, v in headers.items())
|
||||
|
||||
|
||||
class SyncAuthenticationMetadataClientInterceptor(
|
||||
grpc.UnaryUnaryClientInterceptor,
|
||||
grpc.UnaryStreamClientInterceptor,
|
||||
grpc.StreamUnaryClientInterceptor,
|
||||
grpc.StreamStreamClientInterceptor,
|
||||
):
|
||||
"""Synchronous gRPC client interceptor that adds authentication metadata."""
|
||||
|
||||
def _intercept_call_details(self, client_call_details):
|
||||
"""Helper method to add authentication metadata to client call details."""
|
||||
metadata = list(client_call_details.metadata or [])
|
||||
metadata.extend(_get_authentication_metadata_tuple())
|
||||
|
||||
return _ClientCallDetails(
|
||||
method=client_call_details.method,
|
||||
timeout=client_call_details.timeout,
|
||||
metadata=metadata,
|
||||
credentials=client_call_details.credentials,
|
||||
wait_for_ready=getattr(client_call_details, "wait_for_ready", None),
|
||||
compression=getattr(client_call_details, "compression", None),
|
||||
)
|
||||
|
||||
def intercept_unary_unary(self, continuation, client_call_details, request):
|
||||
new_details = self._intercept_call_details(client_call_details)
|
||||
return continuation(new_details, request)
|
||||
|
||||
def intercept_unary_stream(self, continuation, client_call_details, request):
|
||||
new_details = self._intercept_call_details(client_call_details)
|
||||
return continuation(new_details, request)
|
||||
|
||||
def intercept_stream_unary(
|
||||
self, continuation, client_call_details, request_iterator
|
||||
):
|
||||
new_details = self._intercept_call_details(client_call_details)
|
||||
return continuation(new_details, request_iterator)
|
||||
|
||||
def intercept_stream_stream(
|
||||
self, continuation, client_call_details, request_iterator
|
||||
):
|
||||
new_details = self._intercept_call_details(client_call_details)
|
||||
return continuation(new_details, request_iterator)
|
||||
|
||||
|
||||
def _intercept_call_details_async(client_call_details):
|
||||
"""Helper to add authentication metadata to client call details (async version)."""
|
||||
metadata = list(client_call_details.metadata or [])
|
||||
metadata.extend(_get_authentication_metadata_tuple())
|
||||
|
||||
return _ClientCallDetails(
|
||||
method=client_call_details.method,
|
||||
timeout=client_call_details.timeout,
|
||||
metadata=metadata,
|
||||
credentials=client_call_details.credentials,
|
||||
wait_for_ready=getattr(client_call_details, "wait_for_ready", None),
|
||||
compression=getattr(client_call_details, "compression", None),
|
||||
)
|
||||
|
||||
|
||||
# NOTE: gRPC aio's Channel.__init__ uses if-elif chains to categorize interceptors,
|
||||
# so a single class inheriting from multiple interceptor types will only be registered
|
||||
# for the first matching type. We must use separate classes for each RPC type.
|
||||
# See: https://github.com/grpc/grpc/blob/master/src/python/grpcio/grpc/aio/_channel.py
|
||||
|
||||
|
||||
class _AsyncUnaryUnaryAuthInterceptor(aiogrpc.UnaryUnaryClientInterceptor):
|
||||
"""Async unary-unary interceptor that adds authentication metadata."""
|
||||
|
||||
async def intercept_unary_unary(self, continuation, client_call_details, request):
|
||||
new_details = _intercept_call_details_async(client_call_details)
|
||||
return await continuation(new_details, request)
|
||||
|
||||
|
||||
class _AsyncUnaryStreamAuthInterceptor(aiogrpc.UnaryStreamClientInterceptor):
|
||||
"""Async unary-stream interceptor that adds authentication metadata."""
|
||||
|
||||
async def intercept_unary_stream(self, continuation, client_call_details, request):
|
||||
new_details = _intercept_call_details_async(client_call_details)
|
||||
return await continuation(new_details, request)
|
||||
|
||||
|
||||
class _AsyncStreamUnaryAuthInterceptor(aiogrpc.StreamUnaryClientInterceptor):
|
||||
"""Async stream-unary interceptor that adds authentication metadata."""
|
||||
|
||||
async def intercept_stream_unary(
|
||||
self, continuation, client_call_details, request_iterator
|
||||
):
|
||||
new_details = _intercept_call_details_async(client_call_details)
|
||||
return await continuation(new_details, request_iterator)
|
||||
|
||||
|
||||
class _AsyncStreamStreamAuthInterceptor(aiogrpc.StreamStreamClientInterceptor):
|
||||
"""Async stream-stream interceptor that adds authentication metadata."""
|
||||
|
||||
async def intercept_stream_stream(
|
||||
self, continuation, client_call_details, request_iterator
|
||||
):
|
||||
new_details = _intercept_call_details_async(client_call_details)
|
||||
return await continuation(new_details, request_iterator)
|
||||
|
||||
|
||||
def get_async_auth_interceptors():
|
||||
"""Get a list of async authentication interceptors for all RPC types.
|
||||
|
||||
Returns a list of separate interceptor instances, one for each RPC type,
|
||||
because gRPC aio channels only register multi-inheritance interceptors
|
||||
for the first matching type.
|
||||
|
||||
Returns:
|
||||
List of interceptor instances for unary-unary, unary-stream,
|
||||
stream-unary, and stream-stream RPCs.
|
||||
"""
|
||||
return [
|
||||
_AsyncUnaryUnaryAuthInterceptor(),
|
||||
_AsyncUnaryStreamAuthInterceptor(),
|
||||
_AsyncStreamUnaryAuthInterceptor(),
|
||||
_AsyncStreamStreamAuthInterceptor(),
|
||||
]
|
||||
@@ -0,0 +1,232 @@
|
||||
"""gRPC server interceptor for token-based authentication."""
|
||||
|
||||
import logging
|
||||
from typing import Awaitable, Callable
|
||||
|
||||
import grpc
|
||||
from grpc import aio as aiogrpc
|
||||
|
||||
from ray._private.authentication.authentication_constants import (
|
||||
AUTHORIZATION_HEADER_NAME,
|
||||
)
|
||||
from ray._private.authentication.authentication_utils import (
|
||||
is_token_auth_enabled,
|
||||
validate_request_token,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _authenticate_request(metadata: tuple) -> bool:
|
||||
"""Authenticate incoming request. currently only supports token authentication.
|
||||
|
||||
Args:
|
||||
metadata: gRPC metadata tuple of (key, value) pairs
|
||||
Returns:
|
||||
True if authentication succeeds or is not required, False otherwise
|
||||
"""
|
||||
if not is_token_auth_enabled():
|
||||
return True
|
||||
|
||||
# Extract authorization header from metadata
|
||||
auth_header = None
|
||||
for key, value in metadata:
|
||||
if key.lower() == AUTHORIZATION_HEADER_NAME:
|
||||
auth_header = value
|
||||
break
|
||||
|
||||
if not auth_header:
|
||||
logger.warning("Authentication required but no authorization header provided")
|
||||
return False
|
||||
|
||||
# Validate the token format and value
|
||||
# validate_request_token returns bool (True if valid, False otherwise)
|
||||
return validate_request_token(auth_header)
|
||||
|
||||
|
||||
class AsyncAuthenticationServerInterceptor(aiogrpc.ServerInterceptor):
|
||||
"""Async gRPC server interceptor that validates authentication tokens.
|
||||
|
||||
This interceptor checks the "authorization" metadata header for a valid
|
||||
Bearer token when token authentication is enabled via RAY_AUTH_MODE=token.
|
||||
|
||||
If the token is missing or invalid, the request is rejected with UNAUTHENTICATED status.
|
||||
"""
|
||||
|
||||
async def intercept_service(
|
||||
self,
|
||||
continuation: Callable[
|
||||
[grpc.HandlerCallDetails], Awaitable[grpc.RpcMethodHandler]
|
||||
],
|
||||
handler_call_details: grpc.HandlerCallDetails,
|
||||
) -> grpc.RpcMethodHandler:
|
||||
"""Intercept service calls to validate authentication.
|
||||
|
||||
This method is called once per RPC to get the handler. We wrap the handler
|
||||
to validate authentication before executing the actual RPC method.
|
||||
"""
|
||||
# Get the actual handler
|
||||
handler = await continuation(handler_call_details)
|
||||
|
||||
if handler is None:
|
||||
return None
|
||||
|
||||
async def _abort_if_unauthenticated(context):
|
||||
"""Abort the RPC if authentication fails."""
|
||||
if not _authenticate_request(context.invocation_metadata()):
|
||||
await context.abort(
|
||||
grpc.StatusCode.UNAUTHENTICATED,
|
||||
"Invalid or missing authentication token",
|
||||
)
|
||||
|
||||
# Wrap the RPC behavior with authentication check
|
||||
def wrap_unary_response(behavior):
|
||||
"""Wrap a unary response RPC method to validate authentication first."""
|
||||
if behavior is None:
|
||||
return None
|
||||
|
||||
async def wrapped(request_or_iterator, context):
|
||||
await _abort_if_unauthenticated(context)
|
||||
return await behavior(request_or_iterator, context)
|
||||
|
||||
return wrapped
|
||||
|
||||
def wrap_stream_response(behavior):
|
||||
"""Wrap a streaming response RPC method to validate authentication first."""
|
||||
if behavior is None:
|
||||
return None
|
||||
|
||||
async def wrapped(request_or_iterator, context):
|
||||
await _abort_if_unauthenticated(context)
|
||||
async for response in behavior(request_or_iterator, context):
|
||||
yield response
|
||||
|
||||
return wrapped
|
||||
|
||||
# Create a wrapper class that implements RpcMethodHandler interface
|
||||
class AuthenticatedHandler:
|
||||
"""Wrapper handler that validates authentication."""
|
||||
|
||||
def __init__(
|
||||
self, original_handler, unary_wrapper_func, stream_wrapper_func
|
||||
):
|
||||
self._original = original_handler
|
||||
self._wrap_unary = unary_wrapper_func
|
||||
self._wrap_stream = stream_wrapper_func
|
||||
|
||||
@property
|
||||
def request_streaming(self):
|
||||
return self._original.request_streaming
|
||||
|
||||
@property
|
||||
def response_streaming(self):
|
||||
return self._original.response_streaming
|
||||
|
||||
@property
|
||||
def request_deserializer(self):
|
||||
return self._original.request_deserializer
|
||||
|
||||
@property
|
||||
def response_serializer(self):
|
||||
return self._original.response_serializer
|
||||
|
||||
@property
|
||||
def unary_unary(self):
|
||||
return self._wrap_unary(self._original.unary_unary)
|
||||
|
||||
@property
|
||||
def unary_stream(self):
|
||||
return self._wrap_stream(self._original.unary_stream)
|
||||
|
||||
@property
|
||||
def stream_unary(self):
|
||||
return self._wrap_unary(self._original.stream_unary)
|
||||
|
||||
@property
|
||||
def stream_stream(self):
|
||||
return self._wrap_stream(self._original.stream_stream)
|
||||
|
||||
return AuthenticatedHandler(handler, wrap_unary_response, wrap_stream_response)
|
||||
|
||||
|
||||
class SyncAuthenticationServerInterceptor(grpc.ServerInterceptor):
|
||||
"""Synchronous gRPC server interceptor that validates authentication tokens.
|
||||
|
||||
This interceptor checks the "authorization" metadata header for a valid
|
||||
Bearer token when token authentication is enabled via RAY_AUTH_MODE=token.
|
||||
If the token is missing or invalid, the request is rejected with UNAUTHENTICATED status.
|
||||
"""
|
||||
|
||||
def intercept_service(
|
||||
self,
|
||||
continuation: Callable[[grpc.HandlerCallDetails], grpc.RpcMethodHandler],
|
||||
handler_call_details: grpc.HandlerCallDetails,
|
||||
) -> grpc.RpcMethodHandler:
|
||||
"""Intercept service calls to validate authentication.
|
||||
|
||||
This method is called once per RPC to get the handler. We wrap the handler
|
||||
to validate authentication before executing the actual RPC method.
|
||||
"""
|
||||
# Get the actual handler
|
||||
handler = continuation(handler_call_details)
|
||||
|
||||
if handler is None:
|
||||
return None
|
||||
|
||||
# Wrap the RPC behavior with authentication check
|
||||
def wrap_rpc_behavior(behavior):
|
||||
"""Wrap an RPC method to validate authentication first."""
|
||||
if behavior is None:
|
||||
return None
|
||||
|
||||
def wrapped(request_or_iterator, context):
|
||||
if not _authenticate_request(context.invocation_metadata()):
|
||||
context.abort(
|
||||
grpc.StatusCode.UNAUTHENTICATED,
|
||||
"Invalid or missing authentication token",
|
||||
)
|
||||
return behavior(request_or_iterator, context)
|
||||
|
||||
return wrapped
|
||||
|
||||
# Create a wrapper class that implements RpcMethodHandler interface
|
||||
class AuthenticatedHandler:
|
||||
"""Wrapper handler that validates authentication."""
|
||||
|
||||
def __init__(self, original_handler, wrapper_func):
|
||||
self._original = original_handler
|
||||
self._wrap = wrapper_func
|
||||
|
||||
@property
|
||||
def request_streaming(self):
|
||||
return self._original.request_streaming
|
||||
|
||||
@property
|
||||
def response_streaming(self):
|
||||
return self._original.response_streaming
|
||||
|
||||
@property
|
||||
def request_deserializer(self):
|
||||
return self._original.request_deserializer
|
||||
|
||||
@property
|
||||
def response_serializer(self):
|
||||
return self._original.response_serializer
|
||||
|
||||
@property
|
||||
def unary_unary(self):
|
||||
return self._wrap(self._original.unary_unary)
|
||||
|
||||
@property
|
||||
def unary_stream(self):
|
||||
return self._wrap(self._original.unary_stream)
|
||||
|
||||
@property
|
||||
def stream_unary(self):
|
||||
return self._wrap(self._original.stream_unary)
|
||||
|
||||
@property
|
||||
def stream_stream(self):
|
||||
return self._wrap(self._original.stream_stream)
|
||||
|
||||
return AuthenticatedHandler(handler, wrap_rpc_behavior)
|
||||
@@ -0,0 +1,129 @@
|
||||
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
|
||||
Reference in New Issue
Block a user