Allow a callback to populate mock DNS responses.

net/dns/dns_test_util.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_test_util.h"
 
 #include <string>
 
 #include "base/big_endian.h"
 #include "base/bind.h"
+#include "base/callback.h"
 #include "base/location.h"
 #include "base/memory/weak_ptr.h"
 #include "base/single_thread_task_runner.h"
 #include "base/sys_byteorder.h"
 #include "base/threading/thread_task_runner_handle.h"
 #include "net/base/io_buffer.h"
 #include "net/base/net_errors.h"
 #include "net/dns/address_sorter.h"
 #include "net/dns/dns_query.h"
 #include "net/dns/dns_response.h"
@@ skipping 15 matching lines @@
 };
 
 // A DnsTransaction which uses MockDnsClientRuleList to determine the response.
 class MockTransaction : public DnsTransaction,
                         public base::SupportsWeakPtr<MockTransaction> {
  public:
   MockTransaction(const MockDnsClientRuleList& rules,
                   const std::string& hostname,
                   uint16_t qtype,
                   const DnsTransactionFactory::CallbackType& callback)
-      : result_(MockDnsClientRule::FAIL),
-        hostname_(hostname),
+      : hostname_(hostname),
         qtype_(qtype),
         callback_(callback),
         started_(false),
         delayed_(false) {
     // Find the relevant rule which matches |qtype| and prefix of |hostname|.
     for (size_t i = 0; i < rules.size(); ++i) {
       const std::string& prefix = rules[i].prefix;
       if ((rules[i].qtype == qtype) &&
           (hostname.size() >= prefix.size()) &&
           (hostname.compare(0, prefix.size(), prefix) == 0)) {
-        result_ = rules[i].result;
+        response_callback_ = rules[i].response_callback;
         delayed_ = rules[i].delay;
         break;
       }
     }
   }
 
   const std::string& GetHostname() const override { return hostname_; }
 
   uint16_t GetType() const override { return qtype_; }
 
@@ skipping 10 matching lines @@
   void FinishDelayedTransaction() {
     EXPECT_TRUE(delayed_);
     delayed_ = false;
     Finish();
   }
 
   bool delayed() const { return delayed_; }
 
  private:
   void Finish() {
-    switch (result_) {
-      case MockDnsClientRule::EMPTY:
-      case MockDnsClientRule::OK: {
-        std::string qname;
-        DNSDomainFromDot(hostname_, &qname);
-        DnsQuery query(0, qname, qtype_);
+    if (response_callback_.is_null()) {
+      callback_.Run(this, ERR_NAME_NOT_RESOLVED, nullptr);
+      return;
+    }
 
-        DnsResponse response;
-        char* buffer = response.io_buffer()->data();
-        int nbytes = query.io_buffer()->size();
-        memcpy(buffer, query.io_buffer()->data(), nbytes);
-        dns_protocol::Header* header =
-            reinterpret_cast<dns_protocol::Header*>(buffer);
-        header->flags |= dns_protocol::kFlagResponse;
+    std::string qname;
+    DNSDomainFromDot(hostname_, &qname);
+    DnsQuery query(0, qname, qtype_);
 
-        if (MockDnsClientRule::OK == result_) {
-          const uint16_t kPointerToQueryName =
-              static_cast<uint16_t>(0xc000 | sizeof(*header));
+    DnsResponse response;
+    IOBufferWithSize* buffer = response.io_buffer();
+    int query_size = query.io_buffer()->size();
+    CHECK_GE(buffer->size(), query_size);
+    memcpy(buffer->data(), query.io_buffer()->data(), query_size);
+    dns_protocol::Header* header =
+        reinterpret_cast<dns_protocol::Header*>(buffer->data());
+    header->flags |= dns_protocol::kFlagResponse;
 
-          const uint32_t kTTL = 86400;  // One day.
-
-          // Size of RDATA which is a IPv4 or IPv6 address.
-          size_t rdata_size = qtype_ == dns_protocol::kTypeA
-                                  ? IPAddress::kIPv4AddressSize
-                                  : IPAddress::kIPv6AddressSize;
-
-          // 12 is the sum of sizes of the compressed name reference, TYPE,
-          // CLASS, TTL and RDLENGTH.
-          size_t answer_size = 12 + rdata_size;
-
-          // Write answer with loopback IP address.
-          header->ancount = base::HostToNet16(1);
-          base::BigEndianWriter writer(buffer + nbytes, answer_size);
-          writer.WriteU16(kPointerToQueryName);
-          writer.WriteU16(qtype_);
-          writer.WriteU16(dns_protocol::kClassIN);
-          writer.WriteU32(kTTL);
-          writer.WriteU16(static_cast<uint16_t>(rdata_size));
-          if (qtype_ == dns_protocol::kTypeA) {
-            char kIPv4Loopback[] = { 0x7f, 0, 0, 1 };
-            writer.WriteBytes(kIPv4Loopback, sizeof(kIPv4Loopback));
-          } else {
-            char kIPv6Loopback[] = { 0, 0, 0, 0, 0, 0, 0, 0,
-                                     0, 0, 0, 0, 0, 0, 0, 1 };
-            writer.WriteBytes(kIPv6Loopback, sizeof(kIPv6Loopback));
-          }
-          nbytes += answer_size;
-        }
-        EXPECT_TRUE(response.InitParse(nbytes, query));
-        callback_.Run(this, OK, &response);
-      } break;
-      case MockDnsClientRule::FAIL:
-        callback_.Run(this, ERR_NAME_NOT_RESOLVED, NULL);
-        break;
-      case MockDnsClientRule::TIMEOUT:
-        callback_.Run(this, ERR_DNS_TIMED_OUT, NULL);
-        break;
-      default:
-        NOTREACHED();
-        break;
+    base::BigEndianWriter answer_writer(buffer->data() + query_size,
+                                        buffer->size() - query_size);
+    int net_error = response_callback_.Run(header, &answer_writer);
+    if (net_error == OK) {
+      int nbytes = answer_writer.ptr() - buffer->data();
+      EXPECT_TRUE(response.InitParse(nbytes, query));
+      callback_.Run(this, OK, &response);
+    } else {
+      callback_.Run(this, net_error, nullptr);
     }
   }
 
-  MockDnsClientRule::Result result_;
+  MockDnsClientRule::ResponseCallback response_callback_;
   const std::string hostname_;
   const uint16_t qtype_;
   DnsTransactionFactory::CallbackType callback_;
   bool started_;
   bool delayed_;
 };
 
+// Simply returns the |net_error| argument.
+// Useful as a simple callback that does nothing but reports an error.
+int ReturnNetError(int net_error,
+                   dns_protocol::Header* response_header,
+                   base::BigEndianWriter* answer_writer) {
+  CHECK_LE(net_error, 0);
+  return net_error;
+}
+
+// Writes a |qtype| record for the loopback address using |answer_writer|.
+// |qtype| must be |dns_protocol::kTypeA| or |dns_protocol::kTypeAAAA|.
+// Returns net::OK if successful.
+int WriteLoopbackRecordResponse(uint16_t qtype,
+                                dns_protocol::Header* response_header,
+                                base::BigEndianWriter* answer_writer) {
+  const uint16_t kPointerToQueryName =
+      static_cast<uint16_t>(0xc000 | sizeof(*response_header));
+
+  const uint32_t kTTL = 86400;  // One day.
+
+  // Write answer with loopback IP address.
+  response_header->ancount =
+      base::HostToNet16(base::NetToHost16(response_header->ancount) + 1);
+  CHECK(answer_writer->WriteU16(kPointerToQueryName));
+  CHECK(answer_writer->WriteU16(qtype));
+  CHECK(answer_writer->WriteU16(dns_protocol::kClassIN));
+  CHECK(answer_writer->WriteU32(kTTL));
+  if (qtype == dns_protocol::kTypeA) {
+    char kIPv4Loopback[] = {0x7f, 0, 0, 1};
+    CHECK(answer_writer->WriteU16(IPAddress::kIPv4AddressSize));
+    CHECK(answer_writer->WriteBytes(kIPv4Loopback, sizeof(kIPv4Loopback)));
+  } else if (qtype == dns_protocol::kTypeAAAA) {
+    char kIPv6Loopback[] = {0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1};
+    CHECK(answer_writer->WriteU16(IPAddress::kIPv6AddressSize));
+    CHECK(answer_writer->WriteBytes(kIPv6Loopback, sizeof(kIPv6Loopback)));
+  } else {
+    NOTREACHED();
+  }
+
+  return OK;
+}
+
 }  // namespace
 
 // A DnsTransactionFactory which creates MockTransaction.
 class MockTransactionFactory : public DnsTransactionFactory {
  public:
   explicit MockTransactionFactory(const MockDnsClientRuleList& rules)
       : rules_(rules) {}
 
   ~MockTransactionFactory() override {}
 
@@ skipping 19 matching lines @@
     }
   }
 
  private:
   typedef std::vector<base::WeakPtr<MockTransaction> > DelayedTransactionList;
 
   MockDnsClientRuleList rules_;
   DelayedTransactionList delayed_transactions_;
 };
 
+MockDnsClientRule::MockDnsClientRule(const std::string& prefix,
+                                     uint16_t qtype,
+                                     Result result,
+                                     bool delay)
+    : prefix(prefix), qtype(qtype), delay(delay) {
+  switch (result) {
+    case FAIL:
+      response_callback =
+          base::Bind(&ReturnNetError, net::ERR_NAME_NOT_RESOLVED);
+      break;
+    case TIMEOUT:
+      response_callback = base::Bind(&ReturnNetError, net::ERR_DNS_TIMED_OUT);
+      break;
+    case EMPTY:
+      response_callback = base::Bind(&ReturnNetError, net::OK);
+      break;
+    case OK:
+      response_callback = base::Bind(&WriteLoopbackRecordResponse, qtype);
+      break;
+  }
+  CHECK(!response_callback.is_null());
+}
+
+MockDnsClientRule::MockDnsClientRule(const std::string& prefix,
+                                     uint16_t qtype,
+                                     ResponseCallback response_callback,
+                                     bool delay)
+    : response_callback(response_callback),
+      prefix(prefix),
+      qtype(qtype),
+      delay(delay) {}
+
+MockDnsClientRule::MockDnsClientRule(const MockDnsClientRule& o)
+    : response_callback(o.response_callback),
+      prefix(o.prefix),
+      qtype(o.qtype),
+      delay(o.delay) {}
+
+MockDnsClientRule::MockDnsClientRule(MockDnsClientRule&& o)
+    : response_callback(std::move(o.response_callback)),
+      prefix(std::move(o.prefix)),
+      qtype(o.qtype),
+      delay(o.delay) {}
+
+MockDnsClientRule::~MockDnsClientRule() {}
+
 MockDnsClient::MockDnsClient(const DnsConfig& config,
                              const MockDnsClientRuleList& rules)
       : config_(config),
         factory_(new MockTransactionFactory(rules)),
         address_sorter_(new MockAddressSorter()) {
 }
 
 MockDnsClient::~MockDnsClient() {}
 
 void MockDnsClient::SetConfig(const DnsConfig& config) {
@@ skipping 10 matching lines @@
 
 AddressSorter* MockDnsClient::GetAddressSorter() {
   return address_sorter_.get();
 }
 
 void MockDnsClient::CompleteDelayedTransactions() {
   factory_->CompleteDelayedTransactions();
 }
 
 }  // namespace net