Allow LogDnsClient queries to be rate-limited

components/certificate_transparency/log_dns_client_unittest.cc

Index: components/certificate_transparency/log_dns_client_unittest.cc
diff --git a/components/certificate_transparency/log_dns_client_unittest.cc b/components/certificate_transparency/log_dns_client_unittest.cc
index 64223ef51757768878acef01f6fd8431449782c9..e7d35e997a7e1831b3df0dd12819666c2a2a0adf 100644
--- a/components/certificate_transparency/log_dns_client_unittest.cc
+++ b/components/certificate_transparency/log_dns_client_unittest.cc
@@ -9,6 +9,7 @@
 #include <utility>
 #include <vector>
 
+#include "base/memory/ptr_util.h"
 #include "base/message_loop/message_loop.h"
 #include "base/run_loop.h"
 #include "components/certificate_transparency/mock_log_dns_traffic.h"
@@ -120,21 +121,46 @@ class LogDnsClientTest : public ::testing::TestWithParam<net::IoMode> {
     mock_dns_.InitializeDnsConfig();
   }
 
+  std::unique_ptr<LogDnsClient> CreateLogDnsClient(
+      size_t max_concurrent_queries = 0 /* unlimited */) {

[Ryan Sleevi, 2016/09/12 18:12:44]

https://google.github.io/styleguide/cppguide.html#Default_Arguments generally
discourages this pattern.

[Rob Percival, 2016/09/13 14:06:31]

That's why I didn't set a default value for the constructor parameter, but it
seemed more reasonable in test code, which almost never sets a limit on
concurrent queries. Otherwise, I'd either have to declare "size_t
max_concurrent_queries = 0;  // unlimited" in a few different places or just
have a magic number 0 on the end of the various CreateLogDnsClient calls. This
seemed cleaner.

However, I suppose I've ended up with those trailing 0 arguments in a couple of
places anyway so I'm ok with changing this.

+    return base::MakeUnique<LogDnsClient>(mock_dns_.CreateDnsClient(),
+                                          net::BoundNetLog(),
+                                          max_concurrent_queries);
+  }
+
+  void QueryLeafIndexAsync(LogDnsClient* log_client,
+                           base::StringPiece log_domain,
+                           base::StringPiece leaf_hash,
+                           const LogDnsClient::LeafIndexCallback& callback) {
+    log_client->QueryLeafIndex(log_domain, leaf_hash, callback);
+  }
+
+  // Convenience function for calling QueryLeafIndexAsync synchronously.
   void QueryLeafIndex(base::StringPiece log_domain,
                       base::StringPiece leaf_hash,
                       MockLeafIndexCallback* callback) {
-    LogDnsClient log_client(mock_dns_.CreateDnsClient(), net::BoundNetLog());
-    log_client.QueryLeafIndex(log_domain, leaf_hash, callback->AsCallback());
+    auto log_client = CreateLogDnsClient();

[Ryan Sleevi, 2016/09/12 18:12:43]

There's been a variety of threads on this particular usage of auto (e.g.
https://groups.google.com/a/chromium.org/d/topic/chromium-dev/5-Bt3BJzAo0/dis...
https://groups.google.com/a/chromium.org/d/topic/chromium-dev/MVLrzzu3DcM/dis...
)

TL;DR: Prefer to be explicit here

std::unique_ptr<LogDnsClient> log_client(CreateLogDnsClient());

[Rob Percival, 2016/09/13 14:06:31]

Done.

+    QueryLeafIndexAsync(log_client.get(), log_domain, leaf_hash,
+                        callback->AsCallback());
     callback->WaitUntilRun();
   }
 
+  void QueryAuditProofAsync(LogDnsClient* log_client,
+                            base::StringPiece log_domain,
+                            uint64_t leaf_index,
+                            uint64_t tree_size,
+                            const LogDnsClient::AuditProofCallback& callback) {
+    log_client->QueryAuditProof(log_domain, leaf_index, tree_size, callback);
+  }
+
+  // Convenience function for calling QueryAuditProofAsync synchronously.
   void QueryAuditProof(base::StringPiece log_domain,
                        uint64_t leaf_index,
                        uint64_t tree_size,
                        MockAuditProofCallback* callback) {
-    LogDnsClient log_client(mock_dns_.CreateDnsClient(), net::BoundNetLog());
-    log_client.QueryAuditProof(log_domain, leaf_index, tree_size,
-                               callback->AsCallback());
+    auto log_client = CreateLogDnsClient();

[Ryan Sleevi, 2016/09/12 18:12:44]

see above

[Rob Percival, 2016/09/13 14:06:31]

Done.

+    QueryAuditProofAsync(log_client.get(), log_domain, leaf_index, tree_size,
+                         callback->AsCallback());
     callback->WaitUntilRun();
   }
 
@@ -514,7 +540,7 @@ TEST_P(LogDnsClientTest, QueryAuditProofReportsTimeout) {
 TEST_P(LogDnsClientTest, AdoptsLatestDnsConfigIfValid) {
   std::unique_ptr<net::DnsClient> tmp = mock_dns_.CreateDnsClient();
   net::DnsClient* dns_client = tmp.get();
-  LogDnsClient log_client(std::move(tmp), net::BoundNetLog());
+  LogDnsClient log_client(std::move(tmp), net::BoundNetLog(), 0);
 
   // Get the current DNS config, modify it and broadcast the update.
   net::DnsConfig config(*dns_client->GetConfig());
@@ -530,7 +556,7 @@ TEST_P(LogDnsClientTest, AdoptsLatestDnsConfigIfValid) {
 TEST_P(LogDnsClientTest, IgnoresLatestDnsConfigIfInvalid) {
   std::unique_ptr<net::DnsClient> tmp = mock_dns_.CreateDnsClient();
   net::DnsClient* dns_client = tmp.get();
-  LogDnsClient log_client(std::move(tmp), net::BoundNetLog());
+  LogDnsClient log_client(std::move(tmp), net::BoundNetLog(), 0);
 
   // Get the current DNS config, modify it and broadcast the update.
   net::DnsConfig config(*dns_client->GetConfig());
@@ -543,6 +569,214 @@ TEST_P(LogDnsClientTest, IgnoresLatestDnsConfigIfInvalid) {
   EXPECT_THAT(dns_client->GetConfig()->nameservers, Not(IsEmpty()));
 }
 
+TEST_P(LogDnsClientTest, CanPerformLeafIndexQueriesInParallel) {
+  constexpr size_t num_of_parallel_queries = 3;
+  auto log_client = CreateLogDnsClient(num_of_parallel_queries);
+  MockLeafIndexCallback callbacks[num_of_parallel_queries];
+
+  for (size_t i = 0; i < num_of_parallel_queries; ++i) {
+    mock_dns_.ExpectLeafIndexRequestAndResponse(
+        "D4S6DSV2J743QJZEQMH4UYHEYK7KRQ5JIQOCPMFUHZVJNFGHXACA.hash.ct.test.",
+        "123456");
+  }
+
+  for (size_t i = 0; i < num_of_parallel_queries; ++i) {
+    QueryLeafIndexAsync(log_client.get(), "ct.test", kLeafHash,
+                        callbacks[i].AsCallback());
+  }
+
+  for (size_t i = 0; i < num_of_parallel_queries; ++i) {

[Ryan Sleevi, 2016/09/12 18:12:44]

More inline documentation about the tests and what's being tested.

[Rob Percival, 2016/09/13 14:06:31]

Done.

+    MockLeafIndexCallback& callback = callbacks[i];
+    callback.WaitUntilRun();
+
+    SCOPED_TRACE(i);
+    EXPECT_TRUE(callback.called());
+    if (callback.called()) {
+      EXPECT_THAT(callback.net_error(), IsOk());
+      EXPECT_THAT(callback.leaf_index(), 123456);
+    }
+  }
+}
+
+TEST_P(LogDnsClientTest, CanPerformAuditProofQueriesInParallel) {
+  constexpr size_t num_of_parallel_queries = 3;
+  auto log_client = CreateLogDnsClient(num_of_parallel_queries);
+  MockAuditProofCallback callbacks[num_of_parallel_queries];
+  const std::vector<std::string> audit_proof = GetSampleAuditProof(20);
+
+  // It should require 3 queries to collect the entire audit proof, as there is
+  // only space for 7 nodes per UDP packet. Expect each of these queries to

[Ryan Sleevi, 2016/09/12 18:12:44]

What is the source of this magic value (7 nodes per UDP packet). How does it fit
into things like various network fragmentation limits? This seems to open more
questions than it answers, and hints of a magic number somewhere far away that
could invalidate this test.

[Rob Percival, 2016/09/13 14:06:31]

I've now clarified this. That should have said TXT RDATA, not UDP packet.

+  // occur once for each call to QueryAuditProof.
+
+  for (size_t i = 0; i < num_of_parallel_queries; ++i) {
+    mock_dns_.ExpectAuditProofRequestAndResponse(
+        "0.123456.999999.tree.ct.test.", audit_proof.begin(),
+        audit_proof.begin() + 7);
+  }
+
+  for (size_t i = 0; i < num_of_parallel_queries; ++i) {
+    mock_dns_.ExpectAuditProofRequestAndResponse(
+        "7.123456.999999.tree.ct.test.", audit_proof.begin() + 7,
+        audit_proof.begin() + 14);
+  }
+
+  for (size_t i = 0; i < num_of_parallel_queries; ++i) {
+    mock_dns_.ExpectAuditProofRequestAndResponse(
+        "14.123456.999999.tree.ct.test.", audit_proof.begin() + 14,
+        audit_proof.end());
+  }
+
+  // Begin the queries.
+  for (size_t i = 0; i < num_of_parallel_queries; ++i) {
+    QueryAuditProofAsync(log_client.get(), "ct.test", 123456, 999999,
+                         callbacks[i].AsCallback());
+  }
+
+  // Check that they were all successful.
+  for (size_t i = 0; i < num_of_parallel_queries; ++i) {
+    MockAuditProofCallback& callback = callbacks[i];
+    callbacks[i].WaitUntilRun();
+
+    SCOPED_TRACE(i);
+    EXPECT_TRUE(callback.called());
+    if (callback.called()) {
+      EXPECT_THAT(callback.net_error(), IsOk());
+      EXPECT_THAT(callback.proof(), NotNull());
+      if (callback.proof() != nullptr) {
+        EXPECT_THAT(callback.proof()->leaf_index, 123456);
+        // EXPECT_THAT(callback.proof()->tree_size, 999999);
+        EXPECT_THAT(callback.proof()->nodes, audit_proof);
+      }
+    }
+  }
+}
+
+TEST_P(LogDnsClientTest, CanPerformLeafIndexAndAuditProofQueriesInParallel) {
+  constexpr size_t num_of_parallel_queries = 2;
+  auto log_client = CreateLogDnsClient(num_of_parallel_queries);
+  MockLeafIndexCallback leaf_index_callback;
+  MockAuditProofCallback audit_proof_callback;
+  const std::vector<std::string> audit_proof = GetSampleAuditProof(20);
+
+  mock_dns_.ExpectLeafIndexRequestAndResponse(
+      "D4S6DSV2J743QJZEQMH4UYHEYK7KRQ5JIQOCPMFUHZVJNFGHXACA.hash.ct.test.",
+      "123456");
+
+  // It should require 3 queries to collect the entire audit proof, as there is
+  // only space for 7 nodes per UDP packet.
+  mock_dns_.ExpectAuditProofRequestAndResponse("0.123456.999999.tree.ct.test.",
+                                               audit_proof.begin(),
+                                               audit_proof.begin() + 7);
+  mock_dns_.ExpectAuditProofRequestAndResponse("7.123456.999999.tree.ct.test.",
+                                               audit_proof.begin() + 7,
+                                               audit_proof.begin() + 14);
+  mock_dns_.ExpectAuditProofRequestAndResponse("14.123456.999999.tree.ct.test.",
+                                               audit_proof.begin() + 14,
+                                               audit_proof.end());
+
+  QueryLeafIndexAsync(log_client.get(), "ct.test", kLeafHash,
+                      leaf_index_callback.AsCallback());
+  QueryAuditProofAsync(log_client.get(), "ct.test", 123456, 999999,
+                       audit_proof_callback.AsCallback());
+
+  leaf_index_callback.WaitUntilRun();
+  audit_proof_callback.WaitUntilRun();
+
+  EXPECT_TRUE(leaf_index_callback.called());
+  if (leaf_index_callback.called()) {
+    EXPECT_THAT(leaf_index_callback.net_error(), IsOk());
+    EXPECT_THAT(leaf_index_callback.leaf_index(), 123456);
+  }
+
+  EXPECT_TRUE(audit_proof_callback.called());
+  if (audit_proof_callback.called()) {
+    EXPECT_THAT(audit_proof_callback.net_error(), IsOk());
+    EXPECT_THAT(audit_proof_callback.proof(), NotNull());
+    if (audit_proof_callback.proof() != nullptr) {
+      EXPECT_THAT(audit_proof_callback.proof()->leaf_index, 123456);
+      // EXPECT_THAT(audit_proof_callback.proof()->tree_size, 999999);
+      EXPECT_THAT(audit_proof_callback.proof()->nodes, audit_proof);
+    }
+  }
+}
+
+TEST_P(LogDnsClientTest, CanBeThrottledToOneLeafIndexQueryAtATime) {
+  mock_dns_.ExpectLeafIndexRequestAndResponse(
+      "D4S6DSV2J743QJZEQMH4UYHEYK7KRQ5JIQOCPMFUHZVJNFGHXACA.hash.ct.test.",
+      "123456");
+
+  const size_t max_concurrent_queries = 1;
+  auto log_client = CreateLogDnsClient(max_concurrent_queries);
+
+  MockLeafIndexCallback callback1;
+  QueryLeafIndexAsync(log_client.get(), "ct.test", kLeafHash,
+                      callback1.AsCallback());
+  MockLeafIndexCallback callback2;
+  QueryLeafIndexAsync(log_client.get(), "ct.test", kLeafHash,
+                      callback2.AsCallback());
+
+  callback1.WaitUntilRun();
+  callback2.WaitUntilRun();
+
+  EXPECT_TRUE(callback1.called());
+  if (callback1.called()) {
+    EXPECT_THAT(callback1.net_error(), IsOk());
+    EXPECT_THAT(callback1.leaf_index(), 123456);
+  }
+
+  EXPECT_TRUE(callback2.called());
+  if (callback2.called()) {
+    EXPECT_THAT(callback2.net_error(), IsError(net::ERR_TEMPORARILY_THROTTLED));
+    EXPECT_THAT(callback2.leaf_index(), 0);
+  }
+}
+
+TEST_P(LogDnsClientTest, CanBeThrottledToOneAuditProofQueryAtATime) {
+  const std::vector<std::string> audit_proof = GetSampleAuditProof(20);
+
+  // It should require 3 queries to collect the entire audit proof, as there is
+  // only space for 7 nodes per UDP packet.
+  mock_dns_.ExpectAuditProofRequestAndResponse("0.123456.999999.tree.ct.test.",
+                                               audit_proof.begin(),
+                                               audit_proof.begin() + 7);
+  mock_dns_.ExpectAuditProofRequestAndResponse("7.123456.999999.tree.ct.test.",
+                                               audit_proof.begin() + 7,
+                                               audit_proof.begin() + 14);
+  mock_dns_.ExpectAuditProofRequestAndResponse("14.123456.999999.tree.ct.test.",
+                                               audit_proof.begin() + 14,
+                                               audit_proof.end());
+
+  const size_t max_concurrent_queries = 1;
+  auto log_client = CreateLogDnsClient(max_concurrent_queries);
+
+  MockAuditProofCallback callback1;
+  QueryAuditProofAsync(log_client.get(), "ct.test", 123456, 999999,
+                       callback1.AsCallback());
+  MockAuditProofCallback callback2;
+  QueryAuditProofAsync(log_client.get(), "ct.test", 123456, 999999,
+                       callback2.AsCallback());
+
+  callback1.WaitUntilRun();
+  callback2.WaitUntilRun();
+
+  EXPECT_TRUE(callback1.called());
+  if (callback1.called()) {
+    EXPECT_THAT(callback1.net_error(), IsOk());
+    EXPECT_THAT(callback1.proof(), NotNull());
+    if (callback1.proof() != nullptr) {
+      EXPECT_THAT(callback1.proof()->leaf_index, 123456);
+      // EXPECT_THAT(callback1.proof()->tree_size, 999999);
+      EXPECT_THAT(callback1.proof()->nodes, audit_proof);
+    }
+  }
+
+  EXPECT_TRUE(callback2.called());
+  if (callback2.called()) {
+    EXPECT_THAT(callback2.net_error(), IsError(net::ERR_TEMPORARILY_THROTTLED));
+    EXPECT_THAT(callback2.proof(), IsNull());
+  }
+}
+
 INSTANTIATE_TEST_CASE_P(ReadMode,
                         LogDnsClientTest,
                         ::testing::Values(net::IoMode::ASYNC,