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..2773cb50e711f4a80a467de490b3a8fe8892b7a8 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,42 @@ class LogDnsClientTest : public ::testing::TestWithParam<net::IoMode> {
     mock_dns_.InitializeDnsConfig();
   }
 
+  std::unique_ptr<LogDnsClient> CreateLogDnsClient() {
+    return base::MakeUnique<LogDnsClient>(mock_dns_.CreateDnsClient(),
+                                          net::BoundNetLog());
+  }
+
+  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);
+  }
+
   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();
+    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);
+  }
+
   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();
+    QueryAuditProofAsync(log_client.get(), log_domain, leaf_index, tree_size,
+                         callback->AsCallback());
     callback->WaitUntilRun();
   }
 
@@ -543,6 +565,46 @@ TEST_P(LogDnsClientTest, IgnoresLatestDnsConfigIfInvalid) {
   EXPECT_THAT(dns_client->GetConfig()->nameservers, Not(IsEmpty()));
 }
 
+TEST_P(LogDnsClientTest, CanBeThrottledToOneQueryAtATime) {
+  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());
+
+  auto log_client = CreateLogDnsClient();
+  log_client->SetMaxConcurrentQueries(1);
+
+  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();
+
+  ASSERT_TRUE(callback1.called());
+  EXPECT_THAT(callback1.net_error(), IsOk());
+  ASSERT_THAT(callback1.proof(), NotNull());
+  EXPECT_THAT(callback1.proof()->leaf_index, 123456);
+  // EXPECT_THAT(callback1.proof()->tree_size, 999999);
+  EXPECT_THAT(callback1.proof()->nodes, audit_proof);
+
+  ASSERT_TRUE(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,