Refactoring: Create per-connection packet writers in QuicDispatcher.

net/quic/test_tools/quic_test_utils.h

 // Copyright (c) 2012 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.
 //
 // Common utilities for Quic tests
 
 #ifndef NET_QUIC_TEST_TOOLS_QUIC_TEST_UTILS_H_
 #define NET_QUIC_TEST_TOOLS_QUIC_TEST_UTILS_H_
 
 #include <string>
 #include <vector>
 
 #include "base/strings/string_piece.h"
 #include "net/quic/congestion_control/loss_detection_interface.h"
 #include "net/quic/congestion_control/send_algorithm_interface.h"
 #include "net/quic/quic_ack_notifier.h"
 #include "net/quic/quic_client_session_base.h"
 #include "net/quic/quic_connection.h"
+#include "net/quic/quic_dispatcher.h"
 #include "net/quic/quic_framer.h"
+#include "net/quic/quic_per_connection_packet_writer.h"
 #include "net/quic/quic_sent_packet_manager.h"
 #include "net/quic/quic_session.h"
 #include "net/quic/test_tools/mock_clock.h"
 #include "net/quic/test_tools/mock_random.h"
 #include "net/spdy/spdy_framer.h"
 #include "testing/gmock/include/gmock/gmock.h"
 
 namespace net {
 
 namespace test {
@@ skipping 506 matching lines @@
  public:
   MockNetworkChangeVisitor();
   virtual ~MockNetworkChangeVisitor();
 
   MOCK_METHOD1(OnCongestionWindowChange, void(QuicByteCount));
 
  private:
   DISALLOW_COPY_AND_ASSIGN(MockNetworkChangeVisitor);
 };
 
+// Creates per-connection packet writers that register themselves with the
+// TestWriterFactory on each write so that TestWriterFactory::OnPacketSent can
+// be routed to the appropriate QuicConnection.
+class TestWriterFactory : public QuicDispatcher::PacketWriterFactory {
+ public:
+  TestWriterFactory();
+  virtual ~TestWriterFactory();
+
+  virtual QuicPacketWriter* Create(QuicServerPacketWriter* writer,
+                                   QuicConnection* connection) OVERRIDE;
+
+  // Calls OnPacketSent on the last QuicConnection to write through one of the
+  // packet writers created by this factory.
+  void OnPacketSent(WriteResult result);
+
+ private:
+  class PerConnectionPacketWriter : public QuicPerConnectionPacketWriter {
+   public:
+    PerConnectionPacketWriter(TestWriterFactory* factory,
+                              QuicServerPacketWriter* writer,
+                              QuicConnection* connection);
+    virtual ~PerConnectionPacketWriter();
+
+    virtual WriteResult WritePacket(
+        const char* buffer,
+        size_t buf_len,
+        const IPAddressNumber& self_address,
+        const IPEndPoint& peer_address) OVERRIDE;
+
+   private:
+    TestWriterFactory* factory_;
+  };
+
+  // If an asynchronous write is happening and |writer| gets deleted, this
+  // clears the pointer to it to prevent use-after-free.
+  void Unregister(PerConnectionPacketWriter* writer);
+
+  PerConnectionPacketWriter* current_writer_;
+};
+
 }  // namespace test
 }  // namespace net
 
 #endif  // NET_QUIC_TEST_TOOLS_QUIC_TEST_UTILS_H_