Keep pool of pre-connected DNS sockets

net/dns/dns_transaction.cc

 // Copyright (c) 2012 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.
 
 #include "net/dns/dns_transaction.h"
 
 #include <deque>
 #include <string>
 #include <vector>
 
 #include "base/bind.h"
 #include "base/memory/ref_counted.h"
 #include "base/memory/scoped_ptr.h"
 #include "base/memory/scoped_vector.h"
 #include "base/memory/weak_ptr.h"
 #include "base/message_loop.h"
-#include "base/rand_util.h"
 #include "base/stl_util.h"
 #include "base/string_piece.h"
 #include "base/threading/non_thread_safe.h"
 #include "base/timer.h"
 #include "base/values.h"
 #include "net/base/completion_callback.h"
 #include "net/base/dns_util.h"
 #include "net/base/io_buffer.h"
 #include "net/base/ip_endpoint.h"
 #include "net/base/net_errors.h"
 #include "net/base/net_log.h"
 #include "net/dns/dns_protocol.h"
 #include "net/dns/dns_query.h"
 #include "net/dns/dns_response.h"
 #include "net/dns/dns_session.h"
-#include "net/socket/client_socket_factory.h"
 #include "net/udp/datagram_client_socket.h"
 
 namespace net {
 
 namespace {
 
 // Count labels in the fully-qualified name in DNS format.
 int CountLabels(const std::string& name) {
   size_t count = 0;
   for (size_t i = 0; i < name.size() && name[i]; i += name[i] + 1)
@@ skipping 15 matching lines @@
   return dict;
 };
 
 // ----------------------------------------------------------------------------
 
 // A single asynchronous DNS exchange over UDP, which consists of sending out a
 // DNS query, waiting for a response, and returning the response that it
 // matches. Logging is done in the socket and in the outer DnsTransaction.
 class DnsUDPAttempt {
  public:
-  DnsUDPAttempt(scoped_ptr<DatagramClientSocket> socket,
-                const IPEndPoint& server,
+  DnsUDPAttempt(scoped_ptr<DnsSession::SocketLease> socket_lease,
                 scoped_ptr<DnsQuery> query,
                 const CompletionCallback& callback)
       : next_state_(STATE_NONE),
         received_malformed_response_(false),
-        socket_(socket.Pass()),
-        server_(server),
+        socket_lease_(socket_lease.Pass()),
         query_(query.Pass()),
         callback_(callback) {
   }
 
   // Starts the attempt. Returns ERR_IO_PENDING if cannot complete synchronously
   // and calls |callback| upon completion.
   int Start() {
     DCHECK_EQ(STATE_NONE, next_state_);
-    int rv = socket_->Connect(server_);
-    DCHECK_NE(ERR_IO_PENDING, rv);
-    if (rv < 0)
-      return rv;
     next_state_ = STATE_SEND_QUERY;
     return DoLoop(OK);
   }
 
   const DnsQuery* query() const {
     return query_.get();
   }
 
-  const DatagramClientSocket* socket() const {
-    return socket_.get();
+  const BoundNetLog& socket_net_log() const {
+    return socket_lease_->socket()->NetLog();
   }
 
   // Returns the response or NULL if has not received a matching response from
   // the server.
   const DnsResponse* response() const {
     const DnsResponse* resp = response_.get();
     return (resp != NULL && resp->IsValid()) ? resp : NULL;
   }
 
   // Returns a Value representing the received response, along with a reference
   // to the NetLog source source of the UDP socket used.  The request must have
   // completed before this is called.
   Value* NetLogResponseCallback(NetLog::LogLevel /* log_level */) const {
     DCHECK(response_->IsValid());
 
     DictionaryValue* dict = new DictionaryValue();
     dict->SetInteger("rcode", response_->rcode());
     dict->SetInteger("answer_count", response_->answer_count());
-    socket_->NetLog().source().AddToEventParameters(dict);
+    socket_net_log().source().AddToEventParameters(dict);
     return dict;
   }
 
  private:
   enum State {
     STATE_SEND_QUERY,
     STATE_SEND_QUERY_COMPLETE,
     STATE_READ_RESPONSE,
     STATE_READ_RESPONSE_COMPLETE,
     STATE_NONE,
   };
 
+  DatagramClientSocket* socket() {
+    return socket_lease_->socket();
+  }
+
   int DoLoop(int result) {
     CHECK_NE(STATE_NONE, next_state_);
     int rv = result;
     do {
       State state = next_state_;
       next_state_ = STATE_NONE;
       switch (state) {
         case STATE_SEND_QUERY:
           rv = DoSendQuery();
           break;
@@ skipping 13 matching lines @@
     } while (rv != ERR_IO_PENDING && next_state_ != STATE_NONE);
     // If we received a malformed response, and are now waiting for another one,
     // indicate to the transaction that the server might be misbehaving.
     if (rv == ERR_IO_PENDING && received_malformed_response_)
       return ERR_DNS_MALFORMED_RESPONSE;
     return rv;
   }
 
   int DoSendQuery() {
     next_state_ = STATE_SEND_QUERY_COMPLETE;
-    return socket_->Write(query_->io_buffer(),
-                          query_->io_buffer()->size(),
-                          base::Bind(&DnsUDPAttempt::OnIOComplete,
-                                     base::Unretained(this)));
+    return socket()->Write(query_->io_buffer(),
+                              query_->io_buffer()->size(),
+                              base::Bind(&DnsUDPAttempt::OnIOComplete,
+                                         base::Unretained(this)));
   }
 
   int DoSendQueryComplete(int rv) {
     DCHECK_NE(ERR_IO_PENDING, rv);
     if (rv < 0)
       return rv;
 
     // Writing to UDP should not result in a partial datagram.
     if (rv != query_->io_buffer()->size())
       return ERR_MSG_TOO_BIG;
 
     next_state_ = STATE_READ_RESPONSE;
     return OK;
   }
 
   int DoReadResponse() {
     next_state_ = STATE_READ_RESPONSE_COMPLETE;
     response_.reset(new DnsResponse());
-    return socket_->Read(response_->io_buffer(),
-                         response_->io_buffer()->size(),
-                         base::Bind(&DnsUDPAttempt::OnIOComplete,
-                                    base::Unretained(this)));
+    return socket()->Read(response_->io_buffer(),
+                             response_->io_buffer()->size(),
+                             base::Bind(&DnsUDPAttempt::OnIOComplete,
+                                        base::Unretained(this)));
   }
 
   int DoReadResponseComplete(int rv) {
     DCHECK_NE(ERR_IO_PENDING, rv);
     if (rv < 0)
       return rv;
 
     DCHECK(rv);
     if (!response_->InitParse(rv, *query_)) {
       // Other implementations simply ignore mismatched responses. Since each
@@ skipping 19 matching lines @@
 
   void OnIOComplete(int rv) {
     rv = DoLoop(rv);
     if (rv != ERR_IO_PENDING)
       callback_.Run(rv);
   }
 
   State next_state_;
   bool received_malformed_response_;
 
-  scoped_ptr<DatagramClientSocket> socket_;
-  IPEndPoint server_;
+  scoped_ptr<DnsSession::SocketLease> socket_lease_;
   scoped_ptr<DnsQuery> query_;
 
   scoped_ptr<DnsResponse> response_;
 
   CompletionCallback callback_;
 
   DISALLOW_COPY_AND_ASSIGN(DnsUDPAttempt);
 };
 
 // ----------------------------------------------------------------------------
@@ skipping 139 matching lines @@
 
     net_log_.EndEventWithNetErrorCode(NetLog::TYPE_DNS_TRANSACTION, result.rv);
     callback.Run(this, result.rv, response);
   }
 
   // Makes another attempt at the current name, |qnames_.front()|, using the
   // next nameserver.
   AttemptResult MakeAttempt() {
     unsigned attempt_number = attempts_.size();
 
-#if defined(OS_WIN)
-    // Avoid the Windows firewall warning about explicit UDP binding.
-    // TODO(szym): Reuse a pool of pre-bound sockets. http://crbug.com/107413
-    DatagramSocket::BindType bind_type = DatagramSocket::DEFAULT_BIND;
-#else
-    DatagramSocket::BindType bind_type = DatagramSocket::RANDOM_BIND;
-#endif
-
-    scoped_ptr<DatagramClientSocket> socket(
-        session_->socket_factory()->CreateDatagramClientSocket(
-            bind_type,
-            base::Bind(&base::RandInt),
-            net_log_.net_log(),
-            net_log_.source()));
-
     uint16 id = session_->NextQueryId();
     scoped_ptr<DnsQuery> query;
     if (attempts_.empty()) {
       query.reset(new DnsQuery(id, qnames_.front(), qtype_));
     } else {
       query.reset(attempts_[0]->query()->CloneWithNewId(id));
     }
 
-    net_log_.AddEvent(NetLog::TYPE_DNS_TRANSACTION_ATTEMPT,
-                      socket->NetLog().source().ToEventParametersCallback());
-
     const DnsConfig& config = session_->config();
 
     unsigned server_index = first_server_index_ +
         (attempt_number % config.nameservers.size());
 
+    scoped_ptr<DnsSession::SocketLease> lease =
+        session_->AllocateSocket(server_index, net_log_.source());
+
+    bool got_socket = !!lease.get();
+
     DnsUDPAttempt* attempt = new DnsUDPAttempt(
-        socket.Pass(),
-        config.nameservers[server_index],
+        lease.Pass(),
         query.Pass(),
         base::Bind(&DnsTransactionImpl::OnAttemptComplete,
                    base::Unretained(this),
                    attempt_number));
 
     attempts_.push_back(attempt);
 
+    if (!got_socket)
+      return AttemptResult(ERR_CONNECTION_REFUSED, NULL);
+
+    net_log_.AddEvent(
+        NetLog::TYPE_DNS_TRANSACTION_ATTEMPT,
+        attempt->socket_net_log().source().ToEventParametersCallback());
+
     int rv = attempt->Start();
     if (rv == ERR_IO_PENDING) {
       timer_.Stop();
       base::TimeDelta timeout = session_->NextTimeout(attempt_number);
       timer_.Start(FROM_HERE, timeout, this, &DnsTransactionImpl::OnTimeout);
     }
     return AttemptResult(rv, attempt);
   }
 
   // Begins query for the current name. Makes the first attempt.
@@ skipping 49 matching lines @@
           net_log_.EndEventWithNetErrorCode(
               NetLog::TYPE_DNS_TRANSACTION_QUERY, result.rv);
           // Try next suffix.
           qnames_.pop_front();
           if (qnames_.empty()) {
             return AttemptResult(ERR_NAME_NOT_RESOLVED, NULL);
           } else {
             result = StartQuery();
           }
           break;
+        case ERR_CONNECTION_REFUSED:
         case ERR_DNS_TIMED_OUT:
           if (MoreAttemptsAllowed()) {
             result = MakeAttempt();
           } else {
             return result;
           }
           break;
         default:
           // Server failure.
           DCHECK(result.attempt);
@@ skipping 75 matching lines @@
 }  // namespace
 
 // static
 scoped_ptr<DnsTransactionFactory> DnsTransactionFactory::CreateFactory(
     DnsSession* session) {
   return scoped_ptr<DnsTransactionFactory>(
       new DnsTransactionFactoryImpl(session));
 }
 
 }  // namespace net