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 236353)
+++ chrome/browser/extensions/api/cast_channel/cast_socket.cc	(working copy)
@@ -74,14 +74,13 @@
     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();
@@ -151,77 +150,51 @@
   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::Connect(const net::CompletionCallback& callback) {
+  DCHECK(CalledOnValidThread());
+  connect_callback_ = callback;
+  // Post a task so that we can safely run the callback

[Wez, 2013/11/26 02:45:56]

You mean so that if the Connect() completes synchronously then the callback can
be invoked from within the Connect impl?

[Munjal (Google), 2013/11/26 20:06:14]

Actually I removed this PostTask. Upon thinking about it more, it is safe to
invoke callback synchronously in error cases or in cases when net operations
finish synchronously as long as we don't access data members in the callstack
unwind path since the caller can delete this in the callback. This was per
suggestion from Ryan.

[Wez, 2013/11/27 04:15:19]

That makes the code super-fragile to code changes that don't take the
restriction into account, though; even with this CL as-is DoConnectLoop accesses
connect_state_, a member of the object, after running a connection state
function that may have triggered the callback - if the callback were to delete
the CastSocket then that would lead to a use-after-free. 

[Munjal (Google), 2013/12/02 19:22:09]

Several things to keep in mind:
- DoConnectLoop does not access connect_state_ after invoking the connect
callback. Connection callback is only invoked in one place in the file - inside
DoConnectCallback - which is called in one cplace only - outside the do/while
loop in DoConnectLoop.
- Whether or not we post a task, we have to ensure that it is safe for the
caller to delete the socket in the connection callback

+  base::MessageLoop::current()->PostTask(
+      FROM_HERE,
+      base::Bind(&CastSocket::ConnectInternal, AsWeakPtr()));
 }
 
-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) {
+void CastSocket::ConnectInternal() {
   DCHECK(CalledOnValidThread());
-  int result = net::ERR_CONNECTION_FAILED;
   VLOG(1) << "Connect readyState = " << ready_state_;
   if (ready_state_ != READY_STATE_NONE) {

[Wez, 2013/11/26 02:45:56]

Under what circumstances could |ready_state_| not be NONE when we reach here?
Should this be a DCHECK?

[Munjal (Google), 2013/11/26 20:06:14]

When the caller calls Connect twice. In that case we invoke the callback with
error.

[Wez, 2013/11/27 04:15:19]

You mean that JS calls a connect method twice? Or some other Chrome code? If the
latter then DCHECK seems more appropriate.

[Munjal (Google), 2013/12/02 19:22:09]

Yes, if Js calls connect method twice. Of course DCHECK would be fine if it
cannot be done external to Chrome code. But error seems like a better way to
ahndle it since it can be triggered externally. Another way would be just to
ignore all calls to Connect except the first one. Let me know if you prefer
that.

-    callback.Run(result);
+    DoConnectCallback(net::ERR_CONNECTION_FAILED);
     return;
   }
   if (!ParseChannelUrl(url_)) {
     CloseWithError(cast_channel::CHANNEL_ERROR_CONNECT_ERROR);
-    callback.Run(result);
+    DoConnectCallback(net::ERR_CONNECTION_FAILED);
     return;
   }
-  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);
 }
 
 // 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. ConnectInternal 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*
+    ConnectionState state = connect_state_;
+    // All the Do* methods do not set connect_state_ in case of an
+    // error. So set connect_state_ to NONE to figure out if the Do*
     // method changed state or not.
-    next_state_ = CONN_STATE_NONE;
+    connect_state_ = CONN_STATE_NONE;
     switch (state) {
       case CONN_STATE_TCP_CONNECT:
         rv = DoTcpConnect();
@@ -245,46 +218,47 @@
       case CONN_STATE_AUTH_CHALLENGE_REPLY_COMPLETE:
         rv = DoAuthChallengeReplyComplete(rv);
         break;
-
       default:
-        NOTREACHED() << "BUG in CastSocket state machine code";
+        NOTREACHED() << "BUG in CastSocket connection state machine code";
         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;
+  // If there is no pending IO and if we still got out of the loop then
+  // either we are done successfully or there was an error; invoke the
+  // callback in either case.
+  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) {
   VLOG(1) << "DoTcpConnectComplete: " << result;
   if (result == net::OK) {
     // 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;
+    bool success = tcp_socket_->SetKeepAlive(true, kTcpKeepAliveDelaySecs);
+    LOG_IF(WARNING, !success) << "Failed to SetKeepAlive.";
+    connect_state_ = CONN_STATE_SSL_CONNECT;
   }
   return result;
 }
 
 int CastSocket::DoSslConnect() {
   VLOG(1) << "DoSslConnect";
-  next_state_ = CONN_STATE_SSL_CONNECT_COMPLETE;
+  connect_state_ = CONN_STATE_SSL_CONNECT_COMPLETE;
   socket_ = CreateSslSocket();
   return socket_->Connect(
-      base::Bind(&CastSocket::OnConnectComplete, AsWeakPtr()));
+      base::Bind(&CastSocket::DoConnectLoop, AsWeakPtr()));
 }
 
 int CastSocket::DoSslConnectComplete(int result) {
@@ -292,30 +266,40 @@
   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 a challenge message to avoid re-entrancy
+  PostTaskToSendCastMessage(
+      challenge_message,
+      base::Bind(&CastSocket::DoConnectLoop, AsWeakPtr()));
+  // Always return IO_PENDING since we always get the result asynchronously
+  return net::ERR_IO_PENDING;
 }
 
 int CastSocket::DoAuthChallengeSendComplete(int result) {
   VLOG(1) << "DoAuthChallengeSendComplete: " << result;
-  if (result != net::OK)
+  if (result < 0)

[Wez, 2013/11/26 02:45:56]

Didn't we already fix these?

[Munjal (Google), 2013/11/26 20:06:14]

Yes, we did but I refactored that method a bit and now I consistently handle the
result appropriate in Do* method.

     return result;
-  next_state_ = CONN_STATE_AUTH_CHALLENGE_REPLY_COMPLETE;
-  return ReadAuthChallengeReply();
+  connect_state_ = CONN_STATE_AUTH_CHALLENGE_REPLY_COMPLETE;
+  PostTaskToStartReadLoop();
+  // Always return IO_PENDING since we always get the result asynchronously
+  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,20 +307,14 @@
   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;
+  // Start the ReadData loop if not already started.
+  if (result == net::OK)
+    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) {
@@ -353,102 +331,173 @@
 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::PostTaskToSendCastMessage(
+    const CastMessage& message, const net::CompletionCallback& callback) {
+  base::MessageLoop::current()->PostTask(
+      FROM_HERE,
+      base::Bind(&CastSocket::SendCastMessageInternal, AsWeakPtr(),
+      message, callback));

[Wez, 2013/11/26 02:45:56]

This means you're posting a task on every send; why not try the send, and post
the completion callback if it completes synchronously?

[Munjal (Google), 2013/11/26 20:06:14]

We don't post a task for every SendMessage. We only post a task if we are
sending the auth challenge. SendMessage method directly calls
SendMessageInternal. But SendAuthChallenge calls PostTaskToSendMessage.

[Wez, 2013/11/27 04:15:19]

In that case can't SendAuthChallenge do the post directly? Why does it need
bouncing through another helper method?

[Munjal (Google), 2013/12/02 19:22:09]

Done. I thought it might be handy to have this helper in case we need it from
multiple places. But if we do that, we can add it later. Removed now.

+}
+
+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::ClearWriteQueue() {
+  while (!write_queue_.empty())
+    write_queue_.pop();
+}
+
+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())
+    return;
+
+  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_ERROR:
+        rv = DoWriteError(rv);
+        break;
+      default:
+        NOTREACHED() << "BUG in CastSocket write state machine code";
+        break;
+    }
+  } while (!write_queue_.empty() &&
+           rv != net::ERR_IO_PENDING &&
+           write_state_ != WRITE_STATE_NONE);
+}
+
+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;

[Wez, 2013/11/26 02:45:56]

Surely at this point the write is pending?

[Munjal (Google), 2013/11/26 20:06:14]

Yes. But the caller - DoWriteLoop - will get an IO_PENDING as return value if
the Write below does not finish sync'ly and in that case DoWriteLoop will wait
to get a callback before doing DoWriteComplete

[Wez, 2013/11/27 04:15:19]

Right - so the write has _not_ completed - we're in a write-pending state; the
write hasn't completed as far as we're concerned until the completion callback
returns, at which point we return to an about-to-write or nothing-to-write
state?

[Munjal (Google), 2013/12/02 19:22:09]

I think the confusion is in the name - this is the "next" write state to be
executed but ONLY WHEN the current operation completes.

This is how the code in net does things - it sets the state to the next state
but returns IO_PENDING to indicate that that cannot be executed yet.

I think it keeps the number of states manageable.

+
+  return socket_->Write(
       request.io_buffer.get(),
       request.io_buffer->BytesRemaining(),
-      base::Bind(&CastSocket::OnWriteData, AsWeakPtr()));
-
-  if (result != net::ERR_IO_PENDING)
-    OnWriteData(result);
-
-  return result;
+      base::Bind(&CastSocket::DoWriteLoop, AsWeakPtr()));
 }
 
-void CastSocket::OnWriteData(int result) {
-  DCHECK(CalledOnValidThread());
-  VLOG(1) << "OnWriteComplete result = " << result;
-  DCHECK(write_callback_pending_);
+int CastSocket::DoWriteComplete(int result) {
   DCHECK(!write_queue_.empty());
-  write_callback_pending_ = false;
   WriteRequest& request = write_queue_.front();
-  scoped_refptr<net::DrainableIOBuffer> io_buffer = request.io_buffer;
 
-  if (result >= 0) {
+  if (result < 0) {
+    // ERROR: do the write completion callback and get into error state
+    // to close the session
+    request.callback.Run(result);

[Wez, 2013/11/26 02:45:56]

If the write was triggered by connect-time activity then surely you can end up
in DoConnectLoop, called from DoWriteComplete, via DoWriteLoop. Is that safe?

[Munjal (Google), 2013/11/26 20:06:14]

If the Write was triggered by connect-time activity, DOConnectLoop will not be
on the stack since we always post a task to start DoWriteLoop. So we will be
fine.

[Wez, 2013/11/27 04:15:19]

The case I'm asking about is DoWriteLoop -> DoWriteComplete -> DoConnectLoop.
Are you sure that that is a safe nesting, including the case in which the caller
deletes the socket in response to e.g. failed connect?

[Munjal (Google), 2013/12/02 19:22:09]

You are right. Fixed now. We should not access any member variables after
invoking the callback.

+    write_queue_.pop();
+    write_state_ = WRITE_STATE_ERROR;
+    return result;
+  }
+
+  if (result > 0) {
+    scoped_refptr<net::DrainableIOBuffer> io_buffer = request.io_buffer;
+    // Some bytes were successfully written
     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;
+    if (io_buffer->BytesRemaining() == 0) {
+      // A message is sent fully, do the write completion callback
+      request.callback.Run(io_buffer->BytesConsumed());
+      write_queue_.pop();
     }
-    DCHECK_EQ(io_buffer->BytesConsumed(), io_buffer->size());
-    DCHECK_EQ(io_buffer->BytesRemaining(), 0);
-    result = io_buffer->BytesConsumed();
   }
 
-  request.callback.Run(result);
-  write_queue_.pop();
+  write_state_ = WRITE_STATE_WRITE;

[Wez, 2013/11/26 02:45:56]

Under what circumstances do we ever return to WRITE_STATE_NONE?

[Munjal (Google), 2013/11/26 20:06:14]

We never return to WRITE_STATE_NONE. We only start out in that state and we
never return to that state.

+  return net::OK;
+}
 
-  VLOG(1) << "OnWriteComplete size = " << io_buffer->size()
-          << " consumed " << io_buffer->BytesConsumed()
-          << " remaining " << io_buffer->BytesRemaining()
-          << " # requests " << write_queue_.size();
+int CastSocket::DoWriteError(int result) {
+  DCHECK(!write_queue_.empty());
+  DCHECK_LT(result, 0);
+  // TODO(munjal): Consider reporting error to all pending writes
+  ClearWriteQueue();
+  CloseWithError(CHANNEL_ERROR_SOCKET_ERROR);
+  return net::OK;
+}
 
-  if (result < 0) {
-    CloseWithError(CHANNEL_ERROR_SOCKET_ERROR);
-    return;
+void CastSocket::PostTaskToStartReadLoop() {
+  DCHECK(CalledOnValidThread());
+  // Post a task to avoid re-entrancy into DoConnectLoop
+  base::MessageLoop::current()->PostTask(
+      FROM_HERE,
+      base::Bind(&CastSocket::StartReadLoop, AsWeakPtr()));

[Wez, 2013/11/26 02:45:56]

As for read, why do this? Why not start the read directly, and post a task only
if it completes synchronously, to handle the result?

[Munjal (Google), 2013/11/26 20:06:14]

This felt cleaner and more robust since otherwise we have to keep state to
figure out if read completed synchronously or not deep down into the helper
methods so that it can do the right thing. And we have to this in every code
path that does the callback (right now there is only one though).

Let me know what you think.

[Wez, 2013/11/27 04:15:19]

The amount of logic should be the same; what I'm suggesting is:

<do work>
Read(&Callback)
<if synchronous, Post(Callback), else return>

This way Callback is always handled asynchronously, and you avoid the Post() in
the case where Read() completes asynchronously.

[Munjal (Google), 2013/12/02 19:22:09]

If it finishes synchronously then it will invoke the callback already deep
inside hte stack. So the way to detect that would be to add a boolean -
in_read_loop_ - which is set to true here and then checked in
ParseMessageFromBody.

May be the right thin to do is to refactor the call to DoConnectLoop outside of
ParseMessageBody and add one more state for that?

+}
+
+void CastSocket::StartReadLoop() {
+  // If we are in READ_STATE_NONE then get into appropriate
+  // starting state and start the read loop
+  if (read_state_ == READ_STATE_NONE) {
+    read_state_ = READ_STATE_READ;
+    DoReadLoop(net::OK);
   }
+}
 
-  if (!write_queue_.empty())
-    WriteData();
+void CastSocket::DoReadLoop(int result) {
+  DCHECK(CalledOnValidThread());
+  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_ERROR:
+        rv = DoReadError(rv);
+        break;
+      default:
+        NOTREACHED() << "BUG in read state machine";
+        break;
+    }
+  } while (rv != net::ERR_IO_PENDING && read_state_ != READ_STATE_NONE);
 }
 
-int CastSocket::ReadData() {
-  DCHECK(CalledOnValidThread());
-  if (!socket_.get())
-    return net::ERR_FAILED;
-  DCHECK(!read_callback_pending_);
-  read_callback_pending_ = true;
+int CastSocket::DoRead() {
+  read_state_ = READ_STATE_READ_COMPLETE;
   // Figure out if we are reading the header or body, and the remaining bytes.
   uint32 num_bytes_to_read = 0;
   if (header_read_buffer_->RemainingCapacity() > 0) {
@@ -462,57 +511,60 @@
     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(
+  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) {
+int CastSocket::DoReadComplete(int result) {
   DCHECK(CalledOnValidThread());
-  VLOG(1) << "OnReadData result = " << 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;
+    error_state_ = CHANNEL_ERROR_SOCKET_ERROR;
+    read_state_ = READ_STATE_ERROR;
+    return result;
   }
+
   // We read some data.  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);
 
-  bool should_continue = true;
+  bool success = 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();
+    // If we have read a full header, process the contents.
+    success = ProcessHeader();
   } 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();
+    success = ProcessBody();
   }
-  if (should_continue)
-    ReadData();
+
+  if (success)
+    read_state_ = READ_STATE_READ;
+  return net::OK;
 }
 
+int CastSocket::DoReadError(int result) {
+  DCHECK_LT(result, 0);
+  CloseWithError(error_state_);
+  return net::OK;
+}
+
 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);
+    error_state_ = cast_channel::CHANNEL_ERROR_INVALID_MESSAGE;
     return false;
   }
   VLOG(1) << "Parsed header { message_size: " << header.message_size << " }";
@@ -524,7 +576,7 @@
   DCHECK_EQ(static_cast<uint32>(body_read_buffer_->offset()),
             current_message_size_);
   if (!ParseMessageFromBody()) {
-    CloseWithError(cast_channel::CHANNEL_ERROR_INVALID_MESSAGE);
+    error_state_ = cast_channel::CHANNEL_ERROR_INVALID_MESSAGE;
     return false;
   }
   current_message_size_ = 0;
@@ -541,13 +593,14 @@
   CastMessage message_proto;
   if (!message_proto.ParseFromArray(
       body_read_buffer_->StartOfBuffer(),
-      current_message_size_))
+      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);
+    DoConnectLoop(net::OK);
   } else if (delegate_) {
     MessageInfo message;
     if (!CastMessageToMessageInfo(message_proto, &message))