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 243975)
+++ 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);
 
 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_);
 
@@ -147,81 +145,55 @@
   bool result = net::X509Certificate::GetDEREncoded(
      ssl_info.cert->os_cert_handle(), cert);
   if (result)
-    VLOG(1) << "Successfully extracted peer certificate: " << *cert;
+    VLOG(1) << "[URL: " << url_.spec()
+            << "] Successfully extracted peer certificate: " << *cert;
   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 +217,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 +245,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 +263,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 +311,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 +555,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 +649,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 +660,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 +724,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_;