Minor improvements to MockLogDnsTraffic

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 {
 
 namespace {
 
+// This is used for the last mock socket response as a sentinel to prevent
+// trying to read more data than expected.
+const net::MockRead kNoMoreData(net::SYNCHRONOUS, net::ERR_UNEXPECTED, 2);
+
 // Necessary to expose SetDnsConfig for testing.
 class DnsChangeNotifier : public net::NetworkChangeNotifier {
  public:
   static void SetInitialDnsConfig(const net::DnsConfig& config) {
     net::NetworkChangeNotifier::SetInitialDnsConfig(config);
   }
 
   static void SetDnsConfig(const net::DnsConfig& config) {
     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) {
   std::string encoded_qname;
-  EXPECT_TRUE(net::DNSDomainFromDot(qname, &encoded_qname));
+  CHECK(net::DNSDomainFromDot(qname, &encoded_qname));
 
   // 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());
 
   // Header
   net::dns_protocol::Header header = {};
   header.flags = base::HostToNet16(net::dns_protocol::kFlagRD);
   header.qdcount = base::HostToNet16(1);
-  EXPECT_TRUE(writer.WriteBytes(&header, sizeof(header)));
+  CHECK(writer.WriteBytes(&header, sizeof(header)));
   // Query section
-  EXPECT_TRUE(writer.WriteBytes(encoded_qname.data(), encoded_qname.size()));
-  EXPECT_TRUE(writer.WriteU16(net::dns_protocol::kTypeTXT));
-  EXPECT_TRUE(writer.WriteU16(net::dns_protocol::kClassIN));
-  EXPECT_EQ(0, writer.remaining());
+  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;
 }
 
 std::vector<char> CreateDnsTxtResponse(const std::vector<char>& request,
                                        base::StringPiece answer) {
   // 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());
   // Modify the header
   net::dns_protocol::Header* header =
       reinterpret_cast<net::dns_protocol::Header*>(response.data());
   header->ancount = base::HostToNet16(1);
   header->flags |= base::HostToNet16(net::dns_protocol::kFlagResponse);
 
   // Write the answer section
   base::BigEndianWriter writer(response.data() + request.size(),
                                response.size() - request.size());
-  EXPECT_TRUE(writer.WriteU8(0xc0));  // qname is a pointer
-  EXPECT_TRUE(writer.WriteU8(
+  CHECK(writer.WriteU8(0xc0));  // qname is a pointer
+  CHECK(writer.WriteU8(
       sizeof(*header)));  // address of qname (start of query section)
-  EXPECT_TRUE(writer.WriteU16(net::dns_protocol::kTypeTXT));
-  EXPECT_TRUE(writer.WriteU16(net::dns_protocol::kClassIN));
-  EXPECT_TRUE(writer.WriteU32(ttl));
-  EXPECT_TRUE(writer.WriteU16(answer.size()));
-  EXPECT_TRUE(writer.WriteBytes(answer.data(), answer.size()));
-  EXPECT_EQ(0, writer.remaining());
+  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());
 
   return response;
 }
 
 std::vector<char> CreateDnsErrorResponse(const std::vector<char>& request,
                                          uint8_t rcode) {
   std::vector<char> response(request);
   // Modify the header
   net::dns_protocol::Header* header =
       reinterpret_cast<net::dns_protocol::Header*>(response.data());
   header->ancount = base::HostToNet16(1);
   header->flags |= base::HostToNet16(net::dns_protocol::kFlagResponse | rcode);
 
   return response;
 }
 
 }  // namespace
 
-namespace internal {
+// 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 {
+ public:
+  // A socket that expects one write and one read operation.
+  MockSocketData(const std::vector<char>& write, const std::vector<char>& read)
+      : expected_write_payload_(write),
+        expected_read_payload_(read),
+        expected_write_(net::SYNCHRONOUS,
+                        expected_write_payload_.data(),
+                        expected_write_payload_.size(),
+                        0),
+        expected_reads_{net::MockRead(net::ASYNC,
+                                      expected_read_payload_.data(),
+                                      expected_read_payload_.size(),
+                                      1),
+                        kNoMoreData},
+        socket_data_(expected_reads_, 2, &expected_write_, 1) {}
 
-MockSocketData::MockSocketData(const std::vector<char>& write,
-                               const std::vector<char>& read)
-    : expected_write_payload_(write),
-      expected_read_payload_(read),
-      expected_write_(net::SYNCHRONOUS,
-                      expected_write_payload_.data(),
-                      expected_write_payload_.size(),
-                      0),
-      expected_reads_{net::MockRead(net::ASYNC,
-                                    expected_read_payload_.data(),
-                                    expected_read_payload_.size(),
-                                    1),
-                      no_more_data_},
-      socket_data_(expected_reads_, 2, &expected_write_, 1) {}
+  // A socket that expects one write and a read error.
+  MockSocketData(const std::vector<char>& write, net::Error error)
+      : expected_write_payload_(write),
+        expected_write_(net::SYNCHRONOUS,
+                        expected_write_payload_.data(),
+                        expected_write_payload_.size(),
+                        0),
+        expected_reads_{net::MockRead(net::ASYNC, error, 1), kNoMoreData},
+        socket_data_(expected_reads_, 2, &expected_write_, 1) {}
 
-MockSocketData::MockSocketData(const std::vector<char>& write, int net_error)
-    : expected_write_payload_(write),
-      expected_write_(net::SYNCHRONOUS,
-                      expected_write_payload_.data(),
-                      expected_write_payload_.size(),
-                      0),
-      expected_reads_{net::MockRead(net::ASYNC, net_error, 1), no_more_data_},
-      socket_data_(expected_reads_, 2, &expected_write_, 1) {}
+  // A socket that expects one write and no response.
+  explicit MockSocketData(const std::vector<char>& write)
+      : expected_write_payload_(write),
+        expected_write_(net::SYNCHRONOUS,
+                        expected_write_payload_.data(),
+                        expected_write_payload_.size(),
+                        0),
+        expected_reads_{net::MockRead(net::SYNCHRONOUS, net::ERR_IO_PENDING, 1),
+                        kNoMoreData},
+        socket_data_(expected_reads_, 2, &expected_write_, 1) {}
 
-MockSocketData::MockSocketData(const std::vector<char>& write)
-    : expected_write_payload_(write),
-      expected_write_(net::SYNCHRONOUS,
-                      expected_write_payload_.data(),
-                      expected_write_payload_.size(),
-                      0),
-      expected_reads_{net::MockRead(net::SYNCHRONOUS, net::ERR_IO_PENDING, 1),
-                      no_more_data_},
-      socket_data_(expected_reads_, 2, &expected_write_, 1) {}
+  ~MockSocketData() {}
 
-MockSocketData::~MockSocketData() {}
+  void SetWriteMode(net::IoMode mode) { expected_write_.mode = mode; }
+  void SetReadMode(net::IoMode mode) { expected_reads_[0].mode = mode; }
 
-void MockSocketData::AddToFactory(
-    net::MockClientSocketFactory* socket_factory) {
-  socket_factory->AddSocketDataProvider(&socket_data_);
-}
+  void AddToFactory(net::MockClientSocketFactory* socket_factory) {
+    socket_factory->AddSocketDataProvider(&socket_data_);
+  }
 
-const net::MockRead MockSocketData::no_more_data_(net::SYNCHRONOUS,
-                                                  net::ERR_IO_PENDING,
-                                                  2);
+ 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_;
 
-}  // namespace internal
+  // Encapsulates the data that is expected to be written to a socket.
+  net::MockWrite expected_write_;
 
-using internal::MockSocketData;
+  // 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.
+  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,
                                                       uint8_t rcode) {
   std::vector<char> request = CreateDnsTxtRequest(qname);
   EmplaceMockSocketData(CreateDnsTxtRequest(qname),
                         CreateDnsErrorResponse(request, rcode));
 }
 
 void MockLogDnsTraffic::ExpectRequestAndSocketError(base::StringPiece qname,
-                                                    int net_error) {
-  EmplaceMockSocketData(CreateDnsTxtRequest(qname), net_error);
+                                                    net::Error error) {
+  EmplaceMockSocketData(CreateDnsTxtRequest(qname), error);
 }
 
 void MockLogDnsTraffic::ExpectRequestAndTimeout(base::StringPiece qname) {
   EmplaceMockSocketData(CreateDnsTxtRequest(qname));
 
   // Speed up timeout tests.
   SetDnsTimeout(TestTimeouts::tiny_timeout());
 }
 
 void MockLogDnsTraffic::ExpectRequestAndResponse(
@@ skipping 63 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