Fixing a Dr. Memory hang in the remoting_unittests

remoting/host/security_key/remote_security_key_message_reader_impl.cc

Index: remoting/host/security_key/remote_security_key_message_reader_impl.cc
diff --git a/remoting/host/security_key/remote_security_key_message_reader_impl.cc b/remoting/host/security_key/remote_security_key_message_reader_impl.cc
index 55c4061d1af7b7366f307885e3fb5ce0e04ac87f..30a6686ba6b823ded31ab8cd85c82f42561e29b6 100644
--- a/remoting/host/security_key/remote_security_key_message_reader_impl.cc
+++ b/remoting/host/security_key/remote_security_key_message_reader_impl.cc
@@ -64,33 +64,20 @@ void RemoteSecurityKeyMessageReaderImpl::ReadMessage() {
       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;
     }
@@ -110,6 +97,31 @@ void RemoteSecurityKeyMessageReaderImpl::ReadMessage() {
   }
 }
 
+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());