Handle results that are loaded from cached tabs

Resolves brave/brave-browser#17355

browser/brave_browser_main_parts.cc

@@ -14,6 +14,7 @@
 #include "brave/common/pref_names.h"
 #include "brave/components/brave_sync/buildflags/buildflags.h"
 #include "brave/components/brave_sync/features.h"
+#include "brave/components/speedreader/buildflags.h"
 #include "brave/components/tor/buildflags/buildflags.h"
 #include "chrome/common/chrome_features.h"
 #include "components/prefs/pref_service.h"
@@ -22,6 +23,11 @@
 #include "extensions/buildflags/buildflags.h"
 #include "media/base/media_switches.h"
 
+#if BUILDFLAG(ENABLE_SPEEDREADER)
+#include "brave/browser/speedreader/speedreader_extended_info_handler.h"
+#include "components/sessions/content/content_serialized_navigation_driver.h"
+#endif
+
 #if BUILDFLAG(ENABLE_TOR)
 #include <string>
 #include "base/files/file_util.h"
@@ -60,6 +66,19 @@
 #include "extensions/browser/extension_system.h"
 #endif
 
+void BraveBrowserMainParts::PreBrowserStart() {
+#if BUILDFLAG(ENABLE_SPEEDREADER)
+  // Register() must be called after the SerializedNavigationDriver is
+  // initialized, but before any calls to
+  // ContentSerializedNavigationBuilder::ToNavigationEntries()
+  //
+  // TODO(keur): Can we DCHECK the latter condition?
+  DCHECK(sessions::ContentSerializedNavigationDriver::GetInstance());
+  speedreader::SpeedreaderExtendedInfoHandler::Register();
+#endif
+  ChromeBrowserMainParts::PreBrowserStart();
+}
+
 void BraveBrowserMainParts::PostBrowserStart() {
   ChromeBrowserMainParts::PostBrowserStart();
 

browser/brave_browser_main_parts.h

@@ -14,6 +14,7 @@ class BraveBrowserMainParts : public ChromeBrowserMainParts {
   using ChromeBrowserMainParts::ChromeBrowserMainParts;
   ~BraveBrowserMainParts() override = default;
 
+  void PreBrowserStart() override;
   void PostBrowserStart() override;
   void PreShutdown() override;
   void PreProfileInit() override;

browser/speedreader/sources.gni

@@ -10,6 +10,8 @@ brave_browser_speedreader_deps = []
 
 if (enable_speedreader) {
   brave_browser_speedreader_sources += [
+    "//brave/browser/speedreader/speedreader_extended_info_handler.cc",
+    "//brave/browser/speedreader/speedreader_extended_info_handler.h",
     "//brave/browser/speedreader/speedreader_service_factory.cc",
     "//brave/browser/speedreader/speedreader_service_factory.h",
     "//brave/browser/speedreader/speedreader_tab_helper.cc",

browser/speedreader/speedreader_extended_info_handler.cc

@@ -0,0 +1,98 @@
+/* Copyright 2021 The Brave Authors. All rights reserved.
+ * This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this file,
+ * You can obtain one at http://mozilla.org/MPL/2.0/. */
+
+#include "brave/browser/speedreader/speedreader_extended_info_handler.h"
+
+#include <memory>
+
+#include "base/memory/ptr_util.h"
+#include "base/supports_user_data.h"
+#include "components/sessions/content/content_serialized_navigation_driver.h"
+#include "content/public/browser/navigation_entry.h"
+
+namespace {
+
+// This is the key we register in the extended info map. We also use it for the
+// navigation entry user data.
+constexpr char kSpeedreaderKey[] = "speedreader";
+
+constexpr char kPageSavedReaderMode[] = "reader-mode";
+constexpr char kPageSavedSpeedreaderMode[] = "speedreader-mode";
+
+struct SpeedreaderNavigationData : public base::SupportsUserData::Data {
+  explicit SpeedreaderNavigationData(const std::string& value) : value(value) {}
+  std::string value;
+};
+
+}  // namespace
+
+namespace speedreader {
+
+// static
+void SpeedreaderExtendedInfoHandler::Register() {
+  auto* handler = new SpeedreaderExtendedInfoHandler();
+  sessions::ContentSerializedNavigationDriver::GetInstance()
+      ->RegisterExtendedInfoHandler(kSpeedreaderKey, base::WrapUnique(handler));
+}
+
+// static
+void SpeedreaderExtendedInfoHandler::PersistSpeedreaderMode(
+    content::NavigationEntry* entry) {
+  entry->SetUserData(
+      kSpeedreaderKey,
+      std::make_unique<SpeedreaderNavigationData>(kPageSavedSpeedreaderMode));
+}
+
+// static
+void SpeedreaderExtendedInfoHandler::PersistReaderMode(
+    content::NavigationEntry* entry) {
+  entry->SetUserData(
+      kSpeedreaderKey,
+      std::make_unique<SpeedreaderNavigationData>(kPageSavedReaderMode));
+}
+
+// static
+void SpeedreaderExtendedInfoHandler::ClearPersistedData(
+    content::NavigationEntry* entry) {
+  entry->RemoveUserData(kSpeedreaderKey);
+}
+
+// static
+bool SpeedreaderExtendedInfoHandler::IsCachedSpeedreaderMode(
+    content::NavigationEntry* entry) {
+  auto* data = static_cast<SpeedreaderNavigationData*>(
+      entry->GetUserData(kSpeedreaderKey));
+  if (!data) {
+    return false;
+  }
+  return data->value == kPageSavedSpeedreaderMode;
+}
+
+// static
+bool SpeedreaderExtendedInfoHandler::IsCachedReaderMode(
+    content::NavigationEntry* entry) {
+  auto* data = static_cast<SpeedreaderNavigationData*>(
+      entry->GetUserData(kSpeedreaderKey));
+  if (!data) {
+    return false;
+  }
+  return data->value == kPageSavedReaderMode;
+}
+
+std::string SpeedreaderExtendedInfoHandler::GetExtendedInfo(
+    content::NavigationEntry* entry) const {
+  auto* data = static_cast<SpeedreaderNavigationData*>(
+      entry->GetUserData(kSpeedreaderKey));
+  return data ? data->value : std::string();
+}
+
+void SpeedreaderExtendedInfoHandler::RestoreExtendedInfo(
+    const std::string& info,
+    content::NavigationEntry* entry) {
+  entry->SetUserData(kSpeedreaderKey,
+                     std::make_unique<SpeedreaderNavigationData>(info));
+}
+
+}  // namespace speedreader

browser/speedreader/speedreader_extended_info_handler.h

@@ -0,0 +1,49 @@
+/* Copyright 2021 The Brave Authors. All rights reserved.
+ * This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this file,
+ * You can obtain one at http://mozilla.org/MPL/2.0/. */
+
+#ifndef BRAVE_BROWSER_SPEEDREADER_SPEEDREADER_EXTENDED_INFO_HANDLER_H_
+#define BRAVE_BROWSER_SPEEDREADER_SPEEDREADER_EXTENDED_INFO_HANDLER_H_
+
+#include <string>
+
+#include "components/sessions/content/extended_info_handler.h"
+
+namespace content {
+class NavigationEntry;
+}  // namespace content
+
+namespace speedreader {
+
+// This class is meant to persist data to a content::NavigationEntry so that
+// distilled pages will be recognized on a restored session.
+class SpeedreaderExtendedInfoHandler : public sessions::ExtendedInfoHandler {
+ public:
+  // Register the extended info handler.
+  // Calling this more than once will cause a crash.
+  static void Register();
+
+  static void PersistSpeedreaderMode(content::NavigationEntry* entry);
+  static void PersistReaderMode(content::NavigationEntry* entry);
+  static void ClearPersistedData(content::NavigationEntry* entry);
+
+  static bool IsCachedSpeedreaderMode(content::NavigationEntry* entry);
+  static bool IsCachedReaderMode(content::NavigationEntry* entry);
+
+  SpeedreaderExtendedInfoHandler(const SpeedreaderExtendedInfoHandler&) =
+      delete;
+  SpeedreaderExtendedInfoHandler& operator=(
+      const SpeedreaderExtendedInfoHandler&) = delete;
+
+  SpeedreaderExtendedInfoHandler() = default;
+  ~SpeedreaderExtendedInfoHandler() override = default;
+
+  std::string GetExtendedInfo(content::NavigationEntry* entry) const override;
+  void RestoreExtendedInfo(const std::string& info,
+                           content::NavigationEntry* entry) override;
+};
+
+}  // namespace speedreader
+
+#endif  // BRAVE_BROWSER_SPEEDREADER_SPEEDREADER_EXTENDED_INFO_HANDLER_H_

browser/speedreader/speedreader_tab_helper.cc

@@ -5,8 +5,8 @@
 
 #include "brave/browser/speedreader/speedreader_tab_helper.h"
 
-#include "base/notreached.h"
 #include "brave/browser/brave_browser_process.h"
+#include "brave/browser/speedreader/speedreader_extended_info_handler.h"
 #include "brave/browser/speedreader/speedreader_service_factory.h"
 #include "brave/browser/ui/brave_browser_window.h"
 #include "brave/browser/ui/speedreader/speedreader_bubble_view.h"
@@ -18,6 +18,7 @@
 #include "chrome/browser/profiles/profile.h"
 #include "chrome/browser/ui/browser.h"
 #include "chrome/browser/ui/browser_finder.h"
+#include "content/public/browser/navigation_entry.h"
 #include "content/public/browser/navigation_handle.h"
 #include "content/public/browser/web_contents.h"
 
@@ -86,6 +87,28 @@ void SpeedreaderTabHelper::SingleShotSpeedreader() {
   }
 }
 
+bool SpeedreaderTabHelper::MaybeUpdateCachedState(
+    content::NavigationHandle* handle) {
+  auto* entry = handle->GetNavigationEntry();
+  if (!entry) {
+    return false;
+  }
+
+  bool cached = false;
+  if (SpeedreaderExtendedInfoHandler::IsCachedReaderMode(entry)) {
+    distill_state_ = DistillState::kReaderMode;
+    cached = true;
+  } else if (SpeedreaderExtendedInfoHandler::IsCachedSpeedreaderMode(entry)) {
+    distill_state_ = DistillState::kSpeedreaderMode;
+    cached = true;
+  }
+
+  if (!cached) {
+    SpeedreaderExtendedInfoHandler::ClearPersistedData(entry);
+  }
+  return cached;
+}
+
 void SpeedreaderTabHelper::UpdateActiveState(
     content::NavigationHandle* handle) {
   DCHECK(handle);
@@ -125,14 +148,18 @@ void SpeedreaderTabHelper::SetNextRequestState(DistillState state) {
 void SpeedreaderTabHelper::DidStartNavigation(
     content::NavigationHandle* navigation_handle) {
   if (navigation_handle->IsInMainFrame()) {
-    UpdateActiveState(navigation_handle);
+    if (!MaybeUpdateCachedState(navigation_handle)) {
+      UpdateActiveState(navigation_handle);
+    }
   }
 }
 
 void SpeedreaderTabHelper::DidRedirectNavigation(
     content::NavigationHandle* navigation_handle) {
   if (navigation_handle->IsInMainFrame()) {
-    UpdateActiveState(navigation_handle);
+    if (!MaybeUpdateCachedState(navigation_handle)) {
+      UpdateActiveState(navigation_handle);
+    }
   }
 }
 
@@ -174,18 +201,23 @@ void SpeedreaderTabHelper::HideBubble() {
   }
 }
 
-void SpeedreaderTabHelper::OnDistillComplete(bool success) {
-  DCHECK(SpeedreaderTabHelper::PageWantsDistill(distill_state_));
-  if (!success) {
-    distill_state_ = DistillState::kNone;
-    return;
-  }
-
+void SpeedreaderTabHelper::OnDistillComplete() {
   // Perform a state transition
+  auto* entry = web_contents()->GetController().GetLastCommittedEntry();
+  DCHECK(entry);
+  if (!entry) {
+    return;  // not possible?
+  }
   if (distill_state_ == DistillState::kSpeedreaderModePending) {
+    SpeedreaderExtendedInfoHandler::PersistSpeedreaderMode(entry);
     distill_state_ = DistillState::kSpeedreaderMode;
-  } else {  // distill_state_ == DistillState::kReaderModePending
+  } else if (distill_state_ == DistillState::kReaderModePending) {
+    SpeedreaderExtendedInfoHandler::PersistReaderMode(entry);
     distill_state_ = DistillState::kReaderMode;
+  } else {
+    // We got here via an already cached page.
+    DCHECK(distill_state_ == DistillState::kSpeedreaderMode ||
+           distill_state_ == DistillState::kReaderMode);
   }
 }
 

browser/speedreader/speedreader_tab_helper.h

@@ -11,6 +11,7 @@
 #include "content/public/browser/web_contents_user_data.h"
 
 namespace content {
+class NavigationEntry;
 class NavigationHandle;
 class WebContents;
 }  // namespace content
@@ -71,7 +72,9 @@ class SpeedreaderTabHelper
   }
 
   static bool PageWantsDistill(DistillState state) {
-    return state == DistillState::kReaderModePending ||
+    return state == DistillState::kReaderMode ||
+           state == DistillState::kSpeedreaderMode ||
+           state == DistillState::kReaderModePending ||
            state == DistillState::kSpeedreaderModePending;
   }
 
@@ -126,6 +129,7 @@ class SpeedreaderTabHelper
   // committed to the WebContents.
   bool IsEnabledForSite(const GURL& url);
 
+  bool MaybeUpdateCachedState(content::NavigationHandle* handle);
   void UpdateActiveState(content::NavigationHandle* handle);
   void SetNextRequestState(DistillState state);
 
@@ -136,7 +140,7 @@ class SpeedreaderTabHelper
       content::NavigationHandle* navigation_handle) override;
 
   // SpeedreaderResultDelegate:
-  void OnDistillComplete(bool success) override;
+  void OnDistillComplete() override;
 
   bool single_shot_next_request_ =
       false;  // run speedreader once on next page load

components/speedreader/speedreader_result_delegate.h

@@ -13,12 +13,10 @@ class WebContents;
 }  // namespace content
 
 // SpeedreaderResultDelegate is an interface for the speedreader component to
-// notify a tab_helper that the distillation has completed. To meet the
-// requirements of a tab_helper, the interface is scoped to the lifetime of
-// WebContents.
+// notify a tab_helper that the distillation has completed successufully.
 class SpeedreaderResultDelegate {
  public:
-  virtual void OnDistillComplete(bool success) = 0;
+  virtual void OnDistillComplete() = 0;
 };
 
 #endif  // BRAVE_COMPONENTS_SPEEDREADER_SPEEDREADER_RESULT_DELEGATE_H_

components/speedreader/speedreader_throttle.cc

@@ -30,9 +30,7 @@ class SpeedreaderResultWebContentsLifetimeHelper
       SpeedreaderResultDelegate* delegate)
       : delegate_(delegate) {}
 
-  void OnDistillComplete(bool success) {
-    delegate_->OnDistillComplete(success);
-  }
+  void OnDistillComplete() { delegate_->OnDistillComplete(); }
 
   WEB_CONTENTS_USER_DATA_KEY_DECL();
 
@@ -44,7 +42,7 @@ class SpeedreaderResultWebContentsLifetimeHelper
 
 WEB_CONTENTS_USER_DATA_KEY_IMPL(SpeedreaderResultWebContentsLifetimeHelper)
 
-}  // anonymous namespace
+}  // namespace
 
 namespace speedreader {
 // static
@@ -105,12 +103,12 @@ void SpeedReaderThrottle::Resume() {
   delegate_->Resume();
 }
 
-void SpeedReaderThrottle::OnDistillComplete(bool success) {
+void SpeedReaderThrottle::OnDistillComplete() {
   auto* result_helper =
       SpeedreaderResultWebContentsLifetimeHelper::FromWebContents(
           web_contents_);
   if (result_helper) {
-    result_helper->OnDistillComplete(success);
+    result_helper->OnDistillComplete();
   }
 }