1 files changed, 49 insertions, 39 deletions

diff --git a/net/socket_stream/socket_stream.cc b/net/socket_stream/socket_stream.cc
index ab4c3fd..c549fcb 100644
--- a/net/socket_stream/socket_stream.cc
+++ b/net/socket_stream/socket_stream.cc
@@ -97,6 +97,7 @@ SocketStream::SocketStream(const GURL& url, Delegate* delegate)
       proxy_mode_(kDirectConnection),
       proxy_url_(url),
       pac_request_(NULL),
+      connection_(new ClientSocketHandle),
       privacy_mode_(kPrivacyModeDisabled),
       // Unretained() is required; without it, Bind() creates a circular
       // dependency and the SocketStream object will not be freed.
@@ -206,8 +207,10 @@ bool SocketStream::SendData(const char* data, int len) {
       << "The current base::MessageLoop must be TYPE_IO";
   DCHECK_GT(len, 0);
 
-  if (!socket_.get() || !socket_->IsConnected() || next_state_ == STATE_NONE)
+  if (!connection_->socket() ||
+      !connection_->socket()->IsConnected() || next_state_ == STATE_NONE) {
     return false;
+  }
 
   int total_buffered_bytes = len;
   if (current_write_buf_.get()) {
@@ -265,7 +268,7 @@ void SocketStream::RestartWithAuth(const AuthCredentials& credentials) {
   DCHECK_EQ(base::MessageLoop::TYPE_IO, base::MessageLoop::current()->type())
       << "The current base::MessageLoop must be TYPE_IO";
   DCHECK(proxy_auth_controller_.get());
-  if (!socket_.get()) {
+  if (!connection_->socket()) {
     DVLOG(1) << "Socket is closed before restarting with auth.";
     return;
   }
@@ -370,7 +373,7 @@ void SocketStream::Finish(int result) {
 }
 
 int SocketStream::DidEstablishConnection() {
-  if (!socket_.get() || !socket_->IsConnected()) {
+  if (!connection_->socket() || !connection_->socket()->IsConnected()) {
     next_state_ = STATE_CLOSE;
     return ERR_CONNECTION_FAILED;
   }
@@ -731,11 +734,12 @@ int SocketStream::DoTcpConnect(int result) {
   }
   next_state_ = STATE_TCP_CONNECT_COMPLETE;
   DCHECK(factory_);
-  socket_.reset(factory_->CreateTransportClientSocket(addresses_,
-                                                      net_log_.net_log(),
-                                                      net_log_.source()));
+  connection_->SetSocket(
+      factory_->CreateTransportClientSocket(addresses_,
+                                            net_log_.net_log(),
+                                            net_log_.source()));
   metrics_->OnStartConnection();
-  return socket_->Connect(io_callback_);
+  return connection_->socket()->Connect(io_callback_);
 }
 
 int SocketStream::DoTcpConnectComplete(int result) {
@@ -820,7 +824,8 @@ int SocketStream::DoWriteTunnelHeaders() {
   int buf_len = static_cast<int>(tunnel_request_headers_->headers_.size() -
                                  tunnel_request_headers_bytes_sent_);
   DCHECK_GT(buf_len, 0);
-  return socket_->Write(tunnel_request_headers_.get(), buf_len, io_callback_);
+  return connection_->socket()->Write(
+      tunnel_request_headers_.get(), buf_len, io_callback_);
 }
 
 int SocketStream::DoWriteTunnelHeadersComplete(int result) {
@@ -863,7 +868,8 @@ int SocketStream::DoReadTunnelHeaders() {
   tunnel_response_headers_->SetDataOffset(tunnel_response_headers_len_);
   CHECK(tunnel_response_headers_->data());
 
-  return socket_->Read(tunnel_response_headers_.get(), buf_len, io_callback_);
+  return connection_->socket()->Read(
+      tunnel_response_headers_.get(), buf_len, io_callback_);
 }
 
 int SocketStream::DoReadTunnelHeadersComplete(int result) {
@@ -957,16 +963,17 @@ int SocketStream::DoSOCKSConnect() {
   HostResolver::RequestInfo req_info(HostPortPair::FromURL(url_));
 
   DCHECK(!proxy_info_.is_empty());
-  scoped_ptr<ClientSocketHandle> connection(new ClientSocketHandle);
-  connection->set_socket(socket_.release());
+  scoped_ptr<StreamSocket> s;
   if (proxy_info_.proxy_server().scheme() == ProxyServer::SCHEME_SOCKS5) {
-    socket_.reset(new SOCKS5ClientSocket(connection.release(), req_info));
+    s.reset(new SOCKS5ClientSocket(connection_.Pass(), req_info));
   } else {
-    socket_.reset(new SOCKSClientSocket(
-        connection.release(), req_info, context_->host_resolver()));
+    s.reset(new SOCKSClientSocket(
+        connection_.Pass(), req_info, context_->host_resolver()));
   }
+  connection_.reset(new ClientSocketHandle);
+  connection_->SetSocket(s.Pass());
   metrics_->OnCountConnectionType(SocketStreamMetrics::SOCKS_CONNECTION);
-  return socket_->Connect(io_callback_);
+  return connection_->socket()->Connect(io_callback_);
 }
 
 int SocketStream::DoSOCKSConnectComplete(int result) {
@@ -989,16 +996,16 @@ int SocketStream::DoSecureProxyConnect() {
   ssl_context.cert_verifier = context_->cert_verifier();
   ssl_context.transport_security_state = context_->transport_security_state();
   ssl_context.server_bound_cert_service = context_->server_bound_cert_service();
-  scoped_ptr<ClientSocketHandle> connection(new ClientSocketHandle);
-  connection->set_socket(socket_.release());
-  socket_.reset(factory_->CreateSSLClientSocket(
-      connection.release(),
+  scoped_ptr<StreamSocket> socket(factory_->CreateSSLClientSocket(
+      connection_.Pass(),
       proxy_info_.proxy_server().host_port_pair(),
       proxy_ssl_config_,
       ssl_context));
+  connection_.reset(new ClientSocketHandle);
+  connection_->SetSocket(socket.Pass());
   next_state_ = STATE_SECURE_PROXY_CONNECT_COMPLETE;
   metrics_->OnCountConnectionType(SocketStreamMetrics::SECURE_PROXY_CONNECTION);
-  return socket_->Connect(io_callback_);
+  return connection_->socket()->Connect(io_callback_);
 }
 
 int SocketStream::DoSecureProxyConnectComplete(int result) {
@@ -1030,7 +1037,7 @@ int SocketStream::DoSecureProxyHandleCertError(int result) {
 int SocketStream::DoSecureProxyHandleCertErrorComplete(int result) {
   DCHECK_EQ(STATE_NONE, next_state_);
   if (result == OK) {
-    if (!socket_->IsConnectedAndIdle())
+    if (!connection_->socket()->IsConnectedAndIdle())
       return AllowCertErrorForReconnection(&proxy_ssl_config_);
     next_state_ = STATE_GENERATE_PROXY_AUTH_TOKEN;
   } else {
@@ -1045,15 +1052,16 @@ int SocketStream::DoSSLConnect() {
   ssl_context.cert_verifier = context_->cert_verifier();
   ssl_context.transport_security_state = context_->transport_security_state();
   ssl_context.server_bound_cert_service = context_->server_bound_cert_service();
-  scoped_ptr<ClientSocketHandle> connection(new ClientSocketHandle);
-  connection->set_socket(socket_.release());
-  socket_.reset(factory_->CreateSSLClientSocket(connection.release(),
-                                                HostPortPair::FromURL(url_),
-                                                server_ssl_config_,
-                                                ssl_context));
+  scoped_ptr<StreamSocket> socket(
+      factory_->CreateSSLClientSocket(connection_.Pass(),
+                                      HostPortPair::FromURL(url_),
+                                      server_ssl_config_,
+                                      ssl_context));
+  connection_.reset(new ClientSocketHandle);
+  connection_->SetSocket(socket.Pass());
   next_state_ = STATE_SSL_CONNECT_COMPLETE;
   metrics_->OnCountConnectionType(SocketStreamMetrics::SSL_CONNECTION);
-  return socket_->Connect(io_callback_);
+  return connection_->socket()->Connect(io_callback_);
 }
 
 int SocketStream::DoSSLConnectComplete(int result) {
@@ -1089,7 +1097,7 @@ int SocketStream::DoSSLHandleCertErrorComplete(int result) {
   // we should take care of TLS NPN extension here.
 
   if (result == OK) {
-    if (!socket_->IsConnectedAndIdle())
+    if (!connection_->socket()->IsConnectedAndIdle())
       return AllowCertErrorForReconnection(&server_ssl_config_);
     result = DidEstablishConnection();
   } else {
@@ -1103,7 +1111,7 @@ int SocketStream::DoReadWrite(int result) {
     next_state_ = STATE_CLOSE;
     return result;
   }
-  if (!socket_.get() || !socket_->IsConnected()) {
+  if (!connection_->socket() || !connection_->socket()->IsConnected()) {
     next_state_ = STATE_CLOSE;
     return ERR_CONNECTION_CLOSED;
   }
@@ -1112,7 +1120,7 @@ int SocketStream::DoReadWrite(int result) {
   // let's close the socket.
   // We don't care about receiving data after the socket is closed.
   if (closing_ && !current_write_buf_.get() && pending_write_bufs_.empty()) {
-    socket_->Disconnect();
+    connection_->socket()->Disconnect();
     next_state_ = STATE_CLOSE;
     return OK;
   }
@@ -1124,7 +1132,7 @@ int SocketStream::DoReadWrite(int result) {
     if (!read_buf_.get()) {
       // No read pending and server didn't close the socket.
       read_buf_ = new IOBuffer(kReadBufferSize);
-      result = socket_->Read(
+      result = connection_->socket()->Read(
           read_buf_.get(),
           kReadBufferSize,
           base::Bind(&SocketStream::OnReadCompleted, base::Unretained(this)));
@@ -1163,7 +1171,7 @@ int SocketStream::DoReadWrite(int result) {
     pending_write_bufs_.pop_front();
   }
 
-  result = socket_->Write(
+  result = connection_->socket()->Write(
       current_write_buf_.get(),
       current_write_buf_->BytesRemaining(),
       base::Bind(&SocketStream::OnWriteCompleted, base::Unretained(this)));
@@ -1195,10 +1203,10 @@ int SocketStream::HandleCertificateRequest(int result, SSLConfig* ssl_config) {
     return result;
   }
 
-  DCHECK(socket_.get());
+  DCHECK(connection_->socket());
   scoped_refptr<SSLCertRequestInfo> cert_request_info = new SSLCertRequestInfo;
   SSLClientSocket* ssl_socket =
-      static_cast<SSLClientSocket*>(socket_.get());
+      static_cast<SSLClientSocket*>(connection_->socket());
   ssl_socket->GetSSLCertRequestInfo(cert_request_info.get());
 
   HttpTransactionFactory* factory = context_->http_transaction_factory();
@@ -1244,7 +1252,8 @@ int SocketStream::AllowCertErrorForReconnection(SSLConfig* ssl_config) {
   // allowed bad certificates in |ssl_config|.
   // See also net/http/http_network_transaction.cc HandleCertificateError() and
   // RestartIgnoringLastError().
-  SSLClientSocket* ssl_socket = static_cast<SSLClientSocket*>(socket_.get());
+  SSLClientSocket* ssl_socket =
+      static_cast<SSLClientSocket*>(connection_->socket());
   SSLInfo ssl_info;
   ssl_socket->GetSSLInfo(&ssl_info);
   if (ssl_info.cert.get() == NULL ||
@@ -1266,8 +1275,8 @@ int SocketStream::AllowCertErrorForReconnection(SSLConfig* ssl_config) {
   bad_cert.cert_status = ssl_info.cert_status;
   ssl_config->allowed_bad_certs.push_back(bad_cert);
   // Restart connection ignoring the bad certificate.
-  socket_->Disconnect();
-  socket_.reset();
+  connection_->socket()->Disconnect();
+  connection_->SetSocket(scoped_ptr<StreamSocket>());
   next_state_ = STATE_TCP_CONNECT;
   return OK;
 }
@@ -1293,7 +1302,8 @@ void SocketStream::DoRestartWithAuth() {
 
 int SocketStream::HandleCertificateError(int result) {
   DCHECK(IsCertificateError(result));
-  SSLClientSocket* ssl_socket = static_cast<SSLClientSocket*>(socket_.get());
+  SSLClientSocket* ssl_socket =
+      static_cast<SSLClientSocket*>(connection_->socket());
   DCHECK(ssl_socket);
 
   if (!context_)