Use FrameTreeNode ID as frameId in extension APIs

chrome/browser/extensions/api/messaging/extension_message_port.cc

Index: chrome/browser/extensions/api/messaging/extension_message_port.cc
diff --git a/chrome/browser/extensions/api/messaging/extension_message_port.cc b/chrome/browser/extensions/api/messaging/extension_message_port.cc
index 1e4f87f901d6234975f9d3a6dccfcf139172fd3b..af8e58b742fea4882db7b962ef534d9595dd42a0 100644
--- a/chrome/browser/extensions/api/messaging/extension_message_port.cc
+++ b/chrome/browser/extensions/api/messaging/extension_message_port.cc
@@ -5,30 +5,110 @@
 #include "chrome/browser/extensions/api/messaging/extension_message_port.h"
 
 #include "chrome/browser/profiles/profile.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/browser/web_contents_observer.h"
 #include "extensions/browser/extension_host.h"
 #include "extensions/browser/process_manager.h"
+#include "extensions/browser/process_manager_observer.h"
 #include "extensions/common/extension_messages.h"
 #include "extensions/common/manifest_handlers/background_info.h"
 
 namespace extensions {
 
-ExtensionMessagePort::ExtensionMessagePort(content::RenderProcessHost* process,
-                                           int routing_id,
-                                           const std::string& extension_id)
-     : process_(process),
-       routing_id_(routing_id),
-       extension_id_(extension_id),
-       background_host_ptr_(NULL) {
+const char kReceivingEndDoesntExistError[] =
+    "Could not establish connection. Receiving end does not exist.";
+
+// Helper class to detect when frames are destroyed.
+class ExtensionMessagePort::FrameTracker : public content::WebContentsObserver,
+                                           public ProcessManagerObserver {
+ public:
+  // Watch all frames in an extension process.
+  explicit FrameTracker(ExtensionMessagePort* port) : port_(port) {
+    ProcessManager::Get(port_->browser_context_)->AddObserver(this);

[Devlin, 2015/10/30 01:49:39]

nit: prefer scoped observer for usage like this.

[robwu, 2015/11/02 19:08:33]

Done.

+  }
+  // Watch all frames in a tab.
+  // TODO(robwu): Check whether this includes background pages and <webview>s.
+  FrameTracker(ExtensionMessagePort* port, content::WebContents* tab)
+      : content::WebContentsObserver(tab), port_(port) {

[Devlin, 2015/10/30 01:49:38]

This doesn't observe the process manager.  The comment makes it seem like that's
intentional, but having this behavior in the constructor is pretty weird, and
very opaque from any call site.  I think I'd just prefer to have two methods
here, TrackTabFrames and TrackExtensionProcessFrames or something like that. 
Or, if there's a concrete reason to keep the behavior separate and exclusive, 
make two classes.

[robwu, 2015/11/02 19:08:33]

Done.

+  }
+  ~FrameTracker() override {
+    ProcessManager::Get(port_->browser_context_)->RemoveObserver(this);
+  }
+ private:
+  // content::WebContentsObserver overrides:
+  void RenderFrameHostChanged(content::RenderFrameHost* old_host,
+                              content::RenderFrameHost* new_host) override {
+    port_->UnregisterFrame(old_host);
+  }
+  void FrameDeleted(content::RenderFrameHost* render_frame_host) override {
+    port_->UnregisterFrame(render_frame_host);
+  }
+
+  // extensions::ProcessManagerObserver overrides:
+  void OnExtensionFrameUnregistered(
+      const std::string& extension_id,
+      content::RenderFrameHost* render_frame_host) override {
+    if (extension_id == port_->extension_id_)
+      port_->UnregisterFrame(render_frame_host);
+  }
+
+  ExtensionMessagePort* port_;
+  DISALLOW_COPY_AND_ASSIGN(FrameTracker);
+};
+
+ExtensionMessagePort::ExtensionMessagePort(
+    base::WeakPtr<MessageService> message_service,
+    int port_id,
+    const std::string& extension_id,
+    content::RenderProcessHost* extension_process)
+    : weak_message_service_(message_service),
+      port_id_(port_id),
+      extension_id_(extension_id),
+      did_create_port_(false),
+      browser_context_(extension_process->GetBrowserContext()),
+      frames_(ProcessManager::Get(browser_context_)->
+              GetRenderFrameHostsForExtension(extension_id)),
+      extension_process_(extension_process),
+      background_host_ptr_(nullptr),
+      frame_tracker_(new FrameTracker(this)) {
+}
+
+ExtensionMessagePort::ExtensionMessagePort(
+    base::WeakPtr<MessageService> message_service,
+    int port_id,
+    const std::string& extension_id,
+    content::RenderFrameHost* rfh,
+    bool include_child_frames)
+    : weak_message_service_(message_service),
+      port_id_(port_id),
+      extension_id_(extension_id),
+      did_create_port_(false),
+      browser_context_(rfh->GetProcess()->GetBrowserContext()),
+      frames_(),
+      extension_process_(nullptr),
+      background_host_ptr_(nullptr),
+      frame_tracker_() {

[Devlin, 2015/10/30 01:49:39]

This will default-initialize itself; no need for this line.

[robwu, 2015/11/02 19:08:33]

Acknowledged.

+  content::WebContents* tab = content::WebContents::FromRenderFrameHost(rfh);
+  DCHECK(tab);
+  frame_tracker_.reset(new FrameTracker(this, tab));
+  if (include_child_frames) {
+    tab->ForEachFrame(base::Bind(&ExtensionMessagePort::RegisterFrame,
+                                 base::Unretained(this)));
+  } else {
+    RegisterFrame(rfh);
+  }
+}
+
+ExtensionMessagePort::~ExtensionMessagePort() {
+  frame_tracker_.reset();

[Devlin, 2015/10/30 01:49:38]

Why is the explicit destruction needed?

[robwu, 2015/11/02 19:08:33]

It's not needed - removed.

 }
 
 void ExtensionMessagePort::DispatchOnConnect(
-    int dest_port_id,
     const std::string& channel_name,
     scoped_ptr<base::DictionaryValue> source_tab,
     int source_frame_id,
-    int target_tab_id,
-    int target_frame_id,
     int guest_process_id,
     int guest_render_frame_routing_id,
     const std::string& source_extension_id,
@@ -44,32 +124,26 @@ void ExtensionMessagePort::DispatchOnConnect(
   info.target_id = target_extension_id;
   info.source_id = source_extension_id;
   info.source_url = source_url;
-  info.target_tab_id = target_tab_id;
-  info.target_frame_id = target_frame_id;
   info.guest_process_id = guest_process_id;
   info.guest_render_frame_routing_id = guest_render_frame_routing_id;
 
-  process_->Send(new ExtensionMsg_DispatchOnConnect(
-      routing_id_, dest_port_id, channel_name, source, info, tls_channel_id));
+  SendToPort(new ExtensionMsg_DispatchOnConnect(
+      MSG_ROUTING_NONE, port_id_, channel_name, source, info, tls_channel_id));
 }
 
 void ExtensionMessagePort::DispatchOnDisconnect(
-    int source_port_id,
     const std::string& error_message) {
-  process_->Send(new ExtensionMsg_DispatchOnDisconnect(
-      routing_id_, source_port_id, error_message));
+  SendToPort(new ExtensionMsg_DispatchOnDisconnect(
+      MSG_ROUTING_NONE, port_id_, error_message));
 }
 
-void ExtensionMessagePort::DispatchOnMessage(const Message& message,
-                                             int target_port_id) {
-  process_->Send(new ExtensionMsg_DeliverMessage(
-      routing_id_, target_port_id, message));
+void ExtensionMessagePort::DispatchOnMessage(const Message& message) {
+  SendToPort(new ExtensionMsg_DeliverMessage(
+      MSG_ROUTING_NONE, port_id_, message));
 }
 
 void ExtensionMessagePort::IncrementLazyKeepaliveCount() {
-  Profile* profile =
-      Profile::FromBrowserContext(process_->GetBrowserContext());
-  extensions::ProcessManager* pm = ProcessManager::Get(profile);
+  ProcessManager* pm = ProcessManager::Get(browser_context_);
   ExtensionHost* host = pm->GetBackgroundHostForExtension(extension_id_);
   if (host && BackgroundInfo::HasLazyBackgroundPage(host->extension()))
     pm->IncrementLazyKeepaliveCount(host->extension());
@@ -80,16 +154,61 @@ void ExtensionMessagePort::IncrementLazyKeepaliveCount() {
 }
 
 void ExtensionMessagePort::DecrementLazyKeepaliveCount() {
-  Profile* profile =
-      Profile::FromBrowserContext(process_->GetBrowserContext());
-  extensions::ProcessManager* pm = ProcessManager::Get(profile);
+  ProcessManager* pm = ProcessManager::Get(browser_context_);
   ExtensionHost* host = pm->GetBackgroundHostForExtension(extension_id_);
   if (host && host == background_host_ptr_)
     pm->DecrementLazyKeepaliveCount(host->extension());
 }
 
-content::RenderProcessHost* ExtensionMessagePort::GetRenderProcessHost() {
-  return process_;
+void ExtensionMessagePort::OpenPort(int process_id, int routing_id) {
+  // Mark port as active, so we can distinguish abnormal port closure from
+  // explicit port closure via disconnect().
+  did_create_port_ = true;
+}
+
+void ExtensionMessagePort::ClosePort(int process_id, int routing_id) {
+  if (routing_id == MSG_ROUTING_NONE) {
+    // The only non-frame-specific message is the response to an unhandled
+    // onConnect event in the extension process.
+    DCHECK(extension_process_);
+    frames_.clear();
+    CloseChannel();
+    return;
+  }
+  content::RenderFrameHost* rfh =
+      content::RenderFrameHost::FromID(process_id, routing_id);
+  if (rfh)
+    UnregisterFrame(rfh);
+}
+
+void ExtensionMessagePort::CloseChannel() {
+  std::string error_message = did_create_port_ ? std::string() :
+      kReceivingEndDoesntExistError;
+  if (weak_message_service_)
+    weak_message_service_->CloseChannel(port_id_, error_message);
+}
+
+void ExtensionMessagePort::RegisterFrame(content::RenderFrameHost* rfh) {
+  frames_.insert(rfh);
+}
+
+void ExtensionMessagePort::UnregisterFrame(content::RenderFrameHost* rfh) {
+  if (frames_.erase(rfh) != 0 && frames_.empty())
+    CloseChannel();
+}
+
+void ExtensionMessagePort::SendToPort(IPC::Message* msg) {
+  if (extension_process_) {
+    msg->set_routing_id(MSG_ROUTING_CONTROL);
+    extension_process_->Send(msg);
+    return;
+  }
+  for (content::RenderFrameHost* rfh : frames_) {

[Devlin, 2015/10/30 01:49:39]

Hmm... This is a bit confusing, because of the way this is laid out, but I
*think* this is all just equivalent to doing WebContents::SendToAllFrames().  If
|extension_process_| is null, it means we're following the tab, in which case
|frames_| is just the set of all frames, right?  This also seems to imply we
don't actually need |frames_| at all, because it doesn't really seem to be used
in the extension_process_ case.

[robwu, 2015/10/31 00:10:44]

|frames_| is used to avoid unnecessary IPC, since most frames will not be
interested in the port. The flow is this:
1. Browser: constructor + add all frames
2. Browser: DispatchOnConnect
3. Renderer: Acknowledge or reject port
4. Browser: Based on renderer, either ClosePort (=remove frame) or OpenPort
(=mark port as open).
5. Browser: DispatchOnMessage (zero or more times, depending on the extension
logic).
6. Browser: DispatchOnConnect (if all frames rejected, or if the frames
navigated away).

The |extension_process_| is a small optimization: All extension frames reside in
the same extension process, so a single IPC is sufficient to reach all frames.

[Devlin, 2015/10/31 00:19:56]

What about extension iframes in web pages without site isolation?

[robwu, 2015/11/02 19:08:33]

Extension iframes currently don't have any access to extension APIs, except for
those that are already available in content scripts, and origin-specific APIs
such as databases and gUM. Therefore it's totally fine (even desirable for
backwards-compatibility?) that extension frames in web pages are not notified of
extension messages.

+    IPC::Message* msg_copy = new IPC::Message(*msg);
+    msg_copy->set_routing_id(rfh->GetRoutingID());
+    rfh->Send(msg_copy);
+  }
+  delete msg;
 }
 
 }  // namespace extensions