Add an additional ReliableQuicStream* argument to QuicSession::WritevData in an attempt to defend a…

net/quic/quic_chromium_client_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/quic_chromium_client_stream.h"
 
 #include "base/macros.h"
 #include "base/run_loop.h"
 #include "base/strings/string_number_conversions.h"
 #include "net/base/io_buffer.h"
@@ skipping 48 matching lines @@
   QuicCryptoStream* GetCryptoStream() override { return crypto_stream_.get(); }
 
   // From QuicSession.
   MOCK_METHOD3(OnConnectionClosed,
                void(QuicErrorCode error,
                     const std::string& error_details,
                     ConnectionCloseSource source));
   MOCK_METHOD1(CreateIncomingDynamicStream, QuicSpdyStream*(QuicStreamId id));
   MOCK_METHOD1(CreateOutgoingDynamicStream,
                QuicChromiumClientStream*(SpdyPriority priority));
-  MOCK_METHOD5(WritevData,
-               QuicConsumedData(QuicStreamId id,
+  MOCK_METHOD6(WritevData,
+               QuicConsumedData(ReliableQuicStream* stream,
+                                QuicStreamId id,
                                 QuicIOVector data,
                                 QuicStreamOffset offset,
                                 bool fin,
                                 QuicAckListenerInterface*));
   MOCK_METHOD3(SendRstStream,
                void(QuicStreamId stream_id,
                     QuicRstStreamErrorCode error,
                     QuicStreamOffset bytes_written));
 
   MOCK_METHOD2(OnStreamHeaders,
@@ skipping 45 matching lines @@
 };
 
 MockQuicClientSessionBase::MockQuicClientSessionBase(
     QuicConnection* connection,
     QuicClientPushPromiseIndex* push_promise_index)
     : QuicClientSessionBase(connection,
                             push_promise_index,
                             DefaultQuicConfig()) {
   crypto_stream_.reset(new QuicCryptoStream(this));
   Initialize();
-  ON_CALL(*this, WritevData(_, _, _, _, _))
+  ON_CALL(*this, WritevData(_, _, _, _, _, _))
       .WillByDefault(testing::Return(QuicConsumedData(0, false)));
 }
 
 MockQuicClientSessionBase::~MockQuicClientSessionBase() {}
 
 class QuicChromiumClientStreamTest
     : public ::testing::TestWithParam<QuicVersion> {
  public:
   QuicChromiumClientStreamTest()
       : crypto_config_(CryptoTestUtils::ProofVerifierForTesting()),
@@ skipping 280 matching lines @@
   EXPECT_CALL(delegate_, OnClose(QUIC_NO_ERROR));
 }
 
 TEST_P(QuicChromiumClientStreamTest, WriteStreamData) {
   EXPECT_CALL(delegate_, OnClose(QUIC_NO_ERROR));
 
   const char kData1[] = "hello world";
   const size_t kDataLen = arraysize(kData1);
 
   // All data written.
-  EXPECT_CALL(session_, WritevData(stream_->id(), _, _, _, _))
+  EXPECT_CALL(session_, WritevData(stream_, stream_->id(), _, _, _, _))
       .WillOnce(Return(QuicConsumedData(kDataLen, true)));
   TestCompletionCallback callback;
   EXPECT_EQ(OK, stream_->WriteStreamData(base::StringPiece(kData1, kDataLen),
                                          true, callback.callback()));
 }
 
 TEST_P(QuicChromiumClientStreamTest, WriteStreamDataAsync) {
   EXPECT_CALL(delegate_, HasSendHeadersComplete()).Times(AnyNumber());
   EXPECT_CALL(delegate_, OnClose(QUIC_NO_ERROR));
 
   const char kData1[] = "hello world";
   const size_t kDataLen = arraysize(kData1);
 
   // No data written.
-  EXPECT_CALL(session_, WritevData(stream_->id(), _, _, _, _))
+  EXPECT_CALL(session_, WritevData(stream_, stream_->id(), _, _, _, _))
       .WillOnce(Return(QuicConsumedData(0, false)));
   TestCompletionCallback callback;
   EXPECT_EQ(ERR_IO_PENDING,
             stream_->WriteStreamData(base::StringPiece(kData1, kDataLen), true,
                                      callback.callback()));
   ASSERT_FALSE(callback.have_result());
 
   // All data written.
-  EXPECT_CALL(session_, WritevData(stream_->id(), _, _, _, _))
+  EXPECT_CALL(session_, WritevData(stream_, stream_->id(), _, _, _, _))
       .WillOnce(Return(QuicConsumedData(kDataLen, true)));
   stream_->OnCanWrite();
   ASSERT_TRUE(callback.have_result());
   EXPECT_EQ(OK, callback.WaitForResult());
 }
 
 TEST_P(QuicChromiumClientStreamTest, WritevStreamData) {
   EXPECT_CALL(delegate_, OnClose(QUIC_NO_ERROR));
 
   scoped_refptr<StringIOBuffer> buf1(new StringIOBuffer("hello world!"));
   scoped_refptr<StringIOBuffer> buf2(
       new StringIOBuffer("Just a small payload"));
 
   // All data written.
-  EXPECT_CALL(session_, WritevData(stream_->id(), _, _, _, _))
+  EXPECT_CALL(session_, WritevData(stream_, stream_->id(), _, _, _, _))
       .WillOnce(Return(QuicConsumedData(buf1->size(), false)))
       .WillOnce(Return(QuicConsumedData(buf2->size(), true)));
   TestCompletionCallback callback;
   EXPECT_EQ(OK, stream_->WritevStreamData({buf1.get(), buf2.get()},
                                           {buf1->size(), buf2->size()}, true,
                                           callback.callback()));
 }
 
 TEST_P(QuicChromiumClientStreamTest, WritevStreamDataAsync) {
   EXPECT_CALL(delegate_, HasSendHeadersComplete()).Times(AnyNumber());
   EXPECT_CALL(delegate_, OnClose(QUIC_NO_ERROR));
 
   scoped_refptr<StringIOBuffer> buf1(new StringIOBuffer("hello world!"));
   scoped_refptr<StringIOBuffer> buf2(
       new StringIOBuffer("Just a small payload"));
 
   // Only a part of the data is written.
-  EXPECT_CALL(session_, WritevData(stream_->id(), _, _, _, _))
+  EXPECT_CALL(session_, WritevData(stream_, stream_->id(), _, _, _, _))
       // First piece of data is written.
       .WillOnce(Return(QuicConsumedData(buf1->size(), false)))
       // Second piece of data is queued.
       .WillOnce(Return(QuicConsumedData(0, false)));
   TestCompletionCallback callback;
   EXPECT_EQ(ERR_IO_PENDING,
             stream_->WritevStreamData({buf1.get(), buf2.get()},
                                       {buf1->size(), buf2->size()}, true,
                                       callback.callback()));
   ASSERT_FALSE(callback.have_result());
 
   // The second piece of data is written.
-  EXPECT_CALL(session_, WritevData(stream_->id(), _, _, _, _))
+  EXPECT_CALL(session_, WritevData(stream_, stream_->id(), _, _, _, _))
       .WillOnce(Return(QuicConsumedData(buf2->size(), true)));
   stream_->OnCanWrite();
   ASSERT_TRUE(callback.have_result());
   EXPECT_EQ(OK, callback.WaitForResult());
 }
 
 TEST_P(QuicChromiumClientStreamTest, HeadersBeforeDelegate) {
   // We don't use stream_ because we want an incoming server push
   // stream.
   QuicChromiumClientStream* stream = new QuicChromiumClientStream(
@@ skipping 12 matching lines @@
   stream->SetDelegate(&delegate_);
   base::MessageLoop::current()->RunUntilIdle();
 
   // Times(2) because OnClose will be called for stream and stream_.
   EXPECT_CALL(delegate_, OnClose(QUIC_NO_ERROR)).Times(2);
 }
 
 }  // namespace
 }  // namespace test
 }  // namespace net