// Copyright 2023 The Chromium Authors
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#include "chrome/browser/extensions/api/user_scripts/user_scripts_apitest.h"

#include "base/one_shot_event.h"
#include "base/test/bind.h"
#include "base/test/scoped_run_loop_timeout.h"
#include "build/build_config.h"
#include "chrome/browser/extensions/extension_apitest.h"
#include "chrome/browser/extensions/extension_util.h"
#include "chrome/browser/extensions/user_scripts_test_util.h"
#include "chrome/browser/profiles/profile.h"
#include "content/public/browser/web_contents.h"
#include "content/public/test/browser_test.h"
#include "content/public/test/browser_test_utils.h"
#include "content/public/test/test_navigation_observer.h"
#include "extensions/browser/background_script_executor.h"
#include "extensions/browser/extension_file_task_runner.h"
#include "extensions/browser/extension_prefs.h"
#include "extensions/browser/extension_registry.h"
#include "extensions/browser/extension_util.h"
#include "extensions/browser/renderer_startup_helper.h"
#include "extensions/browser/script_injection_tracker.h"
#include "extensions/browser/user_script_manager.h"
#include "extensions/buildflags/buildflags.h"
#include "extensions/common/features/feature_developer_mode_only.h"
#include "extensions/common/user_scripts_allowed_state.h"
#include "extensions/common/utils/content_script_utils.h"
#include "extensions/test/extension_test_message_listener.h"
#include "extensions/test/result_catcher.h"
#include "extensions/test/test_extension_dir.h"
#include "net/dns/mock_host_resolver.h"
#include "testing/gtest/include/gtest/gtest.h"

static_assert(BUILDFLAG(ENABLE_EXTENSIONS_CORE));

namespace extensions {

void UserScriptsAPITest::SetUpOnMainThread() {
  ExtensionApiTest::SetUpOnMainThread();

  host_resolver()->AddRule("*", "127.0.0.1");
  ASSERT_TRUE(StartEmbeddedTestServer());
}

void UserScriptsAPITest::OpenInCurrentTab(const GURL& url) {
  content::WebContents* web_contents = GetActiveWebContents();
  ASSERT_TRUE(web_contents);

  // NavigateToURL() waits for the load to stop and verifies the navigation
  // succeeded.
  ASSERT_TRUE(NavigateToURL(web_contents, url));
  EXPECT_EQ(url, web_contents->GetLastCommittedURL());
}

content::RenderFrameHost* UserScriptsAPITest::OpenInNewTab(const GURL& url) {
  content::TestNavigationObserver nav_observer(url);
  nav_observer.StartWatchingNewWebContents();
  NavigateToURLInNewTab(url);
  auto* web_contents = GetActiveWebContents();
  EXPECT_TRUE(content::WaitForLoadStop(web_contents));
  nav_observer.Wait();

  EXPECT_TRUE(nav_observer.last_navigation_succeeded());
  EXPECT_EQ(url, web_contents->GetLastCommittedURL());

  return web_contents->GetPrimaryMainFrame();
}

content::EvalJsResult UserScriptsAPITest::GetInjectedElements(
    content::RenderFrameHost* host) {
  static constexpr char kGetInjectedScripts[] =
      R"(const divs = document.body.getElementsByTagName('div');
           JSON.stringify(Array.from(divs).map(div => div.id).sort());)";
  return content::EvalJs(host, kGetInjectedScripts);
}

  // Loads the extension and pauses in-between loading and running the tests to
  // enable the userScripts API.
testing::AssertionResult UserScriptsAPITest::RunUserScriptsExtensionTestImpl(
    const base::FilePath& extension_path,
    bool allow_api) {
  // Load the extension.
  ExtensionTestMessageListener test_ready_listener(
      "ready",
      allow_api ? ReplyBehavior::kWillReply : ReplyBehavior::kWontReply);
  ResultCatcher catcher;
  const Extension* extension = LoadExtension(extension_path);
  if (!extension) {
    return testing::AssertionFailure() << "Failed to load extension";
  }

  if (allow_api) {
    // Wait until extension tests are ready to run, then allow the
    // userScripts API, then continue on with the API testing.
    bool extension_ready = test_ready_listener.WaitUntilSatisfied();
    if (!extension_ready) {
      testing::AssertionFailure()
          << "extension did not signal that it was ready after loading";
    }
    user_scripts_test_util::SetUserScriptsAPIAllowed(profile(), extension->id(),
                                                     /*allowed=*/true);
    test_ready_listener.Reply("");
  }

  // Observe each test result.
  {
    base::test::ScopedRunLoopTimeout timeout(
        FROM_HERE, std::nullopt,
        base::BindRepeating(
            [](const base::FilePath& extension_path) {
              return "GetNextResult timeout while "
                     "RunUserScriptsExtensionTest: " +
                     extension_path.MaybeAsASCII();
            },
            extension_path));
    if (!catcher.GetNextResult()) {
      return testing::AssertionFailure() << catcher.message();
    }
  }

  return testing::AssertionSuccess();
  ;
}

testing::AssertionResult UserScriptsAPITest::RunUserScriptsExtensionTest(
    const char* extension_sub_path) {
  const base::FilePath& root_path = test_data_dir_;
  base::FilePath extension_path = root_path.AppendASCII(extension_sub_path);
  return RunUserScriptsExtensionTestImpl(extension_path, /*allow_api=*/true);
}

testing::AssertionResult
UserScriptsAPITest::RunUserScriptsExtensionTestNotAllowed(
    const base::FilePath& extension_path) {
  return RunUserScriptsExtensionTestImpl(extension_path, /*allow_api=*/false);
}

UserScriptsAPITest::UserScriptsAPITest() = default;
UserScriptsAPITest::~UserScriptsAPITest() = default;

IN_PROC_BROWSER_TEST_F(UserScriptsAPITest, GetUserScripts) {
  ASSERT_TRUE(RunUserScriptsExtensionTest("user_scripts/get_scripts"))
      << message_;
}

// TODO(crbug.com/40935741, crbug.com/335421977): Flaky on Linux debug and on
// "Linux ChromiumOS MSan Tests".
#if (BUILDFLAG(IS_LINUX) && !defined(NDEBUG)) || \
    (BUILDFLAG(IS_CHROMEOS) && defined(MEMORY_SANITIZER))
#define MAYBE_RegisterUserScripts DISABLED_RegisterUserScripts
#else
#define MAYBE_RegisterUserScripts RegisterUserScripts
#endif
IN_PROC_BROWSER_TEST_F(UserScriptsAPITest, MAYBE_RegisterUserScripts) {
  ASSERT_TRUE(RunUserScriptsExtensionTest("user_scripts/register")) << message_;
}

IN_PROC_BROWSER_TEST_F(UserScriptsAPITest, UnregisterUserScripts) {
  ASSERT_TRUE(RunUserScriptsExtensionTest("user_scripts/unregister"))
      << message_;
}

IN_PROC_BROWSER_TEST_F(UserScriptsAPITest, UpdateUserScripts) {
  ASSERT_TRUE(RunUserScriptsExtensionTest("user_scripts/update")) << message_;
}

IN_PROC_BROWSER_TEST_F(UserScriptsAPITest, ExecuteUserScripts) {
  ASSERT_TRUE(RunUserScriptsExtensionTest("user_scripts/execute")) << message_;
}

IN_PROC_BROWSER_TEST_F(UserScriptsAPITest, ExecuteUserScripts_Subframes) {
  // Open up two tabs, each with cross-site iframes, one at a.com and one at
  // d.com. In both cases, the cross-site iframes point to b.com and c.com.
  OpenInCurrentTab(
      embedded_test_server()->GetURL("a.com", "/iframe_cross_site.html"));
  OpenInNewTab(
      embedded_test_server()->GetURL("d.com", "/iframe_cross_site.html"));

  ASSERT_TRUE(
      RunUserScriptsExtensionTest("user_scripts/execute_with_subframes"))
      << message_;
}

IN_PROC_BROWSER_TEST_F(UserScriptsAPITest, ExecuteUserScripts_SizeLimit) {
  auto single_scripts_limit_reset =
      script_parsing::CreateScopedMaxScriptLengthForTesting(700u);
  ASSERT_TRUE(RunUserScriptsExtensionTest("user_scripts/execute_size_limit"))
      << message_;
}

// Tests that executing user scripts correctly handles a cross-document
// navigation occurring during the asynchronous file loading phase, preventing
// unauthorized script execution or improper marking of the new process.
IN_PROC_BROWSER_TEST_F(UserScriptsAPITest,
                       ExecuteUserScripts_CrossDocumentNavigationRace) {
  ExtensionTestMessageListener ready_listener("ready",
                                              ReplyBehavior::kWillReply);
  ExtensionTestMessageListener tab_created_listener("tab_created",
                                                    ReplyBehavior::kWillReply);
  ExtensionTestMessageListener execute_called_listener(
      "execute_called", ReplyBehavior::kWillReply);

  ResultCatcher catcher;

  const Extension* extension =
      LoadExtension(test_data_dir_.AppendASCII("user_scripts/navigation_race"));
  ASSERT_TRUE(extension);

  user_scripts_test_util::SetUserScriptsAPIAllowed(profile(), extension->id(),
                                                   /*allowed=*/true);

  ASSERT_TRUE(ready_listener.WaitUntilSatisfied());
  ready_listener.Reply(
      embedded_test_server()->GetURL("a.com", "/empty.html").spec());

  ASSERT_TRUE(tab_created_listener.WaitUntilSatisfied());

  // Block the file task runner.
  base::WaitableEvent block_file_task_runner;
  GetExtensionFileTaskRunner()->PostTask(
      FROM_HERE, base::BindLambdaForTesting([&]() {
        base::ScopedAllowBaseSyncPrimitivesForTesting allow_blocking;
        block_file_task_runner.Wait();
      }));

  tab_created_listener.Reply("");

  ASSERT_TRUE(execute_called_listener.WaitUntilSatisfied());

  // Now userScripts.execute() has been called. The file task is queued.
  // Navigate the tab to c.com, which the extension does not have permission
  // for.
  content::WebContents* web_contents = GetActiveWebContents();
  ASSERT_TRUE(NavigateToURL(
      web_contents, embedded_test_server()->GetURL("c.com", "/empty.html")));

  // Unblock file task runner
  block_file_task_runner.Signal();

  execute_called_listener.Reply("");

  // Wait for the result from the extension. The execute should result in an
  // expected error due to the script being blocked after the navigation.
  EXPECT_TRUE(catcher.GetNextResult()) << catcher.message();

  // Verify that the new document's process was not marked as having run user
  // scripts.
  content::RenderFrameHost* main_frame = web_contents->GetPrimaryMainFrame();
  EXPECT_FALSE(ScriptInjectionTracker::DidProcessRunUserScriptFromExtension(
      *main_frame->GetProcess(), extension->id()));
}

// TODO(crbug.com/335421977): Flaky on "Linux ChromiumOS MSan Tests".
#if BUILDFLAG(IS_CHROMEOS) && defined(MEMORY_SANITIZER)
#define MAYBE_ConfigureWorld DISABLED_ConfigureWorld
#else
#define MAYBE_ConfigureWorld ConfigureWorld
#endif
IN_PROC_BROWSER_TEST_F(UserScriptsAPITest, MAYBE_ConfigureWorld) {
  ASSERT_TRUE(RunUserScriptsExtensionTest("user_scripts/configure_world"))
      << message_;
}

IN_PROC_BROWSER_TEST_F(UserScriptsAPITest,
                       UserScriptInjectionOrderIsAlphabetical) {
  ASSERT_TRUE(RunUserScriptsExtensionTest("user_scripts/injection_order"))
      << message_;
}

IN_PROC_BROWSER_TEST_F(UserScriptsAPITest, GetAndRemoveWorlds) {
  ASSERT_TRUE(RunUserScriptsExtensionTest("user_scripts/get_and_remove_worlds"))
      << message_;
}

// Tests that registered user scripts are disabled when the userScripts API is
// not allowed and are re-enabled if when the API is allowed again.
IN_PROC_BROWSER_TEST_F(UserScriptsAPITest,
                       UserScriptsAreDisabledWhenAPIIsNotAllowed) {
  const Extension* extension =
      LoadExtension(test_data_dir_.AppendASCII("user_scripts/allowed_tests"));
  ASSERT_TRUE(extension);
  user_scripts_test_util::SetUserScriptsAPIAllowed(profile(), extension->id(),
                                                   /*allowed=*/true);

  // Register a user script and a content script.
  EXPECT_EQ("success",
            BackgroundScriptExecutor::ExecuteScript(
                profile(), extension->id(), "registerUserScripts();",
                BackgroundScriptExecutor::ResultCapture::kSendScriptResult));
  EXPECT_EQ("success",
            BackgroundScriptExecutor::ExecuteScript(
                profile(), extension->id(), "registerContentScript();",
                BackgroundScriptExecutor::ResultCapture::kSendScriptResult));

  const GURL url =
      embedded_test_server()->GetURL("example.com", "/simple.html");

  // Open a new tab.
  content::RenderFrameHost* new_tab = OpenInNewTab(url);

  // Since the userScript API is available (as part of this test suite's setup),
  // both the user script and the content script should inject.
  EXPECT_EQ(R"(["content-script","user-script-code","user-script-file"])",
            GetInjectedElements(new_tab));

  user_scripts_test_util::SetUserScriptsAPIAllowed(profile(), extension->id(),
                                                   /*allowed=*/false);

  // Open a new tab. Now, user scripts should be disabled. However, content
  // scripts should still inject.
  new_tab = OpenInNewTab(url);
  EXPECT_EQ(R"(["content-script"])", GetInjectedElements(new_tab));

  user_scripts_test_util::SetUserScriptsAPIAllowed(profile(), extension->id(),
                                                   /*allowed=*/true);

  // Open a new tab. The user script should inject again.
  new_tab = OpenInNewTab(url);
  EXPECT_EQ(R"(["content-script","user-script-code","user-script-file"])",
            GetInjectedElements(new_tab));
}

// Tests that unregisterContentScripts unregisters only content scripts and
// not user scripts.
IN_PROC_BROWSER_TEST_F(UserScriptsAPITest,
                       ScriptingAPIDoesNotAffectUserScripts) {
  ASSERT_TRUE(RunUserScriptsExtensionTest("scripting/dynamic_user_scripts"))
      << message_;
}

// Base test fixture for tests spanning multiple sessions where a custom arg
// is set before the test is run.
class PersistentUserScriptsAPITest : public UserScriptsAPITest {
 public:
  PersistentUserScriptsAPITest() = default;

  // UserScriptsAPITest override.
  void SetUp() override {
    // Initialize the listener object here before calling SetUp. This avoids a
    // race condition where the extension loads (as part of browser startup) and
    // sends a message before a message listener in C++ has been initialized.
    listener_ = std::make_unique<ExtensionTestMessageListener>(
        "ready", ReplyBehavior::kWillReply);

    UserScriptsAPITest::SetUp();
  }

  // Reset listener before the browser gets torn down.
  void TearDownOnMainThread() override {
    listener_.reset();
    UserScriptsAPITest::TearDownOnMainThread();
  }

 protected:
  // Used to wait for results from extension tests. This is initialized before
  // the test is run which avoids a race condition where the extension is loaded
  // (as part of startup) and finishes its tests before the ResultCatcher is
  // created.
  ResultCatcher result_catcher_;

  // Used to wait for the extension to load and send a ready message so the test
  // can reply which the extension waits for to start its testing functions.
  // This ensures that the testing functions will run after the browser has
  // finished initializing.
  std::unique_ptr<ExtensionTestMessageListener> listener_;
};

// Tests that registered user scripts persist across sessions. The test is run
// across three sessions.
IN_PROC_BROWSER_TEST_F(PersistentUserScriptsAPITest,
                       PRE_PRE_PersistentScripts) {
  const Extension* extension = LoadExtension(
      test_data_dir_.AppendASCII("user_scripts/persistent_scripts"));
  ASSERT_TRUE(extension);
  ASSERT_TRUE(listener_->WaitUntilSatisfied());
  user_scripts_test_util::SetUserScriptsAPIAllowed(profile(), extension->id(),
                                                   /*allowed=*/true);
  listener_->Reply(
      testing::UnitTest::GetInstance()->current_test_info()->name());
  EXPECT_TRUE(result_catcher_.GetNextResult()) << result_catcher_.message();
}

IN_PROC_BROWSER_TEST_F(PersistentUserScriptsAPITest, PRE_PersistentScripts) {
  ASSERT_TRUE(listener_->WaitUntilSatisfied());
  listener_->Reply(
      testing::UnitTest::GetInstance()->current_test_info()->name());
  EXPECT_TRUE(result_catcher_.GetNextResult()) << result_catcher_.message();
}

IN_PROC_BROWSER_TEST_F(PersistentUserScriptsAPITest, PersistentScripts) {
  ASSERT_TRUE(listener_->WaitUntilSatisfied());
  listener_->Reply(
      testing::UnitTest::GetInstance()->current_test_info()->name());
  EXPECT_TRUE(result_catcher_.GetNextResult()) << result_catcher_.message();
}

// Tests that the world configuration of a registered user script is persisted
// across sessions. The test is run across three sessions.
IN_PROC_BROWSER_TEST_F(PersistentUserScriptsAPITest,
                       PRE_PRE_PersistentWorldConfiguration) {
  const Extension* extension = LoadExtension(
      test_data_dir_.AppendASCII("user_scripts/persistent_configure_world"));
  ASSERT_TRUE(extension);
  ASSERT_TRUE(listener_->WaitUntilSatisfied());
  user_scripts_test_util::SetUserScriptsAPIAllowed(profile(), extension->id(),
                                                   /*allowed=*/true);
  listener_->Reply(
      testing::UnitTest::GetInstance()->current_test_info()->name());
  EXPECT_TRUE(result_catcher_.GetNextResult()) << result_catcher_.message();
}

IN_PROC_BROWSER_TEST_F(PersistentUserScriptsAPITest,
                       PRE_PersistentWorldConfiguration) {
  ASSERT_TRUE(listener_->WaitUntilSatisfied());
  listener_->Reply(
      testing::UnitTest::GetInstance()->current_test_info()->name());
  EXPECT_TRUE(result_catcher_.GetNextResult()) << result_catcher_.message();
}

IN_PROC_BROWSER_TEST_F(PersistentUserScriptsAPITest,
                       PersistentWorldConfiguration) {
  ASSERT_TRUE(listener_->WaitUntilSatisfied());
  listener_->Reply(
      testing::UnitTest::GetInstance()->current_test_info()->name());
  EXPECT_TRUE(result_catcher_.GetNextResult()) << result_catcher_.message();
}

class UserScriptsAPITestWithoutAPIAllowed : public UserScriptsAPITest {
 public:
  UserScriptsAPITestWithoutAPIAllowed() = default;

  // UserScriptsAPITest override.
  void SetUp() override {
    // Initialize the listener object here before calling SetUp. This avoids a
    // race condition where the extension loads (as part of browser startup) and
    // sends a message before a message listener in C++ has been initialized.
    background_started_listener_ =
        std::make_unique<ExtensionTestMessageListener>("started");

    UserScriptsAPITest::SetUp();
  }

  // Reset listener before the browser gets torn down.
  void TearDown() override {
    background_started_listener_.reset();
    UserScriptsAPITest::TearDown();
  }

 protected:
  std::unique_ptr<ExtensionTestMessageListener> background_started_listener_;
};

// Verifies that the `chrome.userScripts` API is undefined if the API is not
// allowed yet.
IN_PROC_BROWSER_TEST_F(UserScriptsAPITestWithoutAPIAllowed,
                       UserScriptsAPIIsUndefinedWithoutAPIAllowed) {
  static constexpr char kManifest[] =
      R"({
           "name": "user scripts",
           "manifest_version": 3,
           "version": "0.1",
           "background": {"service_worker": "background.js"},
           "permissions": ["userScripts"]
         })";
  static constexpr char kBackgroundJs[] =
      R"(chrome.test.runTests([
           function userScriptsIsUndefined() {
             chrome.test.assertTrue(chrome.userScripts === undefined);
             chrome.test.succeed();
           },
         ]);)";

  TestExtensionDir test_dir;
  test_dir.WriteManifest(kManifest);
  test_dir.WriteFile(FILE_PATH_LITERAL("background.js"), kBackgroundJs);
  ASSERT_TRUE(RunUserScriptsExtensionTestNotAllowed(test_dir.UnpackedPath()))
      << message_;
}

// Tests that registered user scripts are properly ignored when loading
// stored dynamic scripts if the API is not allowed.
// TODO(crbug.com/441364550): Flaky on desktop Android.
#if BUILDFLAG(IS_ANDROID)
#define MAYBE_PRE_UserScriptsDisabledOnStartupIfAPINotAllowed \
  DISABLED_PRE_UserScriptsDisabledOnStartupIfAPINotAllowed
#else
#define MAYBE_PRE_UserScriptsDisabledOnStartupIfAPINotAllowed \
  PRE_UserScriptsDisabledOnStartupIfAPINotAllowed
#endif
IN_PROC_BROWSER_TEST_F(UserScriptsAPITestWithoutAPIAllowed,
                       MAYBE_PRE_UserScriptsDisabledOnStartupIfAPINotAllowed) {
  // Load an extension and register user scripts and a dynamic content script.
  const Extension* extension =
      LoadExtension(test_data_dir_.AppendASCII("user_scripts/allowed_tests"));
  ASSERT_TRUE(extension);
  ASSERT_TRUE(background_started_listener_->WaitUntilSatisfied());
  user_scripts_test_util::SetUserScriptsAPIAllowed(profile(), extension->id(),
                                                   /*allowed=*/true);

  EXPECT_EQ("success",
            BackgroundScriptExecutor::ExecuteScript(
                profile(), extension->id(), "registerUserScripts();",
                BackgroundScriptExecutor::ResultCapture::kSendScriptResult));
  EXPECT_EQ("success",
            BackgroundScriptExecutor::ExecuteScript(
                profile(), extension->id(), "registerContentScript();",
                BackgroundScriptExecutor::ResultCapture::kSendScriptResult));

  const GURL url =
      embedded_test_server()->GetURL("example.com", "/simple.html");

  // To start, all scripts should inject.
  content::RenderFrameHost* new_tab = OpenInNewTab(url);
  EXPECT_EQ(R"(["content-script","user-script-code","user-script-file"])",
            GetInjectedElements(new_tab));

  // Disallow userScript API, and then re-open the browser...
  user_scripts_test_util::SetUserScriptsAPIAllowed(profile(), extension->id(),
                                                   /*allowed=*/false);
}

// TODO(crbug.com/441364550): Flaky on desktop Android.
#if BUILDFLAG(IS_ANDROID)
#define MAYBE_UserScriptsDisabledOnStartupIfAPINotAllowed \
  DISABLED_UserScriptsDisabledOnStartupIfAPINotAllowed
#else
#define MAYBE_UserScriptsDisabledOnStartupIfAPINotAllowed \
  UserScriptsDisabledOnStartupIfAPINotAllowed
#endif
IN_PROC_BROWSER_TEST_F(UserScriptsAPITestWithoutAPIAllowed,
                       MAYBE_UserScriptsDisabledOnStartupIfAPINotAllowed) {
  // Wait until the extension loads so we can get it's ID.
  ASSERT_TRUE(background_started_listener_->WaitUntilSatisfied());

  // Find the extension's ID so we can make some assertions.
  ExtensionId extension_id;
  for (const auto& extension :
       ExtensionRegistry::Get(profile())->enabled_extensions()) {
    if (extension->name() == "Test") {
      extension_id = extension->id();
    }
  }
  ASSERT_TRUE(!extension_id.empty());

  // userScripts should remain disallowed after browser restart.
  EXPECT_FALSE(GetCurrentUserScriptAllowedState(
                   util::GetBrowserContextId(profile()), extension_id)
                   .value_or(false));

  const GURL url =
      embedded_test_server()->GetURL("example.com", "/simple.html");

  // And, to start, only the content script should inject.
  content::RenderFrameHost* new_tab = OpenInNewTab(url);
  EXPECT_EQ(R"(["content-script"])", GetInjectedElements(new_tab));

  user_scripts_test_util::SetUserScriptsAPIAllowed(profile(), extension_id,
                                                   /*allowed=*/true);

  // All scripts should once again inject.
  new_tab = OpenInNewTab(url);
  EXPECT_EQ(R"(["content-script","user-script-code","user-script-file"])",
            GetInjectedElements(new_tab));
}

}  // namespace extensions