[Core] Support token auth in ray Pub-Sub (#58333)

This PR adds token-based authentication support to the
PythonGcsSubscriber, which previously made direct gRPC calls via the
stub without auth. The rest of the pub-sub layer already uses the shared
gRPC infrastructure (GrpcServer, GrpcClient), which supports token
authentication.

---------

Signed-off-by: sampan <[email protected]>
Signed-off-by: Edward Oakes <[email protected]>
Co-authored-by: sampan <[email protected]>
Co-authored-by: Edward Oakes <[email protected]>

Author: 3 people

src/ray/pubsub/python_gcs_subscriber.cc

@@ -22,6 +22,7 @@
 #include <vector>
 
 #include "ray/gcs_rpc_client/rpc_client.h"
+#include "ray/rpc/authentication/authentication_token_loader.h"
 
 namespace ray {
 namespace pubsub {
@@ -51,6 +52,7 @@ Status PythonGcsSubscriber::Subscribe() {
   }
 
   grpc::ClientContext context;
+  SetAuthenticationToken(context);
 
   rpc::GcsSubscriberCommandBatchRequest request;
   request.set_subscriber_id(subscriber_id_);
@@ -78,6 +80,7 @@ Status PythonGcsSubscriber::DoPoll(int64_t timeout_ms, rpc::PubMessage *message)
       return Status::OK();
     }
     current_polling_context_ = std::make_shared<grpc::ClientContext>();
+    SetAuthenticationToken(*current_polling_context_);
     if (timeout_ms != -1) {
       current_polling_context_->set_deadline(std::chrono::system_clock::now() +
                                              std::chrono::milliseconds(timeout_ms));
@@ -173,6 +176,7 @@ Status PythonGcsSubscriber::Close() {
   }
 
   grpc::ClientContext context;
+  SetAuthenticationToken(context);
 
   rpc::GcsSubscriberCommandBatchRequest request;
   request.set_subscriber_id(subscriber_id_);
@@ -195,5 +199,12 @@ int64_t PythonGcsSubscriber::last_batch_size() {
   return last_batch_size_;
 }
 
+void PythonGcsSubscriber::SetAuthenticationToken(grpc::ClientContext &context) {
+  auto auth_token = ray::rpc::AuthenticationTokenLoader::instance().GetToken();
+  if (auth_token.has_value() && !auth_token->empty()) {
+    auth_token->SetMetadata(context);
+  }
+}
+
 }  // namespace pubsub
 }  // namespace ray

src/ray/pubsub/python_gcs_subscriber.h

@@ -80,6 +80,10 @@ class RAY_EXPORT PythonGcsSubscriber {
   std::deque<rpc::PubMessage> queue_ ABSL_GUARDED_BY(mu_);
   bool closed_ ABSL_GUARDED_BY(mu_) = false;
   std::shared_ptr<grpc::ClientContext> current_polling_context_ ABSL_GUARDED_BY(mu_);
+
+  // Set authentication token on a gRPC client context if token-based authentication is
+  // enabled
+  void SetAuthenticationToken(grpc::ClientContext &context);
 };
 
 /// Get the .lines() attribute of a LogBatch as a std::vector

src/ray/pubsub/tests/BUILD.bazel

@@ -40,3 +40,20 @@ ray_cc_test(
         "@com_google_googletest//:gtest_main",
     ],
 )
+
+ray_cc_test(
+    name = "python_gcs_subscriber_auth_test",
+    size = "small",
+    srcs = ["python_gcs_subscriber_auth_test.cc"],
+    tags = ["team:core"],
+    deps = [
+        "//src/ray/common:ray_config",
+        "//src/ray/common:status",
+        "//src/ray/protobuf:gcs_service_cc_grpc",
+        "//src/ray/pubsub:python_gcs_subscriber",
+        "//src/ray/rpc:grpc_server",
+        "//src/ray/rpc/authentication:authentication_token",
+        "//src/ray/rpc/authentication:authentication_token_loader",
+        "@com_google_googletest//:gtest_main",
+    ],
+)