Merge LogDnsClient's QueryLeafIndex and QueryAuditProof methods

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 18bbd0a07c1d14fde6c61f7a26357036aebfdf52..814f002660345ed86af9e743da0efb8facbb392a 100644
--- a/components/certificate_transparency/log_dns_client.cc
+++ b/components/certificate_transparency/log_dns_client.cc
@@ -79,10 +79,16 @@ bool ParseAuditPath(const net::DnsResponse& response,
   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),
@@ -103,52 +109,10 @@ void LogDnsClient::OnDNSChanged() {
 
 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
@@ -156,14 +120,14 @@ void LogDnsClient::QueryLeafIndex(base::StringPiece domain_for_log,
 // 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;
   }
@@ -173,58 +137,102 @@ void LogDnsClient::QueryAuditProof(base::StringPiece domain_for_log,
         FROM_HERE,
         base::Bind(callback, net::Error::ERR_TEMPORARILY_THROTTLED, nullptr));
     return;
   }
 
+  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::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(audit_proof_queries_.begin(), audit_proof_queries_.end(),
+                   [transaction](const Query& query) {
+                     return query.transaction.get() == transaction;
+                   });
+  if (query_iterator == audit_proof_queries_.end()) {
+    NOTREACHED();
+    return;
+  }
+  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, nullptr));
+    return;
+  }
+
   std::unique_ptr<net::ct::MerkleAuditProof> proof(
       new net::ct::MerkleAuditProof);
-  proof->leaf_index = leaf_index;
+
+  if (!ParseLeafIndex(*response, &proof->leaf_index)) {
+    base::ThreadTaskRunnerHandle::Get()->PostTask(
+        FROM_HERE,
+        base::Bind(query.callback, net::ERR_DNS_MALFORMED_RESPONSE, nullptr));
+    return;
+  }
+
+  // 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.

-                       callback);
-}
-
-void LogDnsClient::QueryLeafIndexComplete(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) {
-                     return query.transaction.get() == transaction;
-                   });
-  if (query_iterator == leaf_index_queries_.end()) {
-    NOTREACHED();
-    return;
-  }
-  const Query<LeafIndexCallback> query = std::move(*query_iterator);
-  leaf_index_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));
-    return;
-  }
-
-  uint64_t leaf_index;
-  if (!ParseLeafIndex(*response, &leaf_index)) {
-    base::ThreadTaskRunnerHandle::Get()->PostTask(
-        FROM_HERE,
-        base::Bind(query.callback, net::ERR_DNS_MALFORMED_RESPONSE, 0));
-    return;
-  }
-
-  base::ThreadTaskRunnerHandle::Get()->PostTask(
-      FROM_HERE, base::Bind(query.callback, net::OK, leaf_index));
+                       query.callback);
 }
 
 void LogDnsClient::QueryAuditProofNodes(
     std::unique_ptr<net::ct::MerkleAuditProof> proof,
     base::StringPiece domain_for_log,
@@ -273,19 +281,19 @@ void LogDnsClient::QueryAuditProofNodesComplete(
   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) {
@@ -320,15 +328,12 @@ void LogDnsClient::QueryAuditProofNodesComplete(
       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);