[Core] Support token auth in ray_syncer (#58176)

This PR adds support for token-based authentication in the Ray
bi-directional syncer, for both client and server sides. It also
includes tests to verify the functionality.

---------

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/gcs/BUILD.bazel

@@ -531,6 +531,7 @@ ray_cc_library(
         "//src/ray/raylet_rpc_client:raylet_client_pool",
         "//src/ray/rpc:grpc_server",
         "//src/ray/rpc:metrics_agent_client",
+        "//src/ray/rpc/authentication:authentication_token_loader",
         "//src/ray/util:counter_map",
         "//src/ray/util:exponential_backoff",
         "//src/ray/util:network_util",

src/ray/gcs/gcs_server.cc

@@ -39,6 +39,7 @@
 #include "ray/observability/metric_constants.h"
 #include "ray/pubsub/publisher.h"
 #include "ray/raylet_rpc_client/raylet_client.h"
+#include "ray/rpc/authentication/authentication_token_loader.h"
 #include "ray/stats/stats.h"
 #include "ray/util/network_util.h"
 
@@ -615,7 +616,8 @@ void GcsServer::InitRaySyncer(const GcsInitData &gcs_init_data) {
       syncer::MessageType::RESOURCE_VIEW, nullptr, gcs_resource_manager_.get());
   ray_syncer_->Register(
       syncer::MessageType::COMMANDS, nullptr, gcs_resource_manager_.get());
-  rpc_server_.RegisterService(std::make_unique<syncer::RaySyncerService>(*ray_syncer_));
+  rpc_server_.RegisterService(std::make_unique<syncer::RaySyncerService>(
+      *ray_syncer_, ray::rpc::AuthenticationTokenLoader::instance().GetToken()));
 }
 
 void GcsServer::InitFunctionManager() {

src/ray/ray_syncer/BUILD.bazel

@@ -19,8 +19,11 @@ ray_cc_library(
     ],
     deps = [
         "//src/ray/common:asio",
+        "//src/ray/common:constants",
         "//src/ray/common:id",
         "//src/ray/protobuf:ray_syncer_cc_grpc",
+        "//src/ray/rpc/authentication:authentication_token",
+        "//src/ray/rpc/authentication:authentication_token_loader",
         "@com_github_grpc_grpc//:grpc++",
         "@com_google_absl//absl/container:flat_hash_map",
     ],

src/ray/ray_syncer/ray_syncer.cc

@@ -244,9 +244,17 @@ ServerBidiReactor *RaySyncerService::StartSync(grpc::CallbackServerContext *cont
         }
         RAY_LOG(INFO).WithField(NodeID::FromBinary(node_id)) << "Connection is broken.";
         syncer_.node_state_->RemoveNode(node_id);
-      });
+      },
+      /*auth_token=*/auth_token_);
   RAY_LOG(DEBUG).WithField(NodeID::FromBinary(reactor->GetRemoteNodeID()))
       << "Get connection";
+
+  // If the reactor has already called Finish() (e.g., due to authentication failure),
+  // skip registration. The reactor will clean itself up via OnDone().
+  if (reactor->IsFinished()) {
+    return reactor;
+  }
+
   // Disconnect exiting connection if there is any.
   // This can happen when there is transient network error
   // and the client reconnects.

src/ray/ray_syncer/ray_syncer.h

@@ -19,6 +19,7 @@
 
 #include <memory>
 #include <string>
+#include <utility>
 #include <vector>
 
 #include "absl/container/flat_hash_map.h"
@@ -28,6 +29,7 @@
 #include "ray/common/asio/periodical_runner.h"
 #include "ray/common/id.h"
 #include "ray/ray_syncer/common.h"
+#include "ray/rpc/authentication/authentication_token.h"
 #include "src/ray/protobuf/ray_syncer.grpc.pb.h"
 
 namespace ray::syncer {
@@ -197,14 +199,20 @@ class RaySyncer {
 /// like tree-based one.
 class RaySyncerService : public ray::rpc::syncer::RaySyncer::CallbackService {
  public:
-  explicit RaySyncerService(RaySyncer &syncer) : syncer_(syncer) {}
+  explicit RaySyncerService(
+      RaySyncer &syncer,
+      std::optional<ray::rpc::AuthenticationToken> auth_token = std::nullopt)
+      : syncer_(syncer), auth_token_(std::move(auth_token)) {}
 
   grpc::ServerBidiReactor<RaySyncMessage, RaySyncMessage> *StartSync(
       grpc::CallbackServerContext *context) override;
 
  private:
   // The ray syncer this RPC wrappers of.
   RaySyncer &syncer_;
+  // Authentication token for validation, will be empty if token authentication is
+  // disabled
+  std::optional<ray::rpc::AuthenticationToken> auth_token_;
 };
 
 }  // namespace ray::syncer

src/ray/ray_syncer/ray_syncer_client.cc

@@ -18,6 +18,8 @@
 #include <string>
 #include <utility>
 
+#include "ray/rpc/authentication/authentication_token_loader.h"
+
 namespace ray::syncer {
 
 RayClientBidiReactor::RayClientBidiReactor(
@@ -32,6 +34,11 @@ RayClientBidiReactor::RayClientBidiReactor(
       cleanup_cb_(std::move(cleanup_cb)),
       stub_(std::move(stub)) {
   client_context_.AddMetadata("node_id", NodeID::FromBinary(local_node_id).Hex());
+  // Add authentication token if token authentication is enabled
+  auto auth_token = ray::rpc::AuthenticationTokenLoader::instance().GetToken();
+  if (auth_token.has_value() && !auth_token->empty()) {
+    auth_token->SetMetadata(client_context_);
+  }
   stub_->async()->StartSync(&client_context_, this);
   // Prevent this call from being terminated.
   // Check https://github.com/grpc/proposal/blob/master/L67-cpp-callback-api.md

src/ray/ray_syncer/ray_syncer_server.cc

@@ -17,6 +17,8 @@
 #include <string>
 #include <utility>
 
+#include "ray/common/constants.h"
+
 namespace ray::syncer {
 
 namespace {
@@ -35,13 +37,39 @@ RayServerBidiReactor::RayServerBidiReactor(
     instrumented_io_context &io_context,
     const std::string &local_node_id,
     std::function<void(std::shared_ptr<const RaySyncMessage>)> message_processor,
-    std::function<void(RaySyncerBidiReactor *, bool)> cleanup_cb)
+    std::function<void(RaySyncerBidiReactor *, bool)> cleanup_cb,
+    const std::optional<ray::rpc::AuthenticationToken> &auth_token)
     : RaySyncerBidiReactorBase<ServerBidiReactor>(
           io_context,
           GetNodeIDFromServerContext(server_context),
           std::move(message_processor)),
       cleanup_cb_(std::move(cleanup_cb)),
-      server_context_(server_context) {
+      server_context_(server_context),
+      auth_token_(auth_token) {
+  if (auth_token_.has_value() && !auth_token_->empty()) {
+    // Validate authentication token
+    const auto &metadata = server_context->client_metadata();
+    auto it = metadata.find(kAuthTokenKey);
+    if (it == metadata.end()) {
+      RAY_LOG(WARNING) << "Missing authorization header in syncer connection from node "
+                       << NodeID::FromBinary(GetRemoteNodeID());
+      Finish(grpc::Status(grpc::StatusCode::UNAUTHENTICATED,
+                          "Missing authorization header"));
+      return;
+    }
+
+    const std::string_view header(it->second.data(), it->second.length());
+    ray::rpc::AuthenticationToken provided_token =
+        ray::rpc::AuthenticationToken::FromMetadata(header);
+
+    if (!auth_token_->Equals(provided_token)) {
+      RAY_LOG(WARNING) << "Invalid bearer token in syncer connection from node "
+                       << NodeID::FromBinary(GetRemoteNodeID());
+      Finish(grpc::Status(grpc::StatusCode::UNAUTHENTICATED, "Invalid bearer token"));
+      return;
+    }
+  }
+
   // Send the local node id to the remote
   server_context_->AddInitialMetadata("node_id", NodeID::FromBinary(local_node_id).Hex());
   StartSendInitialMetadata();

src/ray/ray_syncer/ray_syncer_server.h

@@ -16,11 +16,14 @@
 
 #include <gtest/gtest_prod.h>
 
+#include <atomic>
+#include <optional>
 #include <string>
 
 #include "ray/ray_syncer/common.h"
 #include "ray/ray_syncer/ray_syncer_bidi_reactor.h"
 #include "ray/ray_syncer/ray_syncer_bidi_reactor_base.h"
+#include "ray/rpc/authentication/authentication_token.h"
 
 namespace ray::syncer {
 
@@ -35,20 +38,36 @@ class RayServerBidiReactor : public RaySyncerBidiReactorBase<ServerBidiReactor>
       instrumented_io_context &io_context,
       const std::string &local_node_id,
       std::function<void(std::shared_ptr<const RaySyncMessage>)> message_processor,
-      std::function<void(RaySyncerBidiReactor *, bool)> cleanup_cb);
+      std::function<void(RaySyncerBidiReactor *, bool)> cleanup_cb,
+      const std::optional<ray::rpc::AuthenticationToken> &auth_token);
 
   ~RayServerBidiReactor() override = default;
 
+  bool IsFinished() const { return finished_.load(); }
+
  private:
   void DoDisconnect() override;
   void OnCancel() override;
   void OnDone() override;
 
+  void Finish(grpc::Status status) {
+    finished_.store(true);
+    ServerBidiReactor::Finish(status);
+  }
+
   /// Cleanup callback when the call ends.
   const std::function<void(RaySyncerBidiReactor *, bool)> cleanup_cb_;
 
   /// grpc callback context
   grpc::CallbackServerContext *server_context_;
+
+  /// Authentication token for validation, will be empty if token authentication is
+  /// disabled
+  std::optional<ray::rpc::AuthenticationToken> auth_token_;
+
+  /// Track if Finish() has been called to avoid using a reactor that is terminating
+  std::atomic<bool> finished_{false};
+
   FRIEND_TEST(SyncerReactorTest, TestReactorFailure);
 };
 

src/ray/ray_syncer/tests/BUILD.bazel

@@ -13,6 +13,7 @@ ray_cc_test(
         "//src/mock/ray/ray_syncer:mock_ray_syncer",
         "//src/ray/ray_syncer",
         "//src/ray/rpc:grpc_server",
+        "//src/ray/rpc/authentication:authentication_token",
         "//src/ray/util:network_util",
         "//src/ray/util:path_utils",
         "//src/ray/util:raii",