LogDnsClient now returns some errors synchronously

components/certificate_transparency/log_dns_client_unittest.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 <memory>
 #include <string>
 #include <utility>
 #include <vector>
 
 #include "base/format_macros.h"
 #include "base/memory/ptr_util.h"
 #include "base/message_loop/message_loop.h"
 #include "base/run_loop.h"
 #include "base/strings/string_number_conversions.h"
 #include "base/strings/stringprintf.h"
+#include "base/test/test_timeouts.h"
 #include "components/certificate_transparency/mock_log_dns_traffic.h"
 #include "crypto/sha2.h"
 #include "net/base/net_errors.h"
 #include "net/cert/merkle_audit_proof.h"
 #include "net/cert/signed_certificate_timestamp.h"
 #include "net/dns/dns_client.h"
 #include "net/dns/dns_config_service.h"
 #include "net/dns/dns_protocol.h"
 #include "net/log/net_log.h"
 #include "net/test/gtest_util.h"
@@ skipping 35 matching lines @@
 const uint64_t kLeafIndices[] = {0, 1, 2};
 const uint64_t kTreeSizes[] = {100, 10000, 1000000};
 
 // Only 7 audit proof nodes can fit into a DNS response, because they are sent
 // in a TXT RDATA string, which has a maximum size of 255 bytes, and each node
 // is a SHA-256 hash (32 bytes), i.e. (255 / 32) == 7.
 // This means audit proofs consisting of more than 7 nodes require multiple DNS
 // requests to retrieve.
 const size_t kMaxProofNodesPerDnsResponse = 7;
 
+// Returns an example Merkle audit proof containing |length| nodes.
+// The proof cannot be used for cryptographic purposes; it is merely a
+// placeholder.
 std::vector<std::string> GetSampleAuditProof(size_t length) {
   std::vector<std::string> audit_proof(length);
   // Makes each node of the audit proof different, so that tests are able to
   // confirm that the audit proof is reconstructed in the correct order.
   for (size_t i = 0; i < length; ++i) {
     std::string node(crypto::kSHA256Length, '\0');
     // Each node is 32 bytes, with each byte having a different value.
     for (size_t j = 0; j < crypto::kSHA256Length; ++j) {
       node[j] = static_cast<char>((-127 + i + j) % 128);
     }
     audit_proof[i].assign(std::move(node));
   }
 
   return audit_proof;
 }
 
+// MockAuditProofCallback can be used as an AuditProofCallback.
+// It will record the arguments it is invoked with and provides a helpful
+// method for pumping the message loop until it is invoked.
 class MockAuditProofCallback {
  public:
   MockAuditProofCallback() : called_(false) {}
 
   bool called() const { return called_; }
   int net_error() const { return net_error_; }
   const net::ct::MerkleAuditProof* proof() const { return proof_.get(); }
 
+  // Get this callback as an AuditProofCallback.
+  LogDnsClient::AuditProofCallback AsCallback() {
+    return base::Bind(&MockAuditProofCallback::Run, base::Unretained(this));
+  }
+
+  // Wait until either the callback is invoked or the message loop goes idle
+  // (after a specified |timeout|). Returns immediately if the callback has
+  // already been invoked.
+  void WaitUntilRun(base::TimeDelta timeout) {
+    if (called_) {
+      return;
+    }
+
+    // Pump the message loop until the the callback is invoked, which quits the
+    // RunLoop, or a timeout expires and the message loop goes idle.
+    run_loop_.reset(new base::RunLoop());
+    base::Closure quit_closure = run_loop_->QuitWhenIdleClosure();
+    base::ThreadTaskRunnerHandle::Get()->PostDelayedTask(FROM_HERE,
+                                                         quit_closure, timeout);
+    run_loop_->Run();
+    run_loop_.reset();
+  }
+
+ private:
   void Run(int net_error, std::unique_ptr<net::ct::MerkleAuditProof> proof) {
     EXPECT_TRUE(!called_) << "Callback invoked more than once";
     called_ = true;
     net_error_ = net_error;
     proof_ = std::move(proof);
-    run_loop_.Quit();
+    if (run_loop_) {
+      run_loop_->Quit();
+    }
   }
 
-  LogDnsClient::AuditProofCallback AsCallback() {
-    return base::Bind(&MockAuditProofCallback::Run, base::Unretained(this));
-  }
-
-  void WaitUntilRun() { run_loop_.Run(); }
-
- private:
+  // True if the callback has been invoked.
   bool called_;
+  // The arguments that the callback was invoked with.
   int net_error_;
   std::unique_ptr<net::ct::MerkleAuditProof> proof_;
-  base::RunLoop run_loop_;
+  // The RunLoop currently being used to pump the message loop, as a means to
+  // execute this callback.
+  std::unique_ptr<base::RunLoop> run_loop_;
 };
 
 class LogDnsClientTest : public ::testing::TestWithParam<net::IoMode> {
  protected:
   LogDnsClientTest()
       : network_change_notifier_(net::NetworkChangeNotifier::CreateMock()) {
     mock_dns_.SetSocketReadMode(GetParam());
     mock_dns_.InitializeDnsConfig();
   }
 
   std::unique_ptr<LogDnsClient> CreateLogDnsClient(
       size_t kMaxConcurrentQueries) {
     return base::MakeUnique<LogDnsClient>(
         mock_dns_.CreateDnsClient(), net::BoundNetLog(), kMaxConcurrentQueries);
   }
 
-  void QueryAuditProofAsync(LogDnsClient* log_client,
-                            const std::string& log_domain,
-                            const char leaf_hash[crypto::kSHA256Length],
-                            uint64_t tree_size,
-                            const LogDnsClient::AuditProofCallback& callback) {
-    log_client->QueryAuditProof(log_domain, leaf_hash, tree_size, callback);
+  int QueryAuditProofAsync(LogDnsClient* log_client,
+                           const std::string& log_domain,
+                           const char leaf_hash[crypto::kSHA256Length],
+                           uint64_t tree_size,
+                           const LogDnsClient::AuditProofCallback& callback) {
+    return log_client->QueryAuditProof(log_domain, leaf_hash, tree_size,
+                                       callback);
   }
 
-  // Convenience function for calling QueryAuditProofAsync synchronously.
-  void QueryAuditProof(const std::string& log_domain,
-                       const char leaf_hash[crypto::kSHA256Length],
-                       uint64_t tree_size,
-                       MockAuditProofCallback* callback) {
+  // Convenience function for calling ASSERT_THAT(QueryAuditProofAsync
+  // synchronously.
+  int QueryAuditProof(const std::string& log_domain,
+                      const char leaf_hash[crypto::kSHA256Length],
+                      uint64_t tree_size,
+                      MockAuditProofCallback* callback) {
     std::unique_ptr<LogDnsClient> log_client = CreateLogDnsClient(0);
-    QueryAuditProofAsync(log_client.get(), log_domain, leaf_hash, tree_size,
-                         callback->AsCallback());
-    callback->WaitUntilRun();
+    int result = QueryAuditProofAsync(log_client.get(), log_domain, leaf_hash,
+                                      tree_size, callback->AsCallback());
+    if (result == net::ERR_IO_PENDING) {
+      callback->WaitUntilRun(TestTimeouts::action_max_timeout());
+    } else {

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

It is unclear to me why the result from QueryAuditProofAsync indicates that the
callback is not expected to be called.

I suggest an additional flag to the method (or a NULL callback) to indicate this
is indeed the expected behaviour (you could then add a callback created in this
function to ensure it was not called)

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

If LogDnsClient::QueryAuditProof returns ERR_IO_PENDING, it's declaring that
it's continuing asynchronously and the callback will be invoked. Any other
net::Error indicates otherwise, so there's so point waiting for the callback. It
doesn't matter if the test is expecting it to be called because the test will
fail immediately due to an unexpected return value.

+      // The callback isn't expected to be invoked, but pump the message loop
+      // anyway to confirm that it indeed isn't run. Give up as soon as the
+      // message loop is empty (zero timeout).
+      callback->WaitUntilRun(base::TimeDelta());
+    }
+    return result;
   }
 
   // This will be the NetworkChangeNotifier singleton for the duration of the
   // test. It is accessed statically by LogDnsClient.
   std::unique_ptr<net::NetworkChangeNotifier> network_change_notifier_;
   // Queues and handles asynchronous DNS tasks. Indirectly used by LogDnsClient,
   // the underlying net::DnsClient, and NetworkChangeNotifier.
   base::MessageLoopForIO message_loop_;
   // Allows mock DNS sockets to be setup.
   MockLogDnsTraffic mock_dns_;
 };
 
 TEST_P(LogDnsClientTest, QueryAuditProofReportsThatLogDomainDoesNotExist) {
   mock_dns_.ExpectRequestAndErrorResponse(kLeafIndexQnames[0],
                                           net::dns_protocol::kRcodeNXDOMAIN);
 
   MockAuditProofCallback callback;
-  QueryAuditProof("ct.test", kLeafHashes[0], kTreeSizes[0], &callback);
+  ASSERT_THAT(
+      QueryAuditProof("ct.test", kLeafHashes[0], kTreeSizes[0], &callback),
+      IsError(net::ERR_IO_PENDING));
   ASSERT_TRUE(callback.called());
   EXPECT_THAT(callback.net_error(), IsError(net::ERR_NAME_NOT_RESOLVED));
   EXPECT_THAT(callback.proof(), IsNull());
 }
 
 TEST_P(LogDnsClientTest,
        QueryAuditProofReportsServerFailuresDuringLeafIndexRequests) {
   mock_dns_.ExpectRequestAndErrorResponse(kLeafIndexQnames[0],
                                           net::dns_protocol::kRcodeSERVFAIL);
 
   MockAuditProofCallback callback;
-  QueryAuditProof("ct.test", kLeafHashes[0], kTreeSizes[0], &callback);
+  ASSERT_THAT(
+      QueryAuditProof("ct.test", kLeafHashes[0], kTreeSizes[0], &callback),
+      IsError(net::ERR_IO_PENDING));
   ASSERT_TRUE(callback.called());
   EXPECT_THAT(callback.net_error(), IsError(net::ERR_DNS_SERVER_FAILED));
   EXPECT_THAT(callback.proof(), IsNull());
 }
 
 TEST_P(LogDnsClientTest,
        QueryAuditProofReportsServerRefusalsDuringLeafIndexRequests) {
   mock_dns_.ExpectRequestAndErrorResponse(kLeafIndexQnames[0],
                                           net::dns_protocol::kRcodeREFUSED);
 
   MockAuditProofCallback callback;
-  QueryAuditProof("ct.test", kLeafHashes[0], kTreeSizes[0], &callback);
+  ASSERT_THAT(
+      QueryAuditProof("ct.test", kLeafHashes[0], kTreeSizes[0], &callback),
+      IsError(net::ERR_IO_PENDING));
   ASSERT_TRUE(callback.called());
   EXPECT_THAT(callback.net_error(), IsError(net::ERR_DNS_SERVER_FAILED));
   EXPECT_THAT(callback.proof(), IsNull());
 }
 
 TEST_P(LogDnsClientTest,
        QueryAuditProofReportsMalformedResponseIfLeafIndexIsNotNumeric) {
   mock_dns_.ExpectLeafIndexRequestAndResponse(kLeafIndexQnames[0], "foo");
 
   MockAuditProofCallback callback;
-  QueryAuditProof("ct.test", kLeafHashes[0], kTreeSizes[0], &callback);
+  ASSERT_THAT(
+      QueryAuditProof("ct.test", kLeafHashes[0], kTreeSizes[0], &callback),
+      IsError(net::ERR_IO_PENDING));
   ASSERT_TRUE(callback.called());
   EXPECT_THAT(callback.net_error(), IsError(net::ERR_DNS_MALFORMED_RESPONSE));
   EXPECT_THAT(callback.proof(), IsNull());
 }
 
 TEST_P(LogDnsClientTest,
        QueryAuditProofReportsMalformedResponseIfLeafIndexIsFloatingPoint) {
   mock_dns_.ExpectLeafIndexRequestAndResponse(kLeafIndexQnames[0], "123456.0");
 
   MockAuditProofCallback callback;
-  QueryAuditProof("ct.test", kLeafHashes[0], kTreeSizes[0], &callback);
+  ASSERT_THAT(
+      QueryAuditProof("ct.test", kLeafHashes[0], kTreeSizes[0], &callback),
+      IsError(net::ERR_IO_PENDING));
   ASSERT_TRUE(callback.called());
   EXPECT_THAT(callback.net_error(), IsError(net::ERR_DNS_MALFORMED_RESPONSE));
   EXPECT_THAT(callback.proof(), IsNull());
 }
 
 TEST_P(LogDnsClientTest,
        QueryAuditProofReportsMalformedResponseIfLeafIndexIsEmpty) {
   mock_dns_.ExpectLeafIndexRequestAndResponse(kLeafIndexQnames[0], "");
 
   MockAuditProofCallback callback;
-  QueryAuditProof("ct.test", kLeafHashes[0], kTreeSizes[0], &callback);
+  ASSERT_THAT(
+      QueryAuditProof("ct.test", kLeafHashes[0], kTreeSizes[0], &callback),
+      IsError(net::ERR_IO_PENDING));
   ASSERT_TRUE(callback.called());
   EXPECT_THAT(callback.net_error(), IsError(net::ERR_DNS_MALFORMED_RESPONSE));
   EXPECT_THAT(callback.proof(), IsNull());
 }
 
 TEST_P(LogDnsClientTest,
        QueryAuditProofReportsMalformedResponseIfLeafIndexHasNonNumericPrefix) {
   mock_dns_.ExpectLeafIndexRequestAndResponse(kLeafIndexQnames[0], "foo123456");
 
   MockAuditProofCallback callback;
-  QueryAuditProof("ct.test", kLeafHashes[0], kTreeSizes[0], &callback);
+  ASSERT_THAT(
+      QueryAuditProof("ct.test", kLeafHashes[0], kTreeSizes[0], &callback),
+      IsError(net::ERR_IO_PENDING));
   ASSERT_TRUE(callback.called());
   EXPECT_THAT(callback.net_error(), IsError(net::ERR_DNS_MALFORMED_RESPONSE));
   EXPECT_THAT(callback.proof(), IsNull());
 }
 
 TEST_P(LogDnsClientTest,
        QueryAuditProofReportsMalformedResponseIfLeafIndexHasNonNumericSuffix) {
   mock_dns_.ExpectLeafIndexRequestAndResponse(kLeafIndexQnames[0], "123456foo");
 
   MockAuditProofCallback callback;
-  QueryAuditProof("ct.test", kLeafHashes[0], kTreeSizes[0], &callback);
+  ASSERT_THAT(
+      QueryAuditProof("ct.test", kLeafHashes[0], kTreeSizes[0], &callback),
+      IsError(net::ERR_IO_PENDING));
   ASSERT_TRUE(callback.called());
   EXPECT_THAT(callback.net_error(), IsError(net::ERR_DNS_MALFORMED_RESPONSE));
   EXPECT_THAT(callback.proof(), IsNull());
 }
 
 TEST_P(LogDnsClientTest, QueryAuditProofReportsInvalidArgIfLogDomainIsEmpty) {
   MockAuditProofCallback callback;
-  QueryAuditProof("", kLeafHashes[0], kTreeSizes[0], &callback);
-  ASSERT_TRUE(callback.called());
-  EXPECT_THAT(callback.net_error(), IsError(net::ERR_INVALID_ARGUMENT));
-  EXPECT_THAT(callback.proof(), IsNull());
+  ASSERT_THAT(QueryAuditProof("", kLeafHashes[0], kTreeSizes[0], &callback),
+              IsError(net::ERR_INVALID_ARGUMENT));
+  ASSERT_FALSE(callback.called());
 }
 
 TEST_P(LogDnsClientTest, QueryAuditProofReportsInvalidArgIfLeafHashIsInvalid) {
   MockAuditProofCallback callback;
-  QueryAuditProof("ct.test", "foo", kTreeSizes[0], &callback);
-  ASSERT_TRUE(callback.called());
-  EXPECT_THAT(callback.net_error(), IsError(net::ERR_INVALID_ARGUMENT));
-  EXPECT_THAT(callback.proof(), IsNull());
+  ASSERT_THAT(QueryAuditProof("ct.test", "foo", kTreeSizes[0], &callback),
+              IsError(net::ERR_INVALID_ARGUMENT));
+  ASSERT_FALSE(callback.called());
 }
 
 TEST_P(LogDnsClientTest, QueryAuditProofReportsInvalidArgIfLeafHashIsEmpty) {
   MockAuditProofCallback callback;
-  QueryAuditProof("ct.test", "", kTreeSizes[0], &callback);
-  ASSERT_TRUE(callback.called());
-  EXPECT_THAT(callback.net_error(), IsError(net::ERR_INVALID_ARGUMENT));
-  EXPECT_THAT(callback.proof(), IsNull());
+  ASSERT_THAT(QueryAuditProof("ct.test", "", kTreeSizes[0], &callback),
+              IsError(net::ERR_INVALID_ARGUMENT));
+  ASSERT_FALSE(callback.called());
 }
 
 TEST_P(LogDnsClientTest, QueryAuditProofReportsInvalidArgIfLeafHashIsNull) {
   MockAuditProofCallback callback;
-  QueryAuditProof("ct.test", nullptr, kTreeSizes[0], &callback);
-  ASSERT_TRUE(callback.called());
-  EXPECT_THAT(callback.net_error(), IsError(net::ERR_INVALID_ARGUMENT));
-  EXPECT_THAT(callback.proof(), IsNull());
+  ASSERT_THAT(QueryAuditProof("ct.test", nullptr, kTreeSizes[0], &callback),
+              IsError(net::ERR_INVALID_ARGUMENT));
+  ASSERT_FALSE(callback.called());
 }
 
 TEST_P(LogDnsClientTest,
        QueryAuditProofReportsSocketErrorsDuringLeafIndexRequests) {
   mock_dns_.ExpectRequestAndSocketError(kLeafIndexQnames[0],
                                         net::ERR_CONNECTION_REFUSED);
 
   MockAuditProofCallback callback;
-  QueryAuditProof("ct.test", kLeafHashes[0], kTreeSizes[0], &callback);
+  ASSERT_THAT(
+      QueryAuditProof("ct.test", kLeafHashes[0], kTreeSizes[0], &callback),
+      IsError(net::ERR_IO_PENDING));
   ASSERT_TRUE(callback.called());
   EXPECT_THAT(callback.net_error(), IsError(net::ERR_CONNECTION_REFUSED));
   EXPECT_THAT(callback.proof(), IsNull());
 }
 
 TEST_P(LogDnsClientTest,
        QueryAuditProofReportsTimeoutsDuringLeafIndexRequests) {
   mock_dns_.ExpectRequestAndTimeout(kLeafIndexQnames[0]);
 
   MockAuditProofCallback callback;
-  QueryAuditProof("ct.test", kLeafHashes[0], kTreeSizes[0], &callback);
+  ASSERT_THAT(
+      QueryAuditProof("ct.test", kLeafHashes[0], kTreeSizes[0], &callback),
+      IsError(net::ERR_IO_PENDING));
+
+  callback.WaitUntilRun(TestTimeouts::action_max_timeout());
   ASSERT_TRUE(callback.called());
   EXPECT_THAT(callback.net_error(), IsError(net::ERR_DNS_TIMED_OUT));
   EXPECT_THAT(callback.proof(), IsNull());
 }
 
 TEST_P(LogDnsClientTest, QueryAuditProof) {
   const std::vector<std::string> audit_proof = GetSampleAuditProof(20);
 
   // Expect a leaf index query first, to map the leaf hash to a leaf index.
   mock_dns_.ExpectLeafIndexRequestAndResponse(kLeafIndexQnames[0], "123456");
 
   // It takes a number of DNS requests to retrieve the entire |audit_proof|
   // (see |kMaxProofNodesPerDnsResponse|).
   for (size_t nodes_begin = 0; nodes_begin < audit_proof.size();
        nodes_begin += kMaxProofNodesPerDnsResponse) {
     const size_t nodes_end = std::min(
         nodes_begin + kMaxProofNodesPerDnsResponse, audit_proof.size());
 
     mock_dns_.ExpectAuditProofRequestAndResponse(
         base::StringPrintf("%zu.123456.999999.tree.ct.test.", nodes_begin),
         audit_proof.begin() + nodes_begin, audit_proof.begin() + nodes_end);
   }
 
   MockAuditProofCallback callback;
-  QueryAuditProof("ct.test", kLeafHashes[0], 999999, &callback);
+  ASSERT_THAT(QueryAuditProof("ct.test", kLeafHashes[0], 999999, &callback),
+              IsError(net::ERR_IO_PENDING));
   ASSERT_TRUE(callback.called());
   EXPECT_THAT(callback.net_error(), IsOk());
   ASSERT_THAT(callback.proof(), NotNull());
   EXPECT_THAT(callback.proof()->leaf_index, Eq(123456u));
   // TODO(robpercival): Enable this once MerkleAuditProof has tree_size.
   // EXPECT_THAT(callback.proof()->tree_size, Eq(999999));
   EXPECT_THAT(callback.proof()->nodes, Eq(audit_proof));
 }
 
 TEST_P(LogDnsClientTest, QueryAuditProofHandlesResponsesWithShortAuditPaths) {
@@ skipping 16 matching lines @@
                                                audit_proof.begin() + 6,
                                                audit_proof.begin() + 10);
   mock_dns_.ExpectAuditProofRequestAndResponse("10.123456.999999.tree.ct.test.",
                                                audit_proof.begin() + 10,
                                                audit_proof.begin() + 13);
   mock_dns_.ExpectAuditProofRequestAndResponse("13.123456.999999.tree.ct.test.",
                                                audit_proof.begin() + 13,
                                                audit_proof.end());
 
   MockAuditProofCallback callback;
-  QueryAuditProof("ct.test", kLeafHashes[0], 999999, &callback);
+  ASSERT_THAT(QueryAuditProof("ct.test", kLeafHashes[0], 999999, &callback),
+              IsError(net::ERR_IO_PENDING));
   ASSERT_TRUE(callback.called());
   EXPECT_THAT(callback.net_error(), IsOk());
   ASSERT_THAT(callback.proof(), NotNull());
   EXPECT_THAT(callback.proof()->leaf_index, Eq(123456u));
   // TODO(robpercival): Enable this once MerkleAuditProof has tree_size.
   // EXPECT_THAT(callback.proof()->tree_size, Eq(999999));
   EXPECT_THAT(callback.proof()->nodes, Eq(audit_proof));
 }
 
 TEST_P(LogDnsClientTest,
        QueryAuditProofReportsThatAuditProofQnameDoesNotExist) {
   mock_dns_.ExpectLeafIndexRequestAndResponse(kLeafIndexQnames[0], "123456");
   mock_dns_.ExpectRequestAndErrorResponse("0.123456.999999.tree.ct.test.",
                                           net::dns_protocol::kRcodeNXDOMAIN);
 
   MockAuditProofCallback callback;
-  QueryAuditProof("ct.test", kLeafHashes[0], 999999, &callback);
+  ASSERT_THAT(QueryAuditProof("ct.test", kLeafHashes[0], 999999, &callback),
+              IsError(net::ERR_IO_PENDING));
   ASSERT_TRUE(callback.called());
   EXPECT_THAT(callback.net_error(), IsError(net::ERR_NAME_NOT_RESOLVED));
   EXPECT_THAT(callback.proof(), IsNull());
 }
 
 TEST_P(LogDnsClientTest,
-       QueryAuditProofReportsServerFailureDuringAuditProofRequests) {
+       QueryAuditProofReportsServerFailuresDuringAuditProofRequests) {
   mock_dns_.ExpectLeafIndexRequestAndResponse(kLeafIndexQnames[0], "123456");
   mock_dns_.ExpectRequestAndErrorResponse("0.123456.999999.tree.ct.test.",
                                           net::dns_protocol::kRcodeSERVFAIL);
 
   MockAuditProofCallback callback;
-  QueryAuditProof("ct.test", kLeafHashes[0], 999999, &callback);
+  ASSERT_THAT(QueryAuditProof("ct.test", kLeafHashes[0], 999999, &callback),
+              IsError(net::ERR_IO_PENDING));
   ASSERT_TRUE(callback.called());
   EXPECT_THAT(callback.net_error(), IsError(net::ERR_DNS_SERVER_FAILED));
   EXPECT_THAT(callback.proof(), IsNull());
 }
 
 TEST_P(LogDnsClientTest,
-       QueryAuditProofReportsServerRefusalDuringAuditProofRequests) {
+       QueryAuditProofReportsServerRefusalsDuringAuditProofRequests) {
   mock_dns_.ExpectLeafIndexRequestAndResponse(kLeafIndexQnames[0], "123456");
   mock_dns_.ExpectRequestAndErrorResponse("0.123456.999999.tree.ct.test.",
                                           net::dns_protocol::kRcodeREFUSED);
 
   MockAuditProofCallback callback;
-  QueryAuditProof("ct.test", kLeafHashes[0], 999999, &callback);
+  ASSERT_THAT(QueryAuditProof("ct.test", kLeafHashes[0], 999999, &callback),
+              IsError(net::ERR_IO_PENDING));
   ASSERT_TRUE(callback.called());
   EXPECT_THAT(callback.net_error(), IsError(net::ERR_DNS_SERVER_FAILED));
   EXPECT_THAT(callback.proof(), IsNull());
 }
 
 TEST_P(LogDnsClientTest,
        QueryAuditProofReportsResponseMalformedIfNodeTooShort) {
   // node is shorter than a SHA-256 hash (31 vs 32 bytes)
   const std::vector<std::string> audit_proof(1, std::string(31, 'a'));
 
   mock_dns_.ExpectLeafIndexRequestAndResponse(kLeafIndexQnames[0], "123456");
   mock_dns_.ExpectAuditProofRequestAndResponse(
       "0.123456.999999.tree.ct.test.", audit_proof.begin(), audit_proof.end());
 
   MockAuditProofCallback callback;
-  QueryAuditProof("ct.test", kLeafHashes[0], 999999, &callback);
+  ASSERT_THAT(QueryAuditProof("ct.test", kLeafHashes[0], 999999, &callback),
+              IsError(net::ERR_IO_PENDING));
   ASSERT_TRUE(callback.called());
   EXPECT_THAT(callback.net_error(), IsError(net::ERR_DNS_MALFORMED_RESPONSE));
   EXPECT_THAT(callback.proof(), IsNull());
 }
 
 TEST_P(LogDnsClientTest, QueryAuditProofReportsResponseMalformedIfNodeTooLong) {
   // node is longer than a SHA-256 hash (33 vs 32 bytes)
   const std::vector<std::string> audit_proof(1, std::string(33, 'a'));
 
   mock_dns_.ExpectLeafIndexRequestAndResponse(kLeafIndexQnames[0], "123456");
   mock_dns_.ExpectAuditProofRequestAndResponse(
       "0.123456.999999.tree.ct.test.", audit_proof.begin(), audit_proof.end());
 
   MockAuditProofCallback callback;
-  QueryAuditProof("ct.test", kLeafHashes[0], 999999, &callback);
+  ASSERT_THAT(QueryAuditProof("ct.test", kLeafHashes[0], 999999, &callback),
+              IsError(net::ERR_IO_PENDING));
   ASSERT_TRUE(callback.called());
   EXPECT_THAT(callback.net_error(), IsError(net::ERR_DNS_MALFORMED_RESPONSE));
   EXPECT_THAT(callback.proof(), IsNull());
 }
 
 TEST_P(LogDnsClientTest, QueryAuditProofReportsResponseMalformedIfEmpty) {
   const std::vector<std::string> audit_proof;
 
   mock_dns_.ExpectLeafIndexRequestAndResponse(kLeafIndexQnames[0], "123456");
   mock_dns_.ExpectAuditProofRequestAndResponse(
       "0.123456.999999.tree.ct.test.", audit_proof.begin(), audit_proof.end());
 
   MockAuditProofCallback callback;
-  QueryAuditProof("ct.test", kLeafHashes[0], 999999, &callback);
+  ASSERT_THAT(QueryAuditProof("ct.test", kLeafHashes[0], 999999, &callback),
+              IsError(net::ERR_IO_PENDING));
   ASSERT_TRUE(callback.called());
   EXPECT_THAT(callback.net_error(), IsError(net::ERR_DNS_MALFORMED_RESPONSE));
   EXPECT_THAT(callback.proof(), IsNull());
 }
 
 TEST_P(LogDnsClientTest,
        QueryAuditProofReportsInvalidArgIfLeafIndexEqualToTreeSize) {
   mock_dns_.ExpectLeafIndexRequestAndResponse(kLeafIndexQnames[0], "123456");
 
   MockAuditProofCallback callback;
-  QueryAuditProof("ct.test", kLeafHashes[0], 123456, &callback);
+  ASSERT_THAT(QueryAuditProof("ct.test", kLeafHashes[0], 123456, &callback),
+              IsError(net::ERR_IO_PENDING));
   ASSERT_TRUE(callback.called());
   EXPECT_THAT(callback.net_error(), IsError(net::ERR_INVALID_ARGUMENT));
   EXPECT_THAT(callback.proof(), IsNull());
 }
 
 TEST_P(LogDnsClientTest,
        QueryAuditProofReportsInvalidArgIfLeafIndexGreaterThanTreeSize) {
   mock_dns_.ExpectLeafIndexRequestAndResponse(kLeafIndexQnames[0], "999999");
 
   MockAuditProofCallback callback;
-  QueryAuditProof("ct.test", kLeafHashes[0], 123456, &callback);
+  ASSERT_THAT(QueryAuditProof("ct.test", kLeafHashes[0], 123456, &callback),
+              IsError(net::ERR_IO_PENDING));
   ASSERT_TRUE(callback.called());
   EXPECT_THAT(callback.net_error(), IsError(net::ERR_INVALID_ARGUMENT));
   EXPECT_THAT(callback.proof(), IsNull());
 }
 
 TEST_P(LogDnsClientTest,
        QueryAuditProofReportsSocketErrorsDuringAuditProofRequests) {
   mock_dns_.ExpectLeafIndexRequestAndResponse(kLeafIndexQnames[0], "123456");
   mock_dns_.ExpectRequestAndSocketError("0.123456.999999.tree.ct.test.",
                                         net::ERR_CONNECTION_REFUSED);
 
   MockAuditProofCallback callback;
-  QueryAuditProof("ct.test", kLeafHashes[0], 999999, &callback);
+  ASSERT_THAT(QueryAuditProof("ct.test", kLeafHashes[0], 999999, &callback),
+              IsError(net::ERR_IO_PENDING));
   ASSERT_TRUE(callback.called());
   EXPECT_THAT(callback.net_error(), IsError(net::ERR_CONNECTION_REFUSED));
   EXPECT_THAT(callback.proof(), IsNull());
 }
 
 TEST_P(LogDnsClientTest,
        QueryAuditProofReportsTimeoutsDuringAuditProofRequests) {
   mock_dns_.ExpectLeafIndexRequestAndResponse(kLeafIndexQnames[0], "123456");
   mock_dns_.ExpectRequestAndTimeout("0.123456.999999.tree.ct.test.");
 
   MockAuditProofCallback callback;
-  QueryAuditProof("ct.test", kLeafHashes[0], 999999, &callback);
+  ASSERT_THAT(QueryAuditProof("ct.test", kLeafHashes[0], 999999, &callback),
+              IsError(net::ERR_IO_PENDING));
+
+  callback.WaitUntilRun(TestTimeouts::action_max_timeout());
   ASSERT_TRUE(callback.called());
   EXPECT_THAT(callback.net_error(), IsError(net::ERR_DNS_TIMED_OUT));
   EXPECT_THAT(callback.proof(), IsNull());
 }
 
 TEST_P(LogDnsClientTest, AdoptsLatestDnsConfigIfValid) {
   std::unique_ptr<net::DnsClient> tmp = mock_dns_.CreateDnsClient();
   net::DnsClient* dns_client = tmp.get();
   LogDnsClient log_client(std::move(tmp), net::BoundNetLog(), 0);
 
@@ skipping 43 matching lines @@
         base::StringPrintf("%zu.123456.999999.tree.ct.test.", nodes_begin),
         audit_proof.begin() + nodes_begin, audit_proof.begin() + nodes_end);
   }
 
   std::unique_ptr<net::DnsClient> tmp = mock_dns_.CreateDnsClient();
   net::DnsClient* dns_client = tmp.get();
   LogDnsClient log_client(std::move(tmp), net::BoundNetLog(), 0);
 
   // Start query.
   MockAuditProofCallback callback;
-  QueryAuditProofAsync(&log_client, "ct.test", kLeafHashes[0], 999999,
-                       callback.AsCallback());
+  ASSERT_THAT(QueryAuditProofAsync(&log_client, "ct.test", kLeafHashes[0],
+                                   999999, callback.AsCallback()),
+              IsError(net::ERR_IO_PENDING));
 
   // Get the current DNS config, modify it and broadcast the update.
   net::DnsConfig config(*dns_client->GetConfig());
   ASSERT_NE(123, config.attempts);
   config.attempts = 123;
   mock_dns_.SetDnsConfig(config);
 
-  callback.WaitUntilRun();
+  callback.WaitUntilRun(TestTimeouts::action_max_timeout());
   // Check that the DNS changes propogated before the query completed.
   EXPECT_EQ(123, dns_client->GetConfig()->attempts);
 
   ASSERT_TRUE(callback.called());
   EXPECT_THAT(callback.net_error(), IsOk());
   ASSERT_THAT(callback.proof(), NotNull());
   EXPECT_THAT(callback.proof()->leaf_index, Eq(123456u));
   // TODO(robpercival): Enable this once MerkleAuditProof has tree_size.
   // EXPECT_THAT(callback.proof()->tree_size, Eq(999999));
   EXPECT_THAT(callback.proof()->nodes, Eq(audit_proof));
@@ skipping 52 matching lines @@
             base::StringPrintf("%zu.%" PRIu64 ".%" PRIu64 ".tree.ct.test.",
                                start_node, kLeafIndices[query_i],
                                kTreeSizes[query_i]),
             proof.begin() + start_node, proof.begin() + end_node);
       }
     }
   }
 
   // Start the queries.
   for (size_t i = 0; i < kNumOfParallelQueries; ++i) {
-    QueryAuditProofAsync(log_client.get(), "ct.test", kLeafHashes[i],
-                         kTreeSizes[i], callbacks[i].AsCallback());
+    ASSERT_THAT(
+        QueryAuditProofAsync(log_client.get(), "ct.test", kLeafHashes[i],
+                             kTreeSizes[i], callbacks[i].AsCallback()),
+        IsError(net::ERR_IO_PENDING))
+        << "query #" << i;
   }
 
   // Wait for each query to complete and check its results.
   for (size_t i = 0; i < kNumOfParallelQueries; ++i) {
     MockAuditProofCallback& callback = callbacks[i];
-    callbacks[i].WaitUntilRun();
+    callbacks[i].WaitUntilRun(TestTimeouts::action_max_timeout());
 
     SCOPED_TRACE(testing::Message() << "callbacks[" << i << "]");
     ASSERT_TRUE(callback.called());
     EXPECT_THAT(callback.net_error(), IsOk());
     ASSERT_THAT(callback.proof(), NotNull());
     EXPECT_THAT(callback.proof()->leaf_index, Eq(kLeafIndices[i]));
     // TODO(robpercival): Enable this once MerkleAuditProof has tree_size.
     // EXPECT_THAT(callback.proof()->tree_size, kTreeSizes[i]);
     EXPECT_THAT(callback.proof()->nodes, Eq(audit_proofs[i]));
   }
@@ skipping 19 matching lines @@
                                                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 the queries.
+  // Try to start the queries.
   MockAuditProofCallback callback1;
-  QueryAuditProofAsync(log_client.get(), "ct.test", kLeafHashes[0], 999999,
-                       callback1.AsCallback());
+  ASSERT_THAT(QueryAuditProofAsync(log_client.get(), "ct.test", kLeafHashes[0],
+                                   999999, callback1.AsCallback()),
+              IsError(net::ERR_IO_PENDING));
+
   MockAuditProofCallback callback2;
-  QueryAuditProofAsync(log_client.get(), "ct.test", kLeafHashes[1], 999999,
-                       callback2.AsCallback());
+  ASSERT_THAT(QueryAuditProofAsync(log_client.get(), "ct.test", kLeafHashes[1],
+                                   999999, callback2.AsCallback()),
+              IsError(net::ERR_TEMPORARILY_THROTTLED));
 
-  callback1.WaitUntilRun();
-  callback2.WaitUntilRun();
+  // Give the queries a chance to run.
+  callback1.WaitUntilRun(TestTimeouts::action_max_timeout());
+  // This callback shouldn't run, so don't wait much longer for it.
+  callback2.WaitUntilRun(TestTimeouts::tiny_timeout());
 
   // Check that the first query succeeded.
   ASSERT_TRUE(callback1.called());
   EXPECT_THAT(callback1.net_error(), IsOk());
   ASSERT_THAT(callback1.proof(), NotNull());
   EXPECT_THAT(callback1.proof()->leaf_index, Eq(123456u));
   // TODO(robpercival): Enable this once MerkleAuditProof has tree_size.
   // EXPECT_THAT(callback1.proof()->tree_size, Eq(999999));
   EXPECT_THAT(callback1.proof()->nodes, Eq(audit_proof));
 
-  // Check that the second query failed.
-  ASSERT_TRUE(callback2.called());
-  EXPECT_THAT(callback2.net_error(), IsError(net::ERR_TEMPORARILY_THROTTLED));
-  EXPECT_THAT(callback2.proof(), IsNull());
+  // Check that the second query did not run.
+  ASSERT_FALSE(callback2.called());
 }
 
 INSTANTIATE_TEST_CASE_P(ReadMode,
                         LogDnsClientTest,
                         ::testing::Values(net::IoMode::ASYNC,
                                           net::IoMode::SYNCHRONOUS));
 
 }  // namespace
 }  // namespace certificate_transparency