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

net/dns/dns_transaction_unittest.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 "base/bind.h"
 #include "base/memory/scoped_ptr.h"
 #include "base/memory/scoped_vector.h"
 #include "base/rand_util.h"
@@ skipping 13 matching lines @@
 namespace net {
 
 namespace {
 
 std::string DomainFromDot(const base::StringPiece& dotted) {
   std::string out;
   EXPECT_TRUE(DNSDomainFromDot(dotted, &out));
   return out;
 }
 
-// A SocketDataProvider builder for MockUDPClientSocket. Each socket used by a
-// DnsTransaction expects only one write and zero or more reads.
+// A SocketDataProvider builder.
 class DnsSocketData {
  public:
   // The ctor takes parameters for the DnsQuery.
-  DnsSocketData(uint16 id, const char* dotted_name, uint16 qtype, IoMode mode)
+  DnsSocketData(uint16 id,
+                const char* dotted_name,
+                uint16 qtype,
+                IoMode mode,
+                bool use_tcp)
       : query_(new DnsQuery(id, DomainFromDot(dotted_name), qtype)),
-        write_(mode, query_->io_buffer()->data(), query_->io_buffer()->size()) {
+        use_tcp_(use_tcp) {
+    if (use_tcp_) {
+      scoped_ptr<uint16> length(new uint16);
+      *length = base::HostToNet16(query_->io_buffer()->size());
+      writes_.push_back(MockWrite(mode,
+                                  reinterpret_cast<const char*>(length.get()),
+                                  sizeof(uint16)));
+      lengths_.push_back(length.release());
+    }
+    writes_.push_back(MockWrite(mode,
+                                query_->io_buffer()->data(),
+                                query_->io_buffer()->size()));
   }
   ~DnsSocketData() {}
 
   // All responses must be added before GetProvider.
 
-  // Add pre-built DnsResponse.
-  void AddResponse(scoped_ptr<DnsResponse> response, IoMode mode) {
+  // Adds pre-built DnsResponse. |tcp_length| will be used in TCP mode only.
+  void AddResponseWithLength(scoped_ptr<DnsResponse> response, IoMode mode,
+                             uint16 tcp_length) {
     CHECK(!provider_.get());
+    if (use_tcp_) {
+      scoped_ptr<uint16> length(new uint16);
+      *length = base::HostToNet16(tcp_length);
+      reads_.push_back(MockRead(mode,
+                                reinterpret_cast<const char*>(length.get()),
+                                sizeof(uint16)));
+      lengths_.push_back(length.release());
+    }
     reads_.push_back(MockRead(mode,
                               response->io_buffer()->data(),
                               response->io_buffer()->size()));
     responses_.push_back(response.release());
   }
 
+  // Adds pre-built DnsResponse.
+  void AddResponse(scoped_ptr<DnsResponse> response, IoMode mode) {
+    uint16 tcp_length = response->io_buffer()->size();
+    AddResponseWithLength(response.Pass(), mode, tcp_length);
+  }
+
   // Adds pre-built response from |data| buffer.
   void AddResponseData(const uint8* data, size_t length, IoMode mode) {
     CHECK(!provider_.get());
     AddResponse(make_scoped_ptr(
         new DnsResponse(reinterpret_cast<const char*>(data), length, 0)), mode);
   }
 
   // Add no-answer (RCODE only) response matching the query.
   void AddRcode(int rcode, IoMode mode) {
     scoped_ptr<DnsResponse> response(
         new DnsResponse(query_->io_buffer()->data(),
                         query_->io_buffer()->size(),
                         0));
     dns_protocol::Header* header =
         reinterpret_cast<dns_protocol::Header*>(response->io_buffer()->data());
     header->flags |= base::HostToNet16(dns_protocol::kFlagResponse | rcode);
     AddResponse(response.Pass(), mode);
   }
 
   // Build, if needed, and return the SocketDataProvider. No new responses
   // should be added afterwards.
   SocketDataProvider* GetProvider() {
     if (provider_.get())
       return provider_.get();
-    if (reads_.empty()) {
-      // Timeout.
-      provider_.reset(new DelayedSocketData(2, NULL, 0, &write_, 1));
-    } else {
-      // Terminate the reads with ERR_IO_PENDING to prevent overrun.
-      reads_.push_back(MockRead(ASYNC, ERR_IO_PENDING));
-      provider_.reset(new DelayedSocketData(1, &reads_[0], reads_.size(),
-                                            &write_, 1));
+    // Terminate the reads with ERR_IO_PENDING to prevent overrun and default to
+    // timeout.
+    reads_.push_back(MockRead(ASYNC, ERR_IO_PENDING));
+    provider_.reset(new DelayedSocketData(1, &reads_[0], reads_.size(),
+                                          &writes_[0], writes_.size()));
+    if (use_tcp_) {
+      provider_->set_connect_data(MockConnect(reads_[0].mode, OK));
     }
     return provider_.get();
   }
 
   uint16 query_id() const {
     return query_->id();
   }
 
   // Returns true if the expected query was written to the socket.
   bool was_written() const {
     CHECK(provider_.get());
     return provider_->write_index() > 0;
   }
 
  private:
   scoped_ptr<DnsQuery> query_;
+  bool use_tcp_;
+  ScopedVector<uint16> lengths_;
   ScopedVector<DnsResponse> responses_;
-  MockWrite write_;
+  std::vector<MockWrite> writes_;
   std::vector<MockRead> reads_;
   scoped_ptr<DelayedSocketData> provider_;
 
   DISALLOW_COPY_AND_ASSIGN(DnsSocketData);
 };
 
 class TestSocketFactory;
 
 // A variant of MockUDPClientSocket which always fails to Connect.
 class FailingUDPClientSocket : public MockUDPClientSocket {
@@ skipping 205 matching lines @@
     socket_factory_.reset(new TestSocketFactory());
     session_ = new DnsSession(
         config_,
         DnsSocketPool::CreateNull(socket_factory_.get()),
         base::Bind(&DnsTransactionTest::GetNextId, base::Unretained(this)),
         NULL /* NetLog */);
     transaction_factory_ = DnsTransactionFactory::CreateFactory(session_.get());
   }
 
   void AddSocketData(scoped_ptr<DnsSocketData> data) {
+    CHECK(socket_factory_.get());
     transaction_ids_.push_back(data->query_id());
     socket_factory_->AddSocketDataProvider(data->GetProvider());
     socket_data_.push_back(data.release());
   }
 
   // Add expected query for |dotted_name| and |qtype| with |id| and response
   // taken verbatim from |data| of |data_length| bytes. The transaction id in
   // |data| should equal |id|, unless testing mismatched response.
   void AddQueryAndResponse(uint16 id,
                            const char* dotted_name,
                            uint16 qtype,
                            const uint8* response_data,
                            size_t response_length,
-                           IoMode mode) {
+                           IoMode mode,
+                           bool use_tcp) {
     CHECK(socket_factory_.get());
     scoped_ptr<DnsSocketData> data(
-        new DnsSocketData(id, dotted_name, qtype, mode));
+        new DnsSocketData(id, dotted_name, qtype, mode, use_tcp));
     data->AddResponseData(response_data, response_length, mode);
     AddSocketData(data.Pass());
   }
 
   void AddAsyncQueryAndResponse(uint16 id,
                                 const char* dotted_name,
                                 uint16 qtype,
                                 const uint8* data,
                                 size_t data_length) {
-    AddQueryAndResponse(id, dotted_name, qtype, data, data_length, ASYNC);
+    AddQueryAndResponse(id, dotted_name, qtype, data, data_length, ASYNC,
+                        false);
   }
 
   void AddSyncQueryAndResponse(uint16 id,
                                const char* dotted_name,
                                uint16 qtype,
                                const uint8* data,
                                size_t data_length) {
-    AddQueryAndResponse(id, dotted_name, qtype, data, data_length, SYNCHRONOUS);
+    AddQueryAndResponse(id, dotted_name, qtype, data, data_length, SYNCHRONOUS,
+                        false);
   }
 
   // Add expected query of |dotted_name| and |qtype| and no response.
   void AddQueryAndTimeout(const char* dotted_name, uint16 qtype) {
-    CHECK(socket_factory_.get());
     uint16 id = base::RandInt(0, kuint16max);
     scoped_ptr<DnsSocketData> data(
-        new DnsSocketData(id, dotted_name, qtype, ASYNC));
+        new DnsSocketData(id, dotted_name, qtype, ASYNC, false));
     AddSocketData(data.Pass());
   }
 
   // Add expected query of |dotted_name| and |qtype| and matching response with
   // no answer and RCODE set to |rcode|. The id will be generated randomly.
   void AddQueryAndRcode(const char* dotted_name,
                         uint16 qtype,
                         int rcode,
-                        IoMode mode) {
-    CHECK(socket_factory_.get());
+                        IoMode mode,
+                        bool use_tcp) {
     CHECK_NE(dns_protocol::kRcodeNOERROR, rcode);
     uint16 id = base::RandInt(0, kuint16max);
     scoped_ptr<DnsSocketData> data(
-        new DnsSocketData(id, dotted_name, qtype, mode));
+        new DnsSocketData(id, dotted_name, qtype, mode, use_tcp));
     data->AddRcode(rcode, mode);
     AddSocketData(data.Pass());
   }
 
   void AddAsyncQueryAndRcode(const char* dotted_name, uint16 qtype, int rcode) {
-    AddQueryAndRcode(dotted_name, qtype, rcode, ASYNC);
+    AddQueryAndRcode(dotted_name, qtype, rcode, ASYNC, false);
   }
 
   void AddSyncQueryAndRcode(const char* dotted_name, uint16 qtype, int rcode) {
-    AddQueryAndRcode(dotted_name, qtype, rcode, SYNCHRONOUS);
+    AddQueryAndRcode(dotted_name, qtype, rcode, SYNCHRONOUS, false);
   }
 
   // Checks if the sockets were connected in the order matching the indices in
   // |servers|.
   void CheckServerOrder(const unsigned* servers, size_t num_attempts) {
     ASSERT_EQ(num_attempts, socket_factory_->remote_endpoints_.size());
     for (size_t i = 0; i < num_attempts; ++i) {
       EXPECT_EQ(socket_factory_->remote_endpoints_[i],
                 session_->config().nameservers[servers[i]]);
     }
@@ skipping 31 matching lines @@
   ScopedVector<DnsSocketData> socket_data_;
 
   std::deque<int> transaction_ids_;
   scoped_ptr<TestSocketFactory> socket_factory_;
   scoped_refptr<DnsSession> session_;
   scoped_ptr<DnsTransactionFactory> transaction_factory_;
 };
 
 TEST_F(DnsTransactionTest, Lookup) {
   AddAsyncQueryAndResponse(0 /* id */, kT0HostName, kT0Qtype,
-                   kT0ResponseDatagram, arraysize(kT0ResponseDatagram));
+                           kT0ResponseDatagram, arraysize(kT0ResponseDatagram));
 
   TransactionHelper helper0(kT0HostName, kT0Qtype, kT0RecordCount);
   EXPECT_TRUE(helper0.Run(transaction_factory_.get()));
 }
 
 // Concurrent lookup tests assume that DnsTransaction::Start immediately
 // consumes a socket from ClientSocketFactory.
 TEST_F(DnsTransactionTest, ConcurrentLookup) {
   AddAsyncQueryAndResponse(0 /* id */, kT0HostName, kT0Qtype,
                            kT0ResponseDatagram, arraysize(kT0ResponseDatagram));
@@ skipping 25 matching lines @@
   helper0.Cancel();
 
   MessageLoop::current()->RunUntilIdle();
 
   EXPECT_FALSE(helper0.has_completed());
   EXPECT_TRUE(helper1.has_completed());
 }
 
 TEST_F(DnsTransactionTest, DestroyFactory) {
   AddAsyncQueryAndResponse(0 /* id */, kT0HostName, kT0Qtype,
-                   kT0ResponseDatagram, arraysize(kT0ResponseDatagram));
+                           kT0ResponseDatagram, arraysize(kT0ResponseDatagram));
 
   TransactionHelper helper0(kT0HostName, kT0Qtype, kT0RecordCount);
   helper0.StartTransaction(transaction_factory_.get());
 
   // Destroying the client does not affect running requests.
   transaction_factory_.reset(NULL);
 
   MessageLoop::current()->RunUntilIdle();
 
   EXPECT_TRUE(helper0.has_completed());
 }
 
 TEST_F(DnsTransactionTest, CancelFromCallback) {
   AddAsyncQueryAndResponse(0 /* id */, kT0HostName, kT0Qtype,
                            kT0ResponseDatagram, arraysize(kT0ResponseDatagram));
 
   TransactionHelper helper0(kT0HostName, kT0Qtype, kT0RecordCount);
   helper0.set_cancel_in_callback();
   EXPECT_TRUE(helper0.Run(transaction_factory_.get()));
 }
 
 TEST_F(DnsTransactionTest, MismatchedResponseSync) {
   config_.attempts = 2;
   config_.timeout = TestTimeouts::tiny_timeout();
   ConfigureFactory();
 
   // Attempt receives mismatched response followed by valid response.
   scoped_ptr<DnsSocketData> data(
-      new DnsSocketData(0 /* id */, kT0HostName, kT0Qtype, SYNCHRONOUS));
+      new DnsSocketData(0 /* id */, kT0HostName, kT0Qtype, SYNCHRONOUS, false));
   data->AddResponseData(kT1ResponseDatagram,
                         arraysize(kT1ResponseDatagram), SYNCHRONOUS);
   data->AddResponseData(kT0ResponseDatagram,
                         arraysize(kT0ResponseDatagram), SYNCHRONOUS);
   AddSocketData(data.Pass());
 
   TransactionHelper helper0(kT0HostName, kT0Qtype, kT0RecordCount);
   EXPECT_TRUE(helper0.RunUntilDone(transaction_factory_.get()));
 }
 
 TEST_F(DnsTransactionTest, MismatchedResponseAsync) {
   config_.attempts = 2;
   config_.timeout = TestTimeouts::tiny_timeout();
   ConfigureFactory();
 
   // First attempt receives mismatched response followed by valid response.
   // Second attempt times out.
   scoped_ptr<DnsSocketData> data(
-      new DnsSocketData(0 /* id */, kT0HostName, kT0Qtype, ASYNC));
+      new DnsSocketData(0 /* id */, kT0HostName, kT0Qtype, ASYNC, false));
   data->AddResponseData(kT1ResponseDatagram,
                         arraysize(kT1ResponseDatagram), ASYNC);
   data->AddResponseData(kT0ResponseDatagram,
                         arraysize(kT0ResponseDatagram), ASYNC);
   AddSocketData(data.Pass());
   AddQueryAndTimeout(kT0HostName, kT0Qtype);
 
   TransactionHelper helper0(kT0HostName, kT0Qtype, kT0RecordCount);
   EXPECT_TRUE(helper0.RunUntilDone(transaction_factory_.get()));
 }
@@ skipping 270 matching lines @@
 }
 
 TEST_F(DnsTransactionTest, ConnectFailure) {
   socket_factory_->create_failing_sockets_ = true;
   transaction_ids_.push_back(0);  // Needed to make a DnsUDPAttempt.
   TransactionHelper helper0("www.chromium.org", dns_protocol::kTypeA,
                             ERR_CONNECTION_REFUSED);
   EXPECT_TRUE(helper0.Run(transaction_factory_.get()));
 }
 
+TEST_F(DnsTransactionTest, TCPLookup) {
+  AddAsyncQueryAndRcode(kT0HostName, kT0Qtype,
+                        dns_protocol::kRcodeNOERROR | dns_protocol::kFlagTC);
+  AddQueryAndResponse(0 /* id */, kT0HostName, kT0Qtype,
+                      kT0ResponseDatagram, arraysize(kT0ResponseDatagram),
+                      ASYNC, true /* use_tcp */);
+
+  TransactionHelper helper0(kT0HostName, kT0Qtype, kT0RecordCount);
+  EXPECT_TRUE(helper0.Run(transaction_factory_.get()));
+}
+
+TEST_F(DnsTransactionTest, TCPFailure) {
+  AddAsyncQueryAndRcode(kT0HostName, kT0Qtype,
+                        dns_protocol::kRcodeNOERROR | dns_protocol::kFlagTC);
+  AddQueryAndRcode(kT0HostName, kT0Qtype, dns_protocol::kRcodeSERVFAIL,
+                   ASYNC, true /* use_tcp */);
+
+  TransactionHelper helper0(kT0HostName, kT0Qtype, ERR_DNS_SERVER_FAILED);
+  EXPECT_TRUE(helper0.Run(transaction_factory_.get()));
+}
+
+TEST_F(DnsTransactionTest, TCPMalformed) {
+  AddAsyncQueryAndRcode(kT0HostName, kT0Qtype,
+                        dns_protocol::kRcodeNOERROR | dns_protocol::kFlagTC);
+  scoped_ptr<DnsSocketData> data(
+      new DnsSocketData(0 /* id */, kT0HostName, kT0Qtype, ASYNC, true));
+  // Valid response but length too short.
+  data->AddResponseWithLength(
+      make_scoped_ptr(
+          new DnsResponse(reinterpret_cast<const char*>(kT0ResponseDatagram),
+                          arraysize(kT0ResponseDatagram), 0)),
+      ASYNC,
+      static_cast<uint16>(kT0QuerySize - 1));
+  AddSocketData(data.Pass());
+
+  TransactionHelper helper0(kT0HostName, kT0Qtype, ERR_DNS_MALFORMED_RESPONSE);
+  EXPECT_TRUE(helper0.Run(transaction_factory_.get()));
+}
+
+TEST_F(DnsTransactionTest, TCPTimeout) {
+  config_.timeout = TestTimeouts::tiny_timeout();
+  ConfigureFactory();
+  AddAsyncQueryAndRcode(kT0HostName, kT0Qtype,
+                        dns_protocol::kRcodeNOERROR | dns_protocol::kFlagTC);
+  AddSocketData(make_scoped_ptr(
+      new DnsSocketData(1 /* id */, kT0HostName, kT0Qtype, ASYNC, true)));
+
+  TransactionHelper helper0(kT0HostName, kT0Qtype, ERR_DNS_TIMED_OUT);
+  EXPECT_TRUE(helper0.RunUntilDone(transaction_factory_.get()));
+}
+
 }  // namespace
 
 }  // namespace net