Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
81 changes: 64 additions & 17 deletions src/ray/core_worker/experimental_mutable_object_provider.cc
Original file line number Diff line number Diff line change
Expand Up @@ -120,8 +120,9 @@ void MutableObjectProvider::HandleRegisterMutableObject(
RegisterReaderChannel(reader_object_id);
}

void MutableObjectProvider::HandlePushMutableObject(
Status MutableObjectProvider::HandlePushMutableObject(
const rpc::PushMutableObjectRequest &request, rpc::PushMutableObjectReply *reply) {
reply->set_done(false);
LocalReaderInfo info;
const ObjectID writer_object_id = ObjectID::FromBinary(request.writer_object_id());
{
Expand All @@ -130,6 +131,17 @@ void MutableObjectProvider::HandlePushMutableObject(
RAY_CHECK(it != remote_writer_object_to_local_reader_.end());
info = it->second;
}
auto clear_written_so_far = [this, &writer_object_id]() {
absl::MutexLock guard(&written_so_far_lock_);
written_so_far_.erase(writer_object_id);
};
Status channel_status =
object_manager_->GetChannelStatus(info.local_object_id, /*is_reader=*/true);
if (!channel_status.ok()) {
clear_written_so_far();
return channel_status;
}

size_t total_data_size = request.total_data_size();
size_t total_metadata_size = request.total_metadata_size();

Expand All @@ -151,17 +163,23 @@ void MutableObjectProvider::HandlePushMutableObject(
if (tmp_written_so_far == 0u) {

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

security-critical critical

Security Vulnerability (Buffer Overflow / Over-read): The incoming PushMutableObjectRequest parameters (offset, chunk_size, and total_metadata_size) are not validated against the actual sizes of the data and metadata payloads, nor against the total allocated data size. This can lead to out-of-bounds writes (buffer overflow) on the receiver side during memcpy operations, or out-of-bounds reads (buffer over-read) if the payload sizes are smaller than expected.

Please validate these bounds before performing any operations or allocating resources.

  if (offset + chunk_size > total_data_size ||
      tmp_written_so_far + chunk_size > total_data_size ||
      request.data().size() < chunk_size ||
      request.metadata().size() < total_metadata_size) {
    clear_written_so_far();
    return Status::InvalidArgument("Invalid chunk offset, size, accumulated length, or metadata size");
  }
  if (tmp_written_so_far == 0u) {

@Pengkun-ZHU Pengkun-ZHU Oct 28, 2026 •

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I checked upstream code and found this appears to be a real issue.
The size check in MutableObjectManager::WriteAcquire only verifies that the declared data-plus-metadata size fits the receiver’s buffer. It does not verify that each chunk’s offset and length stay within that buffer, or that its payload contains the claimed number of bytes.
For example, an 80-byte declared total fits a 100-byte buffer, but copying 10 bytes at offset 95 still writes past its end.

As it is related to existing code that predates this very PR, I'll leave it for committers to comment.

// We set `metadata` to nullptr since the metadata is at the end of the object, which
// we will not have until the last chunk is received.
RAY_CHECK_OK(object_manager_->WriteAcquire(info.local_object_id,
total_data_size,
/*metadata=*/nullptr,
total_metadata_size,
info.num_readers,
object_backing_store));
Status status = object_manager_->WriteAcquire(info.local_object_id,
total_data_size,
/*metadata=*/nullptr,
total_metadata_size,
info.num_readers,
object_backing_store);
if (!status.ok()) {
clear_written_so_far();
return status;
}
} else {
RAY_CHECK_OK(object_manager_->GetObjectBackingStore(info.local_object_id,
total_data_size,
total_metadata_size,
object_backing_store));
Status status = object_manager_->GetObjectBackingStore(
info.local_object_id, total_data_size, total_metadata_size, object_backing_store);
if (!status.ok()) {
clear_written_so_far();
return status;
}
}
RAY_CHECK(object_backing_store);

Expand All @@ -174,11 +192,16 @@ void MutableObjectProvider::HandlePushMutableObject(
request.metadata().data(),
total_metadata_size);
// The entire object has been written, so call `WriteRelease()`.
RAY_CHECK_OK(object_manager_->WriteRelease(info.local_object_id));
Status status = object_manager_->WriteRelease(info.local_object_id);
if (!status.ok()) {
clear_written_so_far();
return status;
}
reply->set_done(true);
} else {
reply->set_done(false);
}
return Status::OK();
}

Status MutableObjectProvider::WriteAcquire(const ObjectID &object_id,
Expand Down Expand Up @@ -248,24 +271,48 @@ void MutableObjectProvider::PollWriterClosure(
RAY_CHECK(object->GetData());
RAY_CHECK(object->GetMetadata());

std::shared_ptr<size_t> num_replied = std::make_shared<size_t>(0);
struct PushCompletionState {
absl::Mutex mutex;
size_t num_replied ABSL_GUARDED_BY(mutex) = 0;
bool failed ABSL_GUARDED_BY(mutex) = false;
};
auto completion_state = std::make_shared<PushCompletionState>();
for (const auto &reader : *remote_readers) {
reader->PushMutableObject(
writer_object_id,
object->GetData()->Size(),
object->GetMetadata()->Size(),
object->GetData()->Data(),
object->GetMetadata()->Data(),
[this, &io_context, writer_object_id, remote_readers, num_replied](
[this, &io_context, writer_object_id, remote_readers, completion_state](
const Status &push_object_status, const rpc::PushMutableObjectReply &reply) {
*num_replied += 1;
if (!push_object_status.ok()) {
RAY_LOG(ERROR)
<< "Failed to transfer object to a remote node for an object id "
<< writer_object_id << ". It can cause hang.";
<< writer_object_id << ": " << push_object_status;
}

if (*num_replied == remote_readers->size()) {
bool close_writer = false;
bool continue_polling = false;
{
absl::MutexLock guard(&completion_state->mutex);
++completion_state->num_replied;
if (!push_object_status.ok() && !completion_state->failed) {
completion_state->failed = true;
close_writer = true;
}
continue_polling = completion_state->num_replied == remote_readers->size() &&
!completion_state->failed;
}
if (close_writer) {
// TODO(pzhu): Add failure recovery and retries for remote push failures.
Status close_status = object_manager_->SetError(writer_object_id);
if (!close_status.ok()) {
RAY_LOG(ERROR) << "Failed to close mutable object channel "
<< writer_object_id
<< " after a remote push failure: " << close_status;
}
} else if (continue_polling) {
io_context.post(
[this, &io_context, writer_object_id, remote_readers]() {
PollWriterClosure(io_context, writer_object_id, remote_readers);
Expand Down
4 changes: 2 additions & 2 deletions src/ray/core_worker/experimental_mutable_object_provider.h
Original file line number Diff line number Diff line change
Expand Up @@ -50,8 +50,8 @@ class MutableObjectProvider : public MutableObjectProviderInterface {
int64_t num_readers,
const ObjectID &reader_object_id) override;

void HandlePushMutableObject(const rpc::PushMutableObjectRequest &request,
rpc::PushMutableObjectReply *reply) override;
Status HandlePushMutableObject(const rpc::PushMutableObjectRequest &request,
rpc::PushMutableObjectReply *reply) override;

Status WriteAcquire(const ObjectID &object_id,
int64_t data_size,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -59,8 +59,12 @@ class MutableObjectProviderInterface {

/// RPC callback for when a writer pushes a mutable object over the network to a reader
/// on this node.
virtual void HandlePushMutableObject(const rpc::PushMutableObjectRequest &request,
rpc::PushMutableObjectReply *reply) = 0;
/// \param[in] request The request containing the pushed object chunk.
/// \param[out] reply The reply indicating whether the complete object has been
/// received. \return The status of processing this chunk. A closed channel is reported
/// to the sender so it can stop polling the failed channel.
virtual Status HandlePushMutableObject(const rpc::PushMutableObjectRequest &request,
rpc::PushMutableObjectReply *reply) = 0;
Comment thread
Pengkun-ZHU marked this conversation as resolved.
Resolved

/// Acquires a write lock on the object that prevents readers from reading
/// until we are done writing. This is safe for concurrent writers.
Expand Down
6 changes: 4 additions & 2 deletions src/ray/core_worker/fake_mutable_object_provider.h
Original file line number Diff line number Diff line change
Expand Up @@ -40,8 +40,10 @@ class FakeMutableObjectProvider : public MutableObjectProviderInterface {
int64_t num_readers,
const ObjectID &reader_object_id) override {}

void HandlePushMutableObject(const rpc::PushMutableObjectRequest &request,
rpc::PushMutableObjectReply *reply) override {}
Status HandlePushMutableObject(const rpc::PushMutableObjectRequest &request,
rpc::PushMutableObjectReply *reply) override {
return Status::OK();
}

Status WriteAcquire(const ObjectID &object_id,
int64_t data_size,
Expand Down
2 changes: 2 additions & 0 deletions src/ray/core_worker/tests/BUILD.bazel
Original file line number Diff line number Diff line change
Expand Up @@ -221,6 +221,8 @@ ray_cc_test(
"@com_google_absl//absl/functional:bind_front",
"@com_google_absl//absl/random",
"@com_google_absl//absl/strings:str_format",
"@com_google_absl//absl/synchronization",
"@com_google_absl//absl/time",
],
)

Expand Down
Loading
Loading