chore: import upstream snapshot with attribution

python/ray/tests/authentication/BUILD.bazel

@@ -0,0 +1,23 @@
+load("@rules_python//python:defs.bzl", "py_library")
+load("//bazel:python.bzl", "py_test_run_all_subdirectory")
+
+py_library(
+    name = "conftest",
+    srcs = ["conftest.py"],
+)
+
+py_test_run_all_subdirectory(
+    size = "medium",
+    include = glob(["test_*.py"]),
+    exclude = [],
+    extra_srcs = [],
+    tags = [
+        "exclusive",
+        "medium_size_python_tests_shard_0",
+        "team:core",
+    ],
+    deps = [
+        ":conftest",
+        "//:ray_lib",
+    ],
+)

python/ray/tests/authentication/conftest.py

@@ -0,0 +1,144 @@
+import grpc
+import pytest
+from grpc import aio as aiogrpc
+
+from ray._private.authentication.authentication_token_generator import (
+    generate_new_authentication_token,
+)
+from ray._private.authentication_test_utils import (
+    authentication_env_guard,
+    reset_auth_token_state,
+    set_auth_mode,
+    set_env_auth_token,
+)
+from ray._private.grpc_utils import create_grpc_server_with_interceptors
+from ray.core.generated import reporter_pb2, reporter_pb2_grpc
+
+
+class SyncReporterService(reporter_pb2_grpc.ReporterServiceServicer):
+    """Simple synchronous test service for testing auth interceptors."""
+
+    def HealthCheck(self, request, context):
+        """Simple health check endpoint."""
+        return reporter_pb2.HealthCheckReply()
+
+
+class AsyncReporterService(reporter_pb2_grpc.ReporterServiceServicer):
+    """Simple asynchronous test service for testing auth interceptors."""
+
+    async def HealthCheck(self, request, context):
+        """Simple health check endpoint (async version)."""
+        return reporter_pb2.HealthCheckReply()
+
+
+class SyncLogService(reporter_pb2_grpc.LogServiceServicer):
+    """Simple synchronous log service for testing streaming auth interceptors."""
+
+    def StreamLog(self, request, context):
+        """Streaming log endpoint - yields test data chunks."""
+        for i in range(3):
+            yield reporter_pb2.StreamLogReply(data=f"chunk{i}".encode())
+
+
+class AsyncLogService(reporter_pb2_grpc.LogServiceServicer):
+    """Simple asynchronous log service for testing streaming auth interceptors."""
+
+    async def StreamLog(self, request, context):
+        """Streaming log endpoint (async version) - yields test data chunks."""
+        for i in range(3):
+            yield reporter_pb2.StreamLogReply(data=f"chunk{i}".encode())
+
+
+def _create_test_server_base(
+    *,
+    asynchronous: bool,
+    with_auth: bool,
+    reporter_servicer_cls,
+    log_servicer_cls,
+):
+    """Internal helper to create sync or async test server with optional auth."""
+
+    if with_auth:
+        # Auth is enabled - server will use interceptor
+        server = create_grpc_server_with_interceptors(
+            max_workers=None if asynchronous else 10,
+            thread_name_prefix="test_server",
+            options=None,
+            asynchronous=asynchronous,
+        )
+    else:
+        # Auth is disabled - create server without helper (no interceptor)
+        if asynchronous:
+            server = aiogrpc.server(options=None)
+        else:
+            from concurrent import futures
+
+            server = grpc.server(
+                futures.ThreadPoolExecutor(max_workers=10),
+                options=None,
+            )
+
+    # Add test services
+    reporter_servicer = reporter_servicer_cls()
+    reporter_pb2_grpc.add_ReporterServiceServicer_to_server(reporter_servicer, server)
+
+    log_servicer = log_servicer_cls()
+    reporter_pb2_grpc.add_LogServiceServicer_to_server(log_servicer, server)
+
+    # Bind to ephemeral port
+    port = server.add_insecure_port("[::]:0")
+
+    return server, port
+
+
+@pytest.fixture
+def create_sync_test_server():
+    """Factory to create synchronous gRPC test server.
+
+    Returns a function that creates a test server and returns (server, port).
+    The server must be stopped by the caller.
+    """
+
+    def _create(with_auth=True):
+        server, port = _create_test_server_base(
+            asynchronous=False,
+            with_auth=with_auth,
+            reporter_servicer_cls=SyncReporterService,
+            log_servicer_cls=SyncLogService,
+        )
+        server.start()
+        return server, port
+
+    return _create
+
+
+@pytest.fixture
+def create_async_test_server():
+    """Factory to create asynchronous gRPC test server.
+
+    Returns an async function that creates a test server and returns (server, port).
+    The server must be stopped by the caller.
+    """
+
+    async def _create(with_auth=True):
+        server, port = _create_test_server_base(
+            asynchronous=True,
+            with_auth=with_auth,
+            reporter_servicer_cls=AsyncReporterService,
+            log_servicer_cls=AsyncLogService,
+        )
+        await server.start()
+        return server, port
+
+    return _create
+
+
+@pytest.fixture
+def setup_auth_environment():
+    """Set up authentication environment with test token."""
+    test_token = generate_new_authentication_token()
+    with authentication_env_guard():
+        set_auth_mode("token")
+        set_env_auth_token(test_token)
+        reset_auth_token_state()
+        yield test_token

python/ray/tests/authentication/test_async_grpc_interceptors.py

@@ -0,0 +1,221 @@
+import grpc
+import pytest
+from grpc import aio as aiogrpc
+
+from ray._private.authentication.authentication_token_generator import (
+    generate_new_authentication_token,
+)
+from ray._private.authentication_test_utils import (
+    authentication_env_guard,
+    reset_auth_token_state,
+    set_auth_mode,
+    set_env_auth_token,
+)
+from ray._private.grpc_utils import init_grpc_channel
+from ray.core.generated import reporter_pb2, reporter_pb2_grpc
+
+
+@pytest.mark.asyncio
+async def test_async_server_and_client_with_valid_token(create_async_test_server):
+    """Test async server + client with matching token succeeds."""
+    token = generate_new_authentication_token()
+
+    with authentication_env_guard():
+        set_auth_mode("token")
+        set_env_auth_token(token)
+        reset_auth_token_state()
+
+        # Create server with auth enabled
+        server, port = await create_async_test_server(with_auth=True)
+
+        try:
+            # Client with auth interceptor via init_grpc_channel
+            channel = init_grpc_channel(
+                f"localhost:{port}",
+                options=None,
+                asynchronous=True,
+            )
+            stub = reporter_pb2_grpc.ReporterServiceStub(channel)
+            request = reporter_pb2.HealthCheckRequest()
+            response = await stub.HealthCheck(request, timeout=5)
+            assert response is not None
+        finally:
+            await server.stop(grace=1)
+
+
+@pytest.mark.asyncio
+async def test_async_server_and_client_with_invalid_token(create_async_test_server):
+    """Test async server + client with mismatched token fails."""
+    server_token = generate_new_authentication_token()
+    wrong_token = generate_new_authentication_token()
+
+    with authentication_env_guard():
+        # Set up server with server_token
+        set_auth_mode("token")
+        set_env_auth_token(server_token)
+        reset_auth_token_state()
+
+        server, port = await create_async_test_server(with_auth=True)
+
+        try:
+            # Create client channel and manually add wrong token to metadata
+            channel = aiogrpc.insecure_channel(f"localhost:{port}")
+            stub = reporter_pb2_grpc.ReporterServiceStub(channel)
+
+            # Add invalid token to metadata (not using client interceptor)
+            metadata = (("authorization", f"Bearer {wrong_token}"),)
+            request = reporter_pb2.HealthCheckRequest()
+
+            # Should fail with UNAUTHENTICATED
+            with pytest.raises(grpc.RpcError) as exc_info:
+                await stub.HealthCheck(request, metadata=metadata, timeout=5)
+
+            assert exc_info.value.code() == grpc.StatusCode.UNAUTHENTICATED
+        finally:
+            await server.stop(grace=1)
+
+
+@pytest.mark.asyncio
+async def test_async_server_with_auth_client_without_token(create_async_test_server):
+    """Test async server with auth, client without token fails."""
+    token = generate_new_authentication_token()
+
+    with authentication_env_guard():
+        # Set up server with auth enabled
+        set_auth_mode("token")
+        set_env_auth_token(token)
+        reset_auth_token_state()
+
+        server, port = await create_async_test_server(with_auth=True)
+
+        try:
+            # Create channel without auth metadata
+            channel = aiogrpc.insecure_channel(f"localhost:{port}")
+            stub = reporter_pb2_grpc.ReporterServiceStub(channel)
+            request = reporter_pb2.HealthCheckRequest()
+
+            # Should fail with UNAUTHENTICATED (no metadata provided)
+            with pytest.raises(grpc.RpcError) as exc_info:
+                await stub.HealthCheck(request, timeout=5)
+
+            assert exc_info.value.code() == grpc.StatusCode.UNAUTHENTICATED
+        finally:
+            await server.stop(grace=1)
+
+
+@pytest.mark.asyncio
+async def test_async_server_without_auth(create_async_test_server):
+    """Test async server without auth allows unauthenticated requests."""
+    with authentication_env_guard():
+        # Disable auth mode
+        set_auth_mode("disabled")
+        reset_auth_token_state()
+
+        # Create server without auth
+        server, port = await create_async_test_server(with_auth=False)
+
+        try:
+            # Client without auth
+            channel = aiogrpc.insecure_channel(f"localhost:{port}")
+            stub = reporter_pb2_grpc.ReporterServiceStub(channel)
+            request = reporter_pb2.HealthCheckRequest()
+
+            # Should succeed without auth
+            response = await stub.HealthCheck(request, timeout=5)
+            assert response is not None
+        finally:
+            await server.stop(grace=1)
+
+
+@pytest.mark.asyncio
+async def test_async_server_with_auth_disabled_allows_all(create_async_test_server):
+    """Test async server allows requests when auth mode is disabled."""
+    with authentication_env_guard():
+        # Disable auth mode globally
+        set_auth_mode("disabled")
+        reset_auth_token_state()
+
+        # Even though we call create_async_test_server with with_auth=True,
+        # the server won't enforce auth because auth mode is disabled
+        server, port = await create_async_test_server(with_auth=True)
+
+        try:
+            # Client without token
+            channel = aiogrpc.insecure_channel(f"localhost:{port}")
+            stub = reporter_pb2_grpc.ReporterServiceStub(channel)
+            request = reporter_pb2.HealthCheckRequest()
+
+            # Should succeed because auth is disabled
+            response = await stub.HealthCheck(request, timeout=5)
+            assert response is not None
+        finally:
+            await server.stop(grace=1)
+
+
+@pytest.mark.asyncio
+async def test_async_streaming_response_with_valid_token(create_async_test_server):
+    """Test async server streaming response (unary_stream) works with valid token."""
+    token = generate_new_authentication_token()
+
+    with authentication_env_guard():
+        set_auth_mode("token")
+        set_env_auth_token(token)
+        reset_auth_token_state()
+
+        # Create server with auth enabled
+        server, port = await create_async_test_server(with_auth=True)
+
+        try:
+            # Client with auth interceptor via init_grpc_channel
+            channel = init_grpc_channel(
+                f"localhost:{port}",
+                options=None,
+                asynchronous=True,
+            )
+            stub = reporter_pb2_grpc.LogServiceStub(channel)
+            request = reporter_pb2.StreamLogRequest(log_file_name="test.log")
+
+            # Stream the response - this tests the unary_stream RPC path
+            chunks = []
+            async for response in stub.StreamLog(request, timeout=5):
+                chunks.append(response.data)
+
+            # Verify we got all 3 chunks from the test service
+            assert len(chunks) == 3
+            assert chunks == [b"chunk0", b"chunk1", b"chunk2"]
+        finally:
+            await server.stop(grace=1)
+
+
+@pytest.mark.asyncio
+async def test_async_streaming_response_without_token_fails(create_async_test_server):
+    """Test async server streaming response fails without token."""
+    token = generate_new_authentication_token()
+
+    with authentication_env_guard():
+        set_auth_mode("token")
+        set_env_auth_token(token)
+        reset_auth_token_state()
+
+        server, port = await create_async_test_server(with_auth=True)
+
+        try:
+            # Client without auth token
+            channel = aiogrpc.insecure_channel(f"localhost:{port}")
+            stub = reporter_pb2_grpc.LogServiceStub(channel)
+            request = reporter_pb2.StreamLogRequest(log_file_name="test.log")
+
+            # Should fail with UNAUTHENTICATED when trying to iterate
+            with pytest.raises(grpc.RpcError) as exc_info:
+                async for _ in stub.StreamLog(request, timeout=5):
+                    pass
+
+            assert exc_info.value.code() == grpc.StatusCode.UNAUTHENTICATED
+        finally:
+            await server.stop(grace=1)
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-vv", __file__]))

python/ray/tests/authentication/test_sync_grpc_interceptors.py

@@ -0,0 +1,213 @@
+import grpc
+import pytest
+
+from ray._private.authentication.authentication_token_generator import (
+    generate_new_authentication_token,
+)
+from ray._private.authentication_test_utils import (
+    authentication_env_guard,
+    reset_auth_token_state,
+    set_auth_mode,
+    set_env_auth_token,
+)
+from ray._private.grpc_utils import init_grpc_channel
+from ray.core.generated import reporter_pb2, reporter_pb2_grpc
+
+
+def test_sync_server_and_client_with_valid_token(create_sync_test_server):
+    """Test sync server + client with matching token succeeds."""
+    token = generate_new_authentication_token()
+
+    with authentication_env_guard():
+        set_auth_mode("token")
+        set_env_auth_token(token)
+        reset_auth_token_state()
+
+        # Create server with auth enabled
+        server, port = create_sync_test_server(with_auth=True)
+
+        try:
+            # Client with auth interceptor via init_grpc_channel
+            channel = init_grpc_channel(
+                f"localhost:{port}",
+                options=None,
+                asynchronous=False,
+            )
+            stub = reporter_pb2_grpc.ReporterServiceStub(channel)
+            request = reporter_pb2.HealthCheckRequest()
+            response = stub.HealthCheck(request, timeout=5)
+            assert response is not None
+        finally:
+            server.stop(grace=1)
+
+
+def test_sync_server_and_client_with_invalid_token(create_sync_test_server):
+    """Test sync server + client with mismatched token fails."""
+    server_token = generate_new_authentication_token()
+    wrong_token = generate_new_authentication_token()
+
+    with authentication_env_guard():
+        # Set up server with server_token
+        set_auth_mode("token")
+        set_env_auth_token(server_token)
+        reset_auth_token_state()
+
+        server, port = create_sync_test_server(with_auth=True)
+
+        try:
+            # Create client channel and manually add wrong token to metadata
+            channel = grpc.insecure_channel(f"localhost:{port}")
+            stub = reporter_pb2_grpc.ReporterServiceStub(channel)
+
+            # Add invalid token to metadata (not using client interceptor)
+            metadata = (("authorization", f"Bearer {wrong_token}"),)
+            request = reporter_pb2.HealthCheckRequest()
+
+            # Should fail with UNAUTHENTICATED
+            with pytest.raises(grpc.RpcError) as exc_info:
+                stub.HealthCheck(request, metadata=metadata, timeout=5)
+
+            assert exc_info.value.code() == grpc.StatusCode.UNAUTHENTICATED
+        finally:
+            server.stop(grace=1)
+
+
+def test_sync_server_with_auth_client_without_token(create_sync_test_server):
+    """Test server with auth, client without token fails."""
+    token = generate_new_authentication_token()
+
+    with authentication_env_guard():
+        # Set up server with auth enabled
+        set_auth_mode("token")
+        set_env_auth_token(token)
+        reset_auth_token_state()
+
+        server, port = create_sync_test_server(with_auth=True)
+
+        try:
+            # Create channel without auth metadata
+            channel = grpc.insecure_channel(f"localhost:{port}")
+            stub = reporter_pb2_grpc.ReporterServiceStub(channel)
+            request = reporter_pb2.HealthCheckRequest()
+
+            # Should fail with UNAUTHENTICATED (no metadata provided)
+            with pytest.raises(grpc.RpcError) as exc_info:
+                stub.HealthCheck(request, timeout=5)
+
+            assert exc_info.value.code() == grpc.StatusCode.UNAUTHENTICATED
+        finally:
+            server.stop(grace=1)
+
+
+def test_sync_server_without_auth(create_sync_test_server):
+    """Test server without auth allows unauthenticated requests."""
+    with authentication_env_guard():
+        # Disable auth mode
+        set_auth_mode("disabled")
+        reset_auth_token_state()
+
+        # Create server without auth
+        server, port = create_sync_test_server(with_auth=False)
+
+        try:
+            # Client without auth
+            channel = grpc.insecure_channel(f"localhost:{port}")
+            stub = reporter_pb2_grpc.ReporterServiceStub(channel)
+            request = reporter_pb2.HealthCheckRequest()
+
+            # Should succeed without auth
+            response = stub.HealthCheck(request, timeout=5)
+            assert response is not None
+        finally:
+            server.stop(grace=1)
+
+
+def test_sync_server_with_auth_disabled_allows_all(create_sync_test_server):
+    """Test server allows requests when auth mode is disabled."""
+    with authentication_env_guard():
+        # Disable auth mode globally
+        set_auth_mode("disabled")
+        reset_auth_token_state()
+
+        # Even though we call create_sync_test_server with with_auth=True,
+        # the server won't enforce auth because auth mode is disabled
+        server, port = create_sync_test_server(with_auth=True)
+
+        try:
+            # Client without token
+            channel = grpc.insecure_channel(f"localhost:{port}")
+            stub = reporter_pb2_grpc.ReporterServiceStub(channel)
+            request = reporter_pb2.HealthCheckRequest()
+
+            # Should succeed because auth is disabled
+            response = stub.HealthCheck(request, timeout=5)
+            assert response is not None
+        finally:
+            server.stop(grace=1)
+
+
+def test_sync_streaming_response_with_valid_token(create_sync_test_server):
+    """Test sync server streaming response (unary_stream) works with valid token."""
+    token = generate_new_authentication_token()
+
+    with authentication_env_guard():
+        set_auth_mode("token")
+        set_env_auth_token(token)
+        reset_auth_token_state()
+
+        # Create server with auth enabled
+        server, port = create_sync_test_server(with_auth=True)
+
+        try:
+            # Client with auth interceptor via init_grpc_channel
+            channel = init_grpc_channel(
+                f"localhost:{port}",
+                options=None,
+                asynchronous=False,
+            )
+            stub = reporter_pb2_grpc.LogServiceStub(channel)
+            request = reporter_pb2.StreamLogRequest(log_file_name="test.log")
+
+            # Stream the response - this tests the unary_stream RPC path
+            chunks = []
+            for response in stub.StreamLog(request, timeout=5):
+                chunks.append(response.data)
+
+            # Verify we got all 3 chunks from the test service
+            assert len(chunks) == 3
+            assert chunks == [b"chunk0", b"chunk1", b"chunk2"]
+        finally:
+            server.stop(grace=1)
+
+
+def test_sync_streaming_response_without_token_fails(create_sync_test_server):
+    """Test sync server streaming response fails without token."""
+    token = generate_new_authentication_token()
+
+    with authentication_env_guard():
+        set_auth_mode("token")
+        set_env_auth_token(token)
+        reset_auth_token_state()
+
+        server, port = create_sync_test_server(with_auth=True)
+
+        try:
+            # Client without auth token
+            channel = grpc.insecure_channel(f"localhost:{port}")
+            stub = reporter_pb2_grpc.LogServiceStub(channel)
+            request = reporter_pb2.StreamLogRequest(log_file_name="test.log")
+
+            # Should fail with UNAUTHENTICATED when trying to iterate
+            with pytest.raises(grpc.RpcError) as exc_info:
+                for _ in stub.StreamLog(request, timeout=5):
+                    pass
+
+            assert exc_info.value.code() == grpc.StatusCode.UNAUTHENTICATED
+        finally:
+            server.stop(grace=1)
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-vv", __file__]))