[mojo-edk] Add some buffer checks and fix UAF on NodeChannel

mojo/edk/system/message_pipe_dispatcher.cc

 // Copyright 2015 The Chromium Authors. All rights reserved.
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 
 #include "mojo/edk/system/message_pipe_dispatcher.h"
 
 #include <limits>
 #include <memory>
 
 #include "base/macros.h"
@@ skipping 157 matching lines @@
     MojoHandle* handles,
     uint32_t* num_handles,
     MojoReadMessageFlags flags,
     bool read_any_size) {
   // We can't read from a port that's closed or in transit!
   if (port_closed_ || in_transit_)
     return MOJO_RESULT_INVALID_ARGUMENT;
 
   bool no_space = false;
   bool may_discard = flags & MOJO_READ_MESSAGE_FLAG_MAY_DISCARD;
+  bool invalid_message = false;
 
   // Grab a message if the provided handles buffer is large enough. If the input
   // |num_bytes| is provided and |read_any_size| is false, we also ensure
   // that it specifies a size at least as large as the next available payload.
   //
   // If |read_any_size| is true, the input value of |*num_bytes| is ignored.
   // This flag exists to support both new and old API behavior.
 
   ports::ScopedMessage ports_message;
   int rv = node_controller_->node()->GetMessageIf(
       port_,
-      [read_any_size, num_bytes, num_handles, &no_space, &may_discard](
+      [read_any_size, num_bytes, num_handles, &no_space, &may_discard,
+       &invalid_message](
           const ports::Message& next_message) {
         const PortsMessage& message =
             static_cast<const PortsMessage&>(next_message);
-        DCHECK_GE(message.num_payload_bytes(), sizeof(MessageHeader));
+        if (message.num_payload_bytes() < sizeof(MessageHeader)) {
+          invalid_message = true;
+          return true;
+        }
 
         const MessageHeader* header =
             static_cast<const MessageHeader*>(message.payload_bytes());
-        DCHECK_LE(header->header_size, message.num_payload_bytes());
+        if (header->header_size > message.num_payload_bytes()) {
+          invalid_message = true;
+          return true;
+        }
 
         uint32_t bytes_to_read = 0;
         uint32_t bytes_available =
             static_cast<uint32_t>(message.num_payload_bytes()) -
             header->header_size;
         if (num_bytes) {
           bytes_to_read = std::min(*num_bytes, bytes_available);
           *num_bytes = bytes_available;
         }
 
         uint32_t handles_to_read = 0;
         uint32_t handles_available = header->num_dispatchers;
         if (num_handles) {
           handles_to_read = std::min(*num_handles, handles_available);
           *num_handles = handles_available;
         }
 
         if (handles_to_read < handles_available ||
             (!read_any_size && bytes_to_read < bytes_available)) {
           no_space = true;
           return may_discard;
         }
 
         return true;
       },
       &ports_message);
 
+  if (invalid_message)
+    return MOJO_RESULT_UNKNOWN;
+
   if (rv != ports::OK && rv != ports::ERROR_PORT_PEER_CLOSED) {
     if (rv == ports::ERROR_PORT_UNKNOWN ||
         rv == ports::ERROR_PORT_STATE_UNEXPECTED)
       return MOJO_RESULT_INVALID_ARGUMENT;
 
     NOTREACHED();
     return MOJO_RESULT_UNKNOWN;  // TODO: Add a better error code here?
   }
 
   if (no_space) {
@@ skipping 16 matching lines @@
     return MOJO_RESULT_FAILED_PRECONDITION;
   }
 
   // Alright! We have a message and the caller has provided sufficient storage
   // in which to receive it.
 
   std::unique_ptr<PortsMessage> msg(
       static_cast<PortsMessage*>(ports_message.release()));
 
   const MessageHeader* header =
-      static_cast<const MessageHeader*>( msg->payload_bytes());
+      static_cast<const MessageHeader*>(msg->payload_bytes());
   const DispatcherHeader* dispatcher_headers =
       reinterpret_cast<const DispatcherHeader*>(header + 1);
 
-  const char* dispatcher_data = reinterpret_cast<const char*>(
-      dispatcher_headers + header->num_dispatchers);
+  if (header->num_dispatchers > std::numeric_limits<uint16_t>::max())
+    return MOJO_RESULT_UNKNOWN;
 
   // Deserialize dispatchers.
   if (header->num_dispatchers > 0) {
     CHECK(handles);
     std::vector<DispatcherInTransit> dispatchers(header->num_dispatchers);
     size_t data_payload_index = sizeof(MessageHeader) +
         header->num_dispatchers * sizeof(DispatcherHeader);
+    if (data_payload_index > header->header_size)
+      return MOJO_RESULT_UNKNOWN;
+    const char* dispatcher_data = reinterpret_cast<const char*>(
+        dispatcher_headers + header->num_dispatchers);
     size_t port_index = 0;
     size_t platform_handle_index = 0;
     for (size_t i = 0; i < header->num_dispatchers; ++i) {
       const DispatcherHeader& dh = dispatcher_headers[i];
       Type type = static_cast<Type>(dh.type);
 
-      DCHECK_GE(msg->num_payload_bytes(),
-                data_payload_index + dh.num_bytes);
-      DCHECK_GE(msg->num_ports(),
-                port_index + dh.num_ports);
-      DCHECK_GE(msg->num_handles(),
-                platform_handle_index + dh.num_platform_handles);
+      size_t next_payload_index = data_payload_index + dh.num_bytes;
+      if (msg->num_payload_bytes() < next_payload_index ||
+          next_payload_index < data_payload_index) {
+        return MOJO_RESULT_UNKNOWN;
+      }
+
+      size_t next_port_index = port_index + dh.num_ports;
+      if (msg->num_ports() < next_port_index || next_port_index < port_index)
+        return MOJO_RESULT_UNKNOWN;
+
+      size_t next_platform_handle_index =
+          platform_handle_index + dh.num_platform_handles;
+      if (msg->num_handles() < next_platform_handle_index ||
+          next_platform_handle_index < platform_handle_index) {
+        return MOJO_RESULT_UNKNOWN;
+      }
 
       PlatformHandle* out_handles =
           msg->num_handles() ? msg->handles() + platform_handle_index : nullptr;
       dispatchers[i].dispatcher = Dispatcher::Deserialize(
           type, dispatcher_data, dh.num_bytes, msg->ports() + port_index,
           dh.num_ports, out_handles, dh.num_platform_handles);
       if (!dispatchers[i].dispatcher)
         return MOJO_RESULT_UNKNOWN;
 
       dispatcher_data += dh.num_bytes;
-      data_payload_index += dh.num_bytes;
-      port_index += dh.num_ports;
-      platform_handle_index += dh.num_platform_handles;
+      data_payload_index = next_payload_index;
+      port_index = next_port_index;
+      platform_handle_index = next_platform_handle_index;
     }
 
     if (!node_controller_->core()->AddDispatchersFromTransit(dispatchers,
                                                              handles))
       return MOJO_RESULT_UNKNOWN;
   }
 
   CHECK(msg);
   *message = MessageForTransit::WrapPortsMessage(std::move(msg));
   return MOJO_RESULT_OK;
@@ skipping 207 matching lines @@
                << " endpoint " << endpoint_ << " [port=" << port_.name() << "]";
     }
   }
 #endif
 
   awakables_.AwakeForStateChange(GetHandleSignalsStateNoLock());
 }
 
 }  // namespace edk
 }  // namespace mojo