Mojo C++ bindings: introduce an optional array to store transferred interface IDs in messages.

mojo/public/cpp/bindings/lib/message.cc

Index: mojo/public/cpp/bindings/lib/message.cc
diff --git a/mojo/public/cpp/bindings/lib/message.cc b/mojo/public/cpp/bindings/lib/message.cc
index 0018a03bcd7909b9b9657f5d66b6c750b1941485..7da4a97b835dc75b6effb1fe21fb2453d945c0db 100644
--- a/mojo/public/cpp/bindings/lib/message.cc
+++ b/mojo/public/cpp/bindings/lib/message.cc
@@ -16,6 +16,8 @@
 #include "base/logging.h"
 #include "base/strings/stringprintf.h"
 #include "base/threading/thread_local.h"
+#include "mojo/public/cpp/bindings/associated_group_controller.h"
+#include "mojo/public/cpp/bindings/lib/array_internal.h"
 
 namespace mojo {
 
@@ -37,8 +39,10 @@ Message::Message() {
 }
 
 Message::Message(Message&& other)
-    : buffer_(std::move(other.buffer_)), handles_(std::move(other.handles_)) {
-}
+    : buffer_(std::move(other.buffer_)),
+      handles_(std::move(other.handles_)),
+      associated_endpoint_handles_(
+          std::move(other.associated_endpoint_handles_)) {}
 
 Message::~Message() {
   CloseHandles();
@@ -48,12 +52,14 @@ Message& Message::operator=(Message&& other) {
   Reset();
   std::swap(other.buffer_, buffer_);
   std::swap(other.handles_, handles_);
+  std::swap(other.associated_endpoint_handles_, associated_endpoint_handles_);
   return *this;
 }
 
 void Message::Reset() {
   CloseHandles();
   handles_.clear();
+  associated_endpoint_handles_.clear();
   buffer_.reset();
 }
 
@@ -70,7 +76,52 @@ void Message::InitializeFromMojoMessage(ScopedMessageHandle message,
   handles_.swap(*handles);
 }
 
+const uint8_t* Message::payload() const {
+  if (version() < 2)
+    return data() + header()->num_bytes;
+
+  return static_cast<const uint8_t*>(header_v2()->payload.Get());
+}
+
+uint32_t Message::payload_num_bytes() const {
+  DCHECK_GE(data_num_bytes(), header()->num_bytes);
+  size_t num_bytes;
+  if (version() < 2) {
+    num_bytes = data_num_bytes() - header()->num_bytes;
+  } else {
+    auto payload = reinterpret_cast<uintptr_t>(header_v2()->payload.Get());
+    if (!payload) {
+      num_bytes = 0;
+    } else {
+      auto payload_end =
+          reinterpret_cast<uintptr_t>(header_v2()->payload_interface_ids.Get());
+      if (!payload_end)
+        payload_end = reinterpret_cast<uintptr_t>(data() + data_num_bytes());
+      DCHECK_GE(payload_end, payload);
+      num_bytes = payload_end - payload;
+    }
+  }
+  DCHECK_LE(num_bytes, std::numeric_limits<uint32_t>::max());
+  return static_cast<uint32_t>(num_bytes);
+}
+
+uint32_t Message::payload_num_interface_ids() const {
+  auto array_pointer =
+      version() < 2 ? nullptr : header_v2()->payload_interface_ids.Get();
+  return array_pointer ? static_cast<uint32_t>(array_pointer->size()) : 0;
+}
+
+const uint32_t* Message::payload_interface_ids() const {
+  auto array_pointer =
+      version() < 2 ? nullptr : header_v2()->payload_interface_ids.Get();
+  return array_pointer ? array_pointer->storage() : nullptr;
+}
+
 ScopedMessageHandle Message::TakeMojoMessage() {
+  // If there are associated endpoints transferred,
+  // SerializeAssociatedEndpointHandles() must be called before this method.
+  DCHECK(associated_endpoint_handles_.empty());
+
   if (handles_.empty())  // Fast path for the common case: No handles.
     return buffer_->TakeMessage();
 
@@ -114,6 +165,55 @@ void Message::CloseHandles() {
   }
 }
 
+void Message::SerializeAssociatedEndpointHandles(
+    AssociatedGroupController* group_controller) {
+  if (associated_endpoint_handles_.empty())
+    return;
+
+  DCHECK_GE(version(), 2u);
+  DCHECK(header_v2()->payload_interface_ids.is_null());
+
+  size_t size = associated_endpoint_handles_.size();
+  auto data = internal::Array_Data<uint32_t>::New(size, buffer());
+  header_v2()->payload_interface_ids.Set(data);
+
+  for (size_t i = 0; i < size; ++i) {
+    ScopedInterfaceEndpointHandle& handle = associated_endpoint_handles_[i];
+
+    DCHECK(handle.is_valid());
+    DCHECK(!handle.is_local());
+    DCHECK_EQ(group_controller, handle.group_controller());
+    data->storage()[i] = handle.release();
+  }
+  associated_endpoint_handles_.clear();
+}
+
+bool Message::DeserializeAssociatedEndpointHandles(
+    AssociatedGroupController* group_controller) {
+  associated_endpoint_handles_.clear();
+
+  uint32_t num_ids = payload_num_interface_ids();
+  if (num_ids == 0)
+    return true;
+
+  associated_endpoint_handles_.reserve(num_ids);
+  uint32_t* ids = header_v2()->payload_interface_ids.Get()->storage();
+  bool result = true;
+  for (uint32_t i = 0; i < num_ids; ++i) {
+    auto handle = group_controller->CreateLocalEndpointHandle(ids[i]);
+    if (IsValidInterfaceId(ids[i]) && !handle.is_valid()) {
+      // |ids[i]| itself is valid but handle creation failed. In that case, mark
+      // deserialization as failed but continue to deserialize the rest of
+      // handles.
+      result = false;
+    }
+
+    associated_endpoint_handles_.push_back(std::move(handle));
+    ids[i] = kInvalidInterfaceId;
+  }
+  return result;
+}
+
 PassThroughFilter::PassThroughFilter() {}
 
 PassThroughFilter::~PassThroughFilter() {}
@@ -195,6 +295,8 @@ ReportBadMessageCallback GetBadMessageCallback() {
 
 namespace internal {
 
+MessageHeaderV2::MessageHeaderV2() = default;
+
 MessageDispatchContext::MessageDispatchContext(Message* message)
     : outer_context_(current()), message_(message) {
   g_tls_message_dispatch_context.Get().Set(this);