Fix bug in ReliableQuicStream::WritevData() where if the write failed and the connection was closed…

net/quic/reliable_quic_stream_test.cc

 // 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.
 
 #include "net/quic/reliable_quic_stream.h"
 
 #include "net/quic/quic_connection.h"
 #include "net/quic/quic_flags.h"
 #include "net/quic/quic_utils.h"
 #include "net/quic/quic_write_blocked_list.h"
 #include "net/quic/spdy_utils.h"
 #include "net/quic/test_tools/quic_config_peer.h"
 #include "net/quic/test_tools/quic_connection_peer.h"
 #include "net/quic/test_tools/quic_flow_controller_peer.h"
 #include "net/quic/test_tools/quic_session_peer.h"
 #include "net/quic/test_tools/quic_test_utils.h"
 #include "net/quic/test_tools/reliable_quic_stream_peer.h"
 #include "net/test/gtest_util.h"
 #include "testing/gmock/include/gmock/gmock.h"
 #include "testing/gmock_mutant.h"
 
 using base::StringPiece;
 using std::min;
 using std::string;
 using testing::AnyNumber;
 using testing::AtLeast;
 using testing::CreateFunctor;
 using testing::InSequence;
 using testing::Invoke;
+using testing::DoAll;
 using testing::Return;
 using testing::StrictMock;
 using testing::WithArgs;
 using testing::_;
 
 namespace net {
 namespace test {
 namespace {
 
 const char kData1[] = "FooAndBar";
@@ skipping 93 matching lines @@
 
   void set_initial_flow_control_window_bytes(uint32 val) {
     initial_flow_control_window_bytes_ = val;
   }
 
   bool HasWriteBlockedStreams() {
     return write_blocked_list_->HasWriteBlockedCryptoOrHeadersStream() ||
            write_blocked_list_->HasWriteBlockedDataStreams();
   }
 
+  QuicConsumedData CloseStreamOnWriteError(
+      QuicStreamId id,
+      QuicIOVector /*iov*/,
+      QuicStreamOffset /*offset*/,
+      bool /*fin*/,
+      FecProtection /*fec_protection*/,
+      QuicAckListenerInterface* /*ack_notifier_delegate*/) {
+    session_->CloseStream(id);
+    return QuicConsumedData(1, false);
+
+  }
+
  protected:
   MockConnectionHelper helper_;
   MockConnection* connection_;
   scoped_ptr<MockQuicSpdySession> session_;
   TestStream* stream_;
   SpdyHeaderBlock headers_;
   QuicWriteBlockedList* write_blocked_list_;
   uint32 initial_flow_control_window_bytes_;
   QuicTime::Delta zero_;
   QuicVersionVector supported_versions_;
@@ skipping 51 matching lines @@
   Initialize(kShouldProcessData);
 
   // Write no data and a fin.  If we consume nothing we should be write blocked,
   // as the fin was not consumed.
   EXPECT_CALL(*session_, WritevData(kTestStreamId, _, _, _, _, _))
       .WillOnce(Return(QuicConsumedData(0, false)));
   stream_->WriteOrBufferData(StringPiece(), true, nullptr);
   ASSERT_EQ(1u, write_blocked_list_->NumBlockedStreams());
 }
 
+TEST_F(ReliableQuicStreamTest, CloseOnPartialWrite) {
+  Initialize(kShouldProcessData);
+
+  // Write some data and no fin. However, while writing the data
+  // close the stream and verify that MarkConnectionLevelWriteBlocked does not
+  // crash with an unknown stream.
+  EXPECT_CALL(*session_, WritevData(kTestStreamId, _, _, _, _, _))
+      .WillOnce(Invoke(this, &ReliableQuicStreamTest::CloseStreamOnWriteError));
+  stream_->WriteOrBufferData(StringPiece(kData1, 2), false, nullptr);
+  ASSERT_EQ(0u, write_blocked_list_->NumBlockedStreams());
+}
+
 TEST_F(ReliableQuicStreamTest, WriteOrBufferData) {
   Initialize(kShouldProcessData);
 
   EXPECT_FALSE(HasWriteBlockedStreams());
   size_t length = 1 + QuicPacketCreator::StreamFramePacketOverhead(
                           PACKET_8BYTE_CONNECTION_ID, !kIncludeVersion,
                           PACKET_6BYTE_PACKET_NUMBER, 0u, NOT_IN_FEC_GROUP);
   connection_->SetMaxPacketLength(length);
 
   EXPECT_CALL(*session_, WritevData(_, _, _, _, _, _)).WillOnce(
@@ skipping 518 matching lines @@
     EXPECT_TRUE(stream_->HasFinalReceivedByteOffset());
   } else {
     EXPECT_FALSE(stream_->fin_received());
     EXPECT_FALSE(stream_->HasFinalReceivedByteOffset());
   }
 }
 
 }  // namespace
 }  // namespace test
 }  // namespace net