Minor improvements to MockLogDnsTraffic

components/certificate_transparency/mock_log_dns_traffic.cc

Index: components/certificate_transparency/mock_log_dns_traffic.cc
diff --git a/components/certificate_transparency/mock_log_dns_traffic.cc b/components/certificate_transparency/mock_log_dns_traffic.cc
index 22599950dd70a1df0a034adf0d6e230f73dd162b..7a616592caf4b01ebaeced556e79bb4e8c9d352f 100644
--- a/components/certificate_transparency/mock_log_dns_traffic.cc
+++ b/components/certificate_transparency/mock_log_dns_traffic.cc
@@ -21,10 +21,14 @@
 
 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);
@@ -41,11 +45,11 @@ 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
@@ -58,16 +62,16 @@ std::vector<char> CreateDnsTxtRequest(base::StringPiece qname) {
 
   // 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,
@@ -92,19 +96,19 @@ std::vector<char> CreateDnsTxtResponse(const std::vector<char>& request,
   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,
@@ -119,60 +123,85 @@ std::vector<char> CreateDnsErrorResponse(const std::vector<char>& request,
   return response;
 }
 
 }  // namespace
 
-namespace internal {
-
-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) {}
-
-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) {}
-
-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() {}
-
-void MockSocketData::AddToFactory(
-    net::MockClientSocketFactory* socket_factory) {
-  socket_factory->AddSocketDataProvider(&socket_data_);
-}
-
-const net::MockRead MockSocketData::no_more_data_(net::SYNCHRONOUS,
-                                                  net::ERR_IO_PENDING,
-                                                  2);
-
-}  // namespace internal
-
-using internal::MockSocketData;
+// 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) {}
+
+  // 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) {}
+
+  // 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() {}
+
+  void SetWriteMode(net::IoMode mode) { expected_write_.mode = mode; }
+  void SetReadMode(net::IoMode mode) { expected_reads_[0].mode = mode; }
+
+  void AddToFactory(net::MockClientSocketFactory* socket_factory) {
+    socket_factory->AddSocketDataProvider(&socket_data_);
+  }
+
+ 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.
+  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() {}
 
@@ -182,12 +211,12 @@ void MockLogDnsTraffic::ExpectRequestAndErrorResponse(base::StringPiece 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));