Limit the total memory usage for Stream instances

content/browser/byte_stream_unittest.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 "content/browser/byte_stream.h"
 
 #include <deque>
+#include <limits>
 
 #include "base/bind.h"
 #include "base/callback.h"
 #include "base/memory/ref_counted.h"
 #include "base/message_loop/message_loop.h"
 #include "base/test/test_simple_task_runner.h"
 #include "net/base/io_buffer.h"
 #include "testing/gtest/include/gtest/gtest.h"
 
 namespace content {
@@ skipping 91 matching lines @@
 
   // Push a series of IO buffers on; test pushback happening and
   // that it's advisory.
   EXPECT_TRUE(Write(byte_stream_input.get(), 1024));
   EXPECT_TRUE(Write(byte_stream_input.get(), 1024));
   EXPECT_TRUE(Write(byte_stream_input.get(), 1024));
   EXPECT_FALSE(Write(byte_stream_input.get(), 1));
   EXPECT_FALSE(Write(byte_stream_input.get(), 1024));
   // Flush
   byte_stream_input->Close(0);
+  EXPECT_EQ(4 * 1024U + 1U, byte_stream_input->GetTotalBufferedBytes());
   message_loop_.RunUntilIdle();
+  // Data already sent to reader is also counted in.
+  EXPECT_EQ(4 * 1024U + 1U, byte_stream_input->GetTotalBufferedBytes());
 
   // Pull the IO buffers out; do we get the same buffers and do they
   // have the same contents?
   scoped_refptr<net::IOBuffer> output_io_buffer;
   size_t output_length;
   EXPECT_EQ(ByteStreamReader::STREAM_HAS_DATA,
             byte_stream_output->Read(&output_io_buffer, &output_length));
   EXPECT_TRUE(ValidateIOBuffer(output_io_buffer, output_length));
 
   EXPECT_EQ(ByteStreamReader::STREAM_HAS_DATA,
             byte_stream_output->Read(&output_io_buffer, &output_length));
   EXPECT_TRUE(ValidateIOBuffer(output_io_buffer, output_length));
 
   EXPECT_EQ(ByteStreamReader::STREAM_HAS_DATA,
             byte_stream_output->Read(&output_io_buffer, &output_length));
   EXPECT_TRUE(ValidateIOBuffer(output_io_buffer, output_length));
 
   EXPECT_EQ(ByteStreamReader::STREAM_HAS_DATA,
             byte_stream_output->Read(&output_io_buffer, &output_length));
   EXPECT_TRUE(ValidateIOBuffer(output_io_buffer, output_length));
 
   EXPECT_EQ(ByteStreamReader::STREAM_HAS_DATA,
             byte_stream_output->Read(&output_io_buffer, &output_length));
   EXPECT_TRUE(ValidateIOBuffer(output_io_buffer, output_length));
 
   EXPECT_EQ(ByteStreamReader::STREAM_COMPLETE,
             byte_stream_output->Read(&output_io_buffer, &output_length));
+
+  message_loop_.RunUntilIdle();
+  // Reader now knows that all data is read out.
+  EXPECT_EQ(1024U, byte_stream_input->GetTotalBufferedBytes());
 }
 
 // Confirm that Flush() method makes the writer to send written contents to
 // the reader.
 TEST_F(ByteStreamTest, ByteStream_Flush) {
   scoped_ptr<ByteStreamWriter> byte_stream_input;
   scoped_ptr<ByteStreamReader> byte_stream_output;
   CreateByteStream(
       message_loop_.message_loop_proxy(), message_loop_.message_loop_proxy(),
       1024, &byte_stream_input, &byte_stream_output);
@@ skipping 415 matching lines @@
   EXPECT_EQ(ByteStreamReader::STREAM_EMPTY,
             byte_stream_output->Read(&output_io_buffer, &output_length));
 
   byte_stream_input->Close(0);
   message_loop_.RunUntilIdle();
 
   EXPECT_EQ(ByteStreamReader::STREAM_COMPLETE,
             byte_stream_output->Read(&output_io_buffer, &output_length));
 }
 
+TEST_F(ByteStreamTest, ByteStream_WriteOverflow) {
+  scoped_ptr<ByteStreamWriter> byte_stream_input;
+  scoped_ptr<ByteStreamReader> byte_stream_output;
+  CreateByteStream(
+      message_loop_.message_loop_proxy(), message_loop_.message_loop_proxy(),
+      std::numeric_limits<size_t>::max(),
+      &byte_stream_input, &byte_stream_output);
+
+  EXPECT_TRUE(Write(byte_stream_input.get(), 1));
+  // 1 + size_t max -> Overflow.
+  scoped_refptr<net::IOBuffer> empty_io_buffer;
+  EXPECT_FALSE(byte_stream_input->Write(empty_io_buffer,
+                                        std::numeric_limits<size_t>::max()));
+  message_loop_.RunUntilIdle();
+
+  // The first write is below PostToPeer threshold. We shouldn't get anything
+  // from the output.
+  scoped_refptr<net::IOBuffer> output_io_buffer;
+  size_t output_length;
+  EXPECT_EQ(ByteStreamReader::STREAM_EMPTY,
+            byte_stream_output->Read(&output_io_buffer, &output_length));
+}
+
 }  // namespace content