Unify the QUIC dispatcher and make the QuicServer work.

net/quic/quic_server.cc

 // Copyright 2014 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 "net/quic/quic_server.h"
 
 #include <string.h>
 
 #include "net/base/ip_endpoint.h"
 #include "net/base/net_errors.h"
 #include "net/quic/crypto/crypto_handshake.h"
 #include "net/quic/crypto/quic_random.h"
 #include "net/quic/quic_crypto_stream.h"
 #include "net/quic/quic_data_reader.h"
-#include "net/quic/quic_dispatcher.h"
+#include "net/quic/quic_per_connection_packet_writer.h"
 #include "net/quic/quic_protocol.h"
 #include "net/quic/quic_server_packet_writer.h"
+#include "net/tools/quic/quic_dispatcher.h"
 #include "net/udp/udp_server_socket.h"
 
 namespace net {
 
+using tools::QuicDispatcher;
+
 namespace {
 
 const char kSourceAddressTokenSecret[] = "secret";
 
 // Allocate some extra space so we can send an error if the client goes over
 // the limit.
 const int kReadBufferSize = 2 * kMaxPacketSize;
 
+// A packet writer factory which wraps a shared QuicServerPacketWriter
+// inside of a QuicPerConnectionPacketWriter. Instead of checking that
+// the shared_writer is the expected writer, this could instead cast
+// from QuicPacketWriter to QuicServerPacketWriter.
+class CustomPacketWriterFactory : public QuicDispatcher::PacketWriterFactory {
+ public:
+  ~CustomPacketWriterFactory() override {}
+
+  QuicPacketWriter* Create(QuicPacketWriter* writer,
+                           QuicConnection* connection) override {
+    if (writer == nullptr) {
+      LOG(DFATAL) << "shared_writer not initialized";
+      return nullptr;
+    }
+    if (writer != shared_writer_) {
+      LOG(DFATAL) << "writer mismatch";
+      return nullptr;
+    }
+    return new QuicPerConnectionPacketWriter(shared_writer_, connection);
+  }
+
+  void set_shared_writer(QuicServerPacketWriter* shared_writer) {
+    shared_writer_ = shared_writer;
+  }
+
+ private:
+  QuicServerPacketWriter* shared_writer_;  // Not owned.
+};
+
 } // namespace
 
 QuicServer::QuicServer(const QuicConfig& config,
                        const QuicVersionVector& supported_versions)
     : helper_(base::MessageLoop::current()->message_loop_proxy().get(),
               &clock_,
               QuicRandom::GetInstance()),
       config_(config),
       crypto_config_(kSourceAddressTokenSecret, QuicRandom::GetInstance()),
       supported_versions_(supported_versions),
@@ skipping 64 matching lines @@
   rc = socket->GetLocalAddress(&server_address_);
   if (rc < 0) {
     LOG(ERROR) << "GetLocalAddress() failed: " << ErrorToString(rc);
     return rc;
   }
 
   DVLOG(1) << "Listening on " << server_address_.ToString();
 
   socket_.swap(socket);
 
+  CustomPacketWriterFactory* factory = new CustomPacketWriterFactory();
   dispatcher_.reset(
       new QuicDispatcher(config_,
                          crypto_config_,
                          supported_versions_,
-                         new QuicDispatcher::DefaultPacketWriterFactory(),
+                         factory,
                          &helper_));
   QuicServerPacketWriter* writer = new QuicServerPacketWriter(
       socket_.get(),
       dispatcher_.get());
-  dispatcher_->Initialize(writer);
+  factory->set_shared_writer(writer);
+  dispatcher_->InitializeWithWriter(writer);
 
   StartReading();
 
   return OK;
 }
 
 void QuicServer::Shutdown() {
   // Before we shut down the epoll server, give all active sessions a chance to
   // notify clients that they're closing.
   dispatcher_->Shutdown();
@@ skipping 44 matching lines @@
     return;
   }
 
   QuicEncryptedPacket packet(read_buffer_->data(), result, false);
   dispatcher_->ProcessPacket(server_address_, client_address_, packet);
 
   StartReading();
 }
 
 }  // namespace net