Limit the total memory usage for Stream instances

content/browser/byte_stream.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 <set>
 #include <utility>
 
@@ skipping 43 matching lines @@
   void SetPeer(ByteStreamReaderImpl* peer,
                scoped_refptr<base::SequencedTaskRunner> peer_task_runner,
                scoped_refptr<LifetimeFlag> peer_lifetime_flag);
 
   // Overridden from ByteStreamWriter.
   virtual bool Write(scoped_refptr<net::IOBuffer> buffer,
                      size_t byte_count) OVERRIDE;
   virtual void Flush() OVERRIDE;
   virtual void Close(int status) OVERRIDE;
   virtual void RegisterCallback(const base::Closure& source_callback) OVERRIDE;
+  virtual size_t GetTotalBufferedBytes() const OVERRIDE;
 
   // PostTask target from |ByteStreamReaderImpl::MaybeUpdateInput|.
   static void UpdateWindow(scoped_refptr<LifetimeFlag> lifetime_flag,
                            ByteStreamWriterImpl* target,
                            size_t bytes_consumed);
 
  private:
   // Called from UpdateWindow when object existence has been validated.
   void UpdateWindowInternal(size_t bytes_consumed);
 
@@ skipping 128 matching lines @@
     scoped_refptr<LifetimeFlag> peer_lifetime_flag) {
   peer_ = peer;
   peer_task_runner_ = peer_task_runner;
   peer_lifetime_flag_ = peer_lifetime_flag;
 }
 
 bool ByteStreamWriterImpl::Write(
     scoped_refptr<net::IOBuffer> buffer, size_t byte_count) {
   DCHECK(my_task_runner_->RunsTasksOnCurrentThread());
 
+  // Check overflow.
+  //
+  // TODO(tyoshino): Discuss with content/browser/download developer and if
+  // they're fine with, set smaller limit and make it configurable.
+  size_t space_limit = std::numeric_limits<size_t>::max() -
+      GetTotalBufferedBytes();
+  if (byte_count > space_limit) {
+    // TODO(tyoshino): Tell the user that Write() failed.
+    // Ignore input.
+    return false;
+  }
+
   input_contents_.push_back(std::make_pair(buffer, byte_count));
   input_contents_size_ += byte_count;
 
   // Arbitrarily, we buffer to a third of the total size before sending.
   if (input_contents_size_ > total_buffer_size_ / kFractionBufferBeforeSending)
     PostToPeer(false, 0);
 
-  return (input_contents_size_ + output_size_used_ <= total_buffer_size_);
+  return GetTotalBufferedBytes() <= total_buffer_size_;
 }
 
 void ByteStreamWriterImpl::Flush() {
   DCHECK(my_task_runner_->RunsTasksOnCurrentThread());
   if (input_contents_size_ > 0)
     PostToPeer(false, 0);
 }
 
 void ByteStreamWriterImpl::Close(int status) {
   DCHECK(my_task_runner_->RunsTasksOnCurrentThread());
   PostToPeer(true, status);
 }
 
 void ByteStreamWriterImpl::RegisterCallback(
     const base::Closure& source_callback) {
   DCHECK(my_task_runner_->RunsTasksOnCurrentThread());
   space_available_callback_ = source_callback;
 }
 
+size_t ByteStreamWriterImpl::GetTotalBufferedBytes() const {
+  DCHECK(my_task_runner_->RunsTasksOnCurrentThread());
+  // This sum doesn't overflow since Write() fails if this sum is going to
+  // overflow.
+  return input_contents_size_ + output_size_used_;
+}
+
 // static
 void ByteStreamWriterImpl::UpdateWindow(
     scoped_refptr<LifetimeFlag> lifetime_flag, ByteStreamWriterImpl* target,
     size_t bytes_consumed) {
   // If the target object isn't alive anymore, we do nothing.
   if (!lifetime_flag->is_alive) return;
 
   target->UpdateWindowInternal(bytes_consumed);
 }
 
 void ByteStreamWriterImpl::UpdateWindowInternal(size_t bytes_consumed) {
   DCHECK(my_task_runner_->RunsTasksOnCurrentThread());
+
+  bool was_above_limit = GetTotalBufferedBytes() > total_buffer_size_;
+
   DCHECK_GE(output_size_used_, bytes_consumed);
   output_size_used_ -= bytes_consumed;
 
   // Callback if we were above the limit and we're now <= to it.
-  size_t total_known_size_used =
-      input_contents_size_ + output_size_used_;
+  bool no_longer_above_limit = GetTotalBufferedBytes() <= total_buffer_size_;
 
-  if (total_known_size_used <= total_buffer_size_ &&
-      (total_known_size_used + bytes_consumed > total_buffer_size_) &&
+  if (no_longer_above_limit && was_above_limit &&
       !space_available_callback_.is_null())
     space_available_callback_.Run();
 }
 
 void ByteStreamWriterImpl::PostToPeer(bool complete, int status) {
   DCHECK(my_task_runner_->RunsTasksOnCurrentThread());
   // Valid contexts in which to call.
   DCHECK(complete || 0 != input_contents_size_);
 
   scoped_ptr<ContentVector> transfer_buffer;
@@ skipping 162 matching lines @@
   ByteStreamReaderImpl* out = new ByteStreamReaderImpl(
       output_task_runner, output_flag, buffer_size);
 
   in->SetPeer(out, output_task_runner, output_flag);
   out->SetPeer(in, input_task_runner, input_flag);
   input->reset(in);
   output->reset(out);
 }
 
 }  // namespace content