Record queries in DnsUdpTracker

Bug: 1093369
Change-Id: Ia04f5a7f3b44c34aa543919ff7c8a8d6468f1792
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/2238376
Commit-Queue: Eric Orth <ericorth@chromium.org>
Reviewed-by: Dan McArdle <dmcardle@chromium.org>
Cr-Commit-Position: refs/heads/master@{#777249}

net/dns/dns_response.cc

@@ -365,8 +365,13 @@ bool DnsResponse::InitParse(size_t nbytes, const DnsQuery& query) {
     return false;
   }
 
+  // At this point, it has been validated that the response is at least large
+  // enough to read the ID field.
+  id_available_ = true;
+
   // Match the query id.
-  if (base::NetToHost16(header()->id) != query.id())
+  DCHECK(id());
+  if (id().value() != query.id())
     return false;
 
   // Not a response?
@@ -393,6 +398,7 @@ bool DnsResponse::InitParseWithoutQuery(size_t nbytes) {
   if (nbytes < kHeaderSize || nbytes > io_buffer_size_) {
     return false;
   }
+  id_available_ = true;
 
   parser_ = DnsRecordParser(io_buffer_->data(), nbytes, kHeaderSize);
 
@@ -411,6 +417,13 @@ bool DnsResponse::InitParseWithoutQuery(size_t nbytes) {
   return true;
 }
 
+base::Optional<uint16_t> DnsResponse::id() const {
+  if (!id_available_)
+    return base::nullopt;
+
+  return base::NetToHost16(header()->id);
+}
+
 bool DnsResponse::IsValid() const {
   return parser_.IsValid();
 }

net/dns/dns_response.h

@@ -169,6 +169,12 @@ class NET_EXPORT_PRIVATE DnsResponse {
   // if the response is constructed from a raw buffer.
   bool InitParseWithoutQuery(size_t nbytes);
 
+  // Does not require the response to be fully parsed and valid, but will return
+  // nullopt if the ID is unknown. The ID will only be known if the response is
+  // successfully constructed from data or if InitParse...() has been able to
+  // parse at least as far as the ID (not necessarily a fully successful parse).
+  base::Optional<uint16_t> id() const;
+
   // Returns true if response is valid, that is, after successful InitParse, or
   // after successful construction of a new response from data.
   bool IsValid() const;
@@ -220,6 +226,7 @@ class NET_EXPORT_PRIVATE DnsResponse {
   // Iterator constructed after InitParse positioned at the answer section.
   // It is never updated afterwards, so can be used in accessors.
   DnsRecordParser parser_;
+  bool id_available_ = false;
 
   DISALLOW_COPY_AND_ASSIGN(DnsResponse);
 };

net/dns/dns_response_unittest.cc

@@ -4,6 +4,7 @@
 
 #include "net/dns/dns_response.h"
 
+#include <algorithm>
 #include <memory>
 
 #include "base/big_endian.h"
@@ -17,6 +18,7 @@
 #include "net/dns/dns_util.h"
 #include "net/dns/public/dns_protocol.h"
 #include "net/dns/record_rdata.h"
+#include "testing/gmock/include/gmock/gmock.h"
 #include "testing/gtest/include/gtest/gtest.h"
 
 namespace net {
@@ -297,26 +299,32 @@ TEST(DnsResponseTest, InitParse) {
   DnsResponse resp;
   memcpy(resp.io_buffer()->data(), response_data, sizeof(response_data));
 
+  EXPECT_FALSE(resp.id());
+
   // Reject too short.
   EXPECT_FALSE(resp.InitParse(query->io_buffer()->size() - 1, *query));
   EXPECT_FALSE(resp.IsValid());
+  EXPECT_FALSE(resp.id());
 
   // Reject wrong id.
   std::unique_ptr<DnsQuery> other_query = query->CloneWithNewId(0xbeef);
   EXPECT_FALSE(resp.InitParse(sizeof(response_data), *other_query));
   EXPECT_FALSE(resp.IsValid());
+  EXPECT_THAT(resp.id(), testing::Optional(0xcafe));
 
   // Reject wrong question.
   std::unique_ptr<DnsQuery> wrong_query(
       new DnsQuery(0xcafe, qname, dns_protocol::kTypeCNAME));
   EXPECT_FALSE(resp.InitParse(sizeof(response_data), *wrong_query));
   EXPECT_FALSE(resp.IsValid());
+  EXPECT_THAT(resp.id(), testing::Optional(0xcafe));
 
   // Accept matching question.
   EXPECT_TRUE(resp.InitParse(sizeof(response_data), *query));
   EXPECT_TRUE(resp.IsValid());
 
   // Check header access.
+  EXPECT_THAT(resp.id(), testing::Optional(0xcafe));
   EXPECT_EQ(0x8180, resp.flags());
   EXPECT_EQ(0x0, resp.rcode());
   EXPECT_EQ(2u, resp.answer_count());
@@ -384,6 +392,7 @@ TEST(DnsResponseTest, InitParseInvalidFlags) {
 
   EXPECT_FALSE(resp.InitParse(sizeof(response_data), *query));
   EXPECT_FALSE(resp.IsValid());
+  EXPECT_THAT(resp.id(), testing::Optional(0xcafe));
 }
 
 TEST(DnsResponseTest, InitParseWithoutQuery) {
@@ -441,6 +450,7 @@ TEST(DnsResponseTest, InitParseWithoutQueryNoQuestions) {
   EXPECT_TRUE(resp.InitParseWithoutQuery(sizeof(response_data)));
 
   // Check header access.
+  EXPECT_THAT(resp.id(), testing::Optional(0xcafe));
   EXPECT_EQ(0x8180, resp.flags());
   EXPECT_EQ(0x0, resp.rcode());
   EXPECT_EQ(0x1u, resp.answer_count());
@@ -483,6 +493,7 @@ TEST(DnsResponseTest, InitParseWithoutQueryInvalidFlags) {
   memcpy(resp.io_buffer()->data(), response_data, sizeof(response_data));
 
   EXPECT_FALSE(resp.InitParseWithoutQuery(sizeof(response_data)));
+  EXPECT_THAT(resp.id(), testing::Optional(0xcafe));
 }
 
 TEST(DnsResponseTest, InitParseWithoutQueryTwoQuestions) {
@@ -1181,6 +1192,7 @@ TEST(DnsResponseWriteTest, WrittenResponseCanBeParsed) {
                        base::nullopt);
   ASSERT_NE(nullptr, response.io_buffer());
   EXPECT_TRUE(response.IsValid());
+  EXPECT_THAT(response.id(), testing::Optional(0x1234));
   EXPECT_EQ(1u, response.answer_count());
   EXPECT_EQ(1u, response.additional_answer_count());
   auto parser = response.Parser();

net/dns/dns_session.h

@@ -16,6 +16,7 @@
 #include "net/base/rand_callback.h"
 #include "net/dns/dns_config.h"
 #include "net/dns/dns_socket_pool.h"
+#include "net/dns/dns_udp_tracker.h"
 
 namespace net {
 
@@ -62,6 +63,7 @@ class NET_EXPORT_PRIVATE DnsSession : public base::RefCounted<DnsSession> {
              NetLog* net_log);
 
   const DnsConfig& config() const { return config_; }
+  DnsUdpTracker* udp_tracker() { return &udp_tracker_; }
   NetLog* net_log() const { return net_log_; }
 
   // Return the next random query ID.
@@ -98,6 +100,7 @@ class NET_EXPORT_PRIVATE DnsSession : public base::RefCounted<DnsSession> {
 
   const DnsConfig config_;
   std::unique_ptr<DnsSocketPool> socket_pool_;
+  DnsUdpTracker udp_tracker_;
   RandCallback rand_callback_;
   NetLog* net_log_;
 

net/dns/dns_transaction.cc

@@ -23,6 +23,7 @@
 #include "base/memory/weak_ptr.h"
 #include "base/metrics/histogram_functions.h"
 #include "base/metrics/histogram_macros.h"
+#include "base/optional.h"
 #include "base/rand_util.h"
 #include "base/single_thread_task_runner.h"
 #include "base/stl_util.h"
@@ -47,6 +48,7 @@
 #include "net/dns/dns_response.h"
 #include "net/dns/dns_server_iterator.h"
 #include "net/dns/dns_session.h"
+#include "net/dns/dns_udp_tracker.h"
 #include "net/dns/dns_util.h"
 #include "net/dns/public/dns_over_https_server_config.h"
 #include "net/dns/public/dns_protocol.h"
@@ -186,11 +188,13 @@ class DnsUDPAttempt : public DnsAttempt {
  public:
   DnsUDPAttempt(size_t server_index,
                 std::unique_ptr<DnsSession::SocketLease> socket_lease,
-                std::unique_ptr<DnsQuery> query)
+                std::unique_ptr<DnsQuery> query,
+                DnsUdpTracker* udp_tracker)
       : DnsAttempt(server_index),
         next_state_(STATE_NONE),
         socket_lease_(std::move(socket_lease)),
-        query_(std::move(query)) {}
+        query_(std::move(query)),
+        udp_tracker_(udp_tracker) {}
 
   // DnsAttempt methods.
 
@@ -199,6 +203,11 @@ class DnsUDPAttempt : public DnsAttempt {
     callback_ = std::move(callback);
     start_time_ = base::TimeTicks::Now();
     next_state_ = STATE_SEND_QUERY;
+
+    IPEndPoint local_address;
+    if (socket_lease_->socket()->GetLocalAddress(&local_address) == OK)
+      udp_tracker_->RecordQuery(local_address.port(), query_->id());
+
     return DoLoop(OK);
   }
 
@@ -294,7 +303,11 @@ class DnsUDPAttempt : public DnsAttempt {
       return rv;
 
     DCHECK(rv);
-    if (!response_->InitParse(rv, *query_))
+    bool parse_result = response_->InitParse(rv, *query_);
+    if (response_->id())
+      udp_tracker_->RecordResponseId(query_->id(), response_->id().value());
+
+    if (!parse_result)
       return ERR_DNS_MALFORMED_RESPONSE;
     if (response_->flags() & dns_protocol::kFlagTC)
       return ERR_DNS_SERVER_REQUIRES_TCP;
@@ -318,6 +331,10 @@ class DnsUDPAttempt : public DnsAttempt {
   std::unique_ptr<DnsSession::SocketLease> socket_lease_;
   std::unique_ptr<DnsQuery> query_;
 
+  // Should be owned by the DnsSession, to which the transaction should own a
+  // reference.
+  DnsUdpTracker* const udp_tracker_;
+
   std::unique_ptr<DnsResponse> response_;
 
   CompletionOnceCallback callback_;
@@ -1202,8 +1219,9 @@ class DnsTransactionImpl : public DnsTransaction,
 
     bool got_socket = !!lease.get();
 
-    DnsUDPAttempt* attempt = new DnsUDPAttempt(
-        non_doh_server_index, std::move(lease), std::move(query));
+    DnsUDPAttempt* attempt =
+        new DnsUDPAttempt(non_doh_server_index, std::move(lease),
+                          std::move(query), session_->udp_tracker());
 
     attempts_.push_back(base::WrapUnique(attempt));
     ++attempts_count_;