Fixing a Dr. Memory hang in the remoting_unittests

remoting/host/security_key/remote_security_key_message_reader_impl.cc

 // Copyright 2016 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 "remoting/host/security_key/remote_security_key_message_reader_impl.h"
 
 #include <cstdint>
 #include <string>
 #include <utility>
 
@@ skipping 46 matching lines @@
 void RemoteSecurityKeyMessageReaderImpl::ReadMessage() {
   DCHECK(read_task_runner_->RunsTasksOnCurrentThread());
 
   while (true) {
     if (!read_stream_.IsValid()) {
       LOG(ERROR) << "Cannot read from invalid stream.";
       NotifyError();
       return;
     }
 
-    // Read the message header to retrieve the remaining message length.
-    uint32_t total_message_size_bytes;
-    int read_result = read_stream_.ReadAtCurrentPos(
-        reinterpret_cast<char*>(&total_message_size_bytes),
-        SecurityKeyMessage::kHeaderSizeBytes);
-    if (read_result != SecurityKeyMessage::kHeaderSizeBytes) {
-      // 0 means EOF which is normal and should not be logged as an error.
-      if (read_result != 0) {
-        LOG(ERROR) << "Failed to read message header, read returned "
-                   << read_result;
-      }
+    uint32_t message_length_bytes = 0;
+    if (!ReadFromStream(reinterpret_cast<char*>(&message_length_bytes), 4)) {
       NotifyError();
       return;
     }
 
-    if (!SecurityKeyMessage::IsValidMessageSize(total_message_size_bytes)) {
-      LOG(ERROR) << "Message size too large: " << total_message_size_bytes;
+    if (!SecurityKeyMessage::IsValidMessageSize(message_length_bytes)) {
+      LOG(ERROR) << "Message size is invalid: " << message_length_bytes;
       NotifyError();
       return;
     }
 
-    std::string message_data(total_message_size_bytes, '\0');
-    read_result = read_stream_.ReadAtCurrentPos(string_as_array(&message_data),
-                                                total_message_size_bytes);
-    // The static cast is safe as we know the value is smaller than max int.
-    if (read_result != static_cast<int>(total_message_size_bytes)) {
-      LOG(ERROR) << "Failed to read message: " << read_result;
+    std::string message_data(message_length_bytes, '\0');
+    if (!ReadFromStream(string_as_array(&message_data), message_data.size())) {
       NotifyError();
       return;
     }
 
     std::unique_ptr<SecurityKeyMessage> message(new SecurityKeyMessage());
     if (!message->ParseMessage(message_data)) {
       LOG(ERROR) << "Invalid message data received.";
       NotifyError();
       return;
     }
 
     // Notify callback of the new message received.
     main_task_runner_->PostTask(
         FROM_HERE,
         base::Bind(&RemoteSecurityKeyMessageReaderImpl::InvokeMessageCallback,
                    weak_factory_.GetWeakPtr(), base::Passed(&message)));
   }
 }
 
+bool RemoteSecurityKeyMessageReaderImpl::ReadFromStream(
+    char* buffer,
+    uint32_t bytes_to_read) {

[Sergey Ulanov, 2016/04/19 17:36:39]

size_t

[joedow, 2016/04/19 21:07:56]

Done.

+  DCHECK(buffer);
+  DCHECK_GT(bytes_to_read, 0u);
+
+  uint32_t bytes_read = 0;

[Sergey Ulanov, 2016/04/19 17:36:40]

size_t

[joedow, 2016/04/19 21:07:56]

Done.

+  do {
+    int read_result = read_stream_.ReadAtCurrentPos(buffer + bytes_read,

[Sergey Ulanov, 2016/04/19 17:36:40]

Use ReadAtCurrentPosNoBestEffort() here?

[joedow, 2016/04/19 21:07:56]

Done.

+                                                    bytes_to_read - bytes_read);
+    if (read_result < 1) {
+      // 0 means EOF which is normal and should not be logged as an error.
+      if (read_result != 0) {
+        LOG(ERROR) << "Failed to read from stream, ReadAtCurrentPos returned "
+                   << read_result;
+      }
+      return false;
+    }
+    bytes_read += read_result;
+  } while (bytes_read < bytes_to_read);
+  DCHECK_EQ(bytes_read, bytes_to_read);
+
+  return true;
+}
+
 void RemoteSecurityKeyMessageReaderImpl::NotifyError() {
   DCHECK(read_task_runner_->RunsTasksOnCurrentThread());
 
   main_task_runner_->PostTask(
       FROM_HERE,
       base::Bind(&RemoteSecurityKeyMessageReaderImpl::InvokeErrorCallback,
                  weak_factory_.GetWeakPtr()));
 }
 
 void RemoteSecurityKeyMessageReaderImpl::InvokeMessageCallback(
     std::unique_ptr<SecurityKeyMessage> message) {
   DCHECK(main_task_runner_->RunsTasksOnCurrentThread());
   message_callback_.Run(std::move(message));
 }
 
 void RemoteSecurityKeyMessageReaderImpl::InvokeErrorCallback() {
   DCHECK(main_task_runner_->RunsTasksOnCurrentThread());
   error_callback_.Run();
 }
 
 }  // namespace remoting