MockLogDnsTraffic now returns bool instead of CHECKing bad arguments

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>
 
@@ skipping 27 matching lines @@
     net::NetworkChangeNotifier::SetDnsConfig(config);
   }
 };
 
 // Always return min, to simplify testing.
 // This should result in the DNS query ID always being 0.
 int FakeRandInt(int min, int max) {
   return min;
 }
 
-std::vector<char> CreateDnsTxtRequest(base::StringPiece qname) {
+bool CreateDnsTxtRequest(base::StringPiece qname, std::vector<char>* request) {
   std::string encoded_qname;
-  CHECK(net::DNSDomainFromDot(qname, &encoded_qname));
+  if (!net::DNSDomainFromDot(qname, &encoded_qname)) {
+    // |qname| is an invalid domain name.
+    return false;
+  }
 
   // DNS query section:
   // N bytes - qname
   // 2 bytes - record type
   // 2 bytes - record class
   // Total = N + 4 bytes
   const size_t query_section_size = encoded_qname.size() + 4;
 
-  std::vector<char> request(sizeof(net::dns_protocol::Header) +
-                            query_section_size);
-  base::BigEndianWriter writer(request.data(), request.size());
+  request->resize(sizeof(net::dns_protocol::Header) + query_section_size);
+  base::BigEndianWriter writer(request->data(), request->size());
 
   // Header
   net::dns_protocol::Header header = {};
   header.flags = base::HostToNet16(net::dns_protocol::kFlagRD);
   header.qdcount = base::HostToNet16(1);
-  CHECK(writer.WriteBytes(&header, sizeof(header)));
-  // Query section
-  CHECK(writer.WriteBytes(encoded_qname.data(), encoded_qname.size()));
-  CHECK(writer.WriteU16(net::dns_protocol::kTypeTXT));
-  CHECK(writer.WriteU16(net::dns_protocol::kClassIN));
-  CHECK_EQ(0, writer.remaining());
 
-  return request;
+  if (!writer.WriteBytes(&header, sizeof(header)) ||
+      !writer.WriteBytes(encoded_qname.data(), encoded_qname.size()) ||
+      !writer.WriteU16(net::dns_protocol::kTypeTXT) ||
+      !writer.WriteU16(net::dns_protocol::kClassIN)) {
+    return false;
+  }
+
+  if (writer.remaining() != 0) {
+    // Less than the expected amount of data was written.
+    return false;
+  }
+
+  return true;
 }
 
-std::vector<char> CreateDnsTxtResponse(const std::vector<char>& request,
-                                       base::StringPiece answer) {
+bool CreateDnsTxtResponse(const std::vector<char>& request,
+                          base::StringPiece answer,
+                          std::vector<char>* response) {
   // DNS answers section:
   // 2 bytes - qname pointer
   // 2 bytes - record type
   // 2 bytes - record class
   // 4 bytes - time-to-live
   // 2 bytes - size of answer (N)
   // N bytes - answer
   // Total = 12 + N bytes
   const size_t answers_section_size = 12 + answer.size();
   constexpr uint32_t ttl = 86400;  // seconds
 
-  std::vector<char> response(request.size() + answers_section_size);
-  std::copy(request.begin(), request.end(), response.begin());
+  response->resize(request.size() + answers_section_size);
+  std::copy(request.begin(), request.end(), response->begin());
+
   // Modify the header
   net::dns_protocol::Header* header =
-      reinterpret_cast<net::dns_protocol::Header*>(response.data());
+      reinterpret_cast<net::dns_protocol::Header*>(response->data());
   header->ancount = base::HostToNet16(1);
   header->flags |= base::HostToNet16(net::dns_protocol::kFlagResponse);
 
+  // The qname is at the start of the query section (just after the header).
+  const uint8_t qname_ptr = sizeof(*header);
+
   // Write the answer section
-  base::BigEndianWriter writer(response.data() + request.size(),
-                               response.size() - request.size());
-  CHECK(writer.WriteU8(0xc0));  // qname is a pointer
-  CHECK(writer.WriteU8(
-      sizeof(*header)));  // address of qname (start of query section)
-  CHECK(writer.WriteU16(net::dns_protocol::kTypeTXT));
-  CHECK(writer.WriteU16(net::dns_protocol::kClassIN));
-  CHECK(writer.WriteU32(ttl));
-  CHECK(writer.WriteU16(answer.size()));
-  CHECK(writer.WriteBytes(answer.data(), answer.size()));
-  CHECK_EQ(0, writer.remaining());
+  base::BigEndianWriter writer(response->data() + request.size(),
+                               response->size() - request.size());
+  if (!writer.WriteU8(net::dns_protocol::kLabelPointer) ||
+      !writer.WriteU8(qname_ptr) ||
+      !writer.WriteU16(net::dns_protocol::kTypeTXT) ||
+      !writer.WriteU16(net::dns_protocol::kClassIN) || !writer.WriteU32(ttl) ||
+      !writer.WriteU16(answer.size()) ||
+      !writer.WriteBytes(answer.data(), answer.size())) {
+    return false;
+  }
 
-  return response;
+  if (writer.remaining() != 0) {
+    // Less than the expected amount of data was written.
+    return false;
+  }
+
+  return true;
 }
 
-std::vector<char> CreateDnsErrorResponse(const std::vector<char>& request,
-                                         uint8_t rcode) {
-  std::vector<char> response(request);
+bool CreateDnsErrorResponse(const std::vector<char>& request,
+                            uint8_t rcode,
+                            std::vector<char>* response) {
+  *response = request;
   // Modify the header
   net::dns_protocol::Header* header =
-      reinterpret_cast<net::dns_protocol::Header*>(response.data());
+      reinterpret_cast<net::dns_protocol::Header*>(response->data());
   header->ancount = base::HostToNet16(1);
   header->flags |= base::HostToNet16(net::dns_protocol::kFlagResponse | rcode);
-
-  return response;
+  return true;
 }
 
 }  // namespace
 
 // A container for all of the data needed for simulating a socket.
 // This is useful because Mock{Read,Write}, SequencedSocketData and
 // MockClientSocketFactory all do not take ownership of or copy their arguments,
 // so it is necessary to manage the lifetime of those arguments. Wrapping all
 // of that up in a single class simplifies this.
 class MockLogDnsTraffic::MockSocketData {
@@ skipping 46 matching lines @@
  private:
   // This class only supports one write and one read, so just need to store one
   // payload each.
   const std::vector<char> expected_write_payload_;
   const std::vector<char> expected_read_payload_;
 
   // Encapsulates the data that is expected to be written to a socket.
   net::MockWrite expected_write_;
 
   // Encapsulates the data/error that should be returned when reading from a
-  // socket. The second "expected" read is a sentinel that causes socket reads
-  // beyond the first to hang until they timeout. This results in better
-  // test failure messages (rather than a CHECK-fail due to a socket read
-  // overrunning the MockRead array) and behaviour more like a real socket when
-  // an unexpected second socket read occurs.
+  // socket. The second "expected" read is a sentinel (see |kNoMoreData|).
   net::MockRead expected_reads_[2];
 
   // Holds pointers to |expected_write_| and |expected_reads_|. This is what is
   // added to net::MockClientSocketFactory to prepare a mock socket.
   net::SequencedSocketData socket_data_;
 
   DISALLOW_COPY_AND_ASSIGN(MockSocketData);
 };
 
 MockLogDnsTraffic::MockLogDnsTraffic() : socket_read_mode_(net::ASYNC) {}
 
 MockLogDnsTraffic::~MockLogDnsTraffic() {}
 
-void MockLogDnsTraffic::ExpectRequestAndErrorResponse(base::StringPiece qname,
+bool MockLogDnsTraffic::ExpectRequestAndErrorResponse(base::StringPiece qname,
                                                       uint8_t rcode) {
-  std::vector<char> request = CreateDnsTxtRequest(qname);
-  EmplaceMockSocketData(CreateDnsTxtRequest(qname),
-                        CreateDnsErrorResponse(request, rcode));
+  std::vector<char> request;
+  if (!CreateDnsTxtRequest(qname, &request)) {
+    return false;
+  }
+
+  std::vector<char> response;
+  if (!CreateDnsErrorResponse(request, rcode, &response)) {
+    return false;
+  }
+
+  EmplaceMockSocketData(request, response);
+  return true;
 }
 
-void MockLogDnsTraffic::ExpectRequestAndSocketError(base::StringPiece qname,
+bool MockLogDnsTraffic::ExpectRequestAndSocketError(base::StringPiece qname,
                                                     net::Error error) {
-  EmplaceMockSocketData(CreateDnsTxtRequest(qname), error);
+  std::vector<char> request;
+  if (!CreateDnsTxtRequest(qname, &request)) {
+    return false;
+  }
+
+  EmplaceMockSocketData(request, error);
+  return true;
 }
 
-void MockLogDnsTraffic::ExpectRequestAndTimeout(base::StringPiece qname) {
-  EmplaceMockSocketData(CreateDnsTxtRequest(qname));
+bool MockLogDnsTraffic::ExpectRequestAndTimeout(base::StringPiece qname) {
+  std::vector<char> request;
+  if (!CreateDnsTxtRequest(qname, &request)) {
+    return false;
+  }
+
+  EmplaceMockSocketData(request);
 
   // Speed up timeout tests.
   SetDnsTimeout(TestTimeouts::tiny_timeout());
+
+  return true;
 }
 
-void MockLogDnsTraffic::ExpectRequestAndResponse(
+bool 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));
+  std::vector<char> request;
+  if (!CreateDnsTxtRequest(qname, &request)) {
+    return false;
+  }
+
+  std::vector<char> response;
+  if (!CreateDnsTxtResponse(request, answer, &response)) {
+    return false;
+  }
+
+  EmplaceMockSocketData(request, response);
+  return true;
 }
 
-void MockLogDnsTraffic::ExpectLeafIndexRequestAndResponse(
+bool MockLogDnsTraffic::ExpectLeafIndexRequestAndResponse(
     base::StringPiece qname,
     uint64_t leaf_index) {
-  ExpectRequestAndResponse(qname, { base::Uint64ToString(leaf_index) });
+  return ExpectRequestAndResponse(qname, {base::Uint64ToString(leaf_index)});
 }
 
-void MockLogDnsTraffic::ExpectAuditProofRequestAndResponse(
+bool 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());
 
-  ExpectRequestAndResponse(qname, { proof });
+  return 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;
@@ skipping 25 matching lines @@
 }
 
 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