Allow LogDnsClient queries to be rate-limited

components/certificate_transparency/log_dns_client.cc

Index: components/certificate_transparency/log_dns_client.cc
diff --git a/components/certificate_transparency/log_dns_client.cc b/components/certificate_transparency/log_dns_client.cc
index 5fb28b973eb7b431532fa677777810973ccfc2d1..1607d4637d3555faf305ee14c85e27763154afe0 100644
--- a/components/certificate_transparency/log_dns_client.cc
+++ b/components/certificate_transparency/log_dns_client.cc
@@ -82,9 +82,11 @@ bool ParseAuditPath(const net::DnsResponse& response,
 }  // namespace
 
 LogDnsClient::LogDnsClient(std::unique_ptr<net::DnsClient> dns_client,
-                           const net::BoundNetLog& net_log)
+                           const net::BoundNetLog& net_log,
+                           size_t max_concurrent_queries)
     : dns_client_(std::move(dns_client)),
       net_log_(net_log),
+      max_concurrent_queries_(max_concurrent_queries),
       weak_ptr_factory_(this) {
   CHECK(dns_client_);
   net::NetworkChangeNotifier::AddDNSObserver(this);
@@ -112,6 +114,13 @@ void LogDnsClient::QueryLeafIndex(base::StringPiece domain_for_log,
     return;
   }
 
+  if (HasMaxConcurrentQueriesInProgress()) {
+    base::ThreadTaskRunnerHandle::Get()->PostTask(
+        FROM_HERE,
+        base::Bind(callback, net::Error::ERR_TEMPORARILY_THROTTLED, 0));
+    return;
+  }
+
   std::string encoded_leaf_hash =
       base32::Base32Encode(leaf_hash, base32::Base32EncodePolicy::OMIT_PADDING);
   DCHECK_EQ(encoded_leaf_hash.size(), 52u);
@@ -146,7 +155,8 @@ void LogDnsClient::QueryLeafIndex(base::StringPiece domain_for_log,
 // of the code would increase though, as it would need to detect gaps in the
 // audit proof caused by the server not responding with the anticipated number
 // of nodes. Ownership of the proof would need to change, as it would be shared
-// between simultaneous DNS transactions.
+// between simultaneous DNS transactions. Throttling of queries would also need
+// to take into account this increase in parallelism.
 void LogDnsClient::QueryAuditProof(base::StringPiece domain_for_log,
                                    uint64_t leaf_index,
                                    uint64_t tree_size,
@@ -158,6 +168,13 @@ void LogDnsClient::QueryAuditProof(base::StringPiece domain_for_log,
     return;
   }
 
+  if (HasMaxConcurrentQueriesInProgress()) {
+    base::ThreadTaskRunnerHandle::Get()->PostTask(
+        FROM_HERE,
+        base::Bind(callback, net::Error::ERR_TEMPORARILY_THROTTLED, nullptr));
+    return;
+  }
+
   std::unique_ptr<net::ct::MerkleAuditProof> proof(
       new net::ct::MerkleAuditProof);
   proof->leaf_index = leaf_index;
@@ -304,6 +321,14 @@ void LogDnsClient::QueryAuditProofNodesComplete(
       base::Bind(query.callback, net::OK, base::Passed(std::move(proof))));
 }
 
+bool LogDnsClient::HasMaxConcurrentQueriesInProgress() const {
+  const size_t queries_in_progress =
+      leaf_index_queries_.size() + audit_proof_queries_.size();
+
+  return max_concurrent_queries_ != 0 &&
+         queries_in_progress >= max_concurrent_queries_;
+}
+
 void LogDnsClient::UpdateDnsConfig() {
   net::DnsConfig config;
   net::NetworkChangeNotifier::GetDnsConfig(&config);