[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..6c1edc6eca1b097126ca7fd6534d258382360015 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,10 +68,41 @@ 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.

[mmenke, 2012/12/19 16:44:09]

There's no |callback| defined here.

[szym, 2012/12/19 23:35:56]

Changed so that CompletionCallback is passed to Start.

+  virtual int Start() = 0;
+
+  virtual const DnsQuery* query() const = 0;
+
+  // Returns the response or NULL if has not received a matching response from
+  // the server.
+  virtual const DnsResponse* response() const = 0;
+
+  virtual const BoundNetLog& socket_net_log() const = 0;
+
+  virtual unsigned server_index() const = 0;

[mmenke, 2012/12/19 16:44:09]

Virtual methods should use CamelCase.

[mmenke, 2012/12/19 16:44:09]

While these methods do fairly obvious things, still think it's a good idea to
document virtual methods.

[szym, 2012/12/19 23:35:56]

Done. Done.

+
+  // 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) const {
+    DCHECK(response()->IsValid());
+
+    DictionaryValue* dict = new DictionaryValue();
+    dict->SetInteger("rcode", response()->rcode());
+    dict->SetInteger("answer_count", response()->answer_count());
+    socket_net_log().source().AddToEventParameters(dict);
+    return dict;
+  }
+};
+
+class DnsUDPAttempt : public DnsAttempt {
  public:
   DnsUDPAttempt(scoped_ptr<DnsSession::SocketLease> socket_lease,
                 scoped_ptr<DnsQuery> query,
@@ -81,41 +114,28 @@ class DnsUDPAttempt {
         callback_(callback) {
   }
 

[mmenke, 2012/12/19 16:44:09]

nit:  "DnsAttempt:" or "DnsAttempt implementation:"

[szym, 2012/12/19 23:35:56]

Done.

-  // Starts the attempt. Returns ERR_IO_PENDING if cannot complete synchronously
-  // and calls |callback| upon completion.
-  int Start() {
+  virtual int Start() OVERRIDE {
     DCHECK_EQ(STATE_NONE, next_state_);
     start_time_ = base::TimeTicks::Now();
     next_state_ = STATE_SEND_QUERY;
     return DoLoop(OK);
   }
 
-  const DnsQuery* query() const {
+  virtual const DnsQuery* query() 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* response() 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& socket_net_log() 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 server_index() const OVERRIDE {
+    return socket_lease_->server_index();
   }
 
  private:
@@ -247,6 +267,215 @@ class DnsUDPAttempt {
   DISALLOW_COPY_AND_ASSIGN(DnsUDPAttempt);
 };
 
+class DnsTCPAttempt : public DnsAttempt {
+ public:
+  DnsTCPAttempt(scoped_ptr<StreamSocket> socket,
+                scoped_ptr<DnsQuery> query,
+                const CompletionCallback& callback)
+      : next_state_(STATE_NONE),
+        socket_(socket.Pass()),
+        query_(query.Pass()),
+        length_buffer_(new IOBufferWithSize(sizeof(uint16))),
+        response_length_(0),
+        callback_(callback) {
+  }
+

[mmenke, 2012/12/19 16:44:09]

nit:  "DnsAttempt:" or "DnsAttempt implementation:"

[szym, 2012/12/19 23:35:56]

Done.

+  virtual int Start() OVERRIDE {
+    DCHECK_EQ(STATE_NONE, next_state_);
+    start_time_ = base::TimeTicks::Now();
+    // TODO(szym): allow sockets to be pre-connected.
+    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* query() const OVERRIDE {
+    return query_.get();
+  }
+
+  virtual const DnsResponse* response() const OVERRIDE {
+    const DnsResponse* resp = response_.get();
+    return (resp != NULL && resp->IsValid()) ? resp : NULL;
+  }
+
+  virtual const BoundNetLog& socket_net_log() const OVERRIDE {
+    return socket_->NetLog();
+  }
+
+  virtual unsigned server_index() 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) {

[mmenke, 2012/12/19 16:44:09]

Maybe a DCHECK_EQ(STATE_NONE, next_state_);?

[szym, 2012/12/19 23:35:56]

Done.

+      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_);
+    // Allocate more space so that DnsResponse's sanity check passes.

[mmenke, 2012/12/19 16:44:09]

Which check is that?

[szym, 2012/12/19 17:01:17]

size check for |nbytes| in DnsResponse::InitParse.

+    response_.reset(new DnsResponse(response_length_ + 1));

[mmenke, 2012/12/19 16:44:09]

That +1 could trigger overflow.  Also, may want to do a simple sanity check - if
the response length is less than the query length, can just fail.  May be worth
having a test for this, too.

May also want to use a hard-coded size limit, though suppose a 2^16 limit isn't
too many times greater than the largest reasonable size limit.

[mmenke, 2012/12/19 16:44:09]

Speaking of which, we may want to fuzz test this, too, though I think
response_length_ is the only new possible attack surface area we provide, not
present in the UDP version.

[mmenke, 2012/12/19 16:58:46]

Oops - that "test" comment was meant for the second paragraph - a length of
2^16-1.

[szym, 2012/12/19 17:01:17]

I don't think so. Because |response_length_| will be promoted and there will be
no truncation because it is DnsResponse(size_t).
Good idea. We already have a check for that in DnsResponse::InitParse, but this
would allow earlier abort.
I think 2^16 is a reasonable limit, I'd add a comment if it's not obvious that
there is a limit.

[mmenke, 2012/12/19 17:13:51]

You're right - I had been incorrectly assuming promotion would occur at the
point of assignment.
I think we can live without it.

[szym, 2012/12/19 23:35:56]

Done.

+    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;
+
+    CHECK(response());
+    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 +498,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 +551,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.
@@ -433,13 +663,49 @@ class DnsTransactionImpl : public DnsTransaction,
 
     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);
   }
 
+  AttemptResult MakeTCPAttempt(const DnsAttempt* previous_attempt) {
+    DCHECK(previous_attempt);
+    DCHECK(!had_tcp_attempt_);
+
+    scoped_ptr<StreamSocket> socket(
+        session_->AllocateTCPSocket(previous_attempt->server_index(),
+                                    net_log_.source()));
+
+    // TODO(szym): Reuse the same id to help the server?
+    uint16 id = session_->NextQueryId();
+    scoped_ptr<DnsQuery> query(previous_attempt->query()->CloneWithNewId(id));
+
+    net_log_.AddEvent(NetLog::TYPE_DNS_TRANSACTION_TCP_ATTEMPT,
+                      socket->NetLog().source().ToEventParametersCallback());
+
+    // Cancel all other attempts, no point waiting on them.

[mmenke, 2012/12/19 16:44:09]

No point, even if there are multiple servers?

[szym, 2012/12/19 17:01:17]

We don't really anticipate getting substantially different results from
different servers, so if the first response we got says it's truncated (TC),
then we shouldn't hope we will get a non-TC response from the other servers.

The only rationale would be that other servers could be quicker in responding,
but to have any confidence that one server is faster than another we either need
to keep track of their past performance or we run races.

I'd rather keep the logic simple for now and assume we make **at most one** TCP
attempt per DnsTransaction.

+    attempts_.clear();
+
+    DnsTCPAttempt* attempt = new DnsTCPAttempt(
+        socket.Pass(),
+        query.Pass(),
+        base::Bind(&DnsTransactionImpl::OnAttemptComplete,
+                   base::Unretained(this),
+                   attempts_.size()));
+
+    attempts_.push_back(attempt);
+    had_tcp_attempt_ = true;
+
+    int rv = attempt->Start();
+    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 +715,7 @@ class DnsTransactionImpl : public DnsTransaction,
     first_server_index_ = session_->NextFirstServerIndex();
 
     attempts_.clear();
+    had_tcp_attempt_ = false;
     return MakeAttempt();
   }
 
@@ -456,27 +723,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) {
+  void LogResponse(const DnsAttempt* attempt) {
     if (attempt && attempt->response()) {
       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) {
@@ -508,6 +777,9 @@ class DnsTransactionImpl : public DnsTransaction,
             return result;
           }
           break;
+        case ERR_DNS_SERVER_REQUIRES_TCP:
+          result = MakeTCPAttempt(result.attempt);

[mmenke, 2012/12/19 16:44:09]

Should we be remembering this somewhere?  Or might a server just require TCP/IP
for longer responses?  No need to worry about it in this CL.

[szym, 2012/12/19 17:01:17]

This is characteristic to the query more than the server, so we could cache it
in HostCache as a small optimization. However, given that <1% queries need TCP,
I wouldn't bother at this time.

[mmenke, 2012/12/19 17:13:51]

If it's just for the query, I agree that it's not worth it.  Wasn't sure if
there were DNS servers out there that insist on TCP or not.

+          break;
         default:
           // Server failure.
           DCHECK(result.attempt);
@@ -550,7 +822,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_;