[GCM] Check if all of message has arrived before processing

A socket read may only return part of a message. In that case, the connection
handler should wait for the rest of the message to arrive, rather than
immediately attempting to decode the data (which will fail and trigger a
reconnect).

BUG=373056

Review URL: https://codereview.chromium.org/285113003

git-svn-id: svn://svn.chromium.org/chrome/trunk/src@270599 0039d316-1c4b-4281-b951-d872f2087c98

3 files changed, 97 insertions, 8 deletions

diff --git a/google_apis/gcm/base/socket_stream.cc b/google_apis/gcm/base/socket_stream.cc
index 1a0b29d..8c152c6 100644
--- a/google_apis/gcm/base/socket_stream.cc
+++ b/google_apis/gcm/base/socket_stream.cc
@@ -89,8 +89,8 @@ net::Error SocketInputStream::Refresh(const base::Closure& callback,
   DCHECK_GT(byte_limit, 0);
 
   if (byte_limit > read_buffer_->BytesRemaining()) {
-    NOTREACHED() << "Out of buffer space, closing input stream.";
-    CloseStream(net::ERR_UNEXPECTED, base::Closure());
+    LOG(ERROR) << "Out of buffer space, closing input stream.";
+    CloseStream(net::ERR_FILE_TOO_BIG, base::Closure());
     return net::OK;
   }
 
diff --git a/google_apis/gcm/engine/connection_handler_impl.cc b/google_apis/gcm/engine/connection_handler_impl.cc
index d1ef0ee..d47e0f2 100644
--- a/google_apis/gcm/engine/connection_handler_impl.cc
+++ b/google_apis/gcm/engine/connection_handler_impl.cc
@@ -184,9 +184,9 @@ void ConnectionHandlerImpl::WaitForData(ProcessingState state) {
   }
 
   // Used to determine whether a Socket::Read is necessary.
-  int min_bytes_needed = 0;
+  size_t min_bytes_needed = 0;
   // Used to limit the size of the Socket::Read.
-  int max_bytes_needed = 0;
+  size_t max_bytes_needed = 0;
 
   switch(state) {
     case MCS_VERSION_TAG_AND_SIZE:
@@ -214,20 +214,20 @@ void ConnectionHandlerImpl::WaitForData(ProcessingState state) {
   }
   DCHECK_GE(max_bytes_needed, min_bytes_needed);
 
-  int byte_count = input_stream_->UnreadByteCount();
-  if (min_bytes_needed - byte_count > 0 &&
+  size_t unread_byte_count = input_stream_->UnreadByteCount();
+  if (min_bytes_needed > unread_byte_count &&
       input_stream_->Refresh(
           base::Bind(&ConnectionHandlerImpl::WaitForData,
                      weak_ptr_factory_.GetWeakPtr(),
                      state),
-          max_bytes_needed - byte_count) == net::ERR_IO_PENDING) {
+          max_bytes_needed - unread_byte_count) == net::ERR_IO_PENDING) {
     return;
   }
 
   // Check for refresh errors.
   if (input_stream_->GetState() != SocketInputStream::READY) {
     // An error occurred.
-    int last_error = output_stream_->last_error();
+    int last_error = input_stream_->last_error();
     CloseConnection();
     // If the socket stream had an error, plumb it up, else plumb up FAILED.
     if (last_error == net::OK)
@@ -236,6 +236,20 @@ void ConnectionHandlerImpl::WaitForData(ProcessingState state) {
     return;
   }
 
+  // Check whether read is complete, or needs to be continued (
+  // SocketInputStream::Refresh can finish without reading all the data).
+  if (input_stream_->UnreadByteCount() < min_bytes_needed) {
+    DVLOG(1) << "Socket read finished prematurely. Waiting for "
+             << min_bytes_needed - input_stream_->UnreadByteCount()
+             << " more bytes.";
+    base::MessageLoop::current()->PostTask(
+        FROM_HERE,
+        base::Bind(&ConnectionHandlerImpl::WaitForData,
+                   weak_ptr_factory_.GetWeakPtr(),
+                   MCS_PROTO_BYTES));
+    return;
+  }
+
   // Received enough bytes, process them.
   DVLOG(1) << "Processing MCS data: state == " << state;
   switch(state) {
diff --git a/google_apis/gcm/engine/connection_handler_impl_unittest.cc b/google_apis/gcm/engine/connection_handler_impl_unittest.cc
index 32632a8..6b89644 100644
--- a/google_apis/gcm/engine/connection_handler_impl_unittest.cc
+++ b/google_apis/gcm/engine/connection_handler_impl_unittest.cc
@@ -382,6 +382,7 @@ TEST_F(GCMConnectionHandlerImplTest, RecvMsg) {
   WaitForMessage();  // The data message.
   ASSERT_TRUE(received_message.get());
   EXPECT_EQ(data_message_proto, received_message->SerializeAsString());
+  EXPECT_EQ(net::OK, last_error());
 }
 
 // Verify that if two messages arrive at once, they're treated appropriately.
@@ -419,6 +420,7 @@ TEST_F(GCMConnectionHandlerImplTest, Recv2Msgs) {
   WaitForMessage();  // The second data message.
   ASSERT_TRUE(received_message.get());
   EXPECT_EQ(data_message_proto2, received_message->SerializeAsString());
+  EXPECT_EQ(net::OK, last_error());
 }
 
 // Receive a long (>128 bytes) message.
@@ -450,6 +452,46 @@ TEST_F(GCMConnectionHandlerImplTest, RecvLongMsg) {
   WaitForMessage();  // The data message.
   ASSERT_TRUE(received_message.get());
   EXPECT_EQ(data_message_proto, received_message->SerializeAsString());
+  EXPECT_EQ(net::OK, last_error());
+}
+
+// Receive a long (>128 bytes) message in two synchronous parts.
+TEST_F(GCMConnectionHandlerImplTest, RecvLongMsg2Parts) {
+  std::string handshake_request = EncodeHandshakeRequest();
+  WriteList write_list(1, net::MockWrite(net::ASYNC,
+                                         handshake_request.c_str(),
+                                         handshake_request.size()));
+  std::string handshake_response = EncodeHandshakeResponse();
+
+  std::string data_message_proto =
+      BuildDataMessage(kDataMsgFromLong, kDataMsgCategoryLong);
+  std::string data_message_pkt =
+      EncodePacket(kDataMessageStanzaTag, data_message_proto);
+  DCHECK_GT(data_message_pkt.size(), 128U);
+  ReadList read_list;
+  read_list.push_back(net::MockRead(net::ASYNC,
+                                    handshake_response.c_str(),
+                                    handshake_response.size()));
+
+  int bytes_in_first_message = data_message_pkt.size() / 2;
+  read_list.push_back(net::MockRead(net::SYNCHRONOUS,
+                                    data_message_pkt.c_str(),
+                                    bytes_in_first_message));
+  read_list.push_back(net::MockRead(net::SYNCHRONOUS,
+                                    data_message_pkt.c_str() +
+                                        bytes_in_first_message,
+                                    data_message_pkt.size() -
+                                        bytes_in_first_message));
+  BuildSocket(read_list, write_list);
+
+  ScopedMessage received_message;
+  Connect(&received_message);
+  WaitForMessage();  // The login send.
+  WaitForMessage();  // The login response.
+  WaitForMessage();  // The data message.
+  ASSERT_TRUE(received_message.get());
+  EXPECT_EQ(net::OK, last_error());
+  EXPECT_EQ(data_message_proto, received_message->SerializeAsString());
 }
 
 // Receive two long (>128 bytes) message.
@@ -488,6 +530,7 @@ TEST_F(GCMConnectionHandlerImplTest, Recv2LongMsgs) {
   WaitForMessage();  // The second data message.
   ASSERT_TRUE(received_message.get());
   EXPECT_EQ(data_message_proto2, received_message->SerializeAsString());
+  EXPECT_EQ(net::OK, last_error());
 }
 
 // Simulate a message where the end of the data does not arrive in time and the
@@ -630,5 +673,37 @@ TEST_F(GCMConnectionHandlerImplTest, SendMsgSocketDisconnected) {
   EXPECT_EQ(net::ERR_CONNECTION_CLOSED, last_error());
 }
 
+// Receive a message whose size field was corrupted and is larger than the
+// socket's buffer. Should fail gracefully.
+TEST_F(GCMConnectionHandlerImplTest, CorruptedSize) {
+  std::string handshake_request = EncodeHandshakeRequest();
+  WriteList write_list(1, net::MockWrite(net::ASYNC,
+                                         handshake_request.c_str(),
+                                         handshake_request.size()));
+  std::string handshake_response = EncodeHandshakeResponse();
+
+  // Fill a string with 9000 character zero.
+  std::string data_message_proto(9000, '0');
+  std::string data_message_pkt =
+      EncodePacket(kDataMessageStanzaTag, data_message_proto);
+  ReadList read_list;
+  read_list.push_back(net::MockRead(net::ASYNC,
+                                    handshake_response.c_str(),
+                                    handshake_response.size()));
+  read_list.push_back(net::MockRead(net::ASYNC,
+                                    data_message_pkt.c_str(),
+                                    data_message_pkt.size()));
+  BuildSocket(read_list, write_list);
+
+  ScopedMessage received_message;
+  Connect(&received_message);
+  WaitForMessage();  // The login send.
+  WaitForMessage();  // The login response.
+  received_message.reset();
+  WaitForMessage();  // The data message.
+  EXPECT_FALSE(received_message.get());
+  EXPECT_EQ(net::ERR_FILE_TOO_BIG, last_error());
+}
+
 }  // namespace
 }  // namespace gcm