PlzNavigate: Fix for the PhishingDOMFeatureExtractorTest.SubFrames and SubframeRemoval tests

chrome/renderer/safe_browsing/phishing_dom_feature_extractor_browsertest.cc

 // Copyright (c) 2012 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 "chrome/renderer/safe_browsing/phishing_dom_feature_extractor.h"
 
 #include <memory>
 
 #include "base/bind.h"
+#include "base/bind_helpers.h"
 #include "base/callback.h"
 #include "base/memory/ptr_util.h"
 #include "base/memory/weak_ptr.h"
 #include "base/threading/thread_task_runner_handle.h"
 #include "base/time/time.h"
 #include "chrome/renderer/chrome_content_renderer_client.h"
 #include "chrome/renderer/safe_browsing/features.h"
 #include "chrome/renderer/safe_browsing/mock_feature_extractor_clock.h"
 #include "chrome/renderer/safe_browsing/test_utils.h"
 #include "chrome/test/base/chrome_render_view_test.h"
+#include "content/common/frame_messages.h"
+#include "content/common/navigation_params.h"
 #include "content/public/common/content_switches.h"
+#include "content/public/common/resource_response.h"
 #include "content/public/renderer/render_frame.h"
 #include "content/public/test/test_utils.h"
+#include "ipc/ipc_message_macros.h"
 #include "net/base/escape.h"
 #include "net/base/registry_controlled_domains/registry_controlled_domain.h"
 #include "testing/gmock/include/gmock/gmock.h"
 #include "third_party/WebKit/public/platform/URLConversion.h"
 #include "third_party/WebKit/public/platform/WebString.h"
 #include "third_party/WebKit/public/web/WebFrame.h"
 #include "third_party/WebKit/public/web/WebLocalFrame.h"
 #include "third_party/WebKit/public/web/WebScriptSource.h"
 
 using ::testing::DoAll;
@@ skipping 144 matching lines @@
 
   // Helper for the SubframeRemoval test that posts a message to remove
   // the iframe "frame1" from the document.
   void ScheduleRemoveIframe() {
     base::ThreadTaskRunnerHandle::Get()->PostTask(
         FROM_HERE, base::Bind(&PhishingDOMFeatureExtractorTest::RemoveIframe,
                               weak_factory_.GetWeakPtr()));
   }
 
  protected:
+  // Subclasses the ChromeMockRenderThread class for providing functionality
+  // to mock IPCs being sent to the browser from RenderViewTests.
+  class PhishingMockRenderThread : public ChromeMockRenderThread {
+   public:
+    PhishingMockRenderThread()
+        : current_message_routing_id_(-1) {
+    }
+
+   protected:
+    bool OnMessageReceived(const IPC::Message& msg) override {
+      current_message_routing_id_ = msg.routing_id();
+
+      IPC_BEGIN_MESSAGE_MAP(PhishingMockRenderThread, msg)
+        IPC_MESSAGE_HANDLER(FrameHostMsg_BeginNavigation, OnBeginNavigation)
+      IPC_END_MESSAGE_MAP()
+
+      current_message_routing_id_ = -1;
+      return ChromeMockRenderThread::OnMessageReceived(msg);
+    }
+
+    void OnBeginNavigation(
+        const content::CommonNavigationParams& common_params,
+        const content::BeginNavigationParams& begin_params) {
+      std::unique_ptr<IPC::Message> message;
+
+      message.reset(new FrameMsg_CommitNavigation(

[Nathan Parker, 2016/11/08 23:27:30]

I'm not familiar with these IPCs to verify the correctness of this code. I'm
assuming that the test runs and passes as before, yes?

[ananta, 2016/11/08 23:52:05]

Yes. The BeginNavigation IPC is sent by the renderer to initiate a navigation
with PlzNavigate. The browser responds with a CommitNavigation IPC with the
response stream. The renderer proceeds to load the request. Here we shortcircuit
the process and return the mocked IPC back.

+          current_message_routing_id_, content::ResourceResponseHead(),
+          common_params.url, common_params,
+          content::RequestNavigationParams()));
+      content::RenderFrame* frame = content::RenderFrame::FromRoutingID(
+          current_message_routing_id_);
+
+      base::ThreadTaskRunnerHandle::Get()->PostTask(
+          FROM_HERE,
+          base::Bind(
+              base::IgnoreResult(&content::RenderFrame::OnMessageReceived),
+              base::Unretained(frame), *message));
+    }
+
+    int current_message_routing_id_;
+  };
+
   void SetUp() override {
     ChromeRenderViewTest::SetUp();
     extractor_.reset(new TestPhishingDOMFeatureExtractor(&clock_));
   }
 
   void TearDown() override {
     extractor_.reset(nullptr);
     ChromeRenderViewTest::TearDown();
   }
 
   content::ContentRendererClient* CreateContentRendererClient() override {
     ChromeContentRendererClient* client = new TestChromeContentRendererClient();
     InitChromeContentRendererClient(client);
     return client;
   }
 
   void AnotherExtractionDone(bool success) {
     success_ = success;
     message_loop_->QuitClosure().Run();
   }
 
   // Does the actual work of removing the iframe "frame1" from the document.
   void RemoveIframe() {
     blink::WebFrame* main_frame = GetMainFrame();
     ASSERT_TRUE(main_frame);
     main_frame->executeScript(blink::WebString(
         "document.body.removeChild(document.getElementById('frame1'));"));
   }
 
+  ChromeMockRenderThread* CreateMockRenderThread() override {

[Nathan Parker, 2016/11/08 23:27:30]

Add a comment as to which class this is overriding.

[ananta, 2016/11/08 23:52:05]

Done.

+    return new PhishingMockRenderThread();
+  }
+
   MockFeatureExtractorClock clock_;
   bool success_;
   std::unique_ptr<TestPhishingDOMFeatureExtractor> extractor_;
   scoped_refptr<content::MessageLoopRunner> message_loop_;
   base::WeakPtrFactory<PhishingDOMFeatureExtractorTest> weak_factory_;
 };
 
 TEST_F(PhishingDOMFeatureExtractorTest, FormFeatures) {
   // This test doesn't exercise the extraction timing.
   EXPECT_CALL(clock_, Now()).WillRepeatedly(Return(base::TimeTicks::Now()));
@@ skipping 338 matching lines @@
       "<html><head></head><body>"
       "<iframe src=\"" +
       net::EscapeForHTML(iframe1_url.spec()) +
       "\" id=\"frame1\"></iframe>"
       "<form></form></body></html>");
   ExtractFeatures("host.com", html, &features);
   ExpectFeatureMapsAreEqual(features, expected_features);
 }
 
 }  // namespace safe_browsing