PlzNavigate: Get StreamPrivate API to work.

extensions/browser/guest_view/mime_handler_view/mime_handler_stream_manager.cc

 // Copyright 2015 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 "extensions/browser/guest_view/mime_handler_view/mime_handler_stream_manager.h"
 
+#include <map>
 #include "base/memory/ptr_util.h"
 #include "base/memory/singleton.h"
 #include "components/keyed_service/content/browser_context_dependency_manager.h"
 #include "components/keyed_service/content/browser_context_keyed_service_factory.h"
+#include "content/public/browser/navigation_handle.h"
 #include "content/public/browser/render_frame_host.h"
 #include "content/public/browser/render_process_host.h"
+#include "content/public/browser/web_contents.h"
+#include "content/public/common/browser_side_navigation_policy.h"
 #include "extensions/browser/extension_registry.h"
 #include "extensions/browser/extensions_browser_client.h"
 #include "extensions/browser/guest_view/mime_handler_view/mime_handler_view_guest.h"
 
 namespace extensions {
 namespace {
 
 class MimeHandlerStreamManagerFactory
     : public BrowserContextKeyedServiceFactory {
  public:
@@ skipping 39 matching lines @@
       context);
 }
 
 }  // namespace
 
 // A WebContentsObserver that observes for a particular RenderFrameHost either
 // navigating or closing (including by crashing). This is necessary to ensure
 // that streams that aren't claimed by a MimeHandlerViewGuest are not leaked, by
 // aborting the stream if any of those events occurs.
 class MimeHandlerStreamManager::EmbedderObserver
     : public content::WebContentsObserver {

[nasko, 2016/09/21 21:32:18]

Does this object need to track subframes or are is it only restricted to main
frames? If the latter, I think this code can be much simpler.

From what I recall, this is only used for GuestView, which basically has a
WebContents embedded in a different WebContents. Why not get the WebContents as
a param and get the main frame from it. Is there a case where we will get a
subframe routing id here?

[ananta, 2016/09/21 21:59:12]

This codepath also kicks in for iframes.

[nasko, 2016/09/21 22:08:55]

Can you help me understand in what cases that happens?

  public:
   EmbedderObserver(MimeHandlerStreamManager* stream_manager,
+                   const std::string& view_id,
+                   int frame_tree_node_id,
                    int render_process_id,
-                   int render_frame_id,
-                   const std::string& view_id);
+                   int render_frame_id);
 
  private:
   // WebContentsObserver overrides.
   void RenderFrameDeleted(content::RenderFrameHost* render_frame_host) override;
   void RenderProcessGone(base::TerminationStatus status) override;
   void DidStartProvisionalLoadForFrame(
       content::RenderFrameHost* render_frame_host,
       const GURL& validated_url,
       bool is_error_page,
       bool is_iframe_srcdoc) override;
   void WebContentsDestroyed() override;
+  void DidStartNavigation(
+      content::NavigationHandle* navigation_handle) override;
+  void RenderFrameHostChanged(content::RenderFrameHost* old_host,
+                              content::RenderFrameHost* new_host) override;
 
   void AbortStream();
 
   bool IsTrackedRenderFrameHost(content::RenderFrameHost* render_frame_host);
 
   MimeHandlerStreamManager* const stream_manager_;
+  const std::string view_id_;
+  const int frame_tree_node_id_;
   const int render_process_id_;
   const int render_frame_id_;
-  const std::string view_id_;
+  // With PlzNavigate we get an initial provisional load notification for the
+  // URL the mime handler is serving. We don't want to clean up the stream
+  // here. This field helps us track the first load notification. Defaults
+  // to true.
+  bool initial_load_for_frame_;
+  // Used to track whether a RenderFrameHost is swapped with another
+  // RenderFrameHost. We use it to determine when it is safe to clean up the
+  // stream. For e.g if a RFH is swapped with another RFH, and the first one

[nasko, 2016/09/21 21:32:18]

nit: e.g.

[ananta, 2016/09/21 21:59:12]

Comment has changed

+  // goes away, we want to clean up the stream when the second RFH dies.
+  typedef std::map<content::RenderFrameHost*, content::RenderFrameHost*>
+      RenderFrameHostMap;
+  RenderFrameHostMap changed_hosts_map_;
 };
 
 MimeHandlerStreamManager::MimeHandlerStreamManager()
     : extension_registry_observer_(this) {
 }
 
 MimeHandlerStreamManager::~MimeHandlerStreamManager() {
 }
 
 // static
 MimeHandlerStreamManager* MimeHandlerStreamManager::Get(
     content::BrowserContext* context) {
   return MimeHandlerStreamManagerFactory::GetInstance()->Get(context);
 }
 
 void MimeHandlerStreamManager::AddStream(
     const std::string& view_id,
     std::unique_ptr<StreamContainer> stream,
+    int frame_tree_node_id,
     int render_process_id,
     int render_frame_id) {
   streams_by_extension_id_[stream->extension_id()].insert(view_id);
   auto result = streams_.insert(std::make_pair(view_id, std::move(stream)));
   DCHECK(result.second);
   embedder_observers_[view_id] = base::MakeUnique<EmbedderObserver>(
-      this, render_process_id, render_frame_id, view_id);
+      this, view_id, frame_tree_node_id, render_process_id, render_frame_id);
 }
 
 std::unique_ptr<StreamContainer> MimeHandlerStreamManager::ReleaseStream(
     const std::string& view_id) {
   auto stream = streams_.find(view_id);
   if (stream == streams_.end())
     return nullptr;
 
   std::unique_ptr<StreamContainer> result =
       base::WrapUnique(stream->second.release());
@@ skipping 13 matching lines @@
 
   for (const auto& view_id : streams->second) {
     streams_.erase(view_id);
     embedder_observers_.erase(view_id);
   }
   streams_by_extension_id_.erase(streams);
 }
 
 MimeHandlerStreamManager::EmbedderObserver::EmbedderObserver(
     MimeHandlerStreamManager* stream_manager,
+    const std::string& view_id,
+    int frame_tree_node_id,
     int render_process_id,
-    int render_frame_id,
-    const std::string& view_id)
-    : WebContentsObserver(content::WebContents::FromRenderFrameHost(
-          content::RenderFrameHost::FromID(render_process_id,
-                                           render_frame_id))),
-      stream_manager_(stream_manager),
+    int render_frame_id)
+    : stream_manager_(stream_manager),
+      view_id_(view_id),
+      frame_tree_node_id_(frame_tree_node_id),
       render_process_id_(render_process_id),
       render_frame_id_(render_frame_id),
-      view_id_(view_id) {
+      initial_load_for_frame_(true) {
+  content::WebContents* web_contents = nullptr;
+  if (frame_tree_node_id != -1) {
+    web_contents =
+        content::WebContents::FromFrameTreeNodeId(frame_tree_node_id);
+  } else {
+    web_contents = content::WebContents::FromRenderFrameHost(
+        content::RenderFrameHost::FromID(render_process_id, render_frame_id));
+  }
+  content::WebContentsObserver::Observe(web_contents);
 }
 
 void MimeHandlerStreamManager::EmbedderObserver::RenderFrameDeleted(
     content::RenderFrameHost* render_frame_host) {
-  if (!IsTrackedRenderFrameHost(render_frame_host))
+  if (!IsTrackedRenderFrameHost(render_frame_host)) {

[nasko, 2016/09/21 21:32:18]

nit: Why add {}? The rest of the file omits {} on one line if statements.

[ananta, 2016/09/21 21:59:12]

Done.

     return;
-
+  }
+  // If the |render_frame_host| is swapped with another RFH, then we should
+  // cleanup the stream when the other RFH dies,
+  RenderFrameHostMap::iterator index =
+      changed_hosts_map_.find(render_frame_host);
+  if (index != changed_hosts_map_.end()) {
+    changed_hosts_map_.erase(index);
+    return;
+  }
   AbortStream();
 }
 
 void MimeHandlerStreamManager::EmbedderObserver::RenderProcessGone(
     base::TerminationStatus status) {
   AbortStream();
 }
+
 void MimeHandlerStreamManager::EmbedderObserver::
     DidStartProvisionalLoadForFrame(content::RenderFrameHost* render_frame_host,
                                     const GURL& validated_url,
                                     bool is_error_page,
                                     bool is_iframe_srcdoc) {
   if (!IsTrackedRenderFrameHost(render_frame_host))
     return;
+  // With PlzNavigate we get a provisional load notification for the URL we are
+  // serving. We don't want to clean up the stream here.
+  if (initial_load_for_frame_ && content::IsBrowserSideNavigationEnabled()) {
+    initial_load_for_frame_ = false;
+    return;
+  }
+  AbortStream();
+}
 
-  AbortStream();
+void MimeHandlerStreamManager::EmbedderObserver::DidStartNavigation(
+    content::NavigationHandle* navigation_handle) {
+  // If the top level frame is navigating away, clean up the stream.
+  if (navigation_handle->IsInMainFrame())
+    AbortStream();
+}
+
+void MimeHandlerStreamManager::EmbedderObserver::RenderFrameHostChanged(
+    content::RenderFrameHost* old_host,
+    content::RenderFrameHost* new_host) {
+  changed_hosts_map_[old_host] = new_host;

[nasko, 2016/09/21 21:32:18]

Why do we need to keep a map? Can't we always just remember the new host into a
member variable?

[ananta, 2016/09/21 21:59:12]

Thanks. Done

 }
 
 void MimeHandlerStreamManager::EmbedderObserver::WebContentsDestroyed() {
   AbortStream();
 }
 
 void MimeHandlerStreamManager::EmbedderObserver::AbortStream() {
   Observe(nullptr);
   // This will cause the stream to be destroyed.
   stream_manager_->ReleaseStream(view_id_);
 }
 
 bool MimeHandlerStreamManager::EmbedderObserver::IsTrackedRenderFrameHost(
     content::RenderFrameHost* render_frame_host) {
-  return render_frame_host->GetRoutingID() == render_frame_id_ &&
-         render_frame_host->GetProcess()->GetID() == render_process_id_;
+  if (frame_tree_node_id_ != -1) {
+    return (render_frame_host->GetFrameTreeNodeId() == frame_tree_node_id_);
+  } else {
+    DCHECK((render_frame_id_ != -1) && (render_process_id_ != -1));
+    return render_frame_host->GetRoutingID() == render_frame_id_ &&
+           render_frame_host->GetProcess()->GetID() == render_process_id_;
+  }
 }
 
 }  // namespace extensions