Adds LogDnsClient::NotifyWhenNotThrottled

Allows a callback to be registered that will be invoked once the
LogDnsClient has stopped throttling queries.

BUG=624894

Review-Url: https://chromiumcodereview.appspot.com/2380293002
Cr-Commit-Position: refs/heads/master@{#426472}

components/certificate_transparency/log_dns_client.cc

@@ -408,6 +408,11 @@ void LogDnsClient::OnInitialDNSConfigRead() {
   UpdateDnsConfig();
 }
 
+void LogDnsClient::NotifyWhenNotThrottled(const base::Closure& callback) {
+  DCHECK(HasMaxConcurrentQueriesInProgress());
+  not_throttled_callbacks_.push_back(callback);
+}
+
 // |leaf_hash| is not a const-ref to allow callers to std::move that string into
 // the method, avoiding LogDnsClient::AuditProofQuery having to make a copy.
 net::Error LogDnsClient::QueryAuditProof(
@@ -454,6 +459,12 @@ void LogDnsClient::QueryAuditProofComplete(
   audit_proof_queries_.erase(query_iterator);
 
   callback.Run(net_error);
+
+  // Notify interested parties that the next query will not be throttled.
+  std::list<base::Closure> callbacks = std::move(not_throttled_callbacks_);
+  for (const base::Closure& callback : callbacks) {
+    callback.Run();
+  }
 }
 
 bool LogDnsClient::HasMaxConcurrentQueriesInProgress() const {

components/certificate_transparency/log_dns_client.h

@@ -55,6 +55,13 @@ class LogDnsClient : public net::NetworkChangeNotifier::DNSObserver {
   // The DnsClient's config will be updated in response.
   void OnInitialDNSConfigRead() override;
 
+  // Registers a callback to be invoked when the number of concurrent queries
+  // falls below the limit defined by |max_concurrent_queries| (passed to the
+  // constructor of LogDnsClient). This callback will fire once and then be
+  // unregistered. Should only be used if QueryAuditProof() returns
+  // net::ERR_TEMPORARILY_THROTTLED.
+  void NotifyWhenNotThrottled(const base::Closure& callback);
+
   // Queries a CT log to retrieve an audit proof for the leaf with |leaf_hash|.
   // The log is identified by |domain_for_log|, which is the DNS name used as a
   // suffix for all queries.
@@ -112,6 +119,8 @@ class LogDnsClient : public net::NetworkChangeNotifier::DNSObserver {
   std::list<std::unique_ptr<AuditProofQuery>> audit_proof_queries_;
   // The maximum number of queries that can be in flight at one time.
   size_t max_concurrent_queries_;
+  // Callbacks to invoke when the number of concurrent queries is at its limit.
+  std::list<base::Closure> not_throttled_callbacks_;
   // Creates weak_ptrs to this, for callback purposes.
   base::WeakPtrFactory<LogDnsClient> weak_ptr_factory_;
 

components/certificate_transparency/log_dns_client_unittest.cc

@@ -713,6 +713,62 @@ TEST_P(LogDnsClientTest, CanBeThrottledToOneQueryAtATime) {
   EXPECT_THAT(proof3.nodes, Eq(audit_proof));
 }
 
+TEST_P(LogDnsClientTest, NotifiesWhenNoLongerThrottled) {
+  const std::vector<std::string> audit_proof = GetSampleAuditProof(20);
+
+  mock_dns_.ExpectLeafIndexRequestAndResponse(kLeafIndexQnames[0], 123456);
+  mock_dns_.ExpectAuditProofRequestAndResponse("0.123456.999999.tree.ct.test.",
+                                               audit_proof.begin(),
+                                               audit_proof.begin() + 7);
+  mock_dns_.ExpectAuditProofRequestAndResponse("7.123456.999999.tree.ct.test.",
+                                               audit_proof.begin() + 7,
+                                               audit_proof.begin() + 14);
+  mock_dns_.ExpectAuditProofRequestAndResponse("14.123456.999999.tree.ct.test.",
+                                               audit_proof.begin() + 14,
+                                               audit_proof.end());
+
+  const size_t kMaxConcurrentQueries = 1;
+  std::unique_ptr<LogDnsClient> log_client =
+      CreateLogDnsClient(kMaxConcurrentQueries);
+
+  // Start a query.
+  net::ct::MerkleAuditProof proof1;
+  net::TestCompletionCallback proof_callback1;
+  ASSERT_THAT(log_client->QueryAuditProof("ct.test", kLeafHashes[0], 999999,
+                                          &proof1, proof_callback1.callback()),
+              IsError(net::ERR_IO_PENDING));
+
+  net::TestClosure not_throttled_callback;
+  log_client->NotifyWhenNotThrottled(not_throttled_callback.closure());
+
+  ASSERT_THAT(proof_callback1.WaitForResult(), IsOk());
+  not_throttled_callback.WaitForResult();
+
+  // Start another query to check |not_throttled_callback| doesn't fire again.
+  mock_dns_.ExpectLeafIndexRequestAndResponse(kLeafIndexQnames[1], 666);
+  mock_dns_.ExpectAuditProofRequestAndResponse("0.666.999999.tree.ct.test.",
+                                               audit_proof.begin(),
+                                               audit_proof.begin() + 7);
+  mock_dns_.ExpectAuditProofRequestAndResponse("7.666.999999.tree.ct.test.",
+                                               audit_proof.begin() + 7,
+                                               audit_proof.begin() + 14);
+  mock_dns_.ExpectAuditProofRequestAndResponse("14.666.999999.tree.ct.test.",
+                                               audit_proof.begin() + 14,
+                                               audit_proof.end());
+
+  net::ct::MerkleAuditProof proof2;
+  net::TestCompletionCallback proof_callback2;
+  ASSERT_THAT(log_client->QueryAuditProof("ct.test", kLeafHashes[1], 999999,
+                                          &proof2, proof_callback2.callback()),
+              IsError(net::ERR_IO_PENDING));
+
+  // Give the query a chance to run.
+  ASSERT_THAT(proof_callback2.WaitForResult(), IsOk());
+  // Give |not_throttled_callback| a chance to run - it shouldn't though.
+  base::RunLoop().RunUntilIdle();
+  ASSERT_FALSE(not_throttled_callback.have_result());
+}
+
 INSTANTIATE_TEST_CASE_P(ReadMode,
                         LogDnsClientTest,
                         ::testing::Values(net::IoMode::ASYNC,