DnsClient refactoring + features (timeout, suffix search, server rotation).

net/dns/dns_client.cc

Index: net/dns/dns_client.cc
diff --git a/net/dns/dns_client.cc b/net/dns/dns_client.cc
index de60cc33fee98c30b3c7844035b1bbb3443bacf0..fa52745c97b0b6535bdf3c7b1d0fe46d02cc16e0 100644
--- a/net/dns/dns_client.cc
+++ b/net/dns/dns_client.cc
@@ -4,87 +4,522 @@
 
 #include "net/dns/dns_client.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/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 "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/dns/dns_transaction.h"
 #include "net/socket/client_socket_factory.h"
+#include "net/udp/datagram_client_socket.h"
 
 namespace net {
 
-DnsClient::Request::Request(const base::StringPiece& qname,
-                            uint16 qtype,
-                            const RequestCallback& callback)
-    : qname_(qname.data(), qname.size()),
-      qtype_(qtype),
-      callback_(callback) {
+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)
+    ++count;
+  return count;
+}
+
+bool IsIPLiteral(const std::string& hostname) {
+  IPAddressNumber ip;
+  return ParseIPLiteralToNumber(hostname, &ip);
 }
 
-DnsClient::Request::~Request() {}
+class StartParameters : public NetLog::EventParameters {
+ public:
+  StartParameters(const std::string& hostname,
+                  uint16 qtype,
+                  const NetLog::Source& source)
+      : hostname_(hostname), qtype_(qtype), source_(source) {}
+
+  virtual Value* ToValue() const OVERRIDE {
+    DictionaryValue* dict = new DictionaryValue();
+    dict->SetString("hostname", hostname_);
+    dict->SetInteger("query_type", qtype_);
+    dict->Set("source_dependency", source_.ToValue());
+    return dict;
+  }
+
+ private:
+  const std::string hostname_;
+  const uint16 qtype_;
+  const NetLog::Source source_;
+};
 
-// Implementation of DnsClient that uses DnsTransaction to serve requests.
-class DnsClientImpl : public DnsClient {
+class ResponseParameters : public NetLog::EventParameters {
  public:
-  class RequestImpl : public Request {
-   public:
-    RequestImpl(const base::StringPiece& qname,
-                uint16 qtype,
-                const RequestCallback& callback,
-                DnsSession* session,
-                const BoundNetLog& net_log)
-        : Request(qname, qtype, callback),
-          session_(session),
-          net_log_(net_log) {
+  ResponseParameters(int rcode, int answer_count, const NetLog::Source& source)
+      : rcode_(rcode), answer_count_(answer_count), source_(source) {}
+
+  virtual Value* ToValue() const OVERRIDE {
+    DictionaryValue* dict = new DictionaryValue();
+    dict->SetInteger("rcode", rcode_);
+    dict->SetInteger("answer_count", answer_count_);
+    dict->Set("socket_source", source_.ToValue());
+    return dict;
+  }
+
+ private:
+  const int rcode_;
+  const int answer_count_;
+  const NetLog::Source source_;
+};
+
+// ----------------------------------------------------------------------------
+
+// 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,
+                scoped_ptr<DnsQuery> query,
+                const CompletionCallback& callback)
+      : next_state_(STATE_NONE),
+        socket_(socket.Pass()),
+        server_(server),
+        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_);
+    next_state_ = STATE_CONNECT;
+    return DoLoop(OK);
+  }
+
+  const DnsQuery* query() const {
+    return query_.get();
+  }
+
+  const DatagramClientSocket* socket() const {
+    return socket_.get();
+  }
+
+  // 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->is_valid()) ? resp : NULL;
+  }
+
+ private:
+  enum State {
+    STATE_CONNECT,
+    STATE_SEND_QUERY,
+    STATE_SEND_QUERY_COMPLETE,
+    STATE_READ_RESPONSE,
+    STATE_READ_RESPONSE_COMPLETE,
+    STATE_NONE,
+  };
+
+  int DoLoop(int result) {
+    DCHECK_NE(STATE_NONE, next_state_);
+    int rv = result;
+    do {
+      State state = next_state_;
+      next_state_ = STATE_NONE;
+      switch (state) {
+        case STATE_CONNECT:
+          rv = DoConnect();
+          break;
+        case STATE_SEND_QUERY:
+          rv = DoSendQuery();
+          break;
+        case STATE_SEND_QUERY_COMPLETE:
+          rv = DoSendQueryComplete(rv);
+          break;
+        case STATE_READ_RESPONSE:
+          rv = DoReadResponse();
+          break;
+        case STATE_READ_RESPONSE_COMPLETE:
+          rv = DoReadResponseComplete(rv);
+          break;
+        default:
+          NOTREACHED();
+          break;
+      }
+    } while (rv != ERR_IO_PENDING && next_state_ != STATE_NONE);
+
+    return rv;
+  }
+
+  int DoConnect() {
+    next_state_ = STATE_SEND_QUERY;
+    return socket_->Connect(server_);
+  }
+
+  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)));
+  }
+
+  int DoSendQueryComplete(int 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)));
+  }
+
+  int DoReadResponseComplete(int rv) {
+    DCHECK_NE(ERR_IO_PENDING, rv);
+    if (rv < 0)
+      return rv;
+
+    DCHECK(rv);
+    if (!response_->InitParse(rv, *query_))
+      return ERR_DNS_MALFORMED_RESPONSE;
+    if (response_->flags() & dns_protocol::kFlagTC)
+      return ERR_DNS_SERVER_REQUIRES_TCP;
+    if (response_->rcode() != dns_protocol::kRcodeNOERROR &&
+        response_->rcode() != dns_protocol::kRcodeNXDOMAIN) {
+      return ERR_DNS_SERVER_FAILED;
     }
+    if (response_->answer_count() == 0)
+      return ERR_NAME_NOT_RESOLVED;
+
+    return OK;
+  }
+
+  void OnIOComplete(int rv) {
+    rv = DoLoop(rv);
+    if (rv != ERR_IO_PENDING)
+      callback_.Run(rv);
+  }
+
+  State next_state_;
+
+  scoped_ptr<DatagramClientSocket> socket_;
+  IPEndPoint server_;
+  scoped_ptr<DnsQuery> query_;
+
+  scoped_ptr<DnsResponse> response_;
+
+  CompletionCallback callback_;
+
+  DISALLOW_COPY_AND_ASSIGN(DnsUDPAttempt);
+};
+
+// ----------------------------------------------------------------------------
+
+// Implements DnsTransaction. Configuration is supplied by DnsSession.
+// The suffix list is built according to the DnsConfig from the session.
+// The timeout for each DnsUDPAttempt is given by DnsSession::NextTimeout.
+// The first server to attempt on each query is given by
+// DnsSession::NextFirstServerIndex, and the order is round-robin afterwards.
+// Each server is attempted DnsConfig::attempts times.
+class DnsTransactionImpl : public DnsTransaction, public base::NonThreadSafe {
+ public:
+  DnsTransactionImpl(DnsSession* session,
+                     const std::string& hostname,
+                     uint16 qtype,
+                     const DnsTransactionFactory::CallbackType& callback,
+                     const BoundNetLog& source_net_log)
+    : session_(session),
+      hostname_(hostname),
+      qtype_(qtype),
+      callback_(callback),
+      net_log_(BoundNetLog::Make(session->net_log(),
+                                 NetLog::SOURCE_DNS_TRANSACTION)),
+      successful_attempt_(NULL),
+      first_server_index_(0) {
+    DCHECK(session_);
+    DCHECK(!hostname_.empty());
+    DCHECK(!callback_.is_null());
+
+    DCHECK(!IsIPLiteral(hostname_));
 
-    virtual int Start() OVERRIDE {
-      transaction_.reset(new DnsTransaction(
-          session_,
-          qname(),
-          qtype(),
-          base::Bind(&RequestImpl::OnComplete, base::Unretained(this)),
-          net_log_));
-      return transaction_->Start();
+    net_log_.BeginEvent(NetLog::TYPE_DNS_TRANSACTION, make_scoped_refptr(
+        new StartParameters(hostname_, qtype_, source_net_log.source())));
+  }
+
+  virtual ~DnsTransactionImpl() {
+    STLDeleteElements(&attempts_);
+    if (!callback_.is_null()) {
+      net_log_.AddEvent(NetLog::TYPE_CANCELLED, NULL);
+      net_log_.EndEventWithNetErrorCode(NetLog::TYPE_DNS_TRANSACTION,
+                                        ERR_ABORTED);
     }
+  }
+
+  virtual const std::string& GetHostname() const OVERRIDE {
+    DCHECK(CalledOnValidThread());
+    return hostname_;
+  }
+
+  virtual uint16 GetType() const OVERRIDE {
+    DCHECK(CalledOnValidThread());
+    return qtype_;
+  }
 
-    void OnComplete(DnsTransaction* transaction, int rv) {
-      DCHECK_EQ(transaction_.get(), transaction);
-      // TODO(szym):
-      // - handle retransmissions here instead of DnsTransaction
-      // - handle rcode and flags here instead of DnsTransaction
-      // - update RTT in DnsSession
-      // - perform suffix search
-      // - handle DNS over TCP
-      DoCallback(rv, (rv == OK) ? transaction->response() : NULL);
+  virtual int Start() OVERRIDE {
+    int rv = PrepareSearch();
+    if (rv == OK)
+      rv = StartQuery();
+    DCHECK_NE(OK, rv);
+    return rv;
+  }
+
+ private:
+  // Prepares |qnames_| according to the DnsConfig.
+  int PrepareSearch() {
+    const DnsConfig& config = session_->config();
+
+    std::string labelled_hostname;
+    if (!DNSDomainFromDot(hostname_, &labelled_hostname))
+      return ERR_INVALID_ARGUMENT;
+
+    if (hostname_[hostname_.size() - 1] == '.') {
+      // It's a fully-qualified name, no suffix search.
+      qnames_.push_back(labelled_hostname);
+      return OK;
     }
 
-   private:
-    scoped_refptr<DnsSession> session_;
-    BoundNetLog net_log_;
-    scoped_ptr<DnsTransaction> transaction_;
-  };
+    // Set true when |labelled_hostname| is put on the list.
+    bool had_hostname = false;
+
+    int ndots = CountLabels(labelled_hostname) - 1;
+    if (ndots >= config.ndots) {
+      qnames_.push_back(labelled_hostname);
+      had_hostname = true;
+    }
+
+    std::string qname;
+    for (size_t i = 0; i < config.search.size(); ++i) {
+      // Ignore invalid (too long) combinations.
+      if (!DNSDomainFromDot(hostname_ + "." + config.search[i], &qname))
+        continue;
+      if (qname.size() == labelled_hostname.size()) {
+        if (had_hostname)
+          continue;
+        had_hostname = true;
+      }
+      qnames_.push_back(qname);
+    }
+
+    if (!had_hostname)
+      qnames_.push_back(labelled_hostname);
+
+    return OK;
+  }
+
+  void DoCallback(int rv) {
+    if (callback_.is_null())
+      return;
+    DCHECK_NE(ERR_IO_PENDING, rv);
+    DCHECK(rv != OK || successful_attempt_ != NULL);
+
+    DnsTransactionFactory::CallbackType callback = callback_;
+    callback_.Reset();
+    net_log_.EndEventWithNetErrorCode(NetLog::TYPE_DNS_TRANSACTION, rv);
+    callback.Run(this,
+                 rv,
+                 successful_attempt_ ? successful_attempt_->response() : NULL);
+  }
+
+  // Makes another attempt at the current name, |qnames_.front()|, using the
+  // next nameserver.
+  int MakeAttempt() {
+    int attempt_number = attempts_.size();
+
+    scoped_ptr<DatagramClientSocket> socket(
+        session_->socket_factory()->CreateDatagramClientSocket(
+            DatagramSocket::RANDOM_BIND,
+            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, make_scoped_refptr(
+        new NetLogSourceParameter("socket_source", socket->NetLog().source())));
+
+    const DnsConfig& config = session_->config();
+
+    int server_index = first_server_index_ +
+        (attempt_number % config.nameservers.size());
 
-  explicit DnsClientImpl(DnsSession* session) {
+    DnsUDPAttempt* attempt = new DnsUDPAttempt(
+        socket.Pass(),
+        config.nameservers[server_index],
+        query.Pass(),
+        base::Bind(&DnsTransactionImpl::OnAttemptComplete,
+                   base::Unretained(this),
+                   attempt_number));
+
+    base::TimeDelta timeout = session_->NextTimeout(attempt_number);
+    timer_.Start(FROM_HERE, timeout, this, &DnsTransactionImpl::OnTimeout);
+    attempts_.push_back(attempt);
+    return attempt->Start();
+  }
+
+  // Begins query for the current name. Makes the first attempt.
+  int StartQuery() {
+    std::string dotted_qname = DNSDomainToString(qnames_.front());
+    net_log_.BeginEvent(
+        NetLog::TYPE_DNS_TRANSACTION_QUERY,
+        make_scoped_refptr(new NetLogStringParameter("qname", dotted_qname)));
+
+    first_server_index_ = session_->NextFirstServerIndex();
+
+    STLDeleteElements(&attempts_);
+    return MakeAttempt();
+  }
+
+  void OnAttemptComplete(int attempt_number, int rv) {
+    timer_.Stop();
+
+    const DnsUDPAttempt* attempt = attempts_[attempt_number];
+
+    if (attempt->response()) {
+      net_log_.AddEvent(
+          NetLog::TYPE_DNS_TRANSACTION_RESPONSE,
+          make_scoped_refptr(
+              new ResponseParameters(attempt->response()->rcode(),
+                                     attempt->response()->answer_count(),
+                                     attempt->socket()->NetLog().source())));
+    }
+
+    net_log_.EndEventWithNetErrorCode(NetLog::TYPE_DNS_TRANSACTION_QUERY, rv);
+
+    switch (rv) {
+      case ERR_NAME_NOT_RESOLVED:
+        // Try next suffix.
+        qnames_.pop_front();
+        if (qnames_.empty())
+          rv = ERR_NAME_NOT_RESOLVED;
+        else
+          rv = StartQuery();
+        break;
+      case OK:
+        successful_attempt_ = attempt;
+        break;
+      default:
+        // TODO(szym): Some nameservers could fail and we should just ignore
+        // them.
+        break;
+    }
+    if (rv != ERR_IO_PENDING)
+      DoCallback(rv);
+  }
+
+  void OnTimeout() {
+    const DnsConfig& config = session_->config();
+    if (attempts_.size() == config.attempts * config.nameservers.size()) {
+      DoCallback(ERR_DNS_TIMED_OUT);
+      return;
+    }
+    int rv = MakeAttempt();
+    if (rv != ERR_IO_PENDING)
+      DoCallback(rv);
+  }
+
+  scoped_refptr<DnsSession> session_;
+  std::string hostname_;
+  uint16 qtype_;
+  // Set to null once the transaction completes.
+  DnsTransactionFactory::CallbackType callback_;
+
+  BoundNetLog net_log_;
+
+  // Search list of fully-qualified DNS names to query next (in DNS format).
+  std::deque<std::string> qnames_;
+
+  // List of attempts for the current name.
+  std::vector<DnsUDPAttempt*> attempts_;
+  // The member of |attempts_| that succeeded first.
+  const DnsUDPAttempt* successful_attempt_;
+
+  // Index of the first server to try on each search query.
+  int first_server_index_;
+
+  base::OneShotTimer<DnsTransactionImpl> timer_;
+
+  DISALLOW_COPY_AND_ASSIGN(DnsTransactionImpl);
+};
+
+// ----------------------------------------------------------------------------
+
+// Implementation of DnsTransactionFactory that returns instances of
+// DnsTransactionImpl.
+class DnsTransactionFactoryImpl : public DnsTransactionFactory {
+ public:
+  explicit DnsTransactionFactoryImpl(DnsSession* session) {
     session_ = session;
   }
 
-  virtual Request* CreateRequest(
-      const base::StringPiece& qname,
+  virtual scoped_ptr<DnsTransaction> CreateTransaction(
+      const std::string& hostname,
       uint16 qtype,
-      const RequestCallback& callback,
+      const CallbackType& callback,
       const BoundNetLog& source_net_log) OVERRIDE {
-    return new RequestImpl(qname, qtype, callback, session_, source_net_log);
+    return scoped_ptr<DnsTransaction>(new DnsTransactionImpl(session_,
+                                                             hostname,
+                                                             qtype,
+                                                             callback,
+                                                             source_net_log));
   }
 
  private:
   scoped_refptr<DnsSession> session_;
 };
 
+}  // namespace
+
 // static
-DnsClient* DnsClient::CreateClient(DnsSession* session) {
-  return new DnsClientImpl(session);
+scoped_ptr<DnsTransactionFactory> DnsTransactionFactory::CreateFactory(
+    DnsSession* session) {
+  return scoped_ptr<DnsTransactionFactory>(
+      new DnsTransactionFactoryImpl(session));
 }
 
 }  // namespace net