Skip to content

Commit

Permalink
Go back to untrusted webui for this branch
Browse files Browse the repository at this point in the history
  • Loading branch information
petemill committed Sep 11, 2024
1 parent aad1e48 commit c041d30
Show file tree
Hide file tree
Showing 12 changed files with 306 additions and 49 deletions.
19 changes: 19 additions & 0 deletions browser/ai_chat/BUILD.gn
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,25 @@ import("//printing/buildflags/buildflags.gni")

assert(enable_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
Expand Down
103 changes: 103 additions & 0 deletions browser/ai_chat/ai_chat_throttle_unittest.cc
Original file line number Diff line number Diff line change
@@ -0,0 +1,103 @@
/* 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:
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_.InitAndEnableFeature(ai_chat::features::kAIChat);
}

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_;
};

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);
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
10 changes: 10 additions & 0 deletions browser/brave_content_browser_client.cc
Original file line number Diff line number Diff line change
Expand Up @@ -156,6 +156,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"
Expand Down Expand Up @@ -1258,6 +1259,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;
}

Expand Down
74 changes: 37 additions & 37 deletions browser/ui/webui/ai_chat/ai_chat_ui.cc
Original file line number Diff line number Diff line change
Expand Up @@ -14,19 +14,17 @@
#include "brave/components/ai_chat/core/browser/constants.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"
#include "brave/components/ai_chat/core/common/pref_names.h"
#include "brave/components/ai_chat/resources/page/grit/ai_chat_ui_generated_map.h"
#include "brave/components/constants/webui_url_constants.h"
#include "brave/components/l10n/common/localization_util.h"
#include "chrome/browser/profiles/profile.h"
#include "chrome/browser/ui/webui/favicon_source.h"
#include "chrome/browser/ui/webui/webui_util.h"
#include "components/favicon_base/favicon_url_parser.h"
#include "components/grit/brave_components_resources.h"
#include "components/prefs/pref_service.h"
#include "components/user_prefs/user_prefs.h"
#include "content/public/browser/web_contents.h"
#include "content/public/browser/web_ui_controller.h"
#include "content/public/browser/web_ui_data_source.h"
#include "content/public/common/url_constants.h"

Expand Down Expand Up @@ -58,69 +56,69 @@ content::WebContents* GetActiveWebContents(content::BrowserContext* context) {
#endif

AIChatUI::AIChatUI(content::WebUI* web_ui)
: WebUIController(web_ui), profile_(Profile::FromWebUI(web_ui)) {
: ui::UntrustedWebUIController(web_ui),
profile_(Profile::FromWebUI(web_ui)) {
DCHECK(profile_);
DCHECK(profile_->IsRegularProfile());

// Create a URLDataSource and add resources.
content::WebUIDataSource* source = content::WebUIDataSource::CreateAndAdd(
web_ui->GetWebContents()->GetBrowserContext(), kChatUIHost);
content::WebUIDataSource* untrusted_source =
content::WebUIDataSource::CreateAndAdd(
web_ui->GetWebContents()->GetBrowserContext(), kChatUIURL);

webui::SetupWebUIDataSource(
source, base::make_span(kAiChatUiGenerated, kAiChatUiGeneratedSize),
untrusted_source,
base::make_span(kAiChatUiGenerated, kAiChatUiGeneratedSize),
IDR_CHAT_UI_HTML);

source->AddResourcePath("styles.css", IDR_CHAT_UI_CSS);
untrusted_source->AddResourcePath("styles.css", IDR_CHAT_UI_CSS);

for (const auto& str : ai_chat::GetLocalizedStrings()) {
source->AddString(str.name,
brave_l10n::GetLocalizedResourceUTF16String(str.id));
untrusted_source->AddString(
str.name, brave_l10n::GetLocalizedResourceUTF16String(str.id));
}

base::Time last_accepted_disclaimer =
profile_->GetOriginalProfile()->GetPrefs()->GetTime(
ai_chat::prefs::kLastAcceptedDisclaimer);

source->AddBoolean("hasAcceptedAgreement",
!last_accepted_disclaimer.is_null());
source->AddInteger("customModelMaxPageContentLength",
ai_chat::kCustomModelMaxPageContentLength);
source->AddInteger("customModelLongConversationCharLimit",
ai_chat::kCustomModelLongConversationCharLimit);
untrusted_source->AddBoolean("hasAcceptedAgreement",
!last_accepted_disclaimer.is_null());
untrusted_source->AddInteger("customModelMaxPageContentLength",
ai_chat::kCustomModelMaxPageContentLength);
untrusted_source->AddInteger("customModelLongConversationCharLimit",
ai_chat::kCustomModelLongConversationCharLimit);

#if BUILDFLAG(IS_ANDROID) || BUILDFLAG(IS_IOS)
constexpr bool kIsMobile = true;
#else
constexpr bool kIsMobile = false;
#endif

source->AddBoolean("isMobile", kIsMobile);
source->AddBoolean("isHistoryEnabled",
untrusted_source->AddBoolean("isMobile", kIsMobile);
untrusted_source->AddBoolean("isHistoryEnabled",
ai_chat::features::IsAIChatHistoryEnabled());

source->AddBoolean("hasUserDismissedPremiumPrompt",
profile_->GetOriginalProfile()->GetPrefs()->GetBoolean(
ai_chat::prefs::kUserDismissedPremiumPrompt));
untrusted_source->AddBoolean(
"hasUserDismissedPremiumPrompt",
profile_->GetOriginalProfile()->GetPrefs()->GetBoolean(
ai_chat::prefs::kUserDismissedPremiumPrompt));

source->OverrideContentSecurityPolicy(
untrusted_source->OverrideContentSecurityPolicy(
network::mojom::CSPDirectiveName::ScriptSrc,
"script-src 'self' chrome://resources;");
source->OverrideContentSecurityPolicy(
"script-src 'self' chrome-untrusted://resources;");
untrusted_source->OverrideContentSecurityPolicy(
network::mojom::CSPDirectiveName::StyleSrc,
"style-src 'self' 'unsafe-inline' chrome://resources;");
source->OverrideContentSecurityPolicy(
"style-src 'self' 'unsafe-inline' chrome-untrusted://resources;");
untrusted_source->OverrideContentSecurityPolicy(
network::mojom::CSPDirectiveName::ImgSrc,
"img-src 'self' blob: chrome://resources chrome://favicon2;");
source->OverrideContentSecurityPolicy(
"img-src 'self' blob: chrome-untrusted://resources;");
untrusted_source->OverrideContentSecurityPolicy(
network::mojom::CSPDirectiveName::FontSrc,
"font-src 'self' data: chrome://resources;");
"font-src 'self' data: chrome-untrusted://resources;");

source->OverrideContentSecurityPolicy(
untrusted_source->OverrideContentSecurityPolicy(
network::mojom::CSPDirectiveName::TrustedTypes, "trusted-types default;");

content::URLDataSource::Add(
profile_, std::make_unique<FaviconSource>(
profile_, chrome::FaviconUrlFormat::kFavicon2));
}

AIChatUI::~AIChatUI() = default;
Expand Down Expand Up @@ -161,13 +159,15 @@ void AIChatUI::BindInterface(
std::move(receiver));
}

bool AIChatUIConfig::IsWebUIEnabled(content::BrowserContext* browser_context) {
bool UntrustedChatUIConfig::IsWebUIEnabled(
content::BrowserContext* browser_context) {
return ai_chat::IsAIChatEnabled(
user_prefs::UserPrefs::Get(browser_context)) &&
Profile::FromBrowserContext(browser_context)->IsRegularProfile();
}

AIChatUIConfig::AIChatUIConfig()
: DefaultTopChromeWebUIConfig(content::kChromeUIScheme, kChatUIHost) {}
UntrustedChatUIConfig::UntrustedChatUIConfig()
: DefaultTopChromeWebUIConfig(content::kChromeUIUntrustedScheme,
kChatUIHost) {}

WEB_UI_CONTROLLER_TYPE_IMPL(AIChatUI)
10 changes: 6 additions & 4 deletions browser/ui/webui/ai_chat/ai_chat_ui.h
Original file line number Diff line number Diff line change
Expand Up @@ -13,14 +13,16 @@
#include "chrome/browser/ui/webui/top_chrome/top_chrome_web_ui_controller.h"
#include "chrome/browser/ui/webui/top_chrome/top_chrome_webui_config.h"
#include "content/public/browser/web_ui_controller.h"
#include "ui/webui/mojo_web_ui_controller.h"
#include "ui/webui/untrusted_web_ui_controller.h"

namespace content {
class BrowserContext;
}

class Profile;

class AIChatUI : public content::WebUIController {
class AIChatUI : public ui::UntrustedWebUIController {
public:
explicit AIChatUI(content::WebUI* web_ui);
AIChatUI(const AIChatUI&) = delete;
Expand Down Expand Up @@ -49,10 +51,10 @@ class AIChatUI : public content::WebUIController {
WEB_UI_CONTROLLER_TYPE_DECL();
};

class AIChatUIConfig : public DefaultTopChromeWebUIConfig<AIChatUI> {
class UntrustedChatUIConfig : public DefaultTopChromeWebUIConfig<AIChatUI> {
public:
AIChatUIConfig();
~AIChatUIConfig() override = default;
UntrustedChatUIConfig();
~UntrustedChatUIConfig() override = default;

bool IsWebUIEnabled(content::BrowserContext* browser_context) override;
};
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -11,11 +11,17 @@
#include "brave/browser/ui/webui/brave_wallet/market/market_ui.h"
#include "brave/browser/ui/webui/brave_wallet/nft/nft_ui.h"
#include "brave/browser/ui/webui/brave_wallet/trezor/trezor_ui.h"
#include "brave/components/ai_chat/core/common/buildflags/buildflags.h"
#include "brave/components/brave_vpn/common/buildflags/buildflags.h"
#include "brave/components/playlist/common/buildflags/buildflags.h"
#include "build/build_config.h"
#include "content/public/browser/webui_config_map.h"

#if BUILDFLAG(ENABLE_AI_CHAT)
#include "brave/browser/ui/webui/ai_chat/ai_chat_ui.h"
#include "brave/components/ai_chat/core/common/features.h"
#endif // BUILDFLAG(ENABLE_AI_CHAT)

#if BUILDFLAG(ENABLE_BRAVE_VPN)
#if !BUILDFLAG(IS_ANDROID)
#include "brave/browser/ui/webui/brave_vpn/vpn_panel_ui.h"
Expand Down Expand Up @@ -63,4 +69,11 @@ void RegisterChromeUntrustedWebUIConfigs() {
}
#endif // BUILDFLAG(ENABLE_PLAYLIST_WEBUI)
#endif // !BUILDFLAG(IS_ANDROID)

#if BUILDFLAG(ENABLE_AI_CHAT)
if (ai_chat::features::IsAIChatEnabled()) {
content::WebUIConfigMap::GetInstance().AddUntrustedWebUIConfig(
std::make_unique<UntrustedChatUIConfig>());
}
#endif // BUILDFLAG(ENABLE_AI_CHAT)
}
7 changes: 0 additions & 7 deletions chromium_src/chrome/browser/ui/webui/chrome_web_ui_configs.cc
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,6 @@

#include "chrome/browser/ui/webui/chrome_web_ui_configs.h"

#include "brave/components/ai_chat/core/common/buildflags/buildflags.h"
#include "content/public/browser/webui_config_map.h"

#define RegisterChromeWebUIConfigs RegisterChromeWebUIConfigs_ChromiumImpl
Expand All @@ -22,9 +21,6 @@
#include "brave/browser/ui/webui/brave_wallet/wallet_panel_ui.h"
#include "brave/browser/ui/webui/speedreader/speedreader_toolbar_ui.h"
#endif // !BUILDFLAG(IS_ANDROID)
#if BUILDFLAG(ENABLE_AI_CHAT)
#include "brave/browser/ui/webui/ai_chat/ai_chat_ui.h"
#endif // BUILDFLAG(ENABLE_AI_CHAT)

void RegisterChromeWebUIConfigs() {
RegisterChromeWebUIConfigs_ChromiumImpl();
Expand All @@ -39,7 +35,4 @@ void RegisterChromeWebUIConfigs() {
map.AddWebUIConfig(std::make_unique<WalletPanelUIConfig>());
map.AddWebUIConfig(std::make_unique<SpeedreaderToolbarUIConfig>());
#endif // !BUILDFLAG(IS_ANDROID)
#if BUILDFLAG(ENABLE_AI_CHAT)
map.AddWebUIConfig(std::make_unique<AIChatUIConfig>());
#endif
}
2 changes: 2 additions & 0 deletions components/ai_chat/content/browser/BUILD.gn
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down
Loading

0 comments on commit c041d30

Please sign in to comment.