Skip to content

Commit

Permalink
Add a NavigationThrottle to intercept specific URL requests and open Leo
Browse files Browse the repository at this point in the history
Also adds ValidateOpenLeoButtonNonce in PageContentFetcher
  • Loading branch information
yrliou committed Oct 19, 2024
1 parent 589a245 commit 9ef822a
Show file tree
Hide file tree
Showing 35 changed files with 907 additions and 28 deletions.
13 changes: 7 additions & 6 deletions browser/ai_chat/BUILD.gn
Original file line number Diff line number Diff line change
Expand Up @@ -10,12 +10,12 @@ assert(enable_ai_chat)

static_library("ai_chat") {
sources = [
"//brave/browser/ai_chat/ai_chat_service_factory.cc",
"//brave/browser/ai_chat/ai_chat_service_factory.h",
"//brave/browser/ai_chat/ai_chat_settings_helper.cc",
"//brave/browser/ai_chat/ai_chat_settings_helper.h",
"//brave/browser/ai_chat/ai_chat_utils.cc",
"//brave/browser/ai_chat/ai_chat_utils.h",
"ai_chat_service_factory.cc",
"ai_chat_service_factory.h",
"ai_chat_settings_helper.cc",
"ai_chat_settings_helper.h",
"ai_chat_utils.cc",
"ai_chat_utils.h",
]

deps = [
Expand Down Expand Up @@ -73,6 +73,7 @@ source_set("browser_tests") {
sources = [
"//chrome/browser/renderer_context_menu/render_view_context_menu_browsertest_util.cc",
"//chrome/browser/renderer_context_menu/render_view_context_menu_browsertest_util.h",
"ai_chat_brave_search_throttle_browsertest.cc",
"ai_chat_browsertests.cc",
"ai_chat_metrics_browsertest.cc",
"ai_chat_policy_browsertest.cc",
Expand Down
238 changes: 238 additions & 0 deletions browser/ai_chat/ai_chat_brave_search_throttle_browsertest.cc
Original file line number Diff line number Diff line change
@@ -0,0 +1,238 @@
/* Copyright (c) 2024 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_brave_search_throttle.h"

#include <memory>
#include <string>

#include "base/files/file_path.h"
#include "base/location.h"
#include "base/path_service.h"
#include "base/strings/string_util.h"
#include "brave/browser/ui/brave_browser.h"
#include "brave/browser/ui/sidebar/sidebar_controller.h"
#include "brave/browser/ui/sidebar/sidebar_model.h"
#include "brave/components/constants/brave_paths.h"
#include "chrome/browser/ui/browser.h"
#include "chrome/test/base/in_process_browser_test.h"
#include "components/permissions/permission_request_manager.h"
#include "components/permissions/test/mock_permission_prompt_factory.h"
#include "content/public/test/browser_test.h"
#include "content/public/test/browser_test_utils.h"
#include "content/public/test/content_mock_cert_verifier.h"
#include "content/public/test/test_navigation_observer.h"
#include "content/public/test/test_utils.h"
#include "net/base/net_errors.h"
#include "net/dns/mock_host_resolver.h"
#include "net/test/embedded_test_server/embedded_test_server.h"
#include "testing/gtest/include/gtest/gtest.h"

namespace {

constexpr char kBraveSearchHost[] = "search.brave.com";
constexpr char kLeoPath[] = "/leo";
constexpr char kOpenLeoButtonValidPath[] = "/open_leo_button_valid.html";
constexpr char kOpenLeoButtonInvalidPath[] = "/open_leo_button_invalid.html";

} // namespace

class AIChatBraveSearchThrottleBrowserTest : public InProcessBrowserTest {
public:
AIChatBraveSearchThrottleBrowserTest()
: https_server_(net::EmbeddedTestServer::TYPE_HTTPS) {}

void SetUpOnMainThread() override {
InProcessBrowserTest::SetUpOnMainThread();

mock_cert_verifier_.mock_cert_verifier()->set_default_result(net::OK);
host_resolver()->AddRule("*", "127.0.0.1");
content::SetupCrossSiteRedirector(&https_server_);

base::FilePath test_data_dir =
base::PathService::CheckedGet(brave::DIR_TEST_DATA);
test_data_dir = test_data_dir.AppendASCII("leo");
https_server_.ServeFilesFromDirectory(test_data_dir);
ASSERT_TRUE(https_server_.Start());

permissions::PermissionRequestManager* manager =
permissions::PermissionRequestManager::FromWebContents(
GetActiveWebContents());
prompt_factory_ =
std::make_unique<permissions::MockPermissionPromptFactory>(manager);
}

void TearDownOnMainThread() override {
prompt_factory_.reset();
InProcessBrowserTest::TearDownOnMainThread();
}

void SetUpCommandLine(base::CommandLine* command_line) override {
InProcessBrowserTest::SetUpCommandLine(command_line);
mock_cert_verifier_.SetUpCommandLine(command_line);
}

void SetUpInProcessBrowserTestFixture() override {
InProcessBrowserTest::SetUpInProcessBrowserTestFixture();
mock_cert_verifier_.SetUpInProcessBrowserTestFixture();
}

void TearDownInProcessBrowserTestFixture() override {
mock_cert_verifier_.TearDownInProcessBrowserTestFixture();
InProcessBrowserTest::TearDownInProcessBrowserTestFixture();
}

content::WebContents* GetActiveWebContents() {
return browser()->tab_strip_model()->GetActiveWebContents();
}

void ClickOpenLeoButton() {
// Modify the href to have test server port.
std::string script = R"(
var link = document.getElementById('continue-with-leo')
var url = new URL(link.href)
url.port = '$1'
link.href = url.href
link.click()
)";
std::string port = base::NumberToString(https_server_.port());

ASSERT_TRUE(content::ExecJs(
GetActiveWebContents()->GetPrimaryMainFrame(),
base::ReplaceStringPlaceholders(script, {port}, nullptr)));
}

bool IsLeoOpened() {
sidebar::SidebarController* controller =
static_cast<BraveBrowser*>(browser())->sidebar_controller();
auto index = controller->model()->GetIndexOf(
sidebar::SidebarItem::BuiltInItemType::kChatUI);
return index.has_value() && controller->IsActiveIndex(index);
}

void CloseLeoPanel(const base::Location& location) {
SCOPED_TRACE(testing::Message() << location.ToString());
sidebar::SidebarController* controller =
static_cast<BraveBrowser*>(browser())->sidebar_controller();
controller->DeactivateCurrentPanel();
ASSERT_FALSE(IsLeoOpened());
}

void NavigateToTestPage(const base::Location& location,
const std::string& host,
const std::string& path,
int expected_prompt_count) {
SCOPED_TRACE(testing::Message() << location.ToString());
ASSERT_TRUE(content::NavigateToURL(GetActiveWebContents(),
https_server_.GetURL(host, path)));
EXPECT_FALSE(IsLeoOpened());
EXPECT_EQ(expected_prompt_count, prompt_factory_->show_count());
}

void ClickOpenLeoAndCheckLeoOpenedAndNavigationCancelled(
const base::Location& location,
int expected_prompt_count,
bool expected_leo_opened,
const std::string& expected_last_committed_path =
kOpenLeoButtonValidPath) {
SCOPED_TRACE(testing::Message() << location.ToString());
content::TestNavigationObserver observer(
GetActiveWebContents(), net::ERR_ABORTED,
content::MessageLoopRunner::QuitMode::IMMEDIATE,
false /* ignore_uncommitted_navigations */);
ClickOpenLeoButton();
observer.Wait();

EXPECT_EQ(IsLeoOpened(), expected_leo_opened);
EXPECT_EQ(expected_prompt_count, prompt_factory_->show_count());
EXPECT_EQ(observer.last_navigation_url().path_piece(), kLeoPath);
EXPECT_EQ(GetActiveWebContents()->GetLastCommittedURL().path_piece(),
expected_last_committed_path);
}

protected:
net::test_server::EmbeddedTestServer https_server_;
std::unique_ptr<permissions::MockPermissionPromptFactory> prompt_factory_;

private:
content::ContentMockCertVerifier mock_cert_verifier_;
};

IN_PROC_BROWSER_TEST_F(AIChatBraveSearchThrottleBrowserTest,
OpenLeo_AskAndAccept) {
int cur_prompt_count = 0;
prompt_factory_->set_response_type(
permissions::PermissionRequestManager::ACCEPT_ALL);
NavigateToTestPage(FROM_HERE, kBraveSearchHost, kOpenLeoButtonValidPath,
cur_prompt_count);
ClickOpenLeoAndCheckLeoOpenedAndNavigationCancelled(FROM_HERE,
++cur_prompt_count, true);

CloseLeoPanel(FROM_HERE);
ClickOpenLeoAndCheckLeoOpenedAndNavigationCancelled(FROM_HERE,
cur_prompt_count, true);
}

IN_PROC_BROWSER_TEST_F(AIChatBraveSearchThrottleBrowserTest,
OpenLeo_AskAndDeny) {
int cur_prompt_count = 0;
prompt_factory_->set_response_type(
permissions::PermissionRequestManager::DENY_ALL);
NavigateToTestPage(FROM_HERE, kBraveSearchHost, kOpenLeoButtonValidPath,
cur_prompt_count);
ClickOpenLeoAndCheckLeoOpenedAndNavigationCancelled(
FROM_HERE, ++cur_prompt_count, false);

// Clicking a button again to test no new permission prompt should be shown
// when the permission setting is denied.
ClickOpenLeoAndCheckLeoOpenedAndNavigationCancelled(FROM_HERE,
cur_prompt_count, false);
}

IN_PROC_BROWSER_TEST_F(AIChatBraveSearchThrottleBrowserTest,
OpenLeo_AskAndDismiss) {
int cur_prompt_count = 0;
prompt_factory_->set_response_type(
permissions::PermissionRequestManager::DISMISS);
NavigateToTestPage(FROM_HERE, kBraveSearchHost, kOpenLeoButtonValidPath,
cur_prompt_count);
ClickOpenLeoAndCheckLeoOpenedAndNavigationCancelled(
FROM_HERE, ++cur_prompt_count, false);

// Click a button again after dismissing the permission, permission prompt
// should be shown again.
prompt_factory_->set_response_type(
permissions::PermissionRequestManager::ACCEPT_ALL);
NavigateToTestPage(FROM_HERE, kBraveSearchHost, kOpenLeoButtonValidPath,
cur_prompt_count);
ClickOpenLeoAndCheckLeoOpenedAndNavigationCancelled(FROM_HERE,
++cur_prompt_count, true);
}

IN_PROC_BROWSER_TEST_F(AIChatBraveSearchThrottleBrowserTest,
OpenLeo_MismatchedNonce) {
int cur_prompt_count = 0;
NavigateToTestPage(FROM_HERE, kBraveSearchHost, kOpenLeoButtonInvalidPath,
cur_prompt_count);
// No permission prompt should be shown.
ClickOpenLeoAndCheckLeoOpenedAndNavigationCancelled(
FROM_HERE, cur_prompt_count, false, kOpenLeoButtonInvalidPath);
}

IN_PROC_BROWSER_TEST_F(AIChatBraveSearchThrottleBrowserTest,
OpenLeo_NotBraveSearchURL) {
// The behavior should be the same as without the throttle.
NavigateToTestPage(FROM_HERE, "brave.com", kOpenLeoButtonValidPath, 0);
content::TestNavigationObserver observer(GetActiveWebContents());
ClickOpenLeoButton();
observer.Wait();

EXPECT_FALSE(IsLeoOpened());
EXPECT_EQ(0, prompt_factory_->show_count());
EXPECT_TRUE(observer.last_navigation_succeeded());
EXPECT_EQ(observer.last_navigation_url().path_piece(), kLeoPath);
EXPECT_EQ(GetActiveWebContents()->GetLastCommittedURL().path_piece(),
kLeoPath);
}
55 changes: 55 additions & 0 deletions browser/ai_chat/page_content_fetcher_browsertest.cc
Original file line number Diff line number Diff line change
Expand Up @@ -133,6 +133,22 @@ class PageContentFetcherBrowserTest : public InProcessBrowserTest {
run_loop.Run();
}

void ValidateOpenLeoButtonNonce(const base::Location& location,
const std::string& nonce,
bool expected_is_valid) {
SCOPED_TRACE(testing::Message() << location.ToString());
base::RunLoop run_loop;
page_content_fetcher_ =
std::make_unique<ai_chat::PageContentFetcher>(GetActiveWebContents());
page_content_fetcher_->ValidateOpenLeoButtonNonce(
nonce, base::BindLambdaForTesting(
[&run_loop, expected_is_valid](bool is_valid) {
EXPECT_EQ(expected_is_valid, is_valid);
run_loop.Quit();
}));
run_loop.Run();
}

// Handles returning a .patch file if the user is on a github.com pull request
void SetGithubInterceptor() {
GURL expected_patch_url =
Expand Down Expand Up @@ -269,3 +285,42 @@ IN_PROC_BROWSER_TEST_F(PageContentFetcherBrowserTest, GetSearchSummarizerKey) {
GetSearchSummarizerKey(FROM_HERE, expected_result);
}
}

IN_PROC_BROWSER_TEST_F(PageContentFetcherBrowserTest,
ValidateOpenLeoButtonNonce) {
// Test no open Leo button with continue-with-leo ID present.
GURL url = https_server_.GetURL("a.com", "/open_leo_button.html");
NavigateURL(url);
ValidateOpenLeoButtonNonce(FROM_HERE, "5566", false);

// Test valid case.
NavigateURL(url);
ASSERT_TRUE(content::ExecJs(GetActiveWebContents()->GetPrimaryMainFrame(),
"document.getElementById('valid').setAttribute('"
"id', 'continue-with-leo')"));
ValidateOpenLeoButtonNonce(FROM_HERE, "5566", true);

// The pass in nonce should match with continue-with-leo URL in href.
ValidateOpenLeoButtonNonce(FROM_HERE, "7788", false);

// Empty nonce should return false.
ValidateOpenLeoButtonNonce(FROM_HERE, "", false);

// Test invalid cases.
std::vector<std::string> invalid_cases = {
"invalid", "not-a-tag", "no-href", "no-nonce",
"empty-nonce", "empty-nonce2", "empty-nonce3", "empty-nonce4",
"empty-nonce5", "empty-nonce6", "not-https-url", "not-search-url",
"not-open-leo-url"};

for (const auto& invalid_case : invalid_cases) {
SCOPED_TRACE(testing::Message() << "Invalid case: " << invalid_case);
NavigateURL(url);
ASSERT_TRUE(content::ExecJs(GetActiveWebContents()->GetPrimaryMainFrame(),
base::ReplaceStringPlaceholders(
"document.getElementById('$1')."
"setAttribute('id', 'continue-with-leo')",
{invalid_case}, nullptr)));
ValidateOpenLeoButtonNonce(FROM_HERE, "5566", false);
}
}
20 changes: 20 additions & 0 deletions browser/brave_content_browser_client.cc
Original file line number Diff line number Diff line change
Expand Up @@ -156,14 +156,19 @@ using extensions::ChromeContentBrowserClientExtensionsPart;
#endif

#if BUILDFLAG(ENABLE_AI_CHAT)
#include "brave/browser/ai_chat/ai_chat_service_factory.h"
#include "brave/browser/ui/webui/ai_chat/ai_chat_ui.h"
#include "brave/components/ai_chat/content/browser/ai_chat_brave_search_throttle.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"
#include "brave/components/ai_chat/core/common/mojom/page_content_extractor.mojom.h"
#include "brave/components/ai_chat/core/common/mojom/settings_helper.mojom.h"
#if !BUILDFLAG(IS_ANDROID)
#include "brave/browser/ui/ai_chat/ai_chat_brave_search_throttle_delegate_impl.h"
#endif
#if BUILDFLAG(IS_ANDROID)
#include "brave/components/ai_chat/core/browser/android/ai_chat_iap_subscription_android.h"
#endif
Expand Down Expand Up @@ -1267,6 +1272,21 @@ BraveContentBrowserClient::CreateThrottlesForNavigation(
throttles.push_back(std::move(ai_chat_throttle));
}
}

if (Profile::FromBrowserContext(context)->IsRegularProfile()) {
std::unique_ptr<ai_chat::AIChatBraveSearchThrottle::Delegate> delegate;
#if !BUILDFLAG(IS_ANDROID)
delegate =
std::make_unique<ai_chat::AIChatBraveSearchThrottleDelegateImpl>();
#endif

if (auto ai_chat_brave_search_throttle =
ai_chat::AIChatBraveSearchThrottle::MaybeCreateThrottleFor(
std::move(delegate), handle,
ai_chat::AIChatServiceFactory::GetForBrowserContext(context))) {
throttles.push_back(std::move(ai_chat_brave_search_throttle));
}
}
#endif // ENABLE_AI_CHAT

return throttles;
Expand Down
7 changes: 7 additions & 0 deletions browser/ui/BUILD.gn
Original file line number Diff line number Diff line change
Expand Up @@ -108,6 +108,13 @@ source_set("ui") {
"webui/ai_chat/ai_chat_ui_page_handler.h",
]

if (!is_android && !is_ios) {
sources += [
"ai_chat/ai_chat_brave_search_throttle_delegate_impl.cc",
"ai_chat/ai_chat_brave_search_throttle_delegate_impl.h",
]
}

deps += [ "//brave/browser/ai_chat" ]

if (enable_print_preview) {
Expand Down
Loading

0 comments on commit 9ef822a

Please sign in to comment.