[net/dns] Handle TC bit on DNS response in DnsTransaction.

net/dns/dns_transaction.cc

Index: net/dns/dns_transaction.cc
diff --git a/net/dns/dns_transaction.cc b/net/dns/dns_transaction.cc
index a94bba8490e59bc2ebcaa5ffe9be17e955f39def..545cce22a8b7ab610151cc820e90305ae0c6e715 100644
--- a/net/dns/dns_transaction.cc
+++ b/net/dns/dns_transaction.cc
@@ -21,6 +21,7 @@
 #include "base/threading/non_thread_safe.h"
 #include "base/timer.h"
 #include "base/values.h"
+#include "net/base/big_endian.h"
 #include "net/base/completion_callback.h"
 #include "net/base/dns_util.h"
 #include "net/base/io_buffer.h"
@@ -31,6 +32,7 @@
 #include "net/dns/dns_query.h"
 #include "net/dns/dns_response.h"
 #include "net/dns/dns_session.h"
+#include "net/socket/stream_socket.h"
 #include "net/udp/datagram_client_socket.h"
 
 namespace net {
@@ -66,56 +68,77 @@ Value* NetLogStartCallback(const std::string* hostname,
 
 // ----------------------------------------------------------------------------
 
-// A single asynchronous DNS exchange over UDP, which consists of sending out a
+// A single asynchronous DNS exchange, 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 {
+class DnsAttempt {
+ public:
+  virtual ~DnsAttempt() {}
+  // Starts the attempt. Returns ERR_IO_PENDING if cannot complete synchronously
+  // and calls |callback| upon completion.
+  virtual int Start(const CompletionCallback& callback) = 0;
+
+  // Returns the query of this attempt.
+  virtual const DnsQuery* GetQuery() const = 0;
+
+  // Returns the response or NULL if has not received a matching response from
+  // the server.
+  virtual const DnsResponse* GetResponse() const = 0;
+
+  // Returns the net log bound to the source of the socket.
+  virtual const BoundNetLog& GetSocketNetLog() const = 0;
+
+  // Returns the index of the destination server within DnsConfig::nameservers.
+  virtual unsigned GetServerIndex() const = 0;
+
+  // 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(GetResponse()->IsValid());
+
+    DictionaryValue* dict = new DictionaryValue();
+    dict->SetInteger("rcode", GetResponse()->rcode());
+    dict->SetInteger("answer_count", GetResponse()->answer_count());
+    GetSocketNetLog().source().AddToEventParameters(dict);
+    return dict;
+  }
+};
+
+class DnsUDPAttempt : public DnsAttempt {
  public:
   DnsUDPAttempt(scoped_ptr<DnsSession::SocketLease> socket_lease,
-                scoped_ptr<DnsQuery> query,
-                const CompletionCallback& callback)
+                scoped_ptr<DnsQuery> query)
       : next_state_(STATE_NONE),
         received_malformed_response_(false),
         socket_lease_(socket_lease.Pass()),
-        query_(query.Pass()),
-        callback_(callback) {
+        query_(query.Pass()) {
   }
 
-  // Starts the attempt. Returns ERR_IO_PENDING if cannot complete synchronously
-  // and calls |callback| upon completion.
-  int Start() {
+  // DnsAttempt:
+  virtual int Start(const CompletionCallback& callback) OVERRIDE {
     DCHECK_EQ(STATE_NONE, next_state_);
+    callback_ = callback;
     start_time_ = base::TimeTicks::Now();
     next_state_ = STATE_SEND_QUERY;
     return DoLoop(OK);
   }
 
-  const DnsQuery* query() const {
+  virtual const DnsQuery* GetQuery() const OVERRIDE {
     return query_.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 {
+  virtual const DnsResponse* GetResponse() const OVERRIDE {
     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());
+  virtual const BoundNetLog& GetSocketNetLog() const OVERRIDE {
+    return socket_lease_->socket()->NetLog();
+  }
 
-    DictionaryValue* dict = new DictionaryValue();
-    dict->SetInteger("rcode", response_->rcode());
-    dict->SetInteger("answer_count", response_->answer_count());
-    socket_net_log().source().AddToEventParameters(dict);
-    return dict;
+  virtual unsigned GetServerIndex() const OVERRIDE {
+    return socket_lease_->server_index();
   }
 
  private:
@@ -160,6 +183,7 @@ class DnsUDPAttempt {
     if (rv == ERR_IO_PENDING && received_malformed_response_)
       return ERR_DNS_MALFORMED_RESPONSE;
     if (rv == OK) {
+      DCHECK_EQ(STATE_NONE, next_state_);
       DNS_HISTOGRAM("AsyncDNS.UDPAttemptSuccess",
                     base::TimeTicks::Now() - start_time_);
     } else if (rv != ERR_IO_PENDING) {
@@ -223,7 +247,6 @@ class DnsUDPAttempt {
     if (response_->rcode() != dns_protocol::kRcodeNOERROR)
       return ERR_DNS_SERVER_FAILED;
 
-    CHECK(response());
     return OK;
   }
 
@@ -247,6 +270,216 @@ class DnsUDPAttempt {
   DISALLOW_COPY_AND_ASSIGN(DnsUDPAttempt);
 };
 
+class DnsTCPAttempt : public DnsAttempt {
+ public:
+  DnsTCPAttempt(scoped_ptr<StreamSocket> socket,
+                scoped_ptr<DnsQuery> query)
+      : next_state_(STATE_NONE),
+        socket_(socket.Pass()),
+        query_(query.Pass()),
+        length_buffer_(new IOBufferWithSize(sizeof(uint16))),
+        response_length_(0) {
+  }
+
+  // DnsAttempt:
+  virtual int Start(const CompletionCallback& callback) OVERRIDE {
+    DCHECK_EQ(STATE_NONE, next_state_);
+    callback_ = callback;
+    start_time_ = base::TimeTicks::Now();
+    next_state_ = STATE_CONNECT_COMPLETE;
+    int rv = socket_->Connect(base::Bind(&DnsTCPAttempt::OnIOComplete,
+                                         base::Unretained(this)));
+    if (rv == ERR_IO_PENDING)
+      return rv;
+    return DoLoop(rv);
+  }
+
+  virtual const DnsQuery* GetQuery() const OVERRIDE {
+    return query_.get();
+  }
+
+  virtual const DnsResponse* GetResponse() const OVERRIDE {
+    const DnsResponse* resp = response_.get();
+    return (resp != NULL && resp->IsValid()) ? resp : NULL;
+  }
+
+  virtual const BoundNetLog& GetSocketNetLog() const OVERRIDE {
+    return socket_->NetLog();
+  }
+
+  virtual unsigned GetServerIndex() const OVERRIDE {
+    NOTREACHED();
+    return 0;
+  }
+
+ private:
+  enum State {
+    STATE_CONNECT_COMPLETE,
+    STATE_SEND_LENGTH,
+    STATE_SEND_QUERY,
+    STATE_READ_LENGTH,
+    STATE_READ_RESPONSE,
+    STATE_NONE,
+  };
+
+  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_CONNECT_COMPLETE:
+          rv = DoConnectComplete(rv);
+          break;
+        case STATE_SEND_LENGTH:
+          rv = DoSendLength(rv);
+          break;
+        case STATE_SEND_QUERY:
+          rv = DoSendQuery(rv);
+          break;
+        case STATE_READ_LENGTH:
+          rv = DoReadLength(rv);
+          break;
+        case STATE_READ_RESPONSE:
+          rv = DoReadResponse(rv);
+          break;
+        default:
+          NOTREACHED();
+          break;
+      }
+    } while (rv != ERR_IO_PENDING && next_state_ != STATE_NONE);
+    if (rv == OK) {
+      DCHECK_EQ(STATE_NONE, next_state_);
+      DNS_HISTOGRAM("AsyncDNS.TCPAttemptSuccess",
+                    base::TimeTicks::Now() - start_time_);
+    } else if (rv != ERR_IO_PENDING) {
+      DNS_HISTOGRAM("AsyncDNS.TCPAttemptFail",
+                    base::TimeTicks::Now() - start_time_);
+    }
+    return rv;
+  }
+
+  int DoConnectComplete(int rv) {
+    DCHECK_NE(ERR_IO_PENDING, rv);
+    if (rv < 0)
+      return rv;
+
+    WriteBigEndian<uint16>(length_buffer_->data(), query_->io_buffer()->size());
+    buffer_ = new DrainableIOBuffer(length_buffer_, length_buffer_->size());
+    next_state_ = STATE_SEND_LENGTH;
+    return OK;
+  }
+
+  int DoSendLength(int rv) {
+    DCHECK_NE(ERR_IO_PENDING, rv);
+    if (rv < 0)
+      return rv;
+
+    buffer_->DidConsume(rv);
+    if (buffer_->BytesRemaining() > 0) {
+      next_state_ = STATE_SEND_LENGTH;
+      return socket_->Write(buffer_,
+                            buffer_->BytesRemaining(),
+                            base::Bind(&DnsTCPAttempt::OnIOComplete,
+                                       base::Unretained(this)));
+    }
+    buffer_ = new DrainableIOBuffer(query_->io_buffer(),
+                                    query_->io_buffer()->size());
+    next_state_ = STATE_SEND_QUERY;
+    return OK;
+  }
+
+  int DoSendQuery(int rv) {
+    DCHECK_NE(ERR_IO_PENDING, rv);
+    if (rv < 0)
+      return rv;
+
+    buffer_->DidConsume(rv);
+    if (buffer_->BytesRemaining() > 0) {
+      next_state_ = STATE_SEND_QUERY;
+      return socket_->Write(buffer_,
+                            buffer_->BytesRemaining(),
+                            base::Bind(&DnsTCPAttempt::OnIOComplete,
+                                       base::Unretained(this)));
+    }
+    buffer_ = new DrainableIOBuffer(length_buffer_, length_buffer_->size());
+    next_state_ = STATE_READ_LENGTH;
+    return OK;
+  }
+
+  int DoReadLength(int rv) {
+    DCHECK_NE(ERR_IO_PENDING, rv);
+    if (rv < 0)
+      return rv;
+
+    buffer_->DidConsume(rv);
+    if (buffer_->BytesRemaining() > 0) {
+      next_state_ = STATE_READ_LENGTH;
+      return socket_->Read(buffer_,
+                           buffer_->BytesRemaining(),
+                           base::Bind(&DnsTCPAttempt::OnIOComplete,
+                                      base::Unretained(this)));
+    }
+    ReadBigEndian<uint16>(length_buffer_->data(), &response_length_);
+    // Check if advertised response is too short. (Optimization only.)
+    if (response_length_ < query_->io_buffer()->size())
+      return ERR_DNS_MALFORMED_RESPONSE;
+    // Allocate more space so that DnsResponse::InitParse sanity check passes.
+    response_.reset(new DnsResponse(response_length_ + 1));
+    buffer_ = new DrainableIOBuffer(response_->io_buffer(), response_length_);
+    next_state_ = STATE_READ_RESPONSE;
+    return OK;
+  }
+
+  int DoReadResponse(int rv) {
+    DCHECK_NE(ERR_IO_PENDING, rv);
+    if (rv < 0)
+      return rv;
+
+    buffer_->DidConsume(rv);
+    if (buffer_->BytesRemaining() > 0) {
+      next_state_ = STATE_READ_RESPONSE;
+      return socket_->Read(buffer_,
+                           buffer_->BytesRemaining(),
+                           base::Bind(&DnsTCPAttempt::OnIOComplete,
+                                      base::Unretained(this)));
+    }
+    if (!response_->InitParse(buffer_->BytesConsumed(), *query_))
+      return ERR_DNS_MALFORMED_RESPONSE;
+    if (response_->flags() & dns_protocol::kFlagTC)
+      return ERR_UNEXPECTED;
+    // TODO(szym): Frankly, none of these are expected.
+    if (response_->rcode() == dns_protocol::kRcodeNXDOMAIN)
+      return ERR_NAME_NOT_RESOLVED;
+    if (response_->rcode() != dns_protocol::kRcodeNOERROR)
+      return ERR_DNS_SERVER_FAILED;
+
+    return OK;
+  }
+
+  void OnIOComplete(int rv) {
+    rv = DoLoop(rv);
+    if (rv != ERR_IO_PENDING)
+      callback_.Run(rv);
+  }
+
+  State next_state_;
+  base::TimeTicks start_time_;
+
+  scoped_ptr<StreamSocket> socket_;
+  scoped_ptr<DnsQuery> query_;
+  scoped_refptr<IOBufferWithSize> length_buffer_;
+  scoped_refptr<DrainableIOBuffer> buffer_;
+
+  uint16 response_length_;
+  scoped_ptr<DnsResponse> response_;
+
+  CompletionCallback callback_;
+
+  DISALLOW_COPY_AND_ASSIGN(DnsTCPAttempt);
+};
+
 // ----------------------------------------------------------------------------
 
 // Implements DnsTransaction. Configuration is supplied by DnsSession.
@@ -269,6 +502,7 @@ class DnsTransactionImpl : public DnsTransaction,
       qtype_(qtype),
       callback_(callback),
       net_log_(net_log),
+      had_tcp_attempt_(false),
       first_server_index_(0) {
     DCHECK(session_);
     DCHECK(!hostname_.empty());
@@ -321,11 +555,11 @@ class DnsTransactionImpl : public DnsTransaction,
  private:
   // Wrapper for the result of a DnsUDPAttempt.
   struct AttemptResult {
-    AttemptResult(int rv, const DnsUDPAttempt* attempt)
+    AttemptResult(int rv, const DnsAttempt* attempt)
         : rv(rv), attempt(attempt) {}
 
     int rv;
-    const DnsUDPAttempt* attempt;
+    const DnsAttempt* attempt;
   };
 
   // Prepares |qnames_| according to the DnsConfig.
@@ -380,7 +614,7 @@ class DnsTransactionImpl : public DnsTransaction,
     DCHECK(!callback_.is_null());
     DCHECK_NE(ERR_IO_PENDING, result.rv);
     const DnsResponse* response = result.attempt ?
-        result.attempt->response() : NULL;
+        result.attempt->GetResponse() : NULL;
     CHECK(result.rv != OK || response != NULL);
 
     timer_.Stop();
@@ -402,7 +636,7 @@ class DnsTransactionImpl : public DnsTransaction,
     if (attempts_.empty()) {
       query.reset(new DnsQuery(id, qnames_.front(), qtype_));
     } else {
-      query.reset(attempts_[0]->query()->CloneWithNewId(id));
+      query.reset(attempts_[0]->GetQuery()->CloneWithNewId(id));
     }
 
     const DnsConfig& config = session_->config();
@@ -415,12 +649,7 @@ class DnsTransactionImpl : public DnsTransaction,
 
     bool got_socket = !!lease.get();
 
-    DnsUDPAttempt* attempt = new DnsUDPAttempt(
-        lease.Pass(),
-        query.Pass(),
-        base::Bind(&DnsTransactionImpl::OnAttemptComplete,
-                   base::Unretained(this),
-                   attempt_number));
+    DnsUDPAttempt* attempt = new DnsUDPAttempt(lease.Pass(), query.Pass());
 
     attempts_.push_back(attempt);
 
@@ -429,17 +658,56 @@ class DnsTransactionImpl : public DnsTransaction,
 
     net_log_.AddEvent(
         NetLog::TYPE_DNS_TRANSACTION_ATTEMPT,
-        attempt->socket_net_log().source().ToEventParametersCallback());
+        attempt->GetSocketNetLog().source().ToEventParametersCallback());
 
-    int rv = attempt->Start();
+    int rv = attempt->Start(base::Bind(&DnsTransactionImpl::OnAttemptComplete,
+                                       base::Unretained(this),
+                                       attempt_number));
     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);
   }
 
+  AttemptResult MakeTCPAttempt(const DnsAttempt* previous_attempt) {
+    DCHECK(previous_attempt);
+    DCHECK(!had_tcp_attempt_);
+
+    scoped_ptr<StreamSocket> socket(
+        session_->CreateTCPSocket(previous_attempt->GetServerIndex(),
+                                  net_log_.source()));
+
+    // TODO(szym): Reuse the same id to help the server?
+    uint16 id = session_->NextQueryId();
+    scoped_ptr<DnsQuery> query(
+        previous_attempt->GetQuery()->CloneWithNewId(id));
+
+    // Cancel all other attempts, no point waiting on them.
+    attempts_.clear();
+
+    unsigned attempt_number = attempts_.size();
+
+    DnsTCPAttempt* attempt = new DnsTCPAttempt(socket.Pass(), query.Pass());
+
+    attempts_.push_back(attempt);
+    had_tcp_attempt_ = true;
+
+    net_log_.AddEvent(
+        NetLog::TYPE_DNS_TRANSACTION_TCP_ATTEMPT,
+        attempt->GetSocketNetLog().source().ToEventParametersCallback());
+
+    int rv = attempt->Start(base::Bind(&DnsTransactionImpl::OnAttemptComplete,
+                                       base::Unretained(this),
+                                       attempt_number));
+    if (rv == ERR_IO_PENDING) {
+      // Custom timeout for TCP attempt.
+      base::TimeDelta timeout = timer_.GetCurrentDelay() * 2;
+      timer_.Start(FROM_HERE, timeout, this, &DnsTransactionImpl::OnTimeout);
+    }
+    return AttemptResult(rv, attempt);
+  }
+
   // Begins query for the current name. Makes the first attempt.
   AttemptResult StartQuery() {
     std::string dotted_qname = DNSDomainToString(qnames_.front());
@@ -449,6 +717,7 @@ class DnsTransactionImpl : public DnsTransaction,
     first_server_index_ = session_->NextFirstServerIndex();
 
     attempts_.clear();
+    had_tcp_attempt_ = false;
     return MakeAttempt();
   }
 
@@ -456,27 +725,29 @@ class DnsTransactionImpl : public DnsTransaction,
     if (callback_.is_null())
       return;
     DCHECK_LT(attempt_number, attempts_.size());
-    const DnsUDPAttempt* attempt = attempts_[attempt_number];
+    const DnsAttempt* attempt = attempts_[attempt_number];
     AttemptResult result = ProcessAttemptResult(AttemptResult(rv, attempt));
     if (result.rv != ERR_IO_PENDING)
       DoCallback(result);
   }
 
-  void LogResponse(const DnsUDPAttempt* attempt) {
-    if (attempt && attempt->response()) {
+  void LogResponse(const DnsAttempt* attempt) {
+    if (attempt && attempt->GetResponse()) {
       net_log_.AddEvent(
           NetLog::TYPE_DNS_TRANSACTION_RESPONSE,
-          base::Bind(&DnsUDPAttempt::NetLogResponseCallback,
+          base::Bind(&DnsAttempt::NetLogResponseCallback,
                      base::Unretained(attempt)));
     }
   }
 
   bool MoreAttemptsAllowed() const {
+    if (had_tcp_attempt_)
+      return false;
     const DnsConfig& config = session_->config();
     return attempts_.size() < config.attempts * config.nameservers.size();
   }
 
-  // Resolves the result of a DnsUDPAttempt until a terminal result is reached
+  // Resolves the result of a DnsAttempt until a terminal result is reached
   // or it will complete asynchronously (ERR_IO_PENDING).
   AttemptResult ProcessAttemptResult(AttemptResult result) {
     while (result.rv != ERR_IO_PENDING) {
@@ -487,7 +758,7 @@ class DnsTransactionImpl : public DnsTransaction,
           net_log_.EndEventWithNetErrorCode(
               NetLog::TYPE_DNS_TRANSACTION_QUERY, result.rv);
           DCHECK(result.attempt);
-          DCHECK(result.attempt->response());
+          DCHECK(result.attempt->GetResponse());
           return result;
         case ERR_NAME_NOT_RESOLVED:
           net_log_.EndEventWithNetErrorCode(
@@ -508,6 +779,9 @@ class DnsTransactionImpl : public DnsTransaction,
             return result;
           }
           break;
+        case ERR_DNS_SERVER_REQUIRES_TCP:
+          result = MakeTCPAttempt(result.attempt);
+          break;
         default:
           // Server failure.
           DCHECK(result.attempt);
@@ -517,8 +791,10 @@ class DnsTransactionImpl : public DnsTransaction,
           }
           if (MoreAttemptsAllowed()) {
             result = MakeAttempt();
-          } else if (result.rv == ERR_DNS_MALFORMED_RESPONSE) {
-            // Wait until the last attempt times out.
+          } else if (result.rv == ERR_DNS_MALFORMED_RESPONSE &&
+                     !had_tcp_attempt_) {
+            // For UDP only, ignore the response and wait until the last attempt
+            // times out.
             return AttemptResult(ERR_IO_PENDING, NULL);
           } else {
             return AttemptResult(result.rv, NULL);
@@ -550,7 +826,8 @@ class DnsTransactionImpl : public DnsTransaction,
   std::deque<std::string> qnames_;
 
   // List of attempts for the current name.
-  ScopedVector<DnsUDPAttempt> attempts_;
+  ScopedVector<DnsAttempt> attempts_;
+  bool had_tcp_attempt_;
 
   // Index of the first server to try on each search query.
   int first_server_index_;