net/dns/dns_test_util.cc

   1 // Copyright (c) 2012 The Chromium Authors. All rights reserved.
   2 // Use of this source code is governed by a BSD-style license that can be
   3 // found in the LICENSE file.
   4
   5 #include "net/dns/dns_test_util.h"
   6
   7 #include <string>
   8
   9 #include "base/bind.h"
  10 #include "base/memory/weak_ptr.h"
  11 #include "base/message_loop.h"
  12 #include "base/sys_byteorder.h"
  13 #include "net/base/big_endian.h"
  14 #include "net/base/dns_util.h"
  15 #include "net/base/io_buffer.h"
  16 #include "net/base/net_errors.h"
  17 #include "net/dns/address_sorter.h"
  18 #include "net/dns/dns_client.h"
  19 #include "net/dns/dns_config_service.h"
  20 #include "net/dns/dns_protocol.h"
  21 #include "net/dns/dns_query.h"
  22 #include "net/dns/dns_response.h"
  23 #include "net/dns/dns_transaction.h"
  24 #include "testing/gtest/include/gtest/gtest.h"
  25
  26 namespace net {
  27 namespace {
  28
  29 // A DnsTransaction which uses MockDnsClientRuleList to determine the response.
  30 class MockTransaction : public DnsTransaction,
  31                         public base::SupportsWeakPtr<MockTransaction> {
  32  public:
  33   MockTransaction(const MockDnsClientRuleList& rules,
  34                   const std::string& hostname,
  35                   uint16 qtype,
  36                   const DnsTransactionFactory::CallbackType& callback)
  37       : result_(MockDnsClientRule::FAIL_SYNC),
  38         hostname_(hostname),
  39         qtype_(qtype),
  40         callback_(callback),
  41         started_(false) {
  42     // Find the relevant rule which matches |qtype| and prefix of |hostname|.
  43     for (size_t i = 0; i < rules.size(); ++i) {
  44       const std::string& prefix = rules[i].prefix;
  45       if ((rules[i].qtype == qtype) &&
  46           (hostname.size() >= prefix.size()) &&
  47           (hostname.compare(0, prefix.size(), prefix) == 0)) {
  48         result_ = rules[i].result;
  49         break;
  50       }
  51     }
  52   }
  53
  54   virtual const std::string& GetHostname() const OVERRIDE {
  55     return hostname_;
  56   }
  57
  58   virtual uint16 GetType() const OVERRIDE {
  59     return qtype_;
  60   }
  61
  62   virtual int Start() OVERRIDE {
  63     EXPECT_FALSE(started_);
  64     started_ = true;
  65     if (MockDnsClientRule::FAIL_SYNC == result_)
  66       return ERR_NAME_NOT_RESOLVED;
  67     // Using WeakPtr to cleanly cancel when transaction is destroyed.
  68     MessageLoop::current()->PostTask(
  69         FROM_HERE,
  70         base::Bind(&MockTransaction::Finish, AsWeakPtr()));
  71     return ERR_IO_PENDING;
  72   }
  73
  74  private:
  75   void Finish() {
  76     switch (result_) {
  77       case MockDnsClientRule::EMPTY:
  78       case MockDnsClientRule::OK: {
  79         std::string qname;
  80         DNSDomainFromDot(hostname_, &qname);
  81         DnsQuery query(0, qname, qtype_);
  82
  83         DnsResponse response;
  84         char* buffer = response.io_buffer()->data();
  85         int nbytes = query.io_buffer()->size();
  86         memcpy(buffer, query.io_buffer()->data(), nbytes);
  87         dns_protocol::Header* header =
  88             reinterpret_cast<dns_protocol::Header*>(buffer);
  89         header->flags |= dns_protocol::kFlagResponse;
  90
  91         if (MockDnsClientRule::OK == result_) {
  92           const uint16 kPointerToQueryName =
  93               static_cast<uint16>(0xc000 | sizeof(*header));
  94
  95           const uint32 kTTL = 86400;  // One day.
  96
  97           // Size of RDATA which is a IPv4 or IPv6 address.
  98           size_t rdata_size = qtype_ == net::dns_protocol::kTypeA ?
  99                               net::kIPv4AddressSize : net::kIPv6AddressSize;
 100
 101           // 12 is the sum of sizes of the compressed name reference, TYPE,
 102           // CLASS, TTL and RDLENGTH.
 103           size_t answer_size = 12 + rdata_size;
 104
 105           // Write answer with loopback IP address.
 106           header->ancount = base::HostToNet16(1);
 107           BigEndianWriter writer(buffer + nbytes, answer_size);
 108           writer.WriteU16(kPointerToQueryName);
 109           writer.WriteU16(qtype_);
 110           writer.WriteU16(net::dns_protocol::kClassIN);
 111           writer.WriteU32(kTTL);
 112           writer.WriteU16(rdata_size);
 113           if (qtype_ == net::dns_protocol::kTypeA) {
 114             char kIPv4Loopback[] = { 0x7f, 0, 0, 1 };
 115             writer.WriteBytes(kIPv4Loopback, sizeof(kIPv4Loopback));
 116           } else {
 117             char kIPv6Loopback[] = { 0, 0, 0, 0, 0, 0, 0, 0,
 118                                      0, 0, 0, 0, 0, 0, 0, 1 };
 119             writer.WriteBytes(kIPv6Loopback, sizeof(kIPv6Loopback));
 120           }
 121           nbytes += answer_size;
 122         }
 123         EXPECT_TRUE(response.InitParse(nbytes, query));
 124         callback_.Run(this, OK, &response);
 125       } break;
 126       case MockDnsClientRule::FAIL_ASYNC:
 127         callback_.Run(this, ERR_NAME_NOT_RESOLVED, NULL);
 128         break;
 129       case MockDnsClientRule::TIMEOUT:
 130         callback_.Run(this, ERR_DNS_TIMED_OUT, NULL);
 131         break;
 132       default:
 133         NOTREACHED();
 134         break;
 135     }
 136   }
 137
 138   MockDnsClientRule::Result result_;
 139   const std::string hostname_;
 140   const uint16 qtype_;
 141   DnsTransactionFactory::CallbackType callback_;
 142   bool started_;
 143 };
 144
 145
 146 // A DnsTransactionFactory which creates MockTransaction.
 147 class MockTransactionFactory : public DnsTransactionFactory {
 148  public:
 149   explicit MockTransactionFactory(const MockDnsClientRuleList& rules)
 150       : rules_(rules) {}
 151   virtual ~MockTransactionFactory() {}
 152
 153   virtual scoped_ptr<DnsTransaction> CreateTransaction(
 154       const std::string& hostname,
 155       uint16 qtype,
 156       const DnsTransactionFactory::CallbackType& callback,
 157       const BoundNetLog&) OVERRIDE {
 158     return scoped_ptr<DnsTransaction>(
 159         new MockTransaction(rules_, hostname, qtype, callback));
 160   }
 161
 162  private:
 163   MockDnsClientRuleList rules_;
 164 };
 165
 166 class MockAddressSorter : public AddressSorter {
 167  public:
 168   virtual ~MockAddressSorter() {}
 169   virtual void Sort(const AddressList& list,
 170                     const CallbackType& callback) const OVERRIDE {
 171     // Do nothing.
 172     callback.Run(true, list);
 173   }
 174 };
 175
 176 // MockDnsClient provides MockTransactionFactory.
 177 class MockDnsClient : public DnsClient {
 178  public:
 179   MockDnsClient(const DnsConfig& config,
 180                 const MockDnsClientRuleList& rules)
 181       : config_(config), factory_(rules) {}
 182   virtual ~MockDnsClient() {}
 183
 184   virtual void SetConfig(const DnsConfig& config) OVERRIDE {
 185     config_ = config;
 186   }
 187
 188   virtual const DnsConfig* GetConfig() const OVERRIDE {
 189     return config_.IsValid() ? &config_ : NULL;
 190   }
 191
 192   virtual DnsTransactionFactory* GetTransactionFactory() OVERRIDE {
 193     return config_.IsValid() ? &factory_ : NULL;
 194   }
 195
 196   virtual AddressSorter* GetAddressSorter() OVERRIDE {
 197     return &address_sorter_;
 198   }
 199
 200  private:
 201   DnsConfig config_;
 202   MockTransactionFactory factory_;
 203   MockAddressSorter address_sorter_;
 204 };
 205
 206 }  // namespace
 207
 208 // static
 209 scoped_ptr<DnsClient> CreateMockDnsClient(const DnsConfig& config,
 210                                           const MockDnsClientRuleList& rules) {
 211   return scoped_ptr<DnsClient>(new MockDnsClient(config, rules));
 212 }
 213
 214 }  // namespace net