[Core] support token auth in ray client server

python/ray/_private/authentication/authentication_token_setup.py

@@ -67,7 +67,7 @@ def ensure_token_if_auth_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.
+        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:
@@ -79,11 +79,11 @@ def ensure_token_if_auth_enabled(
     if get_authentication_mode() != AuthenticationMode.TOKEN:
         if (
             system_config
-            and "auth_mode" in system_config
-            and system_config["auth_mode"] != "disabled"
+            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."
+                "Set authentication mode can only be set with the `RAY_AUTH_MODE` environment variable, not using the system_config."
             )
         return
 

python/ray/_private/authentication/authentication_utils.py

@@ -15,7 +15,7 @@ def is_token_auth_enabled() -> bool:
     """Check if token authentication is enabled.
 
     Returns:
-        bool: True if auth_mode is set to "token", False otherwise
+        bool: True if AUTH_MODE is set to "token", False otherwise
     """
     if not _RAYLET_AVAILABLE:
         return False

python/ray/_private/authentication/grpc_authentication_client_interceptor.py

@@ -36,7 +36,7 @@ def _get_authentication_metadata_tuple() -> Tuple[Tuple[str, str], ...]:
     return tuple((k, v) for k, v in headers.items())
 
 
-class AuthenticationMetadataClientInterceptor(
+class SyncAuthenticationMetadataClientInterceptor(
     grpc.UnaryUnaryClientInterceptor,
     grpc.UnaryStreamClientInterceptor,
     grpc.StreamUnaryClientInterceptor,

python/ray/_private/authentication/grpc_authentication_server_interceptor.py

@@ -0,0 +1,212 @@
+"""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
+
+        # 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
+
+            async def wrapped(request_or_iterator, context):
+                if not _authenticate_request(context.invocation_metadata()):
+                    await context.abort(
+                        grpc.StatusCode.UNAUTHENTICATED,
+                        "Invalid or missing authentication token",
+                    )
+                return await 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)
+
+
+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)

python/ray/_private/gcs_utils.py

@@ -79,7 +79,7 @@ def create_gcs_channel(address: str, aio=False):
     Returns:
         grpc.Channel or grpc.aio.Channel to GCS
     """
-    from ray._private.utils import init_grpc_channel
+    from ray._private.grpc_utils import init_grpc_channel
 
     return init_grpc_channel(address, options=_GRPC_OPTIONS, asynchronous=aio)