Prevent navigation in Tab if feature flag isn't enabled

browser/ai_chat/BUILD.gn

@@ -44,6 +44,25 @@ static_library("ai_chat") {
   }
 }
 
+source_set("unit_tests") {
+  testonly = true
+  sources = [ "ai_chat_throttle_unittest.cc" ]
+
+  deps = [
+    "//base",
+    "//base/test:test_support",
+    "//brave/components/ai_chat/content/browser",
+    "//brave/components/ai_chat/core/common",
+    "//brave/components/constants",
+    "//chrome/common",
+    "//chrome/test:test_support",
+    "//content/public/browser",
+    "//content/test:test_support",
+    "//testing/gtest",
+    "//url",
+  ]
+}
+
 source_set("browser_tests") {
   if (!is_android) {
     testonly = true

browser/ai_chat/ai_chat_throttle_unittest.cc

@@ -0,0 +1,122 @@
+/* Copyright (c) 2023 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 https://mozilla.org/MPL/2.0/. */
+
+#include <memory>
+
+#include "base/test/scoped_feature_list.h"
+#include "brave/components/ai_chat/content/browser/ai_chat_throttle.h"
+#include "brave/components/ai_chat/core/common/features.h"
+#include "chrome/test/base/testing_browser_process.h"
+#include "chrome/test/base/testing_profile.h"
+#include "chrome/test/base/testing_profile_manager.h"
+#include "content/public/test/browser_task_environment.h"
+#include "content/public/test/mock_navigation_handle.h"
+#include "content/public/test/web_contents_tester.h"
+#include "testing/gtest/include/gtest/gtest.h"
+#include "url/gurl.h"
+
+namespace ai_chat {
+
+namespace {
+constexpr char kTestProfileName[] = "TestProfile";
+}  // namespace
+
+class AiChatThrottleUnitTest : public testing::Test,
+                               public ::testing::WithParamInterface<bool> {
+ public:
+  AiChatThrottleUnitTest() = default;
+  AiChatThrottleUnitTest(const AiChatThrottleUnitTest&) = delete;
+  AiChatThrottleUnitTest& operator=(const AiChatThrottleUnitTest&) = delete;
+  ~AiChatThrottleUnitTest() override = default;
+
+  void SetUp() override {
+    TestingBrowserProcess* browser_process = TestingBrowserProcess::GetGlobal();
+    profile_manager_ = std::make_unique<TestingProfileManager>(browser_process);
+    ASSERT_TRUE(profile_manager_->SetUp());
+    Profile* profile = profile_manager_->CreateTestingProfile(kTestProfileName);
+
+    web_contents_ =
+        content::WebContentsTester::CreateTestWebContents(profile, nullptr);
+
+    features_.InitWithFeatureStates({
+        {ai_chat::features::kAIChat, true},
+        {ai_chat::features::kAIChatHistory, IsAIChatHistoryEnabled()},
+    });
+  }
+
+  bool IsAIChatHistoryEnabled() { return GetParam(); }
+
+  void TearDown() override {
+    web_contents_.reset();
+    profile_manager_->DeleteTestingProfile(kTestProfileName);
+  }
+
+  content::WebContents* web_contents() { return web_contents_.get(); }
+
+ private:
+  content::BrowserTaskEnvironment task_environment_;
+  std::unique_ptr<content::WebContents> web_contents_;
+  std::unique_ptr<TestingProfileManager> profile_manager_;
+  base::test::ScopedFeatureList features_;
+};
+
+INSTANTIATE_TEST_SUITE_P(
+    ,
+    AiChatThrottleUnitTest,
+    ::testing::Bool(),
+    [](const testing::TestParamInfo<AiChatThrottleUnitTest::ParamType>& info) {
+      return base::StringPrintf("History%s",
+                                info.param ? "Enabled" : "Disabled");
+    });
+
+TEST_F(AiChatThrottleUnitTest, CancelNavigationFromTab) {
+  content::MockNavigationHandle test_handle(web_contents());
+
+  test_handle.set_url(GURL("chrome-untrusted://chat"));
+
+#if BUILDFLAG(IS_ANDROID)
+  ui::PageTransition transition = ui::PageTransitionFromInt(
+      ui::PageTransition::PAGE_TRANSITION_FROM_ADDRESS_BAR);
+#else
+  ui::PageTransition transition = ui::PageTransitionFromInt(
+      ui::PageTransition::PAGE_TRANSITION_FROM_ADDRESS_BAR |
+      ui::PageTransition::PAGE_TRANSITION_TYPED);
+#endif
+
+  test_handle.set_page_transition(transition);
+
+  std::unique_ptr<AiChatThrottle> throttle =
+      AiChatThrottle::MaybeCreateThrottleFor(&test_handle);
+
+  if (IsAIChatHistoryEnabled()) {
+    EXPECT_EQ(throttle.get(), nullptr);
+  } else {
+    EXPECT_NE(throttle.get(), nullptr);
+    EXPECT_EQ(content::NavigationThrottle::CANCEL_AND_IGNORE,
+              throttle->WillStartRequest().action());
+  }
+}
+
+TEST_F(AiChatThrottleUnitTest, AllowNavigationFromPanel) {
+  content::MockNavigationHandle test_handle(web_contents());
+
+  test_handle.set_url(GURL("chrome-untrusted://chat"));
+
+#if BUILDFLAG(IS_ANDROID)
+  ui::PageTransition transition =
+      ui::PageTransitionFromInt(ui::PageTransition::PAGE_TRANSITION_FROM_API);
+#else
+  ui::PageTransition transition = ui::PageTransitionFromInt(
+      ui::PageTransition::PAGE_TRANSITION_AUTO_TOPLEVEL);
+#endif
+
+  test_handle.set_page_transition(transition);
+
+  std::unique_ptr<AiChatThrottle> throttle =
+      AiChatThrottle::MaybeCreateThrottleFor(&test_handle);
+  EXPECT_EQ(throttle.get(), nullptr);
+}
+
+}  // namespace ai_chat

browser/brave_content_browser_client.cc

@@ -158,6 +158,7 @@ using extensions::ChromeContentBrowserClientExtensionsPart;
 #if BUILDFLAG(ENABLE_AI_CHAT)
 #include "brave/browser/ui/webui/ai_chat/ai_chat_ui.h"
 #include "brave/components/ai_chat/content/browser/ai_chat_tab_helper.h"
+#include "brave/components/ai_chat/content/browser/ai_chat_throttle.h"
 #include "brave/components/ai_chat/core/browser/utils.h"
 #include "brave/components/ai_chat/core/common/features.h"
 #include "brave/components/ai_chat/core/common/mojom/ai_chat.mojom.h"
@@ -1259,6 +1260,15 @@ BraveContentBrowserClient::CreateThrottlesForNavigation(
   }
 #endif
 
+#if BUILDFLAG(ENABLE_AI_CHAT)
+  if (Profile::FromBrowserContext(context)->IsRegularProfile()) {
+    if (auto ai_chat_throttle =
+            ai_chat::AiChatThrottle::MaybeCreateThrottleFor(handle)) {
+      throttles.push_back(std::move(ai_chat_throttle));
+    }
+  }
+#endif  // ENABLE_AI_CHAT
+
   return throttles;
 }
 

components/ai_chat/content/browser/BUILD.gn

@@ -12,6 +12,8 @@ static_library("browser") {
   sources = [
     "ai_chat_tab_helper.cc",
     "ai_chat_tab_helper.h",
+    "ai_chat_throttle.cc",
+    "ai_chat_throttle.h",
     "model_service_factory.cc",
     "model_service_factory.h",
     "page_content_fetcher.cc",

components/ai_chat/content/browser/ai_chat_throttle.cc

@@ -0,0 +1,91 @@
+/* Copyright (c) 2023 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 https://mozilla.org/MPL/2.0/. */
+
+#include "brave/components/ai_chat/content/browser/ai_chat_throttle.h"
+
+#include <memory>
+
+#include "brave/components/ai_chat/core/browser/utils.h"
+#include "brave/components/ai_chat/core/common/buildflags/buildflags.h"
+#include "brave/components/ai_chat/core/common/features.h"
+#include "brave/components/constants/webui_url_constants.h"
+#include "components/user_prefs/user_prefs.h"
+#include "content/public/browser/browser_context.h"
+#include "content/public/browser/navigation_handle.h"
+#include "content/public/browser/web_contents.h"
+#include "content/public/common/url_constants.h"
+
+namespace ai_chat {
+
+// static
+std::unique_ptr<AiChatThrottle> AiChatThrottle::MaybeCreateThrottleFor(
+    content::NavigationHandle* navigation_handle) {
+  // The AI Chat WebUI won't be enabled if the feature is disabled
+  if (!ai_chat::IsAIChatEnabled(user_prefs::UserPrefs::Get(
+          navigation_handle->GetWebContents()->GetBrowserContext()))) {
+    return nullptr;
+  }
+
+  // We don't need this throttle if the full-page feature is enabled via proxy
+  // of the AIChatHistory feature flag.
+  if (features::IsAIChatHistoryEnabled()) {
+    return nullptr;
+  }
+
+  // We need this throttle to work only for chrome-untrusted://chat page
+  if (!navigation_handle->GetURL().SchemeIs(
+          content::kChromeUIUntrustedScheme) ||
+      navigation_handle->GetURL().host_piece() != kChatUIHost) {
+    return nullptr;
+  }
+
+  // Purpose of this throttle is to forbid loading of chrome-untrusted://chat
+  // in tab.
+  // Parameters check is made different for Android and Desktop because
+  // there are different flags:
+  // --------+---------------------------------+------------------------------
+  //         | Tab                             | Panel
+  // --------+---------------------------------+------------------------------
+  // Android |PAGE_TRANSITION_FROM_ADDRESS_BAR | PAGE_TRANSITION_FROM_API
+  // --------+---------------------------------+------------------------------
+  // Desktop |PAGE_TRANSITION_TYPED|           | PAGE_TRANSITION_AUTO_TOPLEVEL
+  //         |PAGE_TRANSITION_FROM_ADDRESS_BAR |
+  // -------------------------------------------------------------------------
+  //
+  // So for Android the only allowed transition is PAGE_TRANSITION_FROM_API
+  // because it is pretty unique and means the page is loaded in a custom tab
+  // view.
+  // And for the desktop just disallow PAGE_TRANSITION_FROM_ADDRESS_BAR
+  ui::PageTransition transition = navigation_handle->GetPageTransition();
+#if BUILDFLAG(IS_ANDROID)
+  if (ui::PageTransitionTypeIncludingQualifiersIs(
+          transition, ui::PageTransition::PAGE_TRANSITION_FROM_API)) {
+    return nullptr;
+  }
+#else
+  if (!ui::PageTransitionTypeIncludingQualifiersIs(
+          ui::PageTransitionGetQualifier(transition),
+          ui::PageTransition::PAGE_TRANSITION_FROM_ADDRESS_BAR)) {
+    return nullptr;
+  }
+#endif  // BUILDFLAG(IS_ANDROID)
+
+  return std::make_unique<AiChatThrottle>(navigation_handle);
+}
+
+AiChatThrottle::AiChatThrottle(content::NavigationHandle* handle)
+    : content::NavigationThrottle(handle) {}
+
+AiChatThrottle::~AiChatThrottle() {}
+
+AiChatThrottle::ThrottleCheckResult AiChatThrottle::WillStartRequest() {
+  return CANCEL_AND_IGNORE;
+}
+
+const char* AiChatThrottle::GetNameForLogging() {
+  return "AiChatThrottle";
+}
+
+}  // namespace ai_chat

components/ai_chat/content/browser/ai_chat_throttle.h

@@ -0,0 +1,32 @@
+/* Copyright (c) 2023 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 https://mozilla.org/MPL/2.0/. */
+
+#ifndef BRAVE_COMPONENTS_AI_CHAT_CONTENT_BROWSER_AI_CHAT_THROTTLE_H_
+#define BRAVE_COMPONENTS_AI_CHAT_CONTENT_BROWSER_AI_CHAT_THROTTLE_H_
+
+#include <memory>
+
+#include "content/public/browser/navigation_throttle.h"
+
+namespace ai_chat {
+
+// Prevents navigation to certain AI Chat URLs
+class AiChatThrottle : public content::NavigationThrottle {
+ public:
+  explicit AiChatThrottle(content::NavigationHandle* handle);
+  ~AiChatThrottle() override;
+
+  static std::unique_ptr<AiChatThrottle> MaybeCreateThrottleFor(
+      content::NavigationHandle* navigation_handle);
+
+  // content::NavigationThrottle:
+  // ThrottleCheckResult WillProcessResponse() override;
+  ThrottleCheckResult WillStartRequest() override;
+  const char* GetNameForLogging() override;
+};
+
+}  // namespace ai_chat
+
+#endif  // BRAVE_COMPONENTS_AI_CHAT_CONTENT_BROWSER_AI_CHAT_THROTTLE_H_

test/BUILD.gn

@@ -365,6 +365,7 @@ test("brave_unit_tests") {
 
   if (enable_ai_chat) {
     deps += [
+      "//brave/browser/ai_chat:unit_tests",
       "//brave/browser/ui/ai_chat:unit_tests",
       "//brave/components/ai_chat/core/browser:unit_tests",
       "//brave/components/ai_chat/core/common:unit_tests",