Fix WebSocket crash bug.

net/socket_stream/socket_stream.cc

 // Copyright (c) 2010 The Chromium Authors. All rights reserved.
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 //
 // TODO(ukai): code is similar with http_network_transaction.cc.  We should
 //   think about ways to share code, if possible.
 
 #include "net/socket_stream/socket_stream.h"
 
 #include <set>
@@ skipping 48 matching lines @@
       ALLOW_THIS_IN_INITIALIZER_LIST(
           read_callback_(this, &SocketStream::OnReadCompleted)),
       ALLOW_THIS_IN_INITIALIZER_LIST(
           write_callback_(this, &SocketStream::OnWriteCompleted)),
       read_buf_(NULL),
       write_buf_(NULL),
       current_write_buf_(NULL),
       write_buf_offset_(0),
       write_buf_size_(0),
       closing_(false),
+      server_closed_(false),
       metrics_(new SocketStreamMetrics(url)) {
   DCHECK(MessageLoop::current()) <<
       "The current MessageLoop must exist";
   DCHECK_EQ(MessageLoop::TYPE_IO, MessageLoop::current()->type()) <<
       "The current MessageLoop must be TYPE_IO";
   DCHECK(delegate_);
 }
 
 SocketStream::~SocketStream() {
   set_context(NULL);
@@ skipping 102 matching lines @@
   DCHECK(MessageLoop::current()) <<
       "The current MessageLoop must exist";
   DCHECK_EQ(MessageLoop::TYPE_IO, MessageLoop::current()->type()) <<
       "The current MessageLoop must be TYPE_IO";
   // If next_state_ is STATE_NONE, the socket was not opened, or already
   // closed.  So, return immediately.
   // Otherwise, it might call Finish() more than once, so breaks balance
   // of AddRef() and Release() in Connect() and Finish(), respectively.
   if (next_state_ == STATE_NONE)
     return;
-  closing_ = true;
-  // Close asynchronously, so that delegate won't be called
-  // back before returning Close().
   MessageLoop::current()->PostTask(
       FROM_HERE,
-      NewRunnableMethod(this, &SocketStream::DoLoop, OK));
+      NewRunnableMethod(this, &SocketStream::DoClose));
+}
+
+void SocketStream::DoClose() {
+  closing_ = true;
+  // If next_state_ is STATE_TCP_CONNECT, it's waiting other socket establishing
+  // connection.  If next_state_ is STATE_AUTH_REQUIRED, it's waiting for
+  // restarting.  In these states, we'll close the SocketStream now.
+  if (next_state_ == STATE_TCP_CONNECT || next_state_ == STATE_AUTH_REQUIRED) {
+    DoLoop(ERR_ABORTED);
+    return;
+  }
+  // If next_state_ is STATE_READ_WRITE, we'll run DoLoop and close
+  // the SocketStream.
+  // If it's writing now, we should defer the closing after the current
+  // writing is completed.
+  if (next_state_ == STATE_READ_WRITE && !current_write_buf_)
+    DoLoop(ERR_ABORTED);
+
+  // In other next_state_, we'll wait for callback of other APIs, such as
+  // ResolveProxy().
 }
 
 void SocketStream::RestartWithAuth(
     const string16& username, const string16& password) {
   DCHECK(MessageLoop::current()) <<
       "The current MessageLoop must exist";
   DCHECK_EQ(MessageLoop::TYPE_IO, MessageLoop::current()->type()) <<
       "The current MessageLoop must be TYPE_IO";
   DCHECK(auth_handler_.get());
   if (!socket_.get()) {
@@ skipping 116 matching lines @@
   return OK;
 }
 
 void SocketStream::OnIOCompleted(int result) {
   DoLoop(result);
 }
 
 void SocketStream::OnReadCompleted(int result) {
   if (result == 0) {
     // 0 indicates end-of-file, so socket was closed.
-    next_state_ = STATE_CLOSE;
+    // Don't close the socket if it's still writing.
+    server_closed_ = true;
   } else if (result > 0 && read_buf_) {
     result = DidReceiveData(result);
   }
   DoLoop(result);
 }
 
 void SocketStream::OnWriteCompleted(int result) {
   if (result >= 0 && write_buf_) {
     result = DidSendData(result);
   }
@@ skipping 57 matching lines @@
       case STATE_SSL_CONNECT:
         DCHECK_EQ(OK, result);
         result = DoSSLConnect();
         break;
       case STATE_SSL_CONNECT_COMPLETE:
         result = DoSSLConnectComplete(result);
         break;
       case STATE_READ_WRITE:
         result = DoReadWrite(result);
         break;
+      case STATE_AUTH_REQUIRED:
+        NOTREACHED() << "Should not run DoLoop in STATE_AUTH_REQUIRED state.";
+        Finish(result);
+        return;
       case STATE_CLOSE:
         DCHECK_LE(result, OK);
         Finish(result);
         return;
       default:
-        NOTREACHED() << "bad state";
+        NOTREACHED() << "bad state " << state;
         Finish(result);
         return;
     }
     // If the connection is not established yet and had actual errors,
     // close the connection.
     if (state != STATE_READ_WRITE && result < ERR_IO_PENDING) {
       DCHECK_EQ(next_state_, STATE_CLOSE);
       net_log_.EndEvent(NetLog::TYPE_SOCKET_STREAM_CONNECT,
                         new NetLogIntegerParameter("net_error", result));
     }
@@ skipping 407 matching lines @@
   // let's close the socket.
   // We don't care about receiving data after the socket is closed.
   if (closing_ && !write_buf_ && pending_write_bufs_.empty()) {
     socket_->Disconnect();
     next_state_ = STATE_CLOSE;
     return OK;
   }
 
   next_state_ = STATE_READ_WRITE;
 
-  if (!read_buf_) {
-    // No read pending.
-    read_buf_ = new IOBuffer(kReadBufferSize);
-    result = socket_->Read(read_buf_, kReadBufferSize, &read_callback_);
-    if (result > 0) {
-      return DidReceiveData(result);
-    } else if (result == 0) {
-      // 0 indicates end-of-file, so socket was closed.
-      next_state_ = STATE_CLOSE;
-      return ERR_CONNECTION_CLOSED;
+  // If server already closed the socket, we don't try to read.
+  if (!server_closed_) {
+    if (!read_buf_) {
+      // No read pending and server didn't close the socket.
+      read_buf_ = new IOBuffer(kReadBufferSize);
+      result = socket_->Read(read_buf_, kReadBufferSize, &read_callback_);
+      if (result > 0) {
+        return DidReceiveData(result);
+      } else if (result == 0) {
+        // 0 indicates end-of-file, so socket was closed.
+        next_state_ = STATE_CLOSE;
+        server_closed_ = true;
+        return ERR_CONNECTION_CLOSED;
+      }
+      // If read is pending, try write as well.
+      // Otherwise, return the result and do next loop (to close the
+      // connection).
+      if (result != ERR_IO_PENDING) {
+        next_state_ = STATE_CLOSE;
+        server_closed_ = true;
+        return result;
+      }
     }
-    // If read is pending, try write as well.
-    // Otherwise, return the result and do next loop (to close the connection).
-    if (result != ERR_IO_PENDING) {
-      next_state_ = STATE_CLOSE;
-      return result;
-    }
+    // Read is pending.
+    DCHECK(read_buf_);
   }
-  // Read is pending.
-  DCHECK(read_buf_);
 
   if (write_buf_ && !current_write_buf_) {
     // No write pending.
     current_write_buf_ = new DrainableIOBuffer(write_buf_, write_buf_size_);
     current_write_buf_->SetOffset(write_buf_offset_);
     result = socket_->Write(current_write_buf_,
                             current_write_buf_->BytesRemaining(),
                             &write_callback_);
     if (result > 0) {
       return DidSendData(result);
@@ skipping 112 matching lines @@
 
 SSLConfigService* SocketStream::ssl_config_service() const {
   return context_->ssl_config_service();
 }
 
 ProxyService* SocketStream::proxy_service() const {
   return context_->proxy_service();
 }
 
 }  // namespace net