LogDnsClient now rejects responses unless they contain exactly one TXT RDATA string

components/certificate_transparency/mock_log_dns_traffic.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/mock_log_dns_traffic.h"
 
 #include <algorithm>
 #include <numeric>
 #include <vector>
 
 #include "base/big_endian.h"
+#include "base/numerics/safe_conversions.h"
+#include "base/strings/string_number_conversions.h"
 #include "base/sys_byteorder.h"
 #include "base/test/test_timeouts.h"
 #include "net/dns/dns_client.h"
 #include "net/dns/dns_protocol.h"
 #include "net/dns/dns_util.h"
 #include "net/socket/socket_test_util.h"
 #include "testing/gtest/include/gtest/gtest.h"
 
 namespace certificate_transparency {
 
@@ skipping 164 matching lines @@
   EmplaceMockSocketData(CreateDnsTxtRequest(qname), net_error);
 }
 
 void MockLogDnsTraffic::ExpectRequestAndTimeout(base::StringPiece qname) {
   EmplaceMockSocketData(CreateDnsTxtRequest(qname));
 
   // Speed up timeout tests.
   SetDnsTimeout(TestTimeouts::tiny_timeout());
 }
 
+void MockLogDnsTraffic::ExpectRequestAndResponse(
+    base::StringPiece qname,
+    const std::vector<base::StringPiece>& txt_strings) {
+  std::string answer;
+  for (base::StringPiece str : txt_strings) {
+    // The size of the string must precede it. The size must fit into 1 byte.
+    answer.insert(answer.end(), base::checked_cast<uint8_t>(str.size()));
+    str.AppendToString(&answer);
+  }
+
+  std::vector<char> request = CreateDnsTxtRequest(qname);
+  EmplaceMockSocketData(request, CreateDnsTxtResponse(request, answer));
+}
+
 void MockLogDnsTraffic::ExpectLeafIndexRequestAndResponse(
     base::StringPiece qname,
-    base::StringPiece leaf_index) {
-  // Prepend size to leaf_index to create the query answer (rdata)
-  ASSERT_LE(leaf_index.size(), 0xFFul);  // size must fit into a single byte
-  std::string answer = leaf_index.as_string();
-  answer.insert(answer.begin(), static_cast<char>(leaf_index.size()));
-
-  ExpectRequestAndResponse(qname, answer);
+    uint64_t leaf_index) {
+  ExpectRequestAndResponse(qname, { base::Uint64ToString(leaf_index) });
 }
 
 void MockLogDnsTraffic::ExpectAuditProofRequestAndResponse(
     base::StringPiece qname,
     std::vector<std::string>::const_iterator audit_path_start,
     std::vector<std::string>::const_iterator audit_path_end) {
   // Join nodes in the audit path into a single string.
   std::string proof =
       std::accumulate(audit_path_start, audit_path_end, std::string());
 
-  // Prepend size to proof to create the query answer (rdata)
-  ASSERT_LE(proof.size(), 0xFFul);  // size must fit into a single byte
-  proof.insert(proof.begin(), static_cast<char>(proof.size()));
-
-  ExpectRequestAndResponse(qname, proof);
+  ExpectRequestAndResponse(qname, { proof });
 }
 
 void MockLogDnsTraffic::InitializeDnsConfig() {
   net::DnsConfig dns_config;
   // Use an invalid nameserver address. This prevents the tests accidentally
   // sending real DNS queries. The mock sockets don't care that the address
   // is invalid.
   dns_config.nameservers.push_back(net::IPEndPoint());
   // Don't attempt retransmissions - just fail.
   dns_config.attempts = 1;
   // This ensures timeouts are long enough for memory tests.
   dns_config.timeout = TestTimeouts::action_timeout();
   // Simplify testing - don't require random numbers for the source port.
   // This means our FakeRandInt function should only be called to get query
   // IDs.
   dns_config.randomize_ports = false;
 
   DnsChangeNotifier::SetInitialDnsConfig(dns_config);
 }
 
 void MockLogDnsTraffic::SetDnsConfig(const net::DnsConfig& config) {
   DnsChangeNotifier::SetDnsConfig(config);
 }
 
 std::unique_ptr<net::DnsClient> MockLogDnsTraffic::CreateDnsClient() {
   return net::DnsClient::CreateClientForTesting(nullptr, &socket_factory_,
                                                 base::Bind(&FakeRandInt));
 }
 
-void MockLogDnsTraffic::ExpectRequestAndResponse(base::StringPiece qname,
-                                                 base::StringPiece answer) {
-  std::vector<char> request = CreateDnsTxtRequest(qname);
-  EmplaceMockSocketData(request, CreateDnsTxtResponse(request, answer));
-}
-
 template <typename... Args>
 void MockLogDnsTraffic::EmplaceMockSocketData(Args&&... args) {
   mock_socket_data_.emplace_back(
       new MockSocketData(std::forward<Args>(args)...));
   mock_socket_data_.back()->SetReadMode(socket_read_mode_);
   mock_socket_data_.back()->AddToFactory(&socket_factory_);
 }
 
 void MockLogDnsTraffic::SetDnsTimeout(const base::TimeDelta& timeout) {
   net::DnsConfig dns_config;
   DnsChangeNotifier::GetDnsConfig(&dns_config);
   dns_config.timeout = timeout;
   DnsChangeNotifier::SetDnsConfig(dns_config);
 }
 
 }  // namespace certificate_transparency