Merge LogDnsClient's QueryLeafIndex and QueryAuditProof methods

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 "base/bind.h"
 #include "base/format_macros.h"
 #include "base/location.h"
 #include "base/logging.h"
@@ skipping 63 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
 
+struct LogDnsClient::Query
+{
+  std::unique_ptr<net::DnsTransaction> transaction;
+  AuditProofCallback callback;
+};
+
 LogDnsClient::LogDnsClient(std::unique_ptr<net::DnsClient> dns_client,
                            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::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, 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. 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,
+                                   base::StringPiece leaf_hash,
                                    uint64_t tree_size,
                                    const AuditProofCallback& callback) {
-  if (domain_for_log.empty() || leaf_index >= tree_size) {
+  if (domain_for_log.empty() || leaf_hash.size() != crypto::kSHA256Length) {
     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);
+  QueryLeafIndex(domain_for_log, leaf_hash, tree_size, callback);

[Eran Messeri, 2016/09/23 12:59:31]

Nit: domain_for_log and tree_size seem to always be passed around together.
You could store both in the Query struct or have a separate struct for passing
in the two.
Not sure how much it's worth the effort though.

[Rob Percival, 2016/09/27 17:44:17]

Done.

 }
 
-void LogDnsClient::QueryLeafIndexComplete(net::DnsTransaction* transaction,
+void LogDnsClient::QueryLeafIndex(base::StringPiece domain_for_log,
+                                  base::StringPiece leaf_hash,
+                                  uint64_t tree_size,
+                                  const AuditProofCallback& callback) {
+  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, nullptr));
+    return;
+  }
+
+  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(), domain_for_log, tree_size);

[Eran Messeri, 2016/09/23 12:59:31]

Here it seems like the tree_size and domain_for_log could be stored in the Query
- from what I can tell they are bound to the callback just so that they're
available for the follow-up query.

[Rob Percival, 2016/09/27 17:44:17]

Done.

+
+  std::unique_ptr<net::DnsTransaction> dns_transaction =
+      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::QueryLeafIndexComplete(base::StringPiece domain_for_log,
+                                          uint64_t tree_size,
+                                          net::DnsTransaction* transaction,
                                           int net_error,
                                           const net::DnsResponse* response) {
   auto query_iterator =
-      std::find_if(leaf_index_queries_.begin(), leaf_index_queries_.end(),
-                   [transaction](const Query<LeafIndexCallback>& query) {
+      std::find_if(audit_proof_queries_.begin(), audit_proof_queries_.end(),
+                   [transaction](const Query& query) {
                      return query.transaction.get() == transaction;
                    });
-  if (query_iterator == leaf_index_queries_.end()) {
+  if (query_iterator == audit_proof_queries_.end()) {
     NOTREACHED();
     return;
   }
-  const Query<LeafIndexCallback> query = std::move(*query_iterator);
-  leaf_index_queries_.erase(query_iterator);
+  const Query query = std::move(*query_iterator);
+  audit_proof_queries_.erase(query_iterator);
 
   // If we've received no response but no net::error either (shouldn't happen),
   // report the response as invalid.
   if (response == nullptr && net_error == net::OK) {
     net_error = net::ERR_INVALID_RESPONSE;
   }
 
   if (net_error != net::OK) {
     base::ThreadTaskRunnerHandle::Get()->PostTask(
-        FROM_HERE, base::Bind(query.callback, net_error, 0));
+        FROM_HERE, base::Bind(query.callback, net_error, nullptr));
     return;
   }
 
-  uint64_t leaf_index;
-  if (!ParseLeafIndex(*response, &leaf_index)) {
+  std::unique_ptr<net::ct::MerkleAuditProof> proof(
+      new net::ct::MerkleAuditProof);
+
+  if (!ParseLeafIndex(*response, &proof->leaf_index)) {
     base::ThreadTaskRunnerHandle::Get()->PostTask(
         FROM_HERE,
-        base::Bind(query.callback, net::ERR_DNS_MALFORMED_RESPONSE, 0));
+        base::Bind(query.callback, net::ERR_DNS_MALFORMED_RESPONSE, nullptr));
     return;
   }
 
-  base::ThreadTaskRunnerHandle::Get()->PostTask(
-      FROM_HERE, base::Bind(query.callback, net::OK, leaf_index));
+  // Reject leaf index if it is out-of-range.
+  // This indicates either:
+  // a) the wrong tree_size was provided.
+  // b) the wrong leaf hash was provided.
+  // c) there is a bug server-side.
+  // The first two are more likely, so return ERR_INVALID_ARGUMENT.
+  if (proof->leaf_index >= tree_size) {
+    base::ThreadTaskRunnerHandle::Get()->PostTask(
+        FROM_HERE,
+        base::Bind(query.callback, net::ERR_INVALID_ARGUMENT, nullptr));
+    return;
+  }
+
+  // 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,

[Eran Messeri, 2016/09/23 12:59:32]

Document the 0:
tree_size, 0 /* start node index */, ...

[Rob Percival, 2016/09/27 17:44:17]

Done.

+                       query.callback);
 }
 
 void LogDnsClient::QueryAuditProofNodes(
     std::unique_ptr<net::ct::MerkleAuditProof> proof,
     base::StringPiece domain_for_log,
     uint64_t tree_size,
     uint64_t node_index,
     const AuditProofCallback& callback) {
   // Preconditions that should be guaranteed internally by this class.
   DCHECK(proof);
@@ skipping 32 matching lines @@
     uint64_t tree_size,
     net::DnsTransaction* transaction,
     int net_error,
     const net::DnsResponse* response) {
   // Preconditions that should be guaranteed internally by this class.
   DCHECK(proof);
   DCHECK(!domain_for_log.empty());
 
   auto query_iterator =
       std::find_if(audit_proof_queries_.begin(), audit_proof_queries_.end(),
-                   [transaction](const Query<AuditProofCallback>& query) {
+                   [transaction](const Query& query) {
                      return query.transaction.get() == transaction;
                    });
 
   if (query_iterator == audit_proof_queries_.end()) {
     NOTREACHED();
     return;
   }
-  const Query<AuditProofCallback> query = std::move(*query_iterator);
+  const Query query = std::move(*query_iterator);
   audit_proof_queries_.erase(query_iterator);
 
   // If we've received no response but no net::error either (shouldn't happen),
   // report the response as invalid.
   if (response == nullptr && net_error == net::OK) {
     net_error = net::ERR_INVALID_RESPONSE;
   }
 
   if (net_error != net::OK) {
     base::ThreadTaskRunnerHandle::Get()->PostTask(
@@ skipping 18 matching lines @@
                          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_;
+         audit_proof_queries_.size() >= max_concurrent_queries_;
 }
 
 void LogDnsClient::UpdateDnsConfig() {
   net::DnsConfig config;
   net::NetworkChangeNotifier::GetDnsConfig(&config);
   if (config.IsValid())
     dns_client_->SetConfig(config);
 }
 
 }  // namespace certificate_transparency