Refactor CastSocket code for the following:

chrome/browser/extensions/api/cast_channel/cast_socket.cc

Index: chrome/browser/extensions/api/cast_channel/cast_socket.cc
===================================================================
--- chrome/browser/extensions/api/cast_channel/cast_socket.cc	(revision 240465)
+++ chrome/browser/extensions/api/cast_channel/cast_socket.cc	(working copy)
@@ -35,10 +35,6 @@
 const char kCastInsecureScheme[] = "cast";
 const char kCastSecureScheme[] = "casts";
 
-// Size of the message header, in bytes.  Don't use sizeof(MessageHeader)
-// because of alignment; instead, sum the sizeof() for the fields.
-const uint32 kMessageHeaderSize = sizeof(uint32);
-
 // The default keepalive delay.  On Linux, keepalives probes will be sent after
 // the socket is idle for this length of time, and the socket will be closed
 // after 9 failed probes.  So the total idle time before close is 10 *
@@ -64,6 +60,9 @@
 namespace cast_channel {
 
 const uint32 kMaxMessageSize = 65536;
+// Don't use sizeof(MessageHeader) because of alignment; instead, sum the
+// sizeof() for the fields.
+const uint32 kMessageHeaderSize = sizeof(uint32);

[Ryan Sleevi, 2013/12/18 00:02:13]

Why would you take sizeof(uint32)? You're saying it's 32-bit, why bother with
sizeof? :)

[Munjal (Google), 2013/12/19 18:20:21]

The idea is that(originally by mfoltz) kMessageHeaderSize should be sum of
sizeof of various fields in MessageHeader struct. Currently there is only one
uint32 field. But if we add more fields we will update this constant as well.

 
 CastSocket::CastSocket(const std::string& owner_extension_id,
                        const GURL& url,
@@ -74,19 +73,18 @@
     url_(url),
     delegate_(delegate),
     auth_required_(false),
-    error_state_(CHANNEL_ERROR_NONE),
-    ready_state_(READY_STATE_NONE),
-    write_callback_pending_(false),
-    read_callback_pending_(false),
     current_message_size_(0),
     net_log_(net_log),
-    next_state_(CONN_STATE_NONE),
-    in_connect_loop_(false) {
+    connect_state_(CONN_STATE_NONE),
+    write_state_(WRITE_STATE_NONE),
+    read_state_(READ_STATE_NONE),
+    error_state_(CHANNEL_ERROR_NONE),
+    ready_state_(READY_STATE_NONE) {
   DCHECK(net_log_);
   net_log_source_.type = net::NetLog::SOURCE_SOCKET;
   net_log_source_.id = net_log_->NextID();
 
-  // We reuse these buffers for each message.
+  // Reuse these buffers for each message.
   header_read_buffer_ = new net::GrowableIOBuffer();
   header_read_buffer_->SetCapacity(kMessageHeaderSize);
   body_read_buffer_ = new net::GrowableIOBuffer();
@@ -102,15 +100,15 @@
 
 scoped_ptr<net::TCPClientSocket> CastSocket::CreateTcpSocket() {
   net::AddressList addresses(ip_endpoint_);
-  scoped_ptr<net::TCPClientSocket> tcp_socket(
+  return scoped_ptr<net::TCPClientSocket>(
       new net::TCPClientSocket(addresses, net_log_, net_log_source_));
   // Options cannot be set on the TCPClientSocket yet, because the
-  // underlying platform socket will not be created until we Bind()
-  // or Connect() it.
-  return tcp_socket.Pass();
+  // underlying platform socket will not be created until Bind()
+  // or Connect() is called.
 }
 
-scoped_ptr<net::SSLClientSocket> CastSocket::CreateSslSocket() {
+scoped_ptr<net::SSLClientSocket> CastSocket::CreateSslSocket(
+    scoped_ptr<net::StreamSocket> socket) {
   net::SSLConfig ssl_config;
   // If a peer cert was extracted in a previous attempt to connect, then
   // whitelist that cert.
@@ -130,7 +128,7 @@
   context.transport_security_state = transport_security_state_.get();
 
   scoped_ptr<net::ClientSocketHandle> connection(new net::ClientSocketHandle);
-  connection->SetSocket(tcp_socket_.PassAs<net::StreamSocket>());
+  connection->SetSocket(socket.Pass());
   net::HostPortPair host_and_port = net::HostPortPair::FromIPEndPoint(
       ip_endpoint_);
 
@@ -151,77 +149,50 @@
   return result;
 }
 
-int CastSocket::SendAuthChallenge() {
-  CastMessage challenge_message;
-  CreateAuthChallengeMessage(&challenge_message);
-  VLOG(1) << "Sending challenge: " << CastMessageToString(challenge_message);
-  int result = SendMessageInternal(
-      challenge_message,
-      base::Bind(&CastSocket::OnChallengeEvent, AsWeakPtr()));
-  return (result < 0) ? result : net::OK;
+bool CastSocket::VerifyChallengeReply() {
+  return AuthenticateChallengeReply(*challenge_reply_.get(), peer_cert_);
 }
 
-int CastSocket::ReadAuthChallengeReply() {
-  int result = ReadData();
-  return (result < 0) ? result : net::OK;
-}
-
-void CastSocket::OnConnectComplete(int result) {
-  int rv = DoConnectLoop(result);
-  if (rv != net::ERR_IO_PENDING)
-    DoConnectCallback(rv);
-}
-
-void CastSocket::OnChallengeEvent(int result) {
-  // result >= 0 means read or write succeeded synchronously.
-  int rv = DoConnectLoop(result >= 0 ? net::OK : result);
-  if (rv != net::ERR_IO_PENDING)
-    DoConnectCallback(rv);
-}
-
 void CastSocket::Connect(const net::CompletionCallback& callback) {
   DCHECK(CalledOnValidThread());
-  int result = net::ERR_CONNECTION_FAILED;
   VLOG(1) << "Connect readyState = " << ready_state_;
   if (ready_state_ != READY_STATE_NONE) {
-    callback.Run(result);
+    callback.Run(net::ERR_CONNECTION_FAILED);
     return;
   }
   if (!ParseChannelUrl(url_)) {
-    CloseWithError(cast_channel::CHANNEL_ERROR_CONNECT_ERROR);
-    callback.Run(result);
+    callback.Run(net::ERR_CONNECTION_FAILED);
     return;
   }
+
+  ready_state_ = READY_STATE_CONNECTING;
   connect_callback_ = callback;
-  next_state_ = CONN_STATE_TCP_CONNECT;
-  int rv = DoConnectLoop(net::OK);
-  if (rv != net::ERR_IO_PENDING)
-    DoConnectCallback(rv);
+  connect_state_ = CONN_STATE_TCP_CONNECT;
+  DoConnectLoop(net::OK);
 }
 
+void CastSocket::PostTaskToStartConnectLoop(int result) {
+  DCHECK(CalledOnValidThread());
+  base::MessageLoop::current()->PostTask(
+      FROM_HERE,
+      base::Bind(&CastSocket::DoConnectLoop, AsWeakPtr(), result));
+}
+
 // This method performs the state machine transitions for connection flow.
 // There are two entry points to this method:
-// 1. public Connect method: this starts the flow
-// 2. OnConnectComplete: callback method called when an async operation
-//    is done. OnConnectComplete calls this method to continue the state
-//    machine transitions.
-int CastSocket::DoConnectLoop(int result) {
-  // Avoid re-entrancy as a result of synchronous completion.
-  if (in_connect_loop_)
-    return net::ERR_IO_PENDING;
-  in_connect_loop_ = true;
-
+// 1. Connect method: this starts the flow
+// 2. Callback from network operations that finish asynchronously
+void CastSocket::DoConnectLoop(int result) {
   // Network operations can either finish synchronously or asynchronously.
   // This method executes the state machine transitions in a loop so that
   // correct state transitions happen even when network operations finish
   // synchronously.
   int rv = result;
   do {
-    ConnectionState state = next_state_;
-    // All the Do* methods do not set next_state_ in case of an
-    // error. So set next_state_ to NONE to figure out if the Do*
-    // method changed state or not.
-    next_state_ = CONN_STATE_NONE;
+    ConnectionState state = connect_state_;
+    // Default to CONN_STATE_NONE, which breaks the processing loop if any
+    // handler fails to transition to another state to continue processing.
+    connect_state_ = CONN_STATE_NONE;
     switch (state) {
       case CONN_STATE_TCP_CONNECT:
         rv = DoTcpConnect();
@@ -245,27 +216,26 @@
       case CONN_STATE_AUTH_CHALLENGE_REPLY_COMPLETE:
         rv = DoAuthChallengeReplyComplete(rv);
         break;
-
       default:
-        NOTREACHED() << "BUG in CastSocket state machine code";
+        NOTREACHED() << "BUG in connect flow. Unknown state: " << state;
         break;
     }
-  } while (rv != net::ERR_IO_PENDING && next_state_ != CONN_STATE_NONE);
+  } while (rv != net::ERR_IO_PENDING && connect_state_ != CONN_STATE_NONE);
   // Get out of the loop either when:
   // a. A network operation is pending, OR
   // b. The Do* method called did not change state
 
-  in_connect_loop_ = false;
-
-  return rv;
+  // Connect loop is finished: if there is no pending IO invoke the callback.
+  if (rv != net::ERR_IO_PENDING)
+    DoConnectCallback(rv);
 }
 
 int CastSocket::DoTcpConnect() {
   VLOG(1) << "DoTcpConnect";
-  next_state_ = CONN_STATE_TCP_CONNECT_COMPLETE;
+  connect_state_ = CONN_STATE_TCP_CONNECT_COMPLETE;
   tcp_socket_ = CreateTcpSocket();
   return tcp_socket_->Connect(
-      base::Bind(&CastSocket::OnConnectComplete, AsWeakPtr()));
+      base::Bind(&CastSocket::DoConnectLoop, AsWeakPtr()));
 }
 
 int CastSocket::DoTcpConnectComplete(int result) {
@@ -274,17 +244,17 @@
     // Enable TCP protocol-level keep-alive.
     bool result = tcp_socket_->SetKeepAlive(true, kTcpKeepAliveDelaySecs);
     LOG_IF(WARNING, !result) << "Failed to SetKeepAlive.";
-    next_state_ = CONN_STATE_SSL_CONNECT;
+    connect_state_ = CONN_STATE_SSL_CONNECT;
   }
   return result;
 }
 
 int CastSocket::DoSslConnect() {
   VLOG(1) << "DoSslConnect";
-  next_state_ = CONN_STATE_SSL_CONNECT_COMPLETE;
-  socket_ = CreateSslSocket();
+  connect_state_ = CONN_STATE_SSL_CONNECT_COMPLETE;
+  socket_ = CreateSslSocket(tcp_socket_.PassAs<net::StreamSocket>());
   return socket_->Connect(
-      base::Bind(&CastSocket::OnConnectComplete, AsWeakPtr()));
+      base::Bind(&CastSocket::DoConnectLoop, AsWeakPtr()));
 }
 
 int CastSocket::DoSslConnectComplete(int result) {
@@ -292,30 +262,47 @@
   if (result == net::ERR_CERT_AUTHORITY_INVALID &&
              peer_cert_.empty() &&
              ExtractPeerCert(&peer_cert_)) {
-    next_state_ = CONN_STATE_TCP_CONNECT;
+    connect_state_ = CONN_STATE_TCP_CONNECT;
   } else if (result == net::OK && auth_required_) {
-    next_state_ = CONN_STATE_AUTH_CHALLENGE_SEND;
+    connect_state_ = CONN_STATE_AUTH_CHALLENGE_SEND;
   }
   return result;
 }
 
 int CastSocket::DoAuthChallengeSend() {
   VLOG(1) << "DoAuthChallengeSend";
-  next_state_ = CONN_STATE_AUTH_CHALLENGE_SEND_COMPLETE;
-  return SendAuthChallenge();
+  connect_state_ = CONN_STATE_AUTH_CHALLENGE_SEND_COMPLETE;
+  CastMessage challenge_message;
+  CreateAuthChallengeMessage(&challenge_message);
+  VLOG(1) << "Sending challenge: " << CastMessageToString(challenge_message);
+  // Post a task to send auth challenge so that DoWriteLoop is not nested inside
+  // DoConnectLoop. This is not strictly necessary but keeps the write loop
+  // code decoupled from connect loop code.
+  base::MessageLoop::current()->PostTask(
+      FROM_HERE,
+      base::Bind(&CastSocket::SendCastMessageInternal, AsWeakPtr(),
+                 challenge_message,
+                 base::Bind(&CastSocket::DoConnectLoop, AsWeakPtr())));
+  // Always return IO_PENDING since the result is always asynchronous.
+  return net::ERR_IO_PENDING;
 }
 
 int CastSocket::DoAuthChallengeSendComplete(int result) {
   VLOG(1) << "DoAuthChallengeSendComplete: " << result;
-  if (result != net::OK)
+  if (result < 0)
     return result;
-  next_state_ = CONN_STATE_AUTH_CHALLENGE_REPLY_COMPLETE;
-  return ReadAuthChallengeReply();
+  connect_state_ = CONN_STATE_AUTH_CHALLENGE_REPLY_COMPLETE;
+  // Post a task to start read loop so that DoReadLoop is not nested inside
+  // DoConnectLoop. This is not strictly necessary but keeps the read loop
+  // code decoupled from connect loop code.
+  PostTaskToStartReadLoop();
+  // Always return IO_PENDING since the result is always asynchronous.
+  return net::ERR_IO_PENDING;
 }
 
 int CastSocket::DoAuthChallengeReplyComplete(int result) {
   VLOG(1) << "DoAuthChallengeReplyComplete: " << result;
-  if (result != net::OK)
+  if (result < 0)
     return result;
   if (!VerifyChallengeReply())
     return net::ERR_FAILED;
@@ -323,133 +310,238 @@
   return net::OK;
 }
 
-bool CastSocket::VerifyChallengeReply() {
-  return AuthenticateChallengeReply(*challenge_reply_.get(), peer_cert_);
-}
-
 void CastSocket::DoConnectCallback(int result) {
   ready_state_ = (result == net::OK) ? READY_STATE_OPEN : READY_STATE_CLOSED;
   error_state_ = (result == net::OK) ?
       CHANNEL_ERROR_NONE : CHANNEL_ERROR_CONNECT_ERROR;
+  if (result == net::OK)  // Start the read loop
+    PostTaskToStartReadLoop();
   base::ResetAndReturn(&connect_callback_).Run(result);
-  // Start the ReadData loop if not already started.
-  // If auth_required_ is true we would've started a ReadData loop already.
-  // TODO(munjal): This is a bit ugly. Refactor read and write code.
-  if (result == net::OK && !auth_required_)
-    ReadData();
 }
 
 void CastSocket::Close(const net::CompletionCallback& callback) {
   DCHECK(CalledOnValidThread());
   VLOG(1) << "Close ReadyState = " << ready_state_;
-  tcp_socket_.reset(NULL);
-  socket_.reset(NULL);
-  cert_verifier_.reset(NULL);
-  transport_security_state_.reset(NULL);
+  tcp_socket_.reset();
+  socket_.reset();
+  cert_verifier_.reset();
+  transport_security_state_.reset();
   ready_state_ = READY_STATE_CLOSED;
   callback.Run(net::OK);
+  // |callback| can delete |this|
 }
 
 void CastSocket::SendMessage(const MessageInfo& message,
                              const net::CompletionCallback& callback) {
   DCHECK(CalledOnValidThread());
-  VLOG(1) << "Send ReadyState " << ready_state_;
-  int result = net::ERR_FAILED;
   if (ready_state_ != READY_STATE_OPEN) {
-    callback.Run(result);
+    callback.Run(net::ERR_FAILED);
     return;
   }
   CastMessage message_proto;
   if (!MessageInfoToCastMessage(message, &message_proto)) {
-    CloseWithError(cast_channel::CHANNEL_ERROR_INVALID_MESSAGE);
-    // TODO(mfoltz): Do a better job of signaling cast_channel errors to the
-    // caller.
-    callback.Run(net::OK);
+    callback.Run(net::ERR_FAILED);
     return;
   }
-  SendMessageInternal(message_proto, callback);
+
+  SendCastMessageInternal(message_proto, callback);
 }
 
-int CastSocket::SendMessageInternal(const CastMessage& message_proto,
-                                    const net::CompletionCallback& callback) {
+void CastSocket::SendCastMessageInternal(
+    const CastMessage& message,
+    const net::CompletionCallback& callback) {
   WriteRequest write_request(callback);
-  if (!write_request.SetContent(message_proto))
-    return net::ERR_FAILED;
+  if (!write_request.SetContent(message)) {
+    callback.Run(net::ERR_FAILED);
+    return;
+  }
+
   write_queue_.push(write_request);
-  return WriteData();
+  if (write_state_ == WRITE_STATE_NONE) {
+    write_state_ = WRITE_STATE_WRITE;
+    DoWriteLoop(net::OK);
+  }
 }
 
-int CastSocket::WriteData() {
+void CastSocket::DoWriteLoop(int result) {
   DCHECK(CalledOnValidThread());
   VLOG(1) << "WriteData q = " << write_queue_.size();
-  if (write_queue_.empty() || write_callback_pending_)
-    return net::ERR_FAILED;
 
+  if (write_queue_.empty()) {
+    write_state_ = WRITE_STATE_NONE;
+    return;
+  }
+
+  // Network operations can either finish synchronously or asynchronously.
+  // This method executes the state machine transitions in a loop so that
+  // write state transitions happen even when network operations finish
+  // synchronously.
+  int rv = result;
+  do {
+    WriteState state = write_state_;
+    write_state_ = WRITE_STATE_NONE;
+    switch (state) {
+      case WRITE_STATE_WRITE:
+        rv = DoWrite();
+        break;
+      case WRITE_STATE_WRITE_COMPLETE:
+        rv = DoWriteComplete(rv);
+        break;
+      case WRITE_STATE_DO_CALLBACK:
+        rv = DoWriteCallback();
+        break;
+      case WRITE_STATE_ERROR:
+        rv = DoWriteError(rv);
+        break;
+      default:
+        NOTREACHED() << "BUG in write flow. Unknown state: " << state;
+        break;
+    }
+  } while (!write_queue_.empty() &&
+           rv != net::ERR_IO_PENDING &&
+           write_state_ != WRITE_STATE_NONE);
+
+  // If write loop is done because the queue is empty then set write
+  // state to NONE
+  if (write_queue_.empty())
+    write_state_ = WRITE_STATE_NONE;
+
+  // Write loop is done - if the result is ERR_FAILED then close with error.
+  if (rv == net::ERR_FAILED)
+    CloseWithError(error_state_);
+}
+
+int CastSocket::DoWrite() {
+  DCHECK(!write_queue_.empty());
   WriteRequest& request = write_queue_.front();
 
   VLOG(1) << "WriteData byte_count = " << request.io_buffer->size()
-          << " bytes_written " << request.io_buffer->BytesConsumed();
+           << " bytes_written " << request.io_buffer->BytesConsumed();
 
-  write_callback_pending_ = true;
-  int result = socket_->Write(
+  write_state_ = WRITE_STATE_WRITE_COMPLETE;
+
+  return socket_->Write(
       request.io_buffer.get(),
       request.io_buffer->BytesRemaining(),
-      base::Bind(&CastSocket::OnWriteData, AsWeakPtr()));
+      base::Bind(&CastSocket::DoWriteLoop, AsWeakPtr()));
+}
 
-  if (result != net::ERR_IO_PENDING)
-    OnWriteData(result);
+int CastSocket::DoWriteComplete(int result) {
+  DCHECK(!write_queue_.empty());
+  if (result <= 0) {  // NOTE that 0 also indicates an error
+    error_state_ = CHANNEL_ERROR_SOCKET_ERROR;
+    write_state_ = WRITE_STATE_ERROR;
+    return result == 0 ? net::ERR_FAILED : result;
+  }
 
-  return result;
+  // Some bytes were successfully written
+  WriteRequest& request = write_queue_.front();
+  scoped_refptr<net::DrainableIOBuffer> io_buffer = request.io_buffer;
+  io_buffer->DidConsume(result);
+  if (io_buffer->BytesRemaining() == 0)  // Message fully sent
+    write_state_ = WRITE_STATE_DO_CALLBACK;
+  else
+    write_state_ = WRITE_STATE_WRITE;
+
+  return net::OK;
 }
 
-void CastSocket::OnWriteData(int result) {
-  DCHECK(CalledOnValidThread());
-  VLOG(1) << "OnWriteComplete result = " << result;
-  DCHECK(write_callback_pending_);
+int CastSocket::DoWriteCallback() {
   DCHECK(!write_queue_.empty());
-  write_callback_pending_ = false;
   WriteRequest& request = write_queue_.front();
-  scoped_refptr<net::DrainableIOBuffer> io_buffer = request.io_buffer;
+  int bytes_consumed = request.io_buffer->BytesConsumed();
 
-  if (result >= 0) {
-    io_buffer->DidConsume(result);
-    if (io_buffer->BytesRemaining() > 0) {
-      VLOG(1) << "OnWriteComplete size = " << io_buffer->size()
-              << " consumed " << io_buffer->BytesConsumed()
-              << " remaining " << io_buffer->BytesRemaining()
-              << " # requests " << write_queue_.size();
-      WriteData();
-      return;
-    }
-    DCHECK_EQ(io_buffer->BytesConsumed(), io_buffer->size());
-    DCHECK_EQ(io_buffer->BytesRemaining(), 0);
-    result = io_buffer->BytesConsumed();
+  // If inside connection flow, then there should be exaclty one item in
+  // the write queue.
+  if (ready_state_ == READY_STATE_CONNECTING) {
+    write_queue_.pop();
+    DCHECK(write_queue_.empty());
+    PostTaskToStartConnectLoop(bytes_consumed);
+  } else {
+    WriteRequest& request = write_queue_.front();
+    request.callback.Run(bytes_consumed);
+    write_queue_.pop();
   }
+  write_state_ = WRITE_STATE_WRITE;
+  return net::OK;
+}
 
-  request.callback.Run(result);
-  write_queue_.pop();
+int CastSocket::DoWriteError(int result) {
+  DCHECK(!write_queue_.empty());
+  DCHECK_LT(result, 0);
 
-  VLOG(1) << "OnWriteComplete size = " << io_buffer->size()
-          << " consumed " << io_buffer->BytesConsumed()
-          << " remaining " << io_buffer->BytesRemaining()
-          << " # requests " << write_queue_.size();
+  // If inside connection flow, then there should be exactly one item in
+  // the write queue.
+  if (ready_state_ == READY_STATE_CONNECTING) {
+    write_queue_.pop();
+    DCHECK(write_queue_.empty());
+    PostTaskToStartConnectLoop(result);
+    // Connect loop will handle the error. Return net::OK so that write flow
+    // does not try to report error also.
+    return net::OK;
+  }
 
-  if (result < 0) {
-    CloseWithError(CHANNEL_ERROR_SOCKET_ERROR);
-    return;
+  while (!write_queue_.empty()) {
+    WriteRequest& request = write_queue_.front();
+    request.callback.Run(result);
+    write_queue_.pop();
   }
+  return net::ERR_FAILED;
+}
 
-  if (!write_queue_.empty())
-    WriteData();
+void CastSocket::PostTaskToStartReadLoop() {
+  DCHECK(CalledOnValidThread());
+  base::MessageLoop::current()->PostTask(
+      FROM_HERE,
+      base::Bind(&CastSocket::StartReadLoop, AsWeakPtr()));
 }
 
-int CastSocket::ReadData() {
+void CastSocket::StartReadLoop() {
+  // Read loop would have already been started if read state is not NONE
+  if (read_state_ == READ_STATE_NONE) {
+    read_state_ = READ_STATE_READ;
+    DoReadLoop(net::OK);
+  }
+}
+
+void CastSocket::DoReadLoop(int result) {
   DCHECK(CalledOnValidThread());
-  if (!socket_.get())
-    return net::ERR_FAILED;
-  DCHECK(!read_callback_pending_);
-  read_callback_pending_ = true;
-  // Figure out if we are reading the header or body, and the remaining bytes.
+  // Network operations can either finish synchronously or asynchronously.
+  // This method executes the state machine transitions in a loop so that
+  // write state transitions happen even when network operations finish
+  // synchronously.
+  int rv = result;
+  do {
+    ReadState state = read_state_;
+    read_state_ = READ_STATE_NONE;
+
+    switch (state) {
+      case READ_STATE_READ:
+        rv = DoRead();
+        break;
+      case READ_STATE_READ_COMPLETE:
+        rv = DoReadComplete(rv);
+        break;
+      case READ_STATE_DO_CALLBACK:
+        rv = DoReadCallback();
+        break;
+      case READ_STATE_ERROR:
+        rv = DoReadError(rv);
+        break;
+      default:
+        NOTREACHED() << "BUG in read flow. Unknown state: " << state;
+        break;
+    }
+  } while (rv != net::ERR_IO_PENDING && read_state_ != READ_STATE_NONE);
+
+  // Read loop is done - If the result is ERR_FAILED then close with error.
+  if (rv == net::ERR_FAILED)
+    CloseWithError(error_state_);
+}
+
+int CastSocket::DoRead() {
+  read_state_ = READ_STATE_READ_COMPLETE;
+  // Figure out whether to read header or body, and the remaining bytes.
   uint32 num_bytes_to_read = 0;
   if (header_read_buffer_->RemainingCapacity() > 0) {
     current_read_buffer_ = header_read_buffer_;
@@ -462,59 +554,92 @@
     DCHECK_LE(num_bytes_to_read, kMaxMessageSize);
   }
   DCHECK_GT(num_bytes_to_read, 0U);
-  // We read up to num_bytes_to_read into |current_read_buffer_|.
-  int result = socket_->Read(
+
+  // Read up to num_bytes_to_read into |current_read_buffer_|.
+  return socket_->Read(
       current_read_buffer_.get(),
       num_bytes_to_read,
-      base::Bind(&CastSocket::OnReadData, AsWeakPtr()));
-  VLOG(1) << "ReadData result = " << result;
-  if (result > 0) {
-    OnReadData(result);
-  } else if (result != net::ERR_IO_PENDING) {
-    CloseWithError(CHANNEL_ERROR_SOCKET_ERROR);
-  }
-  return result;
+      base::Bind(&CastSocket::DoReadLoop, AsWeakPtr()));
 }
 
-void CastSocket::OnReadData(int result) {
-  DCHECK(CalledOnValidThread());
-  VLOG(1) << "OnReadData result = " << result
+int CastSocket::DoReadComplete(int result) {
+  VLOG(1) << "DoReadDataComplete result = " << result
           << " header offset = " << header_read_buffer_->offset()
           << " body offset = " << body_read_buffer_->offset();
-  read_callback_pending_ = false;
-  if (result <= 0) {
-    CloseWithError(CHANNEL_ERROR_SOCKET_ERROR);
-    return;
+  if (result <= 0) {  // 0 means EOF: the peer closed the socket
+    error_state_ = CHANNEL_ERROR_SOCKET_ERROR;
+    read_state_ = READ_STATE_ERROR;
+    return result == 0 ? net::ERR_FAILED : result;
   }
-  // We read some data.  Move the offset in the current buffer forward.
+
+  // Some data was read.  Move the offset in the current buffer forward.
   DCHECK_LE(current_read_buffer_->offset() + result,
             current_read_buffer_->capacity());
   current_read_buffer_->set_offset(current_read_buffer_->offset() + result);
+  read_state_ = READ_STATE_READ;
 
-  bool should_continue = true;
   if (current_read_buffer_.get() == header_read_buffer_.get() &&
       current_read_buffer_->RemainingCapacity() == 0) {
-  // If we have read a full header, process the contents.
-    should_continue = ProcessHeader();
+    // A full header is read, process the contents.
+    if (!ProcessHeader()) {
+      error_state_ = cast_channel::CHANNEL_ERROR_INVALID_MESSAGE;
+      read_state_ = READ_STATE_ERROR;
+    }
   } else if (current_read_buffer_.get() == body_read_buffer_.get() &&
              static_cast<uint32>(current_read_buffer_->offset()) ==
              current_message_size_) {
-    // If we have read a full body, process the contents.
-    should_continue = ProcessBody();
+    // Full body is read, process the contents.
+    if (ProcessBody()) {
+      read_state_ = READ_STATE_DO_CALLBACK;
+    } else {
+      error_state_ = cast_channel::CHANNEL_ERROR_INVALID_MESSAGE;
+      read_state_ = READ_STATE_ERROR;
+    }
   }
-  if (should_continue)
-    ReadData();
+
+  return net::OK;
 }
 
+int CastSocket::DoReadCallback() {
+  read_state_ = READ_STATE_READ;
+  if (IsAuthMessage(current_message_)) {
+    // An auth message is received, check that connect flow is running.
+    if (ready_state_ == READY_STATE_CONNECTING) {
+      challenge_reply_.reset(new CastMessage(current_message_));
+      PostTaskToStartConnectLoop(net::OK);
+    } else {
+      read_state_ = READ_STATE_ERROR;
+    }
+  } else if (delegate_) {
+    MessageInfo message;
+    if (CastMessageToMessageInfo(current_message_, &message))
+      delegate_->OnMessage(this, message);
+    else
+      read_state_ = READ_STATE_ERROR;
+  }
+  current_message_.Clear();
+  return net::OK;
+}
+
+int CastSocket::DoReadError(int result) {
+  DCHECK_LE(result, 0);
+  // If inside connection flow, then get back to connect loop.
+  if (ready_state_ == READY_STATE_CONNECTING) {
+    PostTaskToStartConnectLoop(result);
+    // does not try to report error also.
+    return net::OK;
+  }
+  return net::ERR_FAILED;
+}
+
 bool CastSocket::ProcessHeader() {
   DCHECK_EQ(static_cast<uint32>(header_read_buffer_->offset()),
             kMessageHeaderSize);
   MessageHeader header;
   MessageHeader::ReadFromIOBuffer(header_read_buffer_.get(), &header);
-  if (header.message_size > kMaxMessageSize) {
-    CloseWithError(cast_channel::CHANNEL_ERROR_INVALID_MESSAGE);
+  if (header.message_size > kMaxMessageSize)
     return false;
-  }
+
   VLOG(1) << "Parsed header { message_size: " << header.message_size << " }";
   current_message_size_ = header.message_size;
   return true;
@@ -523,8 +648,8 @@
 bool CastSocket::ProcessBody() {
   DCHECK_EQ(static_cast<uint32>(body_read_buffer_->offset()),
             current_message_size_);
-  if (!ParseMessageFromBody()) {
-    CloseWithError(cast_channel::CHANNEL_ERROR_INVALID_MESSAGE);
+  if (!current_message_.ParseFromArray(
+      body_read_buffer_->StartOfBuffer(), current_message_size_)) {
     return false;
   }
   current_message_size_ = 0;
@@ -534,29 +659,6 @@
   return true;
 }
 
-bool CastSocket::ParseMessageFromBody() {
-  DCHECK(CalledOnValidThread());
-  DCHECK_EQ(static_cast<uint32>(body_read_buffer_->offset()),
-            current_message_size_);
-  CastMessage message_proto;
-  if (!message_proto.ParseFromArray(
-      body_read_buffer_->StartOfBuffer(),
-      current_message_size_))
-    return false;
-  VLOG(1) << "Parsed message " << CastMessageToString(message_proto);
-  // If the message is an auth message then we handle it internally.
-  if (IsAuthMessage(message_proto)) {
-    challenge_reply_.reset(new CastMessage(message_proto));
-    OnChallengeEvent(net::OK);
-  } else if (delegate_) {
-    MessageInfo message;
-    if (!CastMessageToMessageInfo(message_proto, &message))
-      return false;
-    delegate_->OnMessage(this, message);
-  }
-  return true;
-}
-
 // static
 bool CastSocket::Serialize(const CastMessage& message_proto,
                            std::string* message_data) {
@@ -621,7 +723,6 @@
 };
 
 void CastSocket::FillChannelInfo(ChannelInfo* channel_info) const {
-  DCHECK(CalledOnValidThread());
   channel_info->channel_id = channel_id_;
   channel_info->url = url_.spec();
   channel_info->ready_state = ready_state_;