Deal with frame removal by content scripts

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
index aa0f5d32b9aa35ee7c4cbb85461f79899b169c6e..e1273b8ba77d0ae758328247dd36ea0de361c526 100644
--- a/extensions/renderer/script_injection_manager.cc
+++ b/extensions/renderer/script_injection_manager.cc
@@ -70,7 +70,6 @@ class ScriptInjectionManager::RFOHelper : public content::RenderFrameObserver {
   // RenderFrameObserver implementation.
   bool OnMessageReceived(const IPC::Message& message) override;
   void DidCreateNewDocument() override;
-  void DidCreateDocumentElement() override;
   void DidFailProvisionalLoad(const blink::WebURLError& error) override;
   void DidFinishDocumentLoad() override;
   void DidFinishLoad() override;
@@ -135,10 +134,6 @@ void ScriptInjectionManager::RFOHelper::DidCreateNewDocument() {
     InvalidateAndResetFrame();
 }
 
-void ScriptInjectionManager::RFOHelper::DidCreateDocumentElement() {
-  manager_->StartInjectScripts(render_frame(), UserScript::DOCUMENT_START);
-}
-
 void ScriptInjectionManager::RFOHelper::DidFailProvisionalLoad(
     const blink::WebURLError& error) {
   FrameStatusMap::iterator it = manager_->frame_statuses_.find(render_frame());
@@ -162,7 +157,6 @@ void ScriptInjectionManager::RFOHelper::DidFailProvisionalLoad(
 
 void ScriptInjectionManager::RFOHelper::DidFinishDocumentLoad() {
   DCHECK(content::RenderThread::Get());
-  manager_->StartInjectScripts(render_frame(), UserScript::DOCUMENT_END);
   // We try to run idle in two places: here and DidFinishLoad.
   // DidFinishDocumentLoad() corresponds to completing the document's load,
   // whereas DidFinishLoad corresponds to completing the document and all
@@ -175,6 +169,9 @@ void ScriptInjectionManager::RFOHelper::DidFinishDocumentLoad() {
       base::Bind(&ScriptInjectionManager::RFOHelper::RunIdle,
                  weak_factory_.GetWeakPtr()),
       base::TimeDelta::FromMilliseconds(kScriptIdleTimeoutInMs));
+
+  // Soon after this method returns, AfterDidFinishDocumentLoad will be called,
+  // which injects the DOCUMENT_END scripts.
 }
 
 void ScriptInjectionManager::RFOHelper::DidFinishLoad() {
@@ -254,6 +251,16 @@ ScriptInjectionManager::~ScriptInjectionManager() {
     injection->invalidate_render_frame();
 }
 
+void ScriptInjectionManager::AfterDidCreateDocumentElement(

[Devlin, 2016/02/12 22:21:56]

In general, it kind of makes me sad that we can't centralize the protection for
render frames a little bit more.  We have this class, ScriptInjection, and
ExtensionFrameHelper all kinda-sorta trying to figure it out.  Unfortunately, I
don't immediately have ideas of how to make it better, either.

+    content::RenderFrame* render_frame) {
+  StartInjectScripts(render_frame, UserScript::DOCUMENT_START);
+}
+
+void ScriptInjectionManager::AfterDidFinishDocumentLoad(
+    content::RenderFrame* render_frame) {
+  StartInjectScripts(render_frame, UserScript::DOCUMENT_END);
+}
+
 void ScriptInjectionManager::OnRenderFrameCreated(
     content::RenderFrame* render_frame) {
   rfo_helpers_.push_back(make_scoped_ptr(new RFOHelper(render_frame, this)));