From 496704750768e8710d101861dca321accdf36302 Mon Sep 17 00:00:00 2001 From: Pete Miller Date: Mon, 21 Oct 2024 22:00:09 -0700 Subject: [PATCH] Prevent navigation in Tab if feature flag isn't enabled --- browser/ai_chat/BUILD.gn | 19 +++ browser/ai_chat/ai_chat_throttle_unittest.cc | 122 ++++++++++++++++++ browser/brave_content_browser_client.cc | 10 ++ components/ai_chat/content/browser/BUILD.gn | 2 + .../content/browser/ai_chat_throttle.cc | 91 +++++++++++++ .../content/browser/ai_chat_throttle.h | 32 +++++ test/BUILD.gn | 1 + 7 files changed, 277 insertions(+) create mode 100644 browser/ai_chat/ai_chat_throttle_unittest.cc create mode 100644 components/ai_chat/content/browser/ai_chat_throttle.cc create mode 100644 components/ai_chat/content/browser/ai_chat_throttle.h diff --git a/browser/ai_chat/BUILD.gn b/browser/ai_chat/BUILD.gn index 9930e54e1ebd..f882ac93f149 100644 --- a/browser/ai_chat/BUILD.gn +++ b/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 diff --git a/browser/ai_chat/ai_chat_throttle_unittest.cc b/browser/ai_chat/ai_chat_throttle_unittest.cc new file mode 100644 index 000000000000..75180cd12ba8 --- /dev/null +++ b/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 + +#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 { + 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(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 web_contents_; + std::unique_ptr profile_manager_; + base::test::ScopedFeatureList features_; +}; + +INSTANTIATE_TEST_SUITE_P( + , + AiChatThrottleUnitTest, + ::testing::Bool(), + [](const testing::TestParamInfo& 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 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 throttle = + AiChatThrottle::MaybeCreateThrottleFor(&test_handle); + EXPECT_EQ(throttle.get(), nullptr); +} + +} // namespace ai_chat diff --git a/browser/brave_content_browser_client.cc b/browser/brave_content_browser_client.cc index 4024ef66ffc0..a37795802975 100644 --- a/browser/brave_content_browser_client.cc +++ b/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; } diff --git a/components/ai_chat/content/browser/BUILD.gn b/components/ai_chat/content/browser/BUILD.gn index 158c1ab2ec9f..1e3a609d705f 100644 --- a/components/ai_chat/content/browser/BUILD.gn +++ b/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", diff --git a/components/ai_chat/content/browser/ai_chat_throttle.cc b/components/ai_chat/content/browser/ai_chat_throttle.cc new file mode 100644 index 000000000000..4b38ae22665a --- /dev/null +++ b/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 + +#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::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(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 diff --git a/components/ai_chat/content/browser/ai_chat_throttle.h b/components/ai_chat/content/browser/ai_chat_throttle.h new file mode 100644 index 000000000000..113270fd2e05 --- /dev/null +++ b/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 + +#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 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_ diff --git a/test/BUILD.gn b/test/BUILD.gn index 7feaeef13918..2e1726abcb48 100644 --- a/test/BUILD.gn +++ b/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",