1 files changed, 33 insertions, 13 deletions

diff --git a/blimp/net/tcp_client_transport.cc b/blimp/net/tcp_client_transport.cc
index 9dbb8803..6a6cc50 100644
--- a/blimp/net/tcp_client_transport.cc
+++ b/blimp/net/tcp_client_transport.cc
@@ -9,38 +9,42 @@
 #include "base/memory/scoped_ptr.h"
 #include "base/message_loop/message_loop.h"
 #include "blimp/net/stream_socket_connection.h"
+#include "net/socket/client_socket_factory.h"
 #include "net/socket/stream_socket.h"
 #include "net/socket/tcp_client_socket.h"
 
 namespace blimp {
 
-TCPClientTransport::TCPClientTransport(const net::AddressList& addresses,
+TCPClientTransport::TCPClientTransport(const net::IPEndPoint& ip_endpoint,
                                        net::NetLog* net_log)
-    : addresses_(addresses), net_log_(net_log) {}
+    : ip_endpoint_(ip_endpoint),
+      net_log_(net_log),
+      socket_factory_(net::ClientSocketFactory::GetDefaultFactory()) {}
 
 TCPClientTransport::~TCPClientTransport() {}
 
+void TCPClientTransport::SetClientSocketFactoryForTest(
+    net::ClientSocketFactory* factory) {
+  DCHECK(factory);
+  socket_factory_ = factory;
+}
+
 void TCPClientTransport::Connect(const net::CompletionCallback& callback) {
   DCHECK(!socket_);
   DCHECK(!callback.is_null());
 
-  socket_.reset(
-      new net::TCPClientSocket(addresses_, net_log_, net::NetLog::Source()));
+  connect_callback_ = callback;
+  socket_ = socket_factory_->CreateTransportClientSocket(
+      net::AddressList(ip_endpoint_), net_log_, net::NetLog::Source());
   net::CompletionCallback completion_callback = base::Bind(
       &TCPClientTransport::OnTCPConnectComplete, base::Unretained(this));
 
   int result = socket_->Connect(completion_callback);
   if (result == net::ERR_IO_PENDING) {
-    connect_callback_ = callback;
     return;
   }
 
-  if (result != net::OK) {
-    socket_ = nullptr;
-  }
-
-  base::MessageLoop::current()->PostTask(FROM_HERE,
-                                         base::Bind(callback, result));
+  OnTCPConnectComplete(result);
 }
 
 scoped_ptr<BlimpConnection> TCPClientTransport::TakeConnection() {
@@ -49,17 +53,33 @@ scoped_ptr<BlimpConnection> TCPClientTransport::TakeConnection() {
   return make_scoped_ptr(new StreamSocketConnection(std::move(socket_)));
 }
 
-const std::string TCPClientTransport::GetName() const {
+const char* TCPClientTransport::GetName() const {
   return "TCP";
 }
 
 void TCPClientTransport::OnTCPConnectComplete(int result) {
   DCHECK_NE(net::ERR_IO_PENDING, result);
-  DCHECK(socket_);
+  OnConnectComplete(result);
+}
+
+void TCPClientTransport::OnConnectComplete(int result) {
   if (result != net::OK) {
     socket_ = nullptr;
   }
   base::ResetAndReturn(&connect_callback_).Run(result);
 }
 
+scoped_ptr<net::StreamSocket> TCPClientTransport::TakeSocket() {
+  return std::move(socket_);
+}
+
+void TCPClientTransport::SetSocket(scoped_ptr<net::StreamSocket> socket) {
+  DCHECK(socket);
+  socket_ = std::move(socket);
+}
+
+net::ClientSocketFactory* TCPClientTransport::socket_factory() const {
+  return socket_factory_;
+}
+
 }  // namespace blimp