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/location.h"
@@ skipping 64 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::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);
   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 << ".";

[Ryan Sleevi, 2016/09/16 02:11:03]

As an aside, is ostringstream the right tool here?

https://google.github.io/styleguide/cppguide.html#Streams (which used to
outright forbid them) says "Use them when they're the right tool". But this just
seems to be concatenating strings. Why not just

std::string qname = encoded_leaf_hash;
qname.append(".hash.");
qname.append(domain_for_log);
qname.append(".");

(We don't have StrCat, otherwise that'd totally be the right tool)

[Rob Percival, 2016/09/16 16:28:19]

One of the earlier iterations of this class did pretty much what you're
suggesting, but a review comment suggested changing it. I'm happy to change it
back though :) The comment was just a "nit" anyway.

On the topic of StrCat, when I first wrote this I went looking for StrCat and
was disappointed that it isn't in Chromium :(

[Ryan Sleevi, 2016/09/17 01:51:52]

Any idea what the review was? :)

My inclination is that this isn't an approved use of ostringstream, but I'd be
curious to understand the original nit.

[Rob Percival, 2016/09/20 18:54:42]

The original nit suggested using StringPrintf instead of appending, but this has
slightly different behaviour than appending or ostringstream (and that was
undesirable, though I can't recall exactly what the issue was - I think it
related to handling of NULL). The reviewer was happy with ostringstream as a
compromise. However, according to the tests, Stringprintf is ok now so I'll use
that.

 
   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,
                                  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 126 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