net/quic/test_tools/quic_test_utils.cc

   1 // Copyright (c) 2012 The Chromium Authors. All rights reserved.
   2 // Use of this source code is governed by a BSD-style license that can be
   3 // found in the LICENSE file.
   4
   5 #include "net/quic/test_tools/quic_test_utils.h"
   6
   7 #include "base/sha1.h"
   8 #include "base/stl_util.h"
   9 #include "base/strings/string_number_conversions.h"
  10 #include "net/quic/crypto/crypto_framer.h"
  11 #include "net/quic/crypto/crypto_handshake.h"
  12 #include "net/quic/crypto/crypto_utils.h"
  13 #include "net/quic/crypto/null_encrypter.h"
  14 #include "net/quic/crypto/quic_decrypter.h"
  15 #include "net/quic/crypto/quic_encrypter.h"
  16 #include "net/quic/quic_framer.h"
  17 #include "net/quic/quic_packet_creator.h"
  18 #include "net/quic/quic_utils.h"
  19 #include "net/quic/test_tools/quic_connection_peer.h"
  20 #include "net/spdy/spdy_frame_builder.h"
  21
  22 using base::StringPiece;
  23 using std::max;
  24 using std::min;
  25 using std::string;
  26 using testing::AnyNumber;
  27 using testing::_;
  28
  29 namespace net {
  30 namespace test {
  31 namespace {
  32
  33 // No-op alarm implementation used by MockHelper.
  34 class TestAlarm : public QuicAlarm {
  35  public:
  36   explicit TestAlarm(QuicAlarm::Delegate* delegate)
  37       : QuicAlarm(delegate) {
  38   }
  39
  40   virtual void SetImpl() OVERRIDE {}
  41   virtual void CancelImpl() OVERRIDE {}
  42 };
  43
  44 }  // namespace
  45
  46 QuicAckFrame MakeAckFrame(QuicPacketSequenceNumber largest_observed) {
  47   QuicAckFrame ack;
  48   ack.largest_observed = largest_observed;
  49   ack.entropy_hash = 0;
  50   return ack;
  51 }
  52
  53 QuicAckFrame MakeAckFrameWithNackRanges(
  54     size_t num_nack_ranges, QuicPacketSequenceNumber least_unacked) {
  55   QuicAckFrame ack = MakeAckFrame(2 * num_nack_ranges + least_unacked);
  56   // Add enough missing packets to get num_nack_ranges nack ranges.
  57   for (QuicPacketSequenceNumber i = 1; i < 2 * num_nack_ranges; i += 2) {
  58     ack.missing_packets.insert(least_unacked + i);
  59   }
  60   return ack;
  61 }
  62
  63 SerializedPacket BuildUnsizedDataPacket(QuicFramer* framer,
  64                                         const QuicPacketHeader& header,
  65                                         const QuicFrames& frames) {
  66   const size_t max_plaintext_size = framer->GetMaxPlaintextSize(kMaxPacketSize);
  67   size_t packet_size = GetPacketHeaderSize(header);
  68   for (size_t i = 0; i < frames.size(); ++i) {
  69     DCHECK_LE(packet_size, max_plaintext_size);
  70     bool first_frame = i == 0;
  71     bool last_frame = i == frames.size() - 1;
  72     const size_t frame_size = framer->GetSerializedFrameLength(
  73         frames[i], max_plaintext_size - packet_size, first_frame, last_frame,
  74         header.is_in_fec_group,
  75         header.public_header.sequence_number_length);
  76     DCHECK(frame_size);
  77     packet_size += frame_size;
  78   }
  79   return framer->BuildDataPacket(header, frames, packet_size);
  80 }
  81
  82 uint64 SimpleRandom::RandUint64() {
  83   unsigned char hash[base::kSHA1Length];
  84   base::SHA1HashBytes(reinterpret_cast<unsigned char*>(&seed_), sizeof(seed_),
  85                       hash);
  86   memcpy(&seed_, hash, sizeof(seed_));
  87   return seed_;
  88 }
  89
  90 MockFramerVisitor::MockFramerVisitor() {
  91   // By default, we want to accept packets.
  92   ON_CALL(*this, OnProtocolVersionMismatch(_))
  93       .WillByDefault(testing::Return(false));
  94
  95   // By default, we want to accept packets.
  96   ON_CALL(*this, OnUnauthenticatedHeader(_))
  97       .WillByDefault(testing::Return(true));
  98
  99   ON_CALL(*this, OnUnauthenticatedPublicHeader(_))
 100       .WillByDefault(testing::Return(true));
 101
 102   ON_CALL(*this, OnPacketHeader(_))
 103       .WillByDefault(testing::Return(true));
 104
 105   ON_CALL(*this, OnStreamFrame(_))
 106       .WillByDefault(testing::Return(true));
 107
 108   ON_CALL(*this, OnAckFrame(_))
 109       .WillByDefault(testing::Return(true));
 110
 111   ON_CALL(*this, OnCongestionFeedbackFrame(_))
 112       .WillByDefault(testing::Return(true));
 113
 114   ON_CALL(*this, OnStopWaitingFrame(_))
 115       .WillByDefault(testing::Return(true));
 116
 117   ON_CALL(*this, OnPingFrame(_))
 118       .WillByDefault(testing::Return(true));
 119
 120   ON_CALL(*this, OnRstStreamFrame(_))
 121       .WillByDefault(testing::Return(true));
 122
 123   ON_CALL(*this, OnConnectionCloseFrame(_))
 124       .WillByDefault(testing::Return(true));
 125
 126   ON_CALL(*this, OnGoAwayFrame(_))
 127       .WillByDefault(testing::Return(true));
 128 }
 129
 130 MockFramerVisitor::~MockFramerVisitor() {
 131 }
 132
 133 bool NoOpFramerVisitor::OnProtocolVersionMismatch(QuicVersion version) {
 134   return false;
 135 }
 136
 137 bool NoOpFramerVisitor::OnUnauthenticatedPublicHeader(
 138     const QuicPacketPublicHeader& header) {
 139   return true;
 140 }
 141
 142 bool NoOpFramerVisitor::OnUnauthenticatedHeader(
 143     const QuicPacketHeader& header) {
 144   return true;
 145 }
 146
 147 bool NoOpFramerVisitor::OnPacketHeader(const QuicPacketHeader& header) {
 148   return true;
 149 }
 150
 151 bool NoOpFramerVisitor::OnStreamFrame(const QuicStreamFrame& frame) {
 152   return true;
 153 }
 154
 155 bool NoOpFramerVisitor::OnAckFrame(const QuicAckFrame& frame) {
 156   return true;
 157 }
 158
 159 bool NoOpFramerVisitor::OnCongestionFeedbackFrame(
 160     const QuicCongestionFeedbackFrame& frame) {
 161   return true;
 162 }
 163
 164 bool NoOpFramerVisitor::OnStopWaitingFrame(
 165     const QuicStopWaitingFrame& frame) {
 166   return true;
 167 }
 168
 169 bool NoOpFramerVisitor::OnPingFrame(const QuicPingFrame& frame) {
 170   return true;
 171 }
 172
 173 bool NoOpFramerVisitor::OnRstStreamFrame(
 174     const QuicRstStreamFrame& frame) {
 175   return true;
 176 }
 177
 178 bool NoOpFramerVisitor::OnConnectionCloseFrame(
 179     const QuicConnectionCloseFrame& frame) {
 180   return true;
 181 }
 182
 183 bool NoOpFramerVisitor::OnGoAwayFrame(const QuicGoAwayFrame& frame) {
 184   return true;
 185 }
 186
 187 bool NoOpFramerVisitor::OnWindowUpdateFrame(
 188     const QuicWindowUpdateFrame& frame) {
 189   return true;
 190 }
 191
 192 bool NoOpFramerVisitor::OnBlockedFrame(const QuicBlockedFrame& frame) {
 193   return true;
 194 }
 195
 196 MockConnectionVisitor::MockConnectionVisitor() {
 197 }
 198
 199 MockConnectionVisitor::~MockConnectionVisitor() {
 200 }
 201
 202 MockHelper::MockHelper() {
 203 }
 204
 205 MockHelper::~MockHelper() {
 206 }
 207
 208 const QuicClock* MockHelper::GetClock() const {
 209   return &clock_;
 210 }
 211
 212 QuicRandom* MockHelper::GetRandomGenerator() {
 213   return &random_generator_;
 214 }
 215
 216 QuicAlarm* MockHelper::CreateAlarm(QuicAlarm::Delegate* delegate) {
 217   return new TestAlarm(delegate);
 218 }
 219
 220 void MockHelper::AdvanceTime(QuicTime::Delta delta) {
 221   clock_.AdvanceTime(delta);
 222 }
 223
 224 namespace {
 225 class NiceMockPacketWriterFactory
 226     : public QuicConnection::PacketWriterFactory {
 227  public:
 228   NiceMockPacketWriterFactory() {}
 229   virtual ~NiceMockPacketWriterFactory() {}
 230
 231   virtual QuicPacketWriter* Create(
 232       QuicConnection* /*connection*/) const override {
 233     return new testing::NiceMock<MockPacketWriter>();
 234   }
 235
 236  private:
 237   DISALLOW_COPY_AND_ASSIGN(NiceMockPacketWriterFactory);
 238 };
 239 }  // namespace
 240
 241 MockConnection::MockConnection(bool is_server)
 242     : QuicConnection(kTestConnectionId,
 243                      IPEndPoint(TestPeerIPAddress(), kTestPort),
 244                      new testing::NiceMock<MockHelper>(),
 245                      NiceMockPacketWriterFactory(),
 246                      /* owns_writer= */ true,
 247                      is_server, QuicSupportedVersions()),
 248       helper_(helper()) {
 249 }
 250
 251 MockConnection::MockConnection(IPEndPoint address,
 252                                bool is_server)
 253     : QuicConnection(kTestConnectionId, address,
 254                      new testing::NiceMock<MockHelper>(),
 255                      NiceMockPacketWriterFactory(),
 256                      /* owns_writer= */ true,
 257                      is_server, QuicSupportedVersions()),
 258       helper_(helper()) {
 259 }
 260
 261 MockConnection::MockConnection(QuicConnectionId connection_id,
 262                                bool is_server)
 263     : QuicConnection(connection_id,
 264                      IPEndPoint(TestPeerIPAddress(), kTestPort),
 265                      new testing::NiceMock<MockHelper>(),
 266                      NiceMockPacketWriterFactory(),
 267                      /* owns_writer= */ true,
 268                      is_server, QuicSupportedVersions()),
 269       helper_(helper()) {
 270 }
 271
 272 MockConnection::MockConnection(bool is_server,
 273                                const QuicVersionVector& supported_versions)
 274     : QuicConnection(kTestConnectionId,
 275                      IPEndPoint(TestPeerIPAddress(), kTestPort),
 276                      new testing::NiceMock<MockHelper>(),
 277                      NiceMockPacketWriterFactory(),
 278                      /* owns_writer= */ true,
 279                      is_server, supported_versions),
 280       helper_(helper()) {
 281 }
 282
 283 MockConnection::~MockConnection() {
 284 }
 285
 286 void MockConnection::AdvanceTime(QuicTime::Delta delta) {
 287   static_cast<MockHelper*>(helper())->AdvanceTime(delta);
 288 }
 289
 290 PacketSavingConnection::PacketSavingConnection(bool is_server)
 291     : MockConnection(is_server) {
 292 }
 293
 294 PacketSavingConnection::PacketSavingConnection(
 295     bool is_server,
 296     const QuicVersionVector& supported_versions)
 297     : MockConnection(is_server, supported_versions) {
 298 }
 299
 300 PacketSavingConnection::~PacketSavingConnection() {
 301   STLDeleteElements(&packets_);
 302   STLDeleteElements(&encrypted_packets_);
 303 }
 304
 305 bool PacketSavingConnection::SendOrQueuePacket(
 306     EncryptionLevel level,
 307     const SerializedPacket& packet,
 308     TransmissionType transmission_type) {
 309   packets_.push_back(packet.packet);
 310   QuicEncryptedPacket* encrypted = QuicConnectionPeer::GetFramer(this)->
 311       EncryptPacket(level, packet.sequence_number, *packet.packet);
 312   encrypted_packets_.push_back(encrypted);
 313   return true;
 314 }
 315
 316 MockSession::MockSession(QuicConnection* connection)
 317     : QuicSession(connection, DefaultQuicConfig()) {
 318   InitializeSession();
 319   ON_CALL(*this, WritevData(_, _, _, _, _, _))
 320       .WillByDefault(testing::Return(QuicConsumedData(0, false)));
 321 }
 322
 323 MockSession::~MockSession() {
 324 }
 325
 326 TestSession::TestSession(QuicConnection* connection, const QuicConfig& config)
 327     : QuicSession(connection, config),
 328       crypto_stream_(NULL) {
 329   InitializeSession();
 330 }
 331
 332 TestSession::~TestSession() {}
 333
 334 void TestSession::SetCryptoStream(QuicCryptoStream* stream) {
 335   crypto_stream_ = stream;
 336 }
 337
 338 QuicCryptoStream* TestSession::GetCryptoStream() {
 339   return crypto_stream_;
 340 }
 341
 342 TestClientSession::TestClientSession(QuicConnection* connection,
 343                                      const QuicConfig& config)
 344     : QuicClientSessionBase(connection, config),
 345       crypto_stream_(NULL) {
 346   EXPECT_CALL(*this, OnProofValid(_)).Times(AnyNumber());
 347   InitializeSession();
 348 }
 349
 350 TestClientSession::~TestClientSession() {}
 351
 352 void TestClientSession::SetCryptoStream(QuicCryptoStream* stream) {
 353   crypto_stream_ = stream;
 354 }
 355
 356 QuicCryptoStream* TestClientSession::GetCryptoStream() {
 357   return crypto_stream_;
 358 }
 359
 360 MockPacketWriter::MockPacketWriter() {
 361 }
 362
 363 MockPacketWriter::~MockPacketWriter() {
 364 }
 365
 366 MockSendAlgorithm::MockSendAlgorithm() {
 367 }
 368
 369 MockSendAlgorithm::~MockSendAlgorithm() {
 370 }
 371
 372 MockLossAlgorithm::MockLossAlgorithm() {
 373 }
 374
 375 MockLossAlgorithm::~MockLossAlgorithm() {
 376 }
 377
 378 MockAckNotifierDelegate::MockAckNotifierDelegate() {
 379 }
 380
 381 MockAckNotifierDelegate::~MockAckNotifierDelegate() {
 382 }
 383
 384 MockNetworkChangeVisitor::MockNetworkChangeVisitor() {
 385 }
 386
 387 MockNetworkChangeVisitor::~MockNetworkChangeVisitor() {
 388 }
 389
 390 namespace {
 391
 392 string HexDumpWithMarks(const char* data, int length,
 393                         const bool* marks, int mark_length) {
 394   static const char kHexChars[] = "0123456789abcdef";
 395   static const int kColumns = 4;
 396
 397   const int kSizeLimit = 1024;
 398   if (length > kSizeLimit || mark_length > kSizeLimit) {
 399     LOG(ERROR) << "Only dumping first " << kSizeLimit << " bytes.";
 400     length = min(length, kSizeLimit);
 401     mark_length = min(mark_length, kSizeLimit);
 402   }
 403
 404   string hex;
 405   for (const char* row = data; length > 0;
 406        row += kColumns, length -= kColumns) {
 407     for (const char *p = row; p < row + 4; ++p) {
 408       if (p < row + length) {
 409         const bool mark =
 410             (marks && (p - data) < mark_length && marks[p - data]);
 411         hex += mark ? '*' : ' ';
 412         hex += kHexChars[(*p & 0xf0) >> 4];
 413         hex += kHexChars[*p & 0x0f];
 414         hex += mark ? '*' : ' ';
 415       } else {
 416         hex += "    ";
 417       }
 418     }
 419     hex = hex + "  ";
 420
 421     for (const char *p = row; p < row + 4 && p < row + length; ++p)
 422       hex += (*p >= 0x20 && *p <= 0x7f) ? (*p) : '.';
 423
 424     hex = hex + '\n';
 425   }
 426   return hex;
 427 }
 428
 429 }  // namespace
 430
 431 IPAddressNumber TestPeerIPAddress() { return Loopback4(); }
 432
 433 QuicVersion QuicVersionMax() { return QuicSupportedVersions().front(); }
 434
 435 QuicVersion QuicVersionMin() { return QuicSupportedVersions().back(); }
 436
 437 IPAddressNumber Loopback4() {
 438   IPAddressNumber addr;
 439   CHECK(ParseIPLiteralToNumber("127.0.0.1", &addr));
 440   return addr;
 441 }
 442
 443 IPAddressNumber Loopback6() {
 444   IPAddressNumber addr;
 445   CHECK(ParseIPLiteralToNumber("::1", &addr));
 446   return addr;
 447 }
 448
 449 void GenerateBody(string* body, int length) {
 450   body->clear();
 451   body->reserve(length);
 452   for (int i = 0; i < length; ++i) {
 453     body->append(1, static_cast<char>(32 + i % (126 - 32)));
 454   }
 455 }
 456
 457 QuicEncryptedPacket* ConstructEncryptedPacket(
 458     QuicConnectionId connection_id,
 459     bool version_flag,
 460     bool reset_flag,
 461     QuicPacketSequenceNumber sequence_number,
 462     const string& data) {
 463   QuicPacketHeader header;
 464   header.public_header.connection_id = connection_id;
 465   header.public_header.connection_id_length = PACKET_8BYTE_CONNECTION_ID;
 466   header.public_header.version_flag = version_flag;
 467   header.public_header.reset_flag = reset_flag;
 468   header.public_header.sequence_number_length = PACKET_6BYTE_SEQUENCE_NUMBER;
 469   header.packet_sequence_number = sequence_number;
 470   header.entropy_flag = false;
 471   header.entropy_hash = 0;
 472   header.fec_flag = false;
 473   header.is_in_fec_group = NOT_IN_FEC_GROUP;
 474   header.fec_group = 0;
 475   QuicStreamFrame stream_frame(1, false, 0, MakeIOVector(data));
 476   QuicFrame frame(&stream_frame);
 477   QuicFrames frames;
 478   frames.push_back(frame);
 479   QuicFramer framer(QuicSupportedVersions(), QuicTime::Zero(), false);
 480   scoped_ptr<QuicPacket> packet(
 481       BuildUnsizedDataPacket(&framer, header, frames).packet);
 482   EXPECT_TRUE(packet != NULL);
 483   QuicEncryptedPacket* encrypted = framer.EncryptPacket(ENCRYPTION_NONE,
 484                                                         sequence_number,
 485                                                         *packet);
 486   EXPECT_TRUE(encrypted != NULL);
 487   return encrypted;
 488 }
 489
 490 void CompareCharArraysWithHexError(
 491     const string& description,
 492     const char* actual,
 493     const int actual_len,
 494     const char* expected,
 495     const int expected_len) {
 496   EXPECT_EQ(actual_len, expected_len);
 497   const int min_len = min(actual_len, expected_len);
 498   const int max_len = max(actual_len, expected_len);
 499   scoped_ptr<bool[]> marks(new bool[max_len]);
 500   bool identical = (actual_len == expected_len);
 501   for (int i = 0; i < min_len; ++i) {
 502     if (actual[i] != expected[i]) {
 503       marks[i] = true;
 504       identical = false;
 505     } else {
 506       marks[i] = false;
 507     }
 508   }
 509   for (int i = min_len; i < max_len; ++i) {
 510     marks[i] = true;
 511   }
 512   if (identical) return;
 513   ADD_FAILURE()
 514       << "Description:\n"
 515       << description
 516       << "\n\nExpected:\n"
 517       << HexDumpWithMarks(expected, expected_len, marks.get(), max_len)
 518       << "\nActual:\n"
 519       << HexDumpWithMarks(actual, actual_len, marks.get(), max_len);
 520 }
 521
 522 bool DecodeHexString(const base::StringPiece& hex, std::string* bytes) {
 523   bytes->clear();
 524   if (hex.empty())
 525     return true;
 526   std::vector<uint8> v;
 527   if (!base::HexStringToBytes(hex.as_string(), &v))
 528     return false;
 529   if (!v.empty())
 530     bytes->assign(reinterpret_cast<const char*>(&v[0]), v.size());
 531   return true;
 532 }
 533
 534 static QuicPacket* ConstructPacketFromHandshakeMessage(
 535     QuicConnectionId connection_id,
 536     const CryptoHandshakeMessage& message,
 537     bool should_include_version) {
 538   CryptoFramer crypto_framer;
 539   scoped_ptr<QuicData> data(crypto_framer.ConstructHandshakeMessage(message));
 540   QuicFramer quic_framer(QuicSupportedVersions(), QuicTime::Zero(), false);
 541
 542   QuicPacketHeader header;
 543   header.public_header.connection_id = connection_id;
 544   header.public_header.reset_flag = false;
 545   header.public_header.version_flag = should_include_version;
 546   header.packet_sequence_number = 1;
 547   header.entropy_flag = false;
 548   header.entropy_hash = 0;
 549   header.fec_flag = false;
 550   header.fec_group = 0;
 551
 552   QuicStreamFrame stream_frame(kCryptoStreamId, false, 0,
 553                                MakeIOVector(data->AsStringPiece()));
 554
 555   QuicFrame frame(&stream_frame);
 556   QuicFrames frames;
 557   frames.push_back(frame);
 558   return BuildUnsizedDataPacket(&quic_framer, header, frames).packet;
 559 }
 560
 561 QuicPacket* ConstructHandshakePacket(QuicConnectionId connection_id,
 562                                      QuicTag tag) {
 563   CryptoHandshakeMessage message;
 564   message.set_tag(tag);
 565   return ConstructPacketFromHandshakeMessage(connection_id, message, false);
 566 }
 567
 568 size_t GetPacketLengthForOneStream(
 569     QuicVersion version,
 570     bool include_version,
 571     QuicSequenceNumberLength sequence_number_length,
 572     InFecGroup is_in_fec_group,
 573     size_t* payload_length) {
 574   *payload_length = 1;
 575   const size_t stream_length =
 576       NullEncrypter().GetCiphertextSize(*payload_length) +
 577       QuicPacketCreator::StreamFramePacketOverhead(
 578           PACKET_8BYTE_CONNECTION_ID, include_version,
 579           sequence_number_length, 0u, is_in_fec_group);
 580   const size_t ack_length = NullEncrypter().GetCiphertextSize(
 581       QuicFramer::GetMinAckFrameSize(
 582           sequence_number_length, PACKET_1BYTE_SEQUENCE_NUMBER)) +
 583       GetPacketHeaderSize(PACKET_8BYTE_CONNECTION_ID, include_version,
 584                           sequence_number_length, is_in_fec_group);
 585   if (stream_length < ack_length) {
 586     *payload_length = 1 + ack_length - stream_length;
 587   }
 588
 589   return NullEncrypter().GetCiphertextSize(*payload_length) +
 590       QuicPacketCreator::StreamFramePacketOverhead(
 591           PACKET_8BYTE_CONNECTION_ID, include_version,
 592           sequence_number_length, 0u, is_in_fec_group);
 593 }
 594
 595 TestEntropyCalculator::TestEntropyCalculator() {}
 596
 597 TestEntropyCalculator::~TestEntropyCalculator() {}
 598
 599 QuicPacketEntropyHash TestEntropyCalculator::EntropyHash(
 600     QuicPacketSequenceNumber sequence_number) const {
 601   return 1u;
 602 }
 603
 604 MockEntropyCalculator::MockEntropyCalculator() {}
 605
 606 MockEntropyCalculator::~MockEntropyCalculator() {}
 607
 608 QuicConfig DefaultQuicConfig() {
 609   QuicConfig config;
 610   config.SetDefaults();
 611   config.SetInitialFlowControlWindowToSend(
 612       kInitialSessionFlowControlWindowForTest);
 613   config.SetInitialStreamFlowControlWindowToSend(
 614       kInitialStreamFlowControlWindowForTest);
 615   config.SetInitialSessionFlowControlWindowToSend(
 616       kInitialSessionFlowControlWindowForTest);
 617   return config;
 618 }
 619
 620 QuicVersionVector SupportedVersions(QuicVersion version) {
 621   QuicVersionVector versions;
 622   versions.push_back(version);
 623   return versions;
 624 }
 625
 626 TestWriterFactory::TestWriterFactory() : current_writer_(NULL) {}
 627 TestWriterFactory::~TestWriterFactory() {}
 628
 629 QuicPacketWriter* TestWriterFactory::Create(QuicServerPacketWriter* writer,
 630                                             QuicConnection* connection) {
 631   return new PerConnectionPacketWriter(this, writer, connection);
 632 }
 633
 634 void TestWriterFactory::OnPacketSent(WriteResult result) {
 635   if (current_writer_ != NULL && result.status == WRITE_STATUS_ERROR) {
 636     current_writer_->connection()->OnWriteError(result.error_code);
 637     current_writer_ = NULL;
 638   }
 639 }
 640
 641 void TestWriterFactory::Unregister(PerConnectionPacketWriter* writer) {
 642   if (current_writer_ == writer) {
 643     current_writer_ = NULL;
 644   }
 645 }
 646
 647 TestWriterFactory::PerConnectionPacketWriter::PerConnectionPacketWriter(
 648     TestWriterFactory* factory,
 649     QuicServerPacketWriter* writer,
 650     QuicConnection* connection)
 651     : QuicPerConnectionPacketWriter(writer, connection),
 652       factory_(factory) {
 653 }
 654
 655 TestWriterFactory::PerConnectionPacketWriter::~PerConnectionPacketWriter() {
 656   factory_->Unregister(this);
 657 }
 658
 659 WriteResult TestWriterFactory::PerConnectionPacketWriter::WritePacket(
 660     const char* buffer,
 661     size_t buf_len,
 662     const IPAddressNumber& self_address,
 663     const IPEndPoint& peer_address) {
 664   // A DCHECK(factory_current_writer_ == NULL) would be wrong here -- this class
 665   // may be used in a setting where connection()->OnPacketSent() is called in a
 666   // different way, so TestWriterFactory::OnPacketSent might never be called.
 667   factory_->current_writer_ = this;
 668   return QuicPerConnectionPacketWriter::WritePacket(buffer,
 669                                                     buf_len,
 670                                                     self_address,
 671                                                     peer_address);
 672 }
 673
 674 }  // namespace test
 675 }  // namespace net