Prevent the destruction of WebContents in NavigationThrottle methods

content/browser/frame_host/navigation_handle_impl_unittest.cc

Index: content/browser/frame_host/navigation_handle_impl_unittest.cc
diff --git a/content/browser/frame_host/navigation_handle_impl_unittest.cc b/content/browser/frame_host/navigation_handle_impl_unittest.cc
index af6aced5b310e942202093109dbfcbde51280eef..798bd55c550d5eea7e92fd7c757b6b80df61b077 100644
--- a/content/browser/frame_host/navigation_handle_impl_unittest.cc
+++ b/content/browser/frame_host/navigation_handle_impl_unittest.cc
@@ -3,8 +3,10 @@
 // found in the LICENSE file.
 
 #include "base/macros.h"
+#include "base/command_line.h"
 #include "content/browser/frame_host/navigation_handle_impl.h"
 #include "content/public/browser/navigation_throttle.h"
+#include "content/public/common/content_switches.h"
 #include "content/test/test_render_frame_host.h"
 
 namespace content {
@@ -45,6 +47,33 @@ class TestNavigationThrottle : public NavigationThrottle {
   int will_redirect_calls_;
 };
 
+// A test version of the NavigationThrottle that will destroy its
+// NavigationHandle through by calling |destructor_callback|.
+class DestroyingThrottle : public NavigationThrottle,
+                           public base::SupportsWeakPtr<DestroyingThrottle> {
+ public:
+  DestroyingThrottle(NavigationHandle* handle,
+                     const base::Closure& destructor_callback)
+      : NavigationThrottle(handle),
+        destructor_callback_(destructor_callback) {}
+
+  ~DestroyingThrottle() override {}
+
+  DestroyingThrottle::ThrottleCheckResult WillStartRequest() override {
+    destructor_callback_.Run();
+    return NavigationThrottle::DESTROYED;
+  }
+
+  DestroyingThrottle::ThrottleCheckResult WillRedirectRequest() override {
+    destructor_callback_.Run();
+    return NavigationThrottle::DESTROYED;
+  }
+
+ private:
+  // The callback to run to destroy the NavigationHandle.
+  base::Closure destructor_callback_;
+};
+
 class NavigationHandleImplTest : public RenderViewHostImplTestHarness {
  public:
   NavigationHandleImplTest()
@@ -124,6 +153,18 @@ class NavigationHandleImplTest : public RenderViewHostImplTestHarness {
     return test_throttle;
   }
 
+  // Creates, register and returns a NavigationThrottle that will destroy the
+  // NavigationHandle when WillStartRequest or WillRedirectRequest are called.
+  base::WeakPtr<DestroyingThrottle> AddDestroyingThrottle() {
+    DestroyingThrottle* destroying_throttle = new DestroyingThrottle(
+        test_handle(),
+        base::Bind(&NavigationHandleImplTest::ResetNavigationHandle,
+                   base::Unretained(this)));
+    test_handle()->RegisterThrottleForTesting(
+        scoped_ptr<NavigationThrottle>(destroying_throttle));
+    return destroying_throttle->AsWeakPtr();
+  }
+
  private:
   // The callback provided to NavigationHandleImpl::WillStartRequest and
   // NavigationHandleImpl::WillRedirectRequest during the tests.
@@ -133,6 +174,10 @@ class NavigationHandleImplTest : public RenderViewHostImplTestHarness {
     was_callback_called_ = true;
   }
 
+  void ResetNavigationHandle() {
+    test_handle_.reset();
+  }
+
   scoped_ptr<NavigationHandleImpl> test_handle_;
   bool was_callback_called_;
   NavigationThrottle::ThrottleCheckResult callback_result_;
@@ -387,4 +432,54 @@ TEST_F(NavigationHandleImplTest, CancelThenProceedWillRedirectRequest) {
   EXPECT_EQ(0, proceed_throttle->will_redirect_calls());
 }
 
+// Checks that the NavigationHandle can be safely destroyed during
+// WillStartRequest as long as the NavigationThrottle returns
+// NavigationThrottle::DESTROYED.
+TEST_F(NavigationHandleImplTest, DestroyDuringWillStart) {
+  base::WeakPtr<DestroyingThrottle> destroy_throttle = AddDestroyingThrottle();
+  EXPECT_NE(destroy_throttle, nullptr);
+  EXPECT_NE(test_handle(), nullptr);
+
+  // Simulate WillStartRequest. This will destroy both the Handle and the
+  // Throttle.
+  SimulateWillStartRequest();
+  EXPECT_EQ(destroy_throttle, nullptr);
+  EXPECT_EQ(test_handle(), nullptr);
+
+  if (base::CommandLine::ForCurrentProcess()->HasSwitch(
+          switches::kEnableBrowserSideNavigation)) {
+    return;
+  }
+
+  // In non-PlzNavigate mode, the callback should have been called and the
+  // Navigation should have been cancelled.
+  EXPECT_TRUE(was_callback_called());
+  EXPECT_EQ(NavigationThrottle::CANCEL_AND_IGNORE, callback_result());
+}
+
+// Checks that the NavigationHandle can be safely destroyed during
+// WillRedirectRequest as long as the NavigationThrottle returns
+// NavigationThrottle::DESTROYED.
+TEST_F(NavigationHandleImplTest, DestroyDuringWillRedirect) {
+  base::WeakPtr<DestroyingThrottle> destroy_throttle = AddDestroyingThrottle();
+  EXPECT_NE(destroy_throttle, nullptr);
+  EXPECT_NE(test_handle(), nullptr);
+
+  // Simulate WillStartRequest. This will destroy both the Handle and the
+  // Throttle.
+  SimulateWillRedirectRequest();
+  EXPECT_EQ(destroy_throttle, nullptr);
+  EXPECT_EQ(test_handle(), nullptr);
+
+  if (base::CommandLine::ForCurrentProcess()->HasSwitch(
+          switches::kEnableBrowserSideNavigation)) {
+    return;
+  }
+
+  // In non-PlzNavigate mode, the callback should have been called and the
+  // Navigation should have been cancelled.
+  EXPECT_TRUE(was_callback_called());
+  EXPECT_EQ(NavigationThrottle::CANCEL_AND_IGNORE, callback_result());
+}
+
 }  // namespace content