LogDnsClient now returns some errors synchronously

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 610e5b96ce263f98b196144b5cac0eb5c99b26ab..a8a35974abdb47edc692569d0338f204b7bea167 100644
--- a/components/certificate_transparency/log_dns_client.cc
+++ b/components/certificate_transparency/log_dns_client.cc
@@ -93,179 +93,250 @@ class LogDnsClient::AuditProofQuery {
   // to leave ownership of |dns_client| with LogDnsClient.
   AuditProofQuery(net::DnsClient* dns_client,
                   const std::string& domain_for_log,
                   uint64_t tree_size,
                   const net::BoundNetLog& net_log)
-      : domain_for_log_(domain_for_log),
+      : next_state_(State::NONE),
+        domain_for_log_(domain_for_log),
         tree_size_(tree_size),
         dns_client_(dns_client),
         net_log_(net_log),
         weak_ptr_factory_(this) {
     DCHECK(dns_client_);
     DCHECK(!domain_for_log_.empty());
   }
 
   // Start the query.
-  void Start(base::StringPiece leaf_hash, CompletionCallback callback) {
+  // If a query was already in progress, this will abort it first.

[Eran Messeri, 2016/09/28 13:44:08]

How often do you expect that case to happen?
If not at all (i.e. starting a query twice is a programming error, then DCHECK
that the next_state_ is State::NONE and then you don't have to worry about
cancelling the current_dns_transaction_.

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

It's not the intended usage (there's little reason to reuse an AuditProofQuery)
so sure, I'll just DCHECK on that.

+  int Start(base::StringPiece leaf_hash, CompletionCallback callback) {
+    // Cancel any existing DNS transaction.
     current_dns_transaction_.reset();
+    // Start a new audit proof.
     proof_.reset(new net::ct::MerkleAuditProof);
+    // The leaf hash and callback will be used later.
+    leaf_hash.CopyToString(&leaf_hash_);
     callback_ = callback;
-    QueryLeafIndex(leaf_hash);
+    // The first step in the query is to request the leaf index corresponding to
+    // |leaf_hash| from the CT log.
+    next_state_ = State::REQUEST_LEAF_INDEX;
+    // Begin the state machine.
+    return DoLoop(net::OK);
   }
 
+  // Take the audit proof obtained by the query.
+  // Should only be called once the CompletionCallback is invoked.
   std::unique_ptr<net::ct::MerkleAuditProof> TakeProof() {
     return std::move(proof_);
   }
 
  private:
-  void QueryLeafIndex(base::StringPiece leaf_hash) {
+  enum class State {
+    NONE,
+    REQUEST_LEAF_INDEX,
+    REQUEST_LEAF_INDEX_COMPLETE,
+    REQUEST_AUDIT_PROOF_NODES,
+    REQUEST_AUDIT_PROOF_NODES_COMPLETE,
+  };
+
+  int DoLoop(int net_error) {
+    CHECK_NE(State::NONE, next_state_);
+    do {
+      State state = next_state_;
+      next_state_ = State::NONE;
+      switch (state) {
+        case State::REQUEST_LEAF_INDEX:
+          net_error = RequestLeafIndex();
+          break;
+        case State::REQUEST_LEAF_INDEX_COMPLETE:
+          net_error = RequestLeafIndexComplete(net_error);
+          break;
+        case State::REQUEST_AUDIT_PROOF_NODES:
+          net_error = RequestAuditProofNodes();
+          break;
+        case State::REQUEST_AUDIT_PROOF_NODES_COMPLETE:
+          net_error = RequestAuditProofNodesComplete(net_error);
+          break;
+        case State::NONE:
+          NOTREACHED();
+          break;
+      }
+    } while (net_error != net::ERR_IO_PENDING && next_state_ != State::NONE);
+
+    return net_error;
+  }
+
+  // When a DnsTransaction completes, store the response and resume the state
+  // machine. It is safe to store a pointer to |response| because |transaction|
+  // is kept alive in |current_dns_transaction_|.
+  void OnDnsTransactionComplete(net::DnsTransaction* transaction,
+                                int net_error,
+                                const net::DnsResponse* response) {
+    last_dns_response_ = response;
+    net_error = DoLoop(net_error);
+
+    if (net_error != net::ERR_IO_PENDING) {

[Eran Messeri, 2016/09/28 13:44:08]

Nit: Document why not doing anything in case of net_error being ERR_IO_PENDING
is OK.

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

Done.

+      base::ThreadTaskRunnerHandle::Get()->PostTask(
+          FROM_HERE, base::Bind(callback_, net_error, base::Unretained(this)));
+    }
+  }
+
+  // Requests the leaf index for the CT log entry with |leaf_hash_|.
+  int RequestLeafIndex() {
     std::string encoded_leaf_hash = base32::Base32Encode(
-        leaf_hash, base32::Base32EncodePolicy::OMIT_PADDING);
+        leaf_hash_, base32::Base32EncodePolicy::OMIT_PADDING);
     DCHECK_EQ(encoded_leaf_hash.size(), 52u);
 
-    std::string qname = base::StringPrintf(
-        "%s.hash.%s.", encoded_leaf_hash.c_str(), domain_for_log_.data());
-
     net::DnsTransactionFactory* factory = dns_client_->GetTransactionFactory();
     if (factory == nullptr) {
-      base::ThreadTaskRunnerHandle::Get()->PostTask(
-          FROM_HERE,
-          base::Bind(callback_, net::Error::ERR_NAME_RESOLUTION_FAILED,
-                     base::Unretained(this)));
-      return;
+      return net::ERR_NAME_RESOLUTION_FAILED;
     }
 
+    std::string qname = base::StringPrintf(
+        "%s.hash.%s.", encoded_leaf_hash.c_str(), domain_for_log_.c_str());
+
     net::DnsTransactionFactory::CallbackType transaction_callback =

[Eran Messeri, 2016/09/28 13:44:08]

Inline transaction_callback creation & use into CreateTransaction call.

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

Done.

-        base::Bind(&LogDnsClient::AuditProofQuery::QueryLeafIndexComplete,
+        base::Bind(&LogDnsClient::AuditProofQuery::OnDnsTransactionComplete,
                    weak_ptr_factory_.GetWeakPtr());
 
     current_dns_transaction_ = factory->CreateTransaction(
         qname, net::dns_protocol::kTypeTXT, transaction_callback, net_log_);
-
     current_dns_transaction_->Start();
+
+    next_state_ = State::REQUEST_LEAF_INDEX_COMPLETE;
+    return net::ERR_IO_PENDING;
   }
 
-  void QueryLeafIndexComplete(net::DnsTransaction* transaction,
-                              int net_error,
-                              const net::DnsResponse* response) {
-    // If we've received no response but no net::error either (shouldn't
-    // happen),
+  // Stores the received leaf index in |proof_->leaf_index|.
+  // If successful, the audit proof nodes will be requested next.
+  int RequestLeafIndexComplete(int net_error) {
+    // If we receive no response but no net::error either (shouldn't happen),

[Eran Messeri, 2016/09/28 13:44:08]

"shouldn't happen" implies this is a programming error, right?
If so, it seems a DCHECK for this case is more appropriate.

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

Done.

     // report the response as invalid.
-    if (response == nullptr && net_error == net::OK) {
+    if (last_dns_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(callback_, net_error, base::Unretained(this)));
-      return;
+      return net_error;
     }
 
-    if (!ParseLeafIndex(*response, &proof_->leaf_index)) {
-      base::ThreadTaskRunnerHandle::Get()->PostTask(
-          FROM_HERE, base::Bind(callback_, net::ERR_DNS_MALFORMED_RESPONSE,
-                                base::Unretained(this)));
-      return;
+    if (!ParseLeafIndex(*last_dns_response_, &proof_->leaf_index)) {
+      return net::ERR_DNS_MALFORMED_RESPONSE;
     }
 
     // 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(callback_, net::ERR_INVALID_ARGUMENT,
-                                base::Unretained(this)));
-      return;
+      return net::ERR_INVALID_ARGUMENT;
     }
 
-    // QueryAuditProof for the first batch of audit proof_ nodes (i.e. starting
-    // from 0).
-    QueryAuditProofNodes(0 /* start node index */);
+    next_state_ = State::REQUEST_AUDIT_PROOF_NODES;
+    return net::OK;
   }
 
-  // Queries a CT log to retrieve part of an audit |proof|. The |node_index|
-  // indicates which node of the audit proof/ should be requested. The CT log
-  // may return up to 7 nodes, starting from |node_index| (this is the maximum
-  // that will fit in a DNS UDP packet). The nodes will be appended to
-  // |proof->nodes|.
-  void QueryAuditProofNodes(uint64_t node_index) {
+  // Requests the next batch of audit proof nodes from a CT log.
+  // The index of the first node required is determined by looking at how many
+  // nodes are already in |proof_->nodes|.
+  // The CT log may return up to 7 nodes - this is the maximum allowed by the
+  // CT-over-DNS draft RFC, as a TXT RDATA string can have a maximum length of
+  // 255 bytes and each node is 32 bytes long (a SHA-256 hash).
+  //
+  // The performance of this could be improved by sending all of the expected
+  // requests 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. It would also undermine LogDnsClient's ability
+  // to rate-limit DNS requests.
+  int RequestAuditProofNodes() {
     DCHECK_LT(proof_->leaf_index, tree_size_);
-    DCHECK_LT(node_index, net::ct::CalculateAuditPathLength(proof_->leaf_index,
-                                                            tree_size_));
-
-    std::string qname = base::StringPrintf(
-        "%" PRIu64 ".%" PRIu64 ".%" PRIu64 ".tree.%s.", node_index,
-        proof_->leaf_index, tree_size_, domain_for_log_.data());
+    DCHECK_LT(proof_->nodes.size(), 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,
-                     base::Unretained(this)));
-      return;
+      return net::ERR_NAME_RESOLUTION_FAILED;
     }
 
+    std::string qname = base::StringPrintf(
+        "%" PRIu64 ".%" PRIu64 ".%" PRIu64 ".tree.%s.", proof_->nodes.size(),
+        proof_->leaf_index, tree_size_, domain_for_log_.c_str());
+
     net::DnsTransactionFactory::CallbackType transaction_callback =

[Eran Messeri, 2016/09/28 13:44:08]

Ditto, inline callback creation.

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

Done.

-        base::Bind(&LogDnsClient::AuditProofQuery::QueryAuditProofNodesComplete,
+        base::Bind(&LogDnsClient::AuditProofQuery::OnDnsTransactionComplete,
                    weak_ptr_factory_.GetWeakPtr());
 
     current_dns_transaction_ = factory->CreateTransaction(
         qname, net::dns_protocol::kTypeTXT, transaction_callback, net_log_);
     current_dns_transaction_->Start();
+
+    next_state_ = State::REQUEST_AUDIT_PROOF_NODES_COMPLETE;
+    return net::ERR_IO_PENDING;
   }
 
-  void QueryAuditProofNodesComplete(net::DnsTransaction* transaction,
-                                    int net_error,
-                                    const net::DnsResponse* response) {
-    // If we receive no response but no net::error either (shouldn't happen),
+  // Appends the received audit proof nodes to |proof_->nodes|.
+  // If any nodes are missing, another request will follow this one.
+  int RequestAuditProofNodesComplete(int net_error) {
+    // If we received no response but no net::error either (shouldn't happen),
     // report the response as invalid.
-    if (response == nullptr && net_error == net::OK) {
+    if (last_dns_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(callback_, net_error, base::Unretained(this)));
-      return;
+      return net_error;
     }
 
     const uint64_t audit_path_length =
         net::ct::CalculateAuditPathLength(proof_->leaf_index, tree_size_);
+
     // The calculated |audit_path_length| can't ever be greater than 64, so
     // deriving the amount of memory to reserve from the untrusted |leaf_index|
     // is safe.
     proof_->nodes.reserve(audit_path_length);
 
-    if (!ParseAuditPath(*response, proof_.get())) {
-      base::ThreadTaskRunnerHandle::Get()->PostTask(
-          FROM_HERE, base::Bind(callback_, net::ERR_DNS_MALFORMED_RESPONSE,
-                                base::Unretained(this)));
-      return;
+    if (!ParseAuditPath(*last_dns_response_, proof_.get())) {
+      return net::ERR_DNS_MALFORMED_RESPONSE;
     }
 
-    const uint64_t audit_path_nodes_received = proof_->nodes.size();
-    if (audit_path_nodes_received < audit_path_length) {
-      QueryAuditProofNodes(audit_path_nodes_received);
-      return;
+    // If we don't have all of the proof nodes yet, request more.
+    if (proof_->nodes.size() < audit_path_length) {
+      next_state_ = State::REQUEST_AUDIT_PROOF_NODES;
     }
 
-    base::ThreadTaskRunnerHandle::Get()->PostTask(
-        FROM_HERE, base::Bind(callback_, net::OK, base::Unretained(this)));
+    return net::OK;
   }
 
+  // The next state that this query will enter.
+  State next_state_;
+  // The DNS domain of the CT log that is being queried.
   std::string domain_for_log_;
+  // The Merkle leaf hash of the CT log entry an audit proof is required for.
+  std::string leaf_hash_;
+  // The size of the CT log's tree.

[Eran Messeri, 2016/09/28 13:44:08]

Add: from which the proof is requested. 

   // TODO(robpercival): Remove |tree_size| once |proof_| has a tree_size member.
   uint64_t tree_size_;
+  // The audit proof. It will be null until the query is started and incomplete
+  // until the query is finished.
   std::unique_ptr<net::ct::MerkleAuditProof> proof_;
+  // The callback to invoke when the query is complete.
   CompletionCallback callback_;
+  // The DnsClient to use for sending DNS requests to the CT log.
   net::DnsClient* dns_client_;
+  // The most recent DNS request. Null if no DNS requests have been made.
   std::unique_ptr<net::DnsTransaction> current_dns_transaction_;
+  // The most recent DNS response. Only valid so long as the corresponding DNS
+  // request is stored in |current_dns_transaction_|.
+  const net::DnsResponse* last_dns_response_;
+  // The NetLog that DNS transactions will log to.
   net::BoundNetLog net_log_;
+  // Produces WeakPtrs to |this| for binding callbacks.
   base::WeakPtrFactory<AuditProofQuery> weak_ptr_factory_;
 };
 
 LogDnsClient::LogDnsClient(std::unique_ptr<net::DnsClient> dns_client,
                            const net::BoundNetLog& net_log,
@@ -289,46 +360,32 @@ void LogDnsClient::OnDNSChanged() {
 
 void LogDnsClient::OnInitialDNSConfigRead() {
   UpdateDnsConfig();
 }
 
-// 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(const std::string& domain_for_log,
-                                   base::StringPiece leaf_hash,
-                                   uint64_t tree_size,
-                                   const AuditProofCallback& callback) {
+int LogDnsClient::QueryAuditProof(base::StringPiece domain_for_log,
+                                  base::StringPiece leaf_hash,
+                                  uint64_t tree_size,
+                                  const AuditProofCallback& 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, nullptr));
-    return;
+    return net::ERR_INVALID_ARGUMENT;
   }
 
   if (HasMaxConcurrentQueriesInProgress()) {
-    base::ThreadTaskRunnerHandle::Get()->PostTask(
-        FROM_HERE,
-        base::Bind(callback, net::Error::ERR_TEMPORARILY_THROTTLED, nullptr));
-    return;
+    return net::ERR_TEMPORARILY_THROTTLED;
   }
 
-  audit_proof_queries_.emplace_back(
-      new AuditProofQuery(dns_client_.get(), domain_for_log, tree_size));
+  AuditProofQuery* query = new AuditProofQuery(
+      dns_client_.get(), domain_for_log.as_string(), tree_size, net_log_);
+  // Transfers ownership of |query| to |audit_proof_queries_|.
+  audit_proof_queries_.emplace_back(query);
 
   AuditProofQuery::CompletionCallback internal_callback =
       base::Bind(&LogDnsClient::QueryAuditProofComplete,
                  weak_ptr_factory_.GetWeakPtr(), callback);
 
-  audit_proof_queries_.back()->Start(leaf_hash, internal_callback);
+  return query->Start(leaf_hash, internal_callback);
 }
 
 void LogDnsClient::QueryAuditProofComplete(const AuditProofCallback& callback,
                                            int result,
                                            AuditProofQuery* query) {