Limit the total memory usage for Stream instances

content/browser/streams/stream.cc

 // Copyright (c) 2013 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/streams/stream.h"
 
+#include <limits>
+
 #include "base/bind.h"
 #include "base/location.h"
 #include "base/message_loop/message_loop_proxy.h"
 #include "content/browser/streams/stream_handle_impl.h"
 #include "content/browser/streams/stream_read_observer.h"
 #include "content/browser/streams/stream_registry.h"
 #include "content/browser/streams/stream_write_observer.h"
 #include "net/base/io_buffer.h"
 
 namespace {
 // Start throttling the connection at about 1MB.
 const size_t kDeferSizeThreshold = 40 * 32768;
 }
 
 namespace content {
 
 Stream::Stream(StreamRegistry* registry,
                StreamWriteObserver* write_observer,
                const GURL& url)
     : data_bytes_read_(0),
       can_add_data_(true),
       url_(url),
       data_length_(0),
+      last_total_buffered_bytes_(0),
       registry_(registry),
       read_observer_(NULL),
       write_observer_(write_observer),
       stream_handle_(NULL),
       weak_ptr_factory_(this) {
   CreateByteStream(base::MessageLoopProxy::current(),
                    base::MessageLoopProxy::current(),
                    kDeferSizeThreshold,
                    &writer_,
                    &reader_);
@@ skipping 20 matching lines @@
 void Stream::RemoveReadObserver(StreamReadObserver* observer) {
   DCHECK(observer == read_observer_);
   read_observer_ = NULL;
 }
 
 void Stream::RemoveWriteObserver(StreamWriteObserver* observer) {
   DCHECK(observer == write_observer_);
   write_observer_ = NULL;
 }
 
+void Stream::Abort() {
+  // Clear all buffer. It's safe to clear reader_ here since the same thread
+  // is used for both input and output operation.
+  writer_.reset();
+  reader_.reset();
+  can_add_data_ = false;
+  registry_->UnregisterStream(url());
+}
+
 void Stream::AddData(scoped_refptr<net::IOBuffer> buffer, size_t size) {
+  if (!writer_.get())
+    return;
+
+  size_t current_buffered_bytes = writer_->GetTotalBufferedBytes();
+  if (size > std::numeric_limits<size_t>::max() - current_buffered_bytes) {
+    Abort();
+    return;
+  }
+  size_t new_buffered_bytes = current_buffered_bytes + size;
+  if (!registry_->UpdateMemoryUsage(url(), new_buffered_bytes)) {
+    Abort();
+    return;
+  }
   can_add_data_ = writer_->Write(buffer, size);
+  last_total_buffered_bytes_ = new_buffered_bytes;
 }
 
 void Stream::AddData(const char* data, size_t size) {
   scoped_refptr<net::IOBuffer> io_buffer(new net::IOBuffer(size));
   memcpy(io_buffer->data(), data, size);
-  can_add_data_ = writer_->Write(io_buffer, size);
+  AddData(io_buffer, size);
 }
 
 void Stream::Finalize() {
+  if (!writer_.get())
+    return;
+
   writer_->Close(0);
-  writer_.reset(NULL);
+  writer_.reset();
 
   // Continue asynchronously.
   base::MessageLoopProxy::current()->PostTask(
       FROM_HERE,
       base::Bind(&Stream::OnDataAvailable, weak_ptr_factory_.GetWeakPtr()));
 }
 
 Stream::StreamState Stream::ReadRawData(net::IOBuffer* buf,
                                         int buf_size,
                                         int* bytes_read) {
   DCHECK(buf);
   DCHECK(bytes_read);
 
   *bytes_read = 0;
   if (!data_.get()) {
+    // TODO(tyoshino): Add STREAM_ABORTED type to tell the reader that this
+    // stream is aborted.
+    if (!reader_.get())
+      return STREAM_EMPTY;
+
     data_length_ = 0;
     data_bytes_read_ = 0;
     ByteStreamReader::StreamState state = reader_->Read(&data_, &data_length_);
     switch (state) {
       case ByteStreamReader::STREAM_HAS_DATA:
         break;
       case ByteStreamReader::STREAM_COMPLETE:
         registry_->UnregisterStream(url());
         return STREAM_COMPLETE;
       case ByteStreamReader::STREAM_EMPTY:
@@ skipping 39 matching lines @@
   if (write_observer_)
     write_observer_->OnSpaceAvailable(this);
 }
 
 void Stream::OnDataAvailable() {
   if (read_observer_)
     read_observer_->OnDataAvailable(this);
 }
 
 }  // namespace content