Allow LogDnsClient queries to be rate-limited

components/certificate_transparency/log_dns_client.cc

 // Copyright 2016 The Chromium Authors. All rights reserved.
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 
 #include "components/certificate_transparency/log_dns_client.h"
 
-#include <sstream>
-
 #include "base/bind.h"
+#include "base/format_macros.h"
 #include "base/location.h"
 #include "base/logging.h"
 #include "base/strings/string_number_conversions.h"
 #include "base/strings/string_util.h"
+#include "base/strings/stringprintf.h"
 #include "base/threading/thread_task_runner_handle.h"
 #include "base/time/time.h"
 #include "components/base32/base32.h"
 #include "crypto/sha2.h"
 #include "net/base/net_errors.h"
 #include "net/cert/merkle_audit_proof.h"
 #include "net/dns/dns_client.h"
 #include "net/dns/dns_config_service.h"
 #include "net/dns/dns_protocol.h"
 #include "net/dns/dns_response.h"
@@ skipping 51 matching lines @@
   for (size_t i = 0; i < audit_path.size(); i += crypto::kSHA256Length) {
     proof->nodes.push_back(audit_path.substr(i, crypto::kSHA256Length));
   }
 
   return true;
 }
 
 }  // namespace
 
 LogDnsClient::LogDnsClient(std::unique_ptr<net::DnsClient> dns_client,
-                           const net::NetLogWithSource& net_log)
+                           const net::NetLogWithSource& 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);
   UpdateDnsConfig();
 }
 
 LogDnsClient::~LogDnsClient() {
   net::NetworkChangeNotifier::RemoveDNSObserver(this);
 }
 
 void LogDnsClient::OnDNSChanged() {
   UpdateDnsConfig();
 }
 
 void LogDnsClient::OnInitialDNSConfigRead() {
   UpdateDnsConfig();
 }
 
 void LogDnsClient::QueryLeafIndex(base::StringPiece domain_for_log,
                                   base::StringPiece leaf_hash,
                                   const LeafIndexCallback& callback) {
   if (domain_for_log.empty() || leaf_hash.size() != crypto::kSHA256Length) {
     base::ThreadTaskRunnerHandle::Get()->PostTask(
         FROM_HERE, base::Bind(callback, net::Error::ERR_INVALID_ARGUMENT, 0));
     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);
 
   net::DnsTransactionFactory* factory = dns_client_->GetTransactionFactory();
   if (factory == nullptr) {
     base::ThreadTaskRunnerHandle::Get()->PostTask(
         FROM_HERE,
         base::Bind(callback, net::Error::ERR_NAME_RESOLUTION_FAILED, 0));
     return;
   }
 
-  std::ostringstream qname;
-  qname << encoded_leaf_hash << ".hash." << domain_for_log << ".";
+  std::string qname = base::StringPrintf(
+      "%s.hash.%s.", encoded_leaf_hash.c_str(), domain_for_log.data());
 
   net::DnsTransactionFactory::CallbackType transaction_callback = base::Bind(
       &LogDnsClient::QueryLeafIndexComplete, weak_ptr_factory_.GetWeakPtr());
 
   std::unique_ptr<net::DnsTransaction> dns_transaction =
-      factory->CreateTransaction(qname.str(), net::dns_protocol::kTypeTXT,
+      factory->CreateTransaction(qname, net::dns_protocol::kTypeTXT,
                                  transaction_callback, net_log_);
 
   dns_transaction->Start();
   leaf_index_queries_.push_back({std::move(dns_transaction), callback});
 }
 
 // The performance of this could be improved by sending all of the expected
 // queries up front. Each response can contain a maximum of 7 audit path nodes,
 // so for an audit proof of size 20, it could send 3 queries (for nodes 0-6,
 // 7-13 and 14-19) immediately. Currently, it sends only the first and then,
 // based on the number of nodes received, sends the next query. The complexity
 // 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,
                                    const AuditProofCallback& callback) {
   if (domain_for_log.empty() || leaf_index >= tree_size) {
     base::ThreadTaskRunnerHandle::Get()->PostTask(
         FROM_HERE,
         base::Bind(callback, net::Error::ERR_INVALID_ARGUMENT, nullptr));
     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;
   // TODO(robpercival): Once a "tree_size" field is added to MerkleAuditProof,
   // pass |tree_size| to QueryAuditProofNodes using that.
 
   // Query for the first batch of audit proof nodes (i.e. starting from 0).
   QueryAuditProofNodes(std::move(proof), domain_for_log, tree_size, 0,
                        callback);
 }
@@ skipping 51 matching lines @@
             net::ct::CalculateAuditPathLength(proof->leaf_index, tree_size));
 
   net::DnsTransactionFactory* factory = dns_client_->GetTransactionFactory();
   if (factory == nullptr) {
     base::ThreadTaskRunnerHandle::Get()->PostTask(
         FROM_HERE,
         base::Bind(callback, net::Error::ERR_NAME_RESOLUTION_FAILED, nullptr));
     return;
   }
 
-  std::ostringstream qname;
-  qname << node_index << "." << proof->leaf_index << "." << tree_size
-        << ".tree." << domain_for_log << ".";
+  std::string qname = base::StringPrintf(
+      "%" PRIu64 ".%" PRIu64 ".%" PRIu64 ".tree.%s.", node_index,
+      proof->leaf_index, tree_size, domain_for_log.data());
 
   net::DnsTransactionFactory::CallbackType transaction_callback =
       base::Bind(&LogDnsClient::QueryAuditProofNodesComplete,
                  weak_ptr_factory_.GetWeakPtr(), base::Passed(std::move(proof)),
                  domain_for_log, tree_size);
 
   std::unique_ptr<net::DnsTransaction> dns_transaction =
-      factory->CreateTransaction(qname.str(), net::dns_protocol::kTypeTXT,
+      factory->CreateTransaction(qname, net::dns_protocol::kTypeTXT,
                                  transaction_callback, net_log_);
   dns_transaction->Start();
   audit_proof_queries_.push_back({std::move(dns_transaction), callback});
 }
 
 void LogDnsClient::QueryAuditProofNodesComplete(
     std::unique_ptr<net::ct::MerkleAuditProof> proof,
     base::StringPiece domain_for_log,
     uint64_t tree_size,
     net::DnsTransaction* transaction,
@@ skipping 44 matching lines @@
     QueryAuditProofNodes(std::move(proof), domain_for_log, tree_size,
                          audit_path_nodes_received, query.callback);
     return;
   }
 
   base::ThreadTaskRunnerHandle::Get()->PostTask(
       FROM_HERE,
       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);
   if (config.IsValid())
     dns_client_->SetConfig(config);
 }
 
 }  // namespace certificate_transparency