Allow LogDnsClient queries to be rate-limited

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/memory/ptr_util.h"
 #include "base/message_loop/message_loop.h"
 #include "base/run_loop.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"
@@ skipping 91 matching lines @@
 };
 
 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() {
+    return base::MakeUnique<LogDnsClient>(mock_dns_.CreateDnsClient(),
+                                          net::BoundNetLog());
+  }
+
+  void QueryLeafIndexAsync(LogDnsClient* log_client,
+                           base::StringPiece log_domain,
+                           base::StringPiece leaf_hash,
+                           const LogDnsClient::LeafIndexCallback& callback) {
+    log_client->QueryLeafIndex(log_domain, leaf_hash, callback);
+  }
+
   void QueryLeafIndex(base::StringPiece log_domain,
                       base::StringPiece leaf_hash,
                       MockLeafIndexCallback* callback) {
-    LogDnsClient log_client(mock_dns_.CreateDnsClient(), net::BoundNetLog());
-    log_client.QueryLeafIndex(log_domain, leaf_hash, callback->AsCallback());
+    auto log_client = CreateLogDnsClient();
+    QueryLeafIndexAsync(log_client.get(), log_domain, leaf_hash,
+                        callback->AsCallback());
     callback->WaitUntilRun();
   }
 
+  void QueryAuditProofAsync(LogDnsClient* log_client,
+                            base::StringPiece log_domain,
+                            uint64_t leaf_index,
+                            uint64_t tree_size,
+                            const LogDnsClient::AuditProofCallback& callback) {
+    log_client->QueryAuditProof(log_domain, leaf_index, tree_size, callback);
+  }
+
   void QueryAuditProof(base::StringPiece log_domain,
                        uint64_t leaf_index,
                        uint64_t tree_size,
                        MockAuditProofCallback* callback) {
-    LogDnsClient log_client(mock_dns_.CreateDnsClient(), net::BoundNetLog());
-    log_client.QueryAuditProof(log_domain, leaf_index, tree_size,
-                               callback->AsCallback());
+    auto log_client = CreateLogDnsClient();
+    QueryAuditProofAsync(log_client.get(), log_domain, leaf_index, tree_size,
+                         callback->AsCallback());
     callback->WaitUntilRun();
   }
 
   // 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.
@@ skipping 388 matching lines @@
   net::DnsConfig config(*dns_client->GetConfig());
   ASSERT_THAT(config.nameservers, Not(IsEmpty()));
   config.nameservers.clear();  // Makes config invalid
   mock_dns_.SetDnsConfig(config);
 
   // Let the DNS config change propogate.
   base::RunLoop().RunUntilIdle();
   EXPECT_THAT(dns_client->GetConfig()->nameservers, Not(IsEmpty()));
 }
 
+TEST_P(LogDnsClientTest, CanBeThrottledToOneQueryAtATime) {
+  const std::vector<std::string> audit_proof = GetSampleAuditProof(20);
+
+  // It should require 3 queries to collect the entire audit proof, as there is
+  // only space for 7 nodes per UDP packet.
+  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());
+
+  auto log_client = CreateLogDnsClient();
+  log_client->SetMaxConcurrentQueries(1);
+
+  MockAuditProofCallback callback1;
+  QueryAuditProofAsync(log_client.get(), "ct.test", 123456, 999999,
+                       callback1.AsCallback());
+  MockAuditProofCallback callback2;
+  QueryAuditProofAsync(log_client.get(), "ct.test", 123456, 999999,
+                       callback2.AsCallback());
+
+  callback1.WaitUntilRun();
+  callback2.WaitUntilRun();
+
+  ASSERT_TRUE(callback1.called());
+  EXPECT_THAT(callback1.net_error(), IsOk());
+  ASSERT_THAT(callback1.proof(), NotNull());
+  EXPECT_THAT(callback1.proof()->leaf_index, 123456);
+  // EXPECT_THAT(callback1.proof()->tree_size, 999999);
+  EXPECT_THAT(callback1.proof()->nodes, audit_proof);
+
+  ASSERT_TRUE(callback2.called());
+  EXPECT_THAT(callback2.net_error(), IsError(net::ERR_TEMPORARILY_THROTTLED));
+  EXPECT_THAT(callback2.proof(), IsNull());
+}
+
 INSTANTIATE_TEST_CASE_P(ReadMode,
                         LogDnsClientTest,
                         ::testing::Values(net::IoMode::ASYNC,
                                           net::IoMode::SYNCHRONOUS));
 
 }  // namespace
 }  // namespace certificate_transparency