diff --git a/python/ray/_private/authentication/authentication_token_generator.py b/python/ray/_private/authentication/authentication_token_generator.py index 331584cb5dd4..34fcd1eab8a3 100644 --- a/python/ray/_private/authentication/authentication_token_generator.py +++ b/python/ray/_private/authentication/authentication_token_generator.py @@ -1,6 +1,9 @@ -import uuid +import secrets -# TODO: this is a placeholder for the actual authentication token generator. Will be replaced with a proper implementation. def generate_new_authentication_token() -> str: - return uuid.uuid4().hex + """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) diff --git a/python/ray/tests/authentication/conftest.py b/python/ray/tests/authentication/conftest.py index 4b1e876f4665..4e0e2db17ff6 100644 --- a/python/ray/tests/authentication/conftest.py +++ b/python/ray/tests/authentication/conftest.py @@ -1,9 +1,10 @@ -import uuid - 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, @@ -109,14 +110,9 @@ async def _create(with_auth=True): @pytest.fixture -def test_token(): - """Generate a test authentication token.""" - return uuid.uuid4().hex - - -@pytest.fixture -def setup_auth_environment(test_token): +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) diff --git a/python/ray/tests/authentication/test_async_grpc_interceptors.py b/python/ray/tests/authentication/test_async_grpc_interceptors.py index ca37fcc6ed60..aa9f4ec577c0 100644 --- a/python/ray/tests/authentication/test_async_grpc_interceptors.py +++ b/python/ray/tests/authentication/test_async_grpc_interceptors.py @@ -1,9 +1,10 @@ -import uuid - 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, @@ -17,7 +18,7 @@ @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 = uuid.uuid4().hex + token = generate_new_authentication_token() with authentication_env_guard(): set_auth_mode("token") @@ -45,8 +46,8 @@ async def test_async_server_and_client_with_valid_token(create_async_test_server @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 = uuid.uuid4().hex - wrong_token = uuid.uuid4().hex + server_token = generate_new_authentication_token() + wrong_token = generate_new_authentication_token() with authentication_env_guard(): # Set up server with server_token @@ -77,7 +78,7 @@ async def test_async_server_and_client_with_invalid_token(create_async_test_serv @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 = uuid.uuid4().hex + token = generate_new_authentication_token() with authentication_env_guard(): # Set up server with auth enabled diff --git a/python/ray/tests/authentication/test_sync_grpc_interceptors.py b/python/ray/tests/authentication/test_sync_grpc_interceptors.py index 276b81740e66..73a30182bd23 100644 --- a/python/ray/tests/authentication/test_sync_grpc_interceptors.py +++ b/python/ray/tests/authentication/test_sync_grpc_interceptors.py @@ -1,8 +1,9 @@ -import uuid - 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, @@ -15,7 +16,7 @@ def test_sync_server_and_client_with_valid_token(create_sync_test_server): """Test sync server + client with matching token succeeds.""" - token = uuid.uuid4().hex + token = generate_new_authentication_token() with authentication_env_guard(): set_auth_mode("token") @@ -42,8 +43,8 @@ def test_sync_server_and_client_with_valid_token(create_sync_test_server): def test_sync_server_and_client_with_invalid_token(create_sync_test_server): """Test sync server + client with mismatched token fails.""" - server_token = uuid.uuid4().hex - wrong_token = uuid.uuid4().hex + server_token = generate_new_authentication_token() + wrong_token = generate_new_authentication_token() with authentication_env_guard(): # Set up server with server_token @@ -73,7 +74,7 @@ def test_sync_server_and_client_with_invalid_token(create_sync_test_server): def test_sync_server_with_auth_client_without_token(create_sync_test_server): """Test server with auth, client without token fails.""" - token = uuid.uuid4().hex + token = generate_new_authentication_token() with authentication_env_guard(): # Set up server with auth enabled diff --git a/python/ray/tests/test_token_auth_integration.py b/python/ray/tests/test_token_auth_integration.py index 4a27f1653a11..86f4402233b5 100644 --- a/python/ray/tests/test_token_auth_integration.py +++ b/python/ray/tests/test_token_auth_integration.py @@ -146,7 +146,7 @@ def test_local_cluster_generates_token(): f"Files in {default_token_path.parent}: {list(default_token_path.parent.iterdir()) if default_token_path.parent.exists() else 'directory does not exist'}" ) token = default_token_path.read_text().strip() - assert len(token) == 32 + assert len(token) == 64 assert all(c in "0123456789abcdef" for c in token) # Verify cluster is working @@ -190,7 +190,7 @@ def test_cluster_token_authentication(tokens_match, setup_cluster_with_token_aut if tokens_match: client_token = cluster_token # Same token - should succeed else: - client_token = "b" * 32 # Different token - should fail + client_token = "b" * 64 # Different token - should fail set_env_auth_token(client_token) reset_auth_token_state() @@ -263,7 +263,7 @@ def test_ray_start_without_token_raises_error(is_head, request): def test_ray_start_head_with_token_succeeds(): """Test that ray start --head succeeds when token auth is enabled with a valid token.""" # Set up environment with token auth and a valid token - test_token = "a" * 32 + test_token = "a" * 64 env = os.environ.copy() env["RAY_AUTH_TOKEN"] = test_token env["RAY_AUTH_MODE"] = "token" @@ -326,7 +326,7 @@ def test_ray_start_address_with_token(token_match, setup_cluster_with_token_auth env["RAY_AUTH_TOKEN"] = cluster_token expect_success = True else: - env["RAY_AUTH_TOKEN"] = "b" * 32 + env["RAY_AUTH_TOKEN"] = "b" * 64 expect_success = False # Start worker node @@ -422,7 +422,7 @@ def job_finished(): @pytest.mark.parametrize("use_generate", [True, False]) def test_get_auth_token_cli(use_generate): """Test ray get-auth-token CLI command.""" - test_token = "a" * 32 + test_token = "a" * 64 with authentication_env_guard(): set_auth_mode("token") @@ -452,7 +452,7 @@ def test_get_auth_token_cli(use_generate): # Verify token is printed to stdout token = result.stdout.strip() - assert len(token) == 32, f"Token should be 32 chars, got {len(token)}" + assert len(token) == 64, token assert all(c in "0123456789abcdef" for c in token), "Token should be hex" if not use_generate: @@ -497,7 +497,7 @@ def test_get_auth_token_cli_no_token_no_generate(): ) def test_get_auth_token_cli_piping(): """Test that ray get-auth-token output can be piped.""" - test_token = "b" * 32 + test_token = "b" * 64 with authentication_env_guard(): set_auth_mode("token") @@ -516,9 +516,8 @@ def test_get_auth_token_cli_piping(): ) assert result.returncode == 0 - # Should be 32 chars (no newline with nl=False) char_count = int(result.stdout.strip()) - assert char_count == 32, f"Expected 32 chars (no newline), got {char_count}" + assert char_count == 64, f"Expected 64 chars (no newline), got {char_count}" if __name__ == "__main__":