Refactor renderer-side script injection

extensions/renderer/script_injection_manager.cc

Index: extensions/renderer/script_injection_manager.cc
diff --git a/extensions/renderer/script_injection_manager.cc b/extensions/renderer/script_injection_manager.cc
new file mode 100644
index 0000000000000000000000000000000000000000..bfb0156128398269bf0a26c4448e2cdc65770295
--- /dev/null
+++ b/extensions/renderer/script_injection_manager.cc
@@ -0,0 +1,351 @@
+// Copyright 2014 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/renderer/script_injection_manager.h"
+
+#include "base/bind.h"
+#include "base/memory/weak_ptr.h"
+#include "base/metrics/histogram.h"
+#include "base/values.h"
+#include "content/public/renderer/render_view.h"
+#include "content/public/renderer/render_view_observer.h"
+#include "extensions/common/extension.h"
+#include "extensions/common/extension_messages.h"
+#include "extensions/common/extension_set.h"
+#include "extensions/renderer/extension_helper.h"
+#include "extensions/renderer/programmatic_script_injection.h"
+#include "extensions/renderer/script_context.h"
+#include "ipc/ipc_message_macros.h"
+#include "third_party/WebKit/public/web/WebFrame.h"
+#include "third_party/WebKit/public/web/WebLocalFrame.h"
+#include "third_party/WebKit/public/web/WebView.h"
+#include "url/gurl.h"
+
+namespace extensions {
+
+namespace {
+
+// The length of time to wait after the DOM is complete to try and run user
+// scripts.
+const int kScriptIdleTimeoutInMs = 200;
+
+// Log information about a given script run.
+void LogScriptsRun(blink::WebFrame* frame,
+                   UserScript::RunLocation location,
+                   const ScriptInjection::ScriptsRunInfo& info) {
+  // Notify the browser if any extensions are now executing scripts.
+  if (!info.executing_scripts.empty()) {
+    content::RenderView* render_view =
+        content::RenderView::FromWebView(frame->view());
+    render_view->Send(new ExtensionHostMsg_ContentScriptsExecuting(
+        render_view->GetRoutingID(),
+        info.executing_scripts,
+        render_view->GetPageId(),
+        ScriptContext::GetDataSourceURLForFrame(frame)));
+  }
+
+  switch (location) {
+    case UserScript::DOCUMENT_START:
+      UMA_HISTOGRAM_COUNTS_100("Extensions.InjectStart_CssCount", info.num_css);
+      UMA_HISTOGRAM_COUNTS_100("Extensions.InjectStart_ScriptCount",
+                               info.num_js);
+      if (info.num_css || info.num_js)
+        UMA_HISTOGRAM_TIMES("Extensions.InjectStart_Time",
+                            info.timer.Elapsed());
+      break;
+    case UserScript::DOCUMENT_END:
+      UMA_HISTOGRAM_COUNTS_100("Extensions.InjectEnd_ScriptCount", info.num_js);
+      if (info.num_js)
+        UMA_HISTOGRAM_TIMES("Extensions.InjectEnd_Time", info.timer.Elapsed());
+      break;
+    case UserScript::DOCUMENT_IDLE:
+      UMA_HISTOGRAM_COUNTS_100("Extensions.InjectIdle_ScriptCount",
+                               info.num_js);
+      if (info.num_js)
+        UMA_HISTOGRAM_TIMES("Extensions.InjectIdle_Time", info.timer.Elapsed());
+      break;
+    case UserScript::RUN_DEFERRED:
+      // TODO(rdevlin.cronin): Add histograms.
+      break;
+    case UserScript::UNDEFINED:
+    case UserScript::RUN_LOCATION_LAST:
+      NOTREACHED();
+  }
+}
+
+}  // namespace
+
+class ScriptInjectionManager::RVOHelper : public content::RenderViewObserver {
+ public:
+  RVOHelper(content::RenderView* render_view, ScriptInjectionManager* manager);
+  virtual ~RVOHelper();
+
+ private:
+  // RenderViewObserver implementation.
+  virtual bool OnMessageReceived(const IPC::Message& message) OVERRIDE;
+  virtual void DidCreateDocumentElement(blink::WebLocalFrame* frame) OVERRIDE;
+  virtual void DidFinishDocumentLoad(blink::WebLocalFrame* frame) OVERRIDE;
+  virtual void DidFinishLoad(blink::WebLocalFrame* frame) OVERRIDE;
+  virtual void DidStartProvisionalLoad(blink::WebLocalFrame* frame) OVERRIDE;
+  virtual void FrameDetached(blink::WebFrame* frame) OVERRIDE;
+  virtual void OnDestruct() OVERRIDE;
+
+  virtual void OnExecuteCode(const ExtensionMsg_ExecuteCode_Params& params);
+  virtual void OnPermitScriptInjection(int64 request_id);
+
+  // Tells the ScriptInjectionManager to run tasks associated with
+  // document_idle.
+  void RunIdle(blink::WebFrame* frame);
+
+  ScriptInjectionManager* manager_;
+  base::WeakPtrFactory<RVOHelper> weak_factory_;
+};
+
+ScriptInjectionManager::RVOHelper::RVOHelper(
+    content::RenderView* render_view,
+    ScriptInjectionManager* manager)
+    : content::RenderViewObserver(render_view),
+      manager_(manager),
+      weak_factory_(this) {
+}
+
+ScriptInjectionManager::RVOHelper::~RVOHelper() {
+}
+
+bool ScriptInjectionManager::RVOHelper::OnMessageReceived(
+    const IPC::Message& message) {
+  bool handled = true;
+  IPC_BEGIN_MESSAGE_MAP(ScriptInjectionManager::RVOHelper, message)
+    IPC_MESSAGE_HANDLER(ExtensionMsg_ExecuteCode, OnExecuteCode)
+    IPC_MESSAGE_HANDLER(ExtensionMsg_PermitScriptInjection,
+                        OnPermitScriptInjection)
+    IPC_MESSAGE_UNHANDLED(handled = false)
+  IPC_END_MESSAGE_MAP()
+  return handled;
+}
+
+void ScriptInjectionManager::RVOHelper::DidCreateDocumentElement(
+    blink::WebLocalFrame* frame) {
+  manager_->InjectScripts(frame, UserScript::DOCUMENT_START);
+}
+
+void ScriptInjectionManager::RVOHelper::DidFinishDocumentLoad(
+    blink::WebLocalFrame* frame) {
+  manager_->InjectScripts(frame, UserScript::DOCUMENT_END);
+  base::MessageLoop::current()->PostDelayedTask(
+      FROM_HERE,
+      base::Bind(&ScriptInjectionManager::RVOHelper::RunIdle,
+                 weak_factory_.GetWeakPtr(),
+                 frame),
+      base::TimeDelta::FromMilliseconds(kScriptIdleTimeoutInMs));
+}
+
+void ScriptInjectionManager::RVOHelper::DidFinishLoad(
+    blink::WebLocalFrame* frame) {
+  // Ensure that running scripts does not keep any progress UI running.
+  base::MessageLoop::current()->PostTask(
+      FROM_HERE,
+      base::Bind(&ScriptInjectionManager::RVOHelper::RunIdle,
+                 weak_factory_.GetWeakPtr(),
+                 frame));
+}
+
+void ScriptInjectionManager::RVOHelper::DidStartProvisionalLoad(
+    blink::WebLocalFrame* frame) {
+  manager_->InvalidateForFrame(frame);
+}
+
+void ScriptInjectionManager::RVOHelper::FrameDetached(blink::WebFrame* frame) {
+  manager_->InvalidateForFrame(frame);
+}
+
+void ScriptInjectionManager::RVOHelper::OnDestruct() {
+  manager_->RemoveObserver(this);
+}
+
+void ScriptInjectionManager::RVOHelper::OnExecuteCode(
+    const ExtensionMsg_ExecuteCode_Params& params) {
+  manager_->HandleExecuteCode(params, render_view());
+}
+
+void ScriptInjectionManager::RVOHelper::OnPermitScriptInjection(
+    int64 request_id) {
+  manager_->HandlePermitScriptInjection(request_id);
+}
+
+void ScriptInjectionManager::RVOHelper::RunIdle(blink::WebFrame* frame) {
+  manager_->InjectScripts(frame, UserScript::DOCUMENT_IDLE);
+}
+
+ScriptInjectionManager::ScriptInjectionManager(
+    const ExtensionSet* extensions,
+    UserScriptSet* user_script_set)
+    : extensions_(extensions),
+      user_script_set_(user_script_set),
+      user_script_set_observer_(this) {
+  user_script_set_observer_.Add(user_script_set_);
+}
+
+ScriptInjectionManager::~ScriptInjectionManager() {
+}
+
+void ScriptInjectionManager::OnRenderViewCreated(
+    content::RenderView* render_view) {
+  rvo_helpers_.push_back(new RVOHelper(render_view, this));
+}
+
+void ScriptInjectionManager::OnUserScriptsUpdated(
+    const std::set<std::string>& changed_extensions,
+    const std::vector<UserScript*>& scripts) {
+  for (ScopedVector<ScriptInjection>::iterator iter =
+           pending_injections_.begin();
+       iter != pending_injections_.end();) {
+    if (changed_extensions.count((*iter)->extension_id()) > 0)
+      iter = pending_injections_.erase(iter);
+    else
+      ++iter;
+  }
+}
+
+void ScriptInjectionManager::RemoveObserver(RVOHelper* helper) {
+  for (ScopedVector<RVOHelper>::iterator iter = rvo_helpers_.begin();
+       iter != rvo_helpers_.end();
+       ++iter) {
+    if (*iter == helper) {
+      rvo_helpers_.erase(iter);
+      break;
+    }
+  }
+}
+
+void ScriptInjectionManager::InvalidateForFrame(blink::WebFrame* frame) {
+  for (ScopedVector<ScriptInjection>::iterator iter =
+           pending_injections_.begin();
+       iter != pending_injections_.end();) {
+    if ((*iter)->web_frame() == frame)
+      iter = pending_injections_.erase(iter);
+    else
+      ++iter;
+  }
+
+  frame_statuses_.erase(frame);
+}
+
+void ScriptInjectionManager::InjectScripts(
+    blink::WebFrame* frame, UserScript::RunLocation run_location) {
+  FrameStatusMap::iterator iter = frame_statuses_.find(frame);
+  // Check if the frame is already in our map.
+  if (iter == frame_statuses_.end()) {
+    // If the frame isn't in our map, and the run location isn't document_start,
+    // then abort. This can happen in two ways:
+    // 1. We just received a delayed idle run for a frame which is invalidated.
+    //    Obviously, we don't want to run.
+    // 2. We somehow received a document_end or document_idle notification
+    //    without ever receiving a document_start. We don't want to run because
+    //    extensions may have requirements that scripts running at start have
+    //    run by the time scripts run at idle. Better to just not run.
+    if (run_location != UserScript::DOCUMENT_START)
+      return;
+
+    // Otherwise, add a new entry to the map.
+    frame_statuses_[frame] = UserScript::DOCUMENT_START;
+  } else {  // Already in the map.
+    // If we've already run the given location (happens in the case of idle
+    // since we notify in two places), return.
+    if (iter->second == run_location)
+      return;
+
+    // Otherwise, update the frame status.
+    iter->second = run_location;
+  }
+
+  // Inject any scripts that were waiting for the right run location.
+  ScriptInjection::ScriptsRunInfo scripts_run_info;
+  for (ScopedVector<ScriptInjection>::iterator iter =
+           pending_injections_.begin();
+       iter != pending_injections_.end();) {
+    if ((*iter)->web_frame() == frame &&
+        (*iter)->TryToInject(run_location,
+                             extensions_->GetByID((*iter)->extension_id()),
+                             &scripts_run_info)) {
+      iter = pending_injections_.erase(iter);
+    } else {
+      ++iter;
+    }
+  }
+
+  // Try to inject any user scripts that should run for this location. If they
+  // don't complete their injection (for example, waiting for a permission
+  // response) then they will be added to |pending_injections_|.
+  ScopedVector<ScriptInjection> user_script_injections;
+  int tab_id = ExtensionHelper::Get(content::RenderView::FromWebView(
+                                        frame->top()->view()))->tab_id();
+  user_script_set_->GetInjections(
+      &user_script_injections, frame, tab_id, run_location);
+  for (ScopedVector<ScriptInjection>::iterator iter =
+           user_script_injections.begin();
+       iter != user_script_injections.end();) {
+    scoped_ptr<ScriptInjection> injection(*iter);
+    iter = user_script_injections.weak_erase(iter);
+    if (!injection->TryToInject(run_location,
+                                extensions_->GetByID(injection->extension_id()),
+                                &scripts_run_info)) {
+      pending_injections_.push_back(injection.release());
+    }
+  }
+
+  LogScriptsRun(frame, run_location, scripts_run_info);
+}
+
+void ScriptInjectionManager::HandleExecuteCode(
+    const ExtensionMsg_ExecuteCode_Params& params,
+    content::RenderView* render_view) {
+  blink::WebFrame* main_frame = render_view->GetWebView()->mainFrame();
+  if (!main_frame) {
+    render_view->Send(
+        new ExtensionHostMsg_ExecuteCodeFinished(render_view->GetRoutingID(),
+                                                 params.request_id,
+                                                 "No main frame",
+                                                 -1,
+                                                 GURL(std::string()),
+                                                 base::ListValue()));
+    return;
+  }
+
+  scoped_ptr<ScriptInjection> injection(new ProgrammaticScriptInjection(
+      main_frame, params, ExtensionHelper::Get(render_view)->tab_id()));
+
+  ScriptInjection::ScriptsRunInfo scripts_run_info;
+  FrameStatusMap::const_iterator iter = frame_statuses_.find(main_frame);
+  if (!injection->TryToInject(
+          iter == frame_statuses_.end() ? UserScript::UNDEFINED : iter->second,
+          extensions_->GetByID(injection->extension_id()),
+          &scripts_run_info)) {
+    pending_injections_.push_back(injection.release());
+  }
+}
+
+void ScriptInjectionManager::HandlePermitScriptInjection(int request_id) {
+  ScopedVector<ScriptInjection>::iterator iter =
+      pending_injections_.begin();
+  for (; iter != pending_injections_.end(); ++iter) {
+    if ((*iter)->request_id() == request_id)
+      break;
+  }
+  if (iter == pending_injections_.end())
+    return;
+
+  scoped_ptr<ScriptInjection> injection(*iter);
+  pending_injections_.weak_erase(iter);
+
+  ScriptInjection::ScriptsRunInfo scripts_run_info;
+  if (injection->OnPermissionGranted(extensions_->GetByID(
+                                         injection->extension_id()),
+                                     &scripts_run_info)) {
+    LogScriptsRun(
+        injection->web_frame(), UserScript::RUN_DEFERRED, scripts_run_info);
+  }
+}
+
+}  // namespace extensions