Relax ad-hoc assumptions on InterceptingResourceHandler

content/browser/loader/intercepting_resource_handler.cc

 // Copyright 2016 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/loader/intercepting_resource_handler.h"
 
 #include "base/logging.h"
 #include "base/strings/string_util.h"
 #include "content/public/common/resource_response.h"
 #include "net/base/io_buffer.h"
 
 namespace content {
 
 InterceptingResourceHandler::InterceptingResourceHandler(
     std::unique_ptr<ResourceHandler> next_handler,
     net::URLRequest* request)
-    : LayeredResourceHandler(request, std::move(next_handler)),
-      state_(State::STARTING),
-      first_read_buffer_size_(0) {}
+    : LayeredResourceHandler(request, std::move(next_handler)) {
+  next_handler_->SetController(this);
+}
 
 InterceptingResourceHandler::~InterceptingResourceHandler() {}
 
+void InterceptingResourceHandler::SetController(
+    ResourceController* controller) {
+  if (state_ == State::PASS_THROUGH)
+    return LayeredResourceHandler::SetController(controller);
+  ResourceHandler::SetController(controller);
+}
+
 bool InterceptingResourceHandler::OnResponseStarted(ResourceResponse* response,
                                                     bool* defer) {
   // If there's no need to switch handlers, just start acting as a blind
   // pass-through ResourceHandler.
   if (!new_handler_) {
-    state_ = State::DONE;
-    first_read_buffer_ = nullptr;
+    state_ = State::PASS_THROUGH;
+    next_handler_->SetController(controller());
     return next_handler_->OnResponseStarted(response, defer);
   }
 
+  DCHECK_EQ(state_, State::STARTING);
   // Otherwise, switch handlers. First, inform the original ResourceHandler
   // that this will be handled entirely by the new ResourceHandler.
-  // TODO(clamy): We will probably need to check the return values of these for
-  // PlzNavigate.
   bool defer_ignored = false;
-  next_handler_->OnResponseStarted(response, &defer_ignored);
+  if (!next_handler_->OnResponseStarted(response, &defer_ignored))
+    return false;
 
   // Although deferring OnResponseStarted is legal, the only downstream handler
   // which does so is CrossSiteResourceHandler. Cross-site transitions should
   // not trigger when switching handlers.
   DCHECK(!defer_ignored);
 
-  // Make a copy of the data in the first read buffer. Despite not having been
-  // informed of any data being stored in first_read_buffer_, the
-  // MimeSniffingResourceHandler has read the data, it's just holding it back.
-  // This data should be passed to the alternate ResourceHandler and not to to
-  // the current ResourceHandler.
-  // TODO(clamy): see if doing the copy should be moved to the
-  // MimeSniffingResourceHandler.
-  if (first_read_buffer_) {
-    first_read_buffer_copy_ = new net::IOBuffer(first_read_buffer_size_);
-    memcpy(first_read_buffer_copy_->data(), first_read_buffer_->data(),
-           first_read_buffer_size_);
-  }
-
-  // Send the payload to the old handler.
-  SendPayloadToOldHandler();
-  first_read_buffer_ = nullptr;
-
-  // The original ResourceHandler is now no longer needed, so replace it with
-  // the new one, before sending the response to the new one.
-  next_handler_ = std::move(new_handler_);
-
-  state_ =
-      first_read_buffer_copy_ ? State::WAITING_FOR_BUFFER_COPY : State::DONE;
-
-  return next_handler_->OnResponseStarted(response, defer);
+  // TODO(yhirano): Retaining ownership from a raw pointer is bad.
+  response_ = response;
+  state_ = State::SENDING_PAYLOAD_TO_OLD_HANDLER;
+  return DoLoop(defer);
 }
 
 bool InterceptingResourceHandler::OnWillRead(scoped_refptr<net::IOBuffer>* buf,
                                              int* buf_size,
                                              int min_size) {
-  if (state_ == State::DONE)
+  if (state_ == State::PASS_THROUGH)
     return next_handler_->OnWillRead(buf, buf_size, min_size);
 
   DCHECK_EQ(State::STARTING, state_);
   DCHECK_EQ(-1, min_size);
 
   if (!next_handler_->OnWillRead(buf, buf_size, min_size))
     return false;
 
   first_read_buffer_ = *buf;
   first_read_buffer_size_ = *buf_size;
+  first_read_buffer_double_ = new net::IOBuffer(static_cast<size_t>(*buf_size));
+  *buf = first_read_buffer_double_;

[mmenke, 2016/10/10 19:40:15]

This code is optimized for the case where we switch to a new handler, at the
cost of doing an extra copy in the case when we don't switch to a new handler. 
I think that 90+% of the time, we don't switch handlers?  Would it be better to
do two extra copies in the case we switch handlers, and 0 when we don't?

[yhirano, 2016/10/11 11:06:40]

My understanding is that we need this copy for correctness. When the following
functions are called in order,

 1. OnWillRead
 2. UseNewHandler
 3. OnResponseStarted

we need to call old handler's OnReadCompleted, but the buffer we exposed to the
client at OnWillRead should not be used because the client hasn't finished
writing to the buffer yet. I wanted to remove [1], but you think the performance
loss is unacceptable, right?

1:
https://cs.chromium.org/chromium/src/content/browser/loader/intercepting_reso...

[mmenke, 2016/10/11 18:04:27]

You're absolutely correct - I mean, the buffer has actually been written to, by
that point, but we haven't been informed about it.  And even if we had, we
really should not modify it before OnReadCompleted is called on the intermediate
ResourceHandler, between the sniffer and the interceptor.

It's not the performance cost of removing [1], it's that it messes up ordering
in which we do other things.

I can't think of a better way to make everything sane than this current code, so
let's keep this as-is.

   return true;
 }
 
 bool InterceptingResourceHandler::OnReadCompleted(int bytes_read, bool* defer) {
-  DCHECK(bytes_read >= 0);
-  if (state_ == State::DONE)
+  DCHECK_GE(bytes_read, 0);
+  if (state_ == State::PASS_THROUGH) {
+    if (first_read_buffer_double_) {
+      // |first_read_buffer_double_| was allocated and the user wrote data to
+      // the buffer, but switching has not been done after all.
+      memcpy(first_read_buffer_->data(), first_read_buffer_double_->data(),
+             bytes_read);
+      first_read_buffer_ = nullptr;
+      first_read_buffer_double_ = nullptr;
+    }
     return next_handler_->OnReadCompleted(bytes_read, defer);
-
-  DCHECK_EQ(State::WAITING_FOR_BUFFER_COPY, state_);
-  state_ = State::DONE;
-
-  // Copy the data from the first read to the new ResourceHandler.
-  scoped_refptr<net::IOBuffer> buf;
-  int buf_len = 0;
-  if (!next_handler_->OnWillRead(&buf, &buf_len, bytes_read))
-    return false;
-
-  CHECK(buf_len >= bytes_read);
-  CHECK_GE(first_read_buffer_size_, static_cast<size_t>(bytes_read));
-  memcpy(buf->data(), first_read_buffer_copy_->data(), bytes_read);
-
-  first_read_buffer_copy_ = nullptr;
-
-  // TODO(clamy): Add a unit test to check that the deferral value is properly
-  // passed to the caller.
-  return next_handler_->OnReadCompleted(bytes_read, defer);
+  }
+
+  DCHECK_EQ(State::WAITING_FOR_ON_READ_COMPLETED, state_);
+  first_read_buffer_bytes_read_ = bytes_read;
+  state_ = State::SENDING_BUFFER_TO_NEW_HANDLER;
+  return DoLoop(defer);
+}
+
+void InterceptingResourceHandler::OnResponseCompleted(
+    const net::URLRequestStatus& status,
+    bool* defer) {
+  switch (state_) {
+    case State::PASS_THROUGH:
+      LayeredResourceHandler::OnResponseCompleted(status, defer);
+      return;
+    case State::NOTIFYING_ON_RESPONSE_COMPLETED_TO_NEW_HANDLER:
+    case State::DONE:
+      NOTREACHED();
+      return;
+    default:
+      break;
+  }
+
+  next_handler_->OnResponseCompleted(status, defer);
+  DCHECK(!*defer);
+
+  state_ = State::NOTIFYING_ON_RESPONSE_COMPLETED_TO_NEW_HANDLER;
+  if (new_handler_)
+    new_handler_->OnResponseCompleted(status, defer);

[mmenke, 2016/10/11 21:34:35]

This seems not quite right.  If we're in
State::NOTIFYING_ON_RESPONSE_STARTED_TO_NEW_HANDLER or
State::WAITING_FOR_ON_READ_COMPLETED, we'll never tell any handler about the
completed response.

Seems like the safest thing to do here is:

void InterceptingResourceHandler::OnResponseCompleted(...) {
  if (state_ != State::PASS_THROUGH) {
    // |this| either is in the middle of transferring control to a handler,
    // or has not yet been notified of the response starting.
    // In the first case, just swap in the new handler and tell it of the
    // failure.  In the second case, just switch to pass through mode.

    if (new_handler_)
      next_handler_ = std::move(new_handler_);

    state_ = State::PASS_THROUGH;
    first_read_buffer_double_ = nullptr;
    next_handler_->SetController(controller());
  }

  DCHECK(!new_handler_);

  return LayeredController::OnResponseCompleted(...);
}

This would also let us get rid of two states.

Also, add tests for this?

[yhirano, 2016/10/12 10:36:21]

When |state_| is STARTING and UseNewHandler has been called, there are two
handlers but the suggested code calls only one OnResponseCompleted. Is it the
expected behavior?

When |state_| is one of NOTIFYING_ON_RESPONSE_STARTED_TO_NEW_HANDLER,
WAITING_FOR_ON_READ_COMPLETED and SENDING_BUFFER_TO_NEW_HANDLER, |next_handler_|
here is actually the new handler and |new_handler_| is null. In that case the
suggested code works as same as PS8 (except for |defer| handling, which I fixed
at PS10), I guess?

Switching the controller is a great idea, Thank you!

+  if (*defer)
+    return;
+  DoLoop(defer);
+}
+
+void InterceptingResourceHandler::Cancel() {

[mmenke, 2016/10/11 21:34:35]

In all of these 4, DCHECK_NE(state_, State::PASS_THROUGH);?  Could also DCHECK
it's not STARTING.

[yhirano, 2016/10/12 10:36:21]

Done.

The state can be STARTING because uninteresting method calls can trigger the
operations.

+  controller()->Cancel();
+}
+
+void InterceptingResourceHandler::CancelAndIgnore() {
+  controller()->CancelAndIgnore();
+}
+
+void InterceptingResourceHandler::CancelWithError(int error_code) {
+  controller()->CancelWithError(error_code);
+}
+
+void InterceptingResourceHandler::Resume() {
+  if (state_ == State::STARTING ||
+      state_ == State::WAITING_FOR_ON_READ_COMPLETED) {
+    // Uninteresting Resume: just delegate to the original resource controller.
+    controller()->Resume();
+    return;
+  }
+  bool defer = false;
+  if (!DoLoop(&defer)) {
+    controller()->Cancel();
+    return;
+  }
+
+  if (!defer)
+    controller()->Resume();
 }
 
 void InterceptingResourceHandler::UseNewHandler(
     std::unique_ptr<ResourceHandler> new_handler,
     const std::string& payload_for_old_handler) {
   new_handler_ = std::move(new_handler);
-  new_handler_->SetController(controller());
+  new_handler_->SetController(this);
   payload_for_old_handler_ = payload_for_old_handler;
 }
 
-void InterceptingResourceHandler::SendPayloadToOldHandler() {
+bool InterceptingResourceHandler::SendPayloadToOldHandler(bool* defer) {
+  DCHECK_EQ(State::SENDING_PAYLOAD_TO_OLD_HANDLER, state_);
+  while (payload_bytes_written_ < payload_for_old_handler_.size()) {
+    scoped_refptr<net::IOBuffer> buffer;
+    int size = 0;
+    if (first_read_buffer_) {
+      buffer = first_read_buffer_;
+      size = first_read_buffer_size_;
+
+      first_read_buffer_ = nullptr;
+      first_read_buffer_size_ = 0;
+    } else {
+      if (!next_handler_->OnWillRead(&buffer, &size, -1))
+        return false;
+    }
+
+    size = std::min(size, static_cast<int>(payload_for_old_handler_.size() -
+                                           payload_bytes_written_));
+    memcpy(buffer->data(),
+           payload_for_old_handler_.data() + payload_bytes_written_, size);
+    if (!next_handler_->OnReadCompleted(size, defer))
+      return false;
+    payload_bytes_written_ += size;
+    if (*defer)
+      return true;
+  }

[mmenke, 2016/10/10 19:40:15]

nit:  Add blank line here?

[yhirano, 2016/10/12 10:36:21]

Done.

   bool defer_ignored = false;
+  net::URLRequestStatus status = {net::URLRequestStatus::SUCCESS, 0};

[mmenke, 2016/10/10 19:40:16]

Just URLRequestStatus?  Or use the constructor explicitly? 
net::URLRequestStatus status(net::URLRequestStatus::SUCCESS, 0);

[yhirano, 2016/10/12 10:36:21]

Done.

   if (payload_for_old_handler_.empty()) {
     // If there is no payload, just finalize the request on the old handler.
-    net::URLRequestStatus status(net::URLRequestStatus::CANCELED,
-                                 net::ERR_ABORTED);
-    next_handler_->OnResponseCompleted(status, &defer_ignored);
-    DCHECK(!defer_ignored);
-    return;
-  }
-
-  // Ensure the old ResourceHandler has a buffer that can store the payload.
-  scoped_refptr<net::IOBuffer> buf;
-  int size = 0;
-  if (first_read_buffer_) {
-    // The first read buffer can be reused. The data inside it has been copied
-    // before calling this function, so it can safely be overriden.
-    buf = first_read_buffer_;
-    size = first_read_buffer_size_;
-  }
-
-  // If there is no first read buffer, ask the old ResourceHandler to create a
-  // buffer that can contain payload.
-  if (!buf)
-    next_handler_->OnWillRead(&buf, &size, -1);
-
-  DCHECK(buf);
-  CHECK_GE(size, static_cast<int>(payload_for_old_handler_.length()));
-  memcpy(buf->data(), payload_for_old_handler_.c_str(),
-         payload_for_old_handler_.length());
-  next_handler_->OnReadCompleted(payload_for_old_handler_.length(),
-                                 &defer_ignored);
-  payload_for_old_handler_ = std::string();
-  DCHECK(!defer_ignored);
-
-  // Finalize the request.
-  net::URLRequestStatus status(net::URLRequestStatus::SUCCESS, 0);
+    status = net::URLRequestStatus(net::URLRequestStatus::CANCELED,
+                                   net::ERR_ABORTED);

[mmenke, 2016/10/10 19:40:16]

net::URLREQuestStatus::FromError(ERR_ABORTED).  We want to get rid of
URLRequestStatus::CANCELED, some day.

[yhirano, 2016/10/12 10:36:21]

Done.

+  }
   next_handler_->OnResponseCompleted(status, &defer_ignored);

[mmenke, 2016/10/10 19:40:16]

Could trivially add support for this deferring the operation - don't think it's
worth a new test.

[yhirano, 2016/10/12 10:36:21]

Let me confirm what you mean by "trivially" - we need to add a new state to stop
the controller from disposing the request. We may need another state for
OnresponseCompleted. Do you agree with that?

[mmenke, 2016/10/12 15:04:01]

Exactly that.  We'd need one more state, but don't think it's a big deal.

[yhirano, 2016/10/13 11:36:21]

Done.

   DCHECK(!defer_ignored);
+
+  next_handler_ = std::move(new_handler_);
+  state_ = State::NOTIFYING_ON_RESPONSE_STARTED_TO_NEW_HANDLER;
+  if (!next_handler_->OnResponseStarted(response_.get(), defer))
+    return false;
+  if (*defer)
+    return true;
+  return DoLoop(defer);

[mmenke, 2016/10/10 19:40:16]

This method is only called from DoLoop, so the above 5 lines should just be:

return next_handler_->OnResponseStarted(response_.get(), defer);

[yhirano, 2016/10/12 10:36:21]

Done.

+}
+
+bool InterceptingResourceHandler::SendFirstReadBufferToNewHandler(bool* defer) {
+  DCHECK_EQ(state_, State::SENDING_BUFFER_TO_NEW_HANDLER);
+
+  while (first_read_buffer_bytes_written_ < first_read_buffer_bytes_read_) {
+    scoped_refptr<net::IOBuffer> buf;
+    int size = 0;
+    if (!next_handler_->OnWillRead(&buf, &size, -1))
+      return false;
+    size = std::min(size, static_cast<int>(first_read_buffer_bytes_read_ -
+                                           first_read_buffer_bytes_written_));
+    memcpy(buf->data(),
+           first_read_buffer_double_->data() + first_read_buffer_bytes_written_,
+           size);
+    if (!next_handler_->OnReadCompleted(size, defer))
+      return false;
+    first_read_buffer_bytes_written_ += size;
+    if (*defer)
+      return true;
+  }
+
+  state_ = State::PASS_THROUGH;
+  first_read_buffer_double_ = nullptr;
+  next_handler_->SetController(controller());
+  return true;
+}
+
+bool InterceptingResourceHandler::DoLoop(bool* defer) {
+  bool result = true;
+  do {
+    switch (state_) {
+      case State::STARTING:
+      case State::WAITING_FOR_ON_READ_COMPLETED:
+      case State::PASS_THROUGH:
+      case State::DONE:
+        NOTREACHED();
+        break;
+      case State::NOTIFYING_ON_RESPONSE_STARTED_TO_NEW_HANDLER:
+        if (first_read_buffer_double_) {
+          // OnWillRead has been called, so copying the data from
+          // |first_read_buffer_double_| to |first_read_buffer_| will be needed
+          // when OnReadCompleted is called.
+          state_ = State::WAITING_FOR_ON_READ_COMPLETED;
+        } else {
+          // OnWillRead has not been called, so no special handling will be
+          // needed from now on.
+          state_ = State::PASS_THROUGH;
+          next_handler_->SetController(controller());
+        }
+        break;
+      case State::SENDING_PAYLOAD_TO_OLD_HANDLER:
+        result = SendPayloadToOldHandler(defer);
+        break;
+      case State::SENDING_BUFFER_TO_NEW_HANDLER:
+        result = SendFirstReadBufferToNewHandler(defer);
+        break;
+      case State::NOTIFYING_ON_RESPONSE_COMPLETED_TO_NEW_HANDLER:
+        state_ = State::DONE;
+        new_handler_ = nullptr;
+        break;
+    }
+  } while (result && !*defer && state_ != State::STARTING &&

[mmenke, 2016/10/11 21:34:35]

I don't think State::STARTING is allowed here?  We never enter the DoLoop with
that state, and we never return to that state, right?

[yhirano, 2016/10/12 10:36:21]

You're right, thanks. Done.

+           state_ != State::WAITING_FOR_ON_READ_COMPLETED &&
+           state_ != State::PASS_THROUGH && state_ != State::DONE);
+  return result;
 }
 
 }  // namespace content