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

#include "extensions/browser/mime_handler/mime_handler_registry.h"

#include "base/containers/span.h"
#include "base/feature_list.h"
#include "base/json/values_util.h"
#include "base/strings/stringprintf.h"
#include "base/test/scoped_feature_list.h"
#include "base/time/time.h"
#include "components/version_info/channel.h"
#include "extensions/browser/extension_pref_names.h"
#include "extensions/browser/extension_prefs.h"
#include "extensions/browser/extension_registry.h"
#include "extensions/browser/extensions_test.h"
#include "extensions/browser/unloaded_extension_reason.h"
#include "extensions/common/constants.h"
#include "extensions/common/extension.h"
#include "extensions/common/extension_builder.h"
#include "extensions/common/extension_features.h"
#include "extensions/common/features/feature_channel.h"
#include "extensions/common/manifest_handlers/mime_types_handler.h"
#include "testing/gtest/include/gtest/gtest.h"

namespace extensions {
namespace {

constexpr char kPdfMimeType[] = "application/pdf";
constexpr char kViewerUrl[] = "viewer.html";

// Expected handler registration for `channel`: the ApiMimeHandler feature
// state (default or explicit override) with the parser's dev-channel
// fallback. Computed at runtime so a change to the flag's default state
// moves these expectations together with the parser gate. Must be called
// while the ScopedFeatureList under test is active.
bool ExpectHandlerRegistered(version_info::Channel channel) {
  return base::FeatureList::IsEnabled(extensions_features::kApiMimeHandler) ||
         channel <= version_info::Channel::DEV;
}

class MimeHandlerRegistryTest : public ExtensionsTest {
 protected:
  void SetUp() override {
    feature_list_.InitAndEnableFeature(extensions_features::kApiMimeHandler);
    ExtensionsTest::SetUp();
    ASSERT_TRUE(registry());
  }

  // Builds an extension with dict-format mime_types_handler.
  scoped_refptr<const Extension> CreateMimeHandlerExtension(
      const std::string& name,
      const std::string& mime_type,
      const std::string& handler_url) {
    std::string json = base::StringPrintf(
        R"("mime_types_handler": {"%s": {"handler_url": "%s"}})",
        mime_type.c_str(), handler_url.c_str());
    return ExtensionBuilder(name).AddJSON(json).Build();
  }

  // Builds an allowlisted dict-format mime handler whose extension ID is
  // `id_from_allowlist` (must be one of `kMIMETypeHandlersAllowlist`).
  scoped_refptr<const Extension> CreateAllowlistedHandler(
      const ExtensionId& id_from_allowlist,
      const std::string& mime_type) {
    std::string json = base::StringPrintf(
        R"("mime_types_handler": {"%s": {"handler_url": "viewer.html"}})",
        mime_type.c_str());
    return ExtensionBuilder(id_from_allowlist)
        .SetID(id_from_allowlist)
        .AddJSON(json)
        .Build();
  }

  // Overrides the persisted first-install time for `ext` so tests can
  // control the sort key used by the registry.
  void SetFirstInstallTime(const Extension* ext, base::Time t) {
    ExtensionPrefs::Get(browser_context())
        ->UpdateExtensionPref(ext->id(), kPrefFirstInstallTime,
                              base::TimeToValue(t));
  }

  void LoadExtension(const Extension* extension) {
    ExtensionRegistry::Get(browser_context())
        ->AddEnabled(base::WrapRefCounted(extension));
    ExtensionRegistry::Get(browser_context())->TriggerOnLoaded(extension);
  }

  void UnloadExtension(const Extension* extension) {
    ExtensionRegistry::Get(browser_context())->RemoveEnabled(extension->id());
    ExtensionRegistry::Get(browser_context())
        ->TriggerOnUnloaded(extension, UnloadedExtensionReason::UNINSTALL);
  }

  MimeHandlerRegistry* registry() {
    return MimeHandlerRegistry::Get(browser_context());
  }

 private:
  // Channel "trunk" is represented as UNKNOWN in the feature system.
  // Required because the dict-format mime_types_handler manifest key
  // is gated to channel "trunk".
  ScopedCurrentChannel channel_{version_info::Channel::UNKNOWN};
  base::test::ScopedFeatureList feature_list_;
};

TEST_F(MimeHandlerRegistryTest, RegisterAndLookup) {
  auto ext =
      CreateMimeHandlerExtension("PDF Handler", kPdfMimeType, kViewerUrl);
  LoadExtension(ext.get());

  EXPECT_EQ(ext->id(), registry()->GetHandlerForMimeType(kPdfMimeType));
}

TEST_F(MimeHandlerRegistryTest, UnregisterRemovesMapping) {
  auto ext =
      CreateMimeHandlerExtension("PDF Handler", kPdfMimeType, kViewerUrl);
  LoadExtension(ext.get());
  UnloadExtension(ext.get());

  EXPECT_TRUE(registry()->GetHandlerForMimeType(kPdfMimeType).empty());
}

TEST_F(MimeHandlerRegistryTest, MostRecentlyInstalledWins) {
  auto ext_a = CreateMimeHandlerExtension("A", kPdfMimeType, "a.html");
  auto ext_b = CreateMimeHandlerExtension("B", kPdfMimeType, "b.html");
  auto ext_c = CreateMimeHandlerExtension("C", kPdfMimeType, "c.html");

  const base::Time t0 = base::Time::Now();
  SetFirstInstallTime(ext_a.get(), t0);
  SetFirstInstallTime(ext_b.get(), t0 + base::Hours(2));
  SetFirstInstallTime(ext_c.get(), t0 + base::Hours(1));

  // Load in a different order than install order. ext_b (newest install)
  // is loaded second, proving neither first-loaded nor last-loaded wins.
  LoadExtension(ext_a.get());
  LoadExtension(ext_b.get());
  LoadExtension(ext_c.get());

  EXPECT_EQ(ext_b->id(), registry()->GetHandlerForMimeType(kPdfMimeType));

  // Unload ext_b; ext_c (next-newest install) becomes active.
  UnloadExtension(ext_b.get());
  EXPECT_EQ(ext_c->id(), registry()->GetHandlerForMimeType(kPdfMimeType));

  // Unload ext_c; ext_a (oldest install) becomes active.
  UnloadExtension(ext_c.get());
  EXPECT_EQ(ext_a->id(), registry()->GetHandlerForMimeType(kPdfMimeType));
}

TEST_F(MimeHandlerRegistryTest, NoHandlerReturnsEmpty) {
  EXPECT_TRUE(registry()->GetHandlerForMimeType(kPdfMimeType).empty());
}

TEST_F(MimeHandlerRegistryTest, PublicHandlerSupersedesAllowlisted) {
  // Allowlisted PDF handler, installed AFTER the public handler.
  auto allowlisted =
      CreateAllowlistedHandler(extension_misc::kPdfExtensionId, kPdfMimeType);
  auto public_ext =
      CreateMimeHandlerExtension("PublicPdf", kPdfMimeType, kViewerUrl);

  const base::Time t0 = base::Time::Now();
  SetFirstInstallTime(public_ext.get(), t0);
  SetFirstInstallTime(allowlisted.get(), t0 + base::Hours(1));

  LoadExtension(public_ext.get());
  LoadExtension(allowlisted.get());

  // Public handlers supersede allowlisted ones, regardless of install time.
  EXPECT_EQ(public_ext->id(), registry()->GetHandlerForMimeType(kPdfMimeType));
}

TEST_F(MimeHandlerRegistryTest, AllowlistArrayIndexBreaksTies) {
  // PDF is at array index 0; QuickOffice is at a higher index.
  // Higher array index wins, even when installed EARLIER than PDF.
  auto pdf =
      CreateAllowlistedHandler(extension_misc::kPdfExtensionId, kPdfMimeType);
  auto qoffice = CreateAllowlistedHandler(
      extension_misc::kQuickOfficeExtensionId, kPdfMimeType);

  const base::Time t0 = base::Time::Now();
  SetFirstInstallTime(qoffice.get(), t0);
  SetFirstInstallTime(pdf.get(), t0 + base::Hours(1));

  LoadExtension(qoffice.get());
  LoadExtension(pdf.get());

  // Expected: qoffice wins by array index, even though pdf is newer.
  EXPECT_EQ(qoffice->id(), registry()->GetHandlerForMimeType(kPdfMimeType));
}

TEST_F(MimeHandlerRegistryTest, GetHandlersByMimeTypeReturnsAllTypes) {
  // Two distinct MIME types. QuickOffice is allowlisted, so it is free to
  // register any MIME type (including a non-public one).
  constexpr char kDocMimeType[] = "application/msword";

  auto pdf_handler =
      CreateMimeHandlerExtension("PdfHandler", kPdfMimeType, "a.html");
  auto doc_handler = CreateAllowlistedHandler(
      extension_misc::kQuickOfficeExtensionId, kDocMimeType);

  LoadExtension(pdf_handler.get());
  LoadExtension(doc_handler.get());

  auto all = registry()->GetHandlersByMimeType();
  ASSERT_EQ(all.size(), 2u);

  auto pdf_it = all.find(kPdfMimeType);
  ASSERT_NE(pdf_it, all.end());
  ASSERT_EQ(pdf_it->second.size(), 1u);
  EXPECT_EQ(pdf_it->second.front(), pdf_handler->id());

  auto doc_it = all.find(kDocMimeType);
  ASSERT_NE(doc_it, all.end());
  ASSERT_EQ(doc_it->second.size(), 1u);
  EXPECT_EQ(doc_it->second.front(), doc_handler->id());
}

TEST_F(MimeHandlerRegistryTest, GetHandlersForMimeTypeIsOrdered) {
  // Exercise all three precedence rules in one ordered span:
  //   * public_new and public_old are public handlers; public beats
  //     allowlisted, and newest-install-time wins among public.
  //   * qoffice and pdf are both allowlisted; qoffice has a higher
  //     index in `kMIMETypeHandlersAllowlist` and must beat pdf.
  auto public_new =
      CreateMimeHandlerExtension("PublicNew", kPdfMimeType, "b.html");
  auto public_old =
      CreateMimeHandlerExtension("PublicOld", kPdfMimeType, "a.html");
  auto qoffice = CreateAllowlistedHandler(
      extension_misc::kQuickOfficeExtensionId, kPdfMimeType);
  auto pdf =
      CreateAllowlistedHandler(extension_misc::kPdfExtensionId, kPdfMimeType);

  const base::Time t0 = base::Time::Now();
  SetFirstInstallTime(public_old.get(), t0);
  SetFirstInstallTime(public_new.get(), t0 + base::Hours(1));
  // Install times for allowlisted handlers do not affect precedence;
  // set deterministic values anyway.
  SetFirstInstallTime(qoffice.get(), t0 + base::Hours(2));
  SetFirstInstallTime(pdf.get(), t0 + base::Hours(3));

  LoadExtension(qoffice.get());
  LoadExtension(pdf.get());
  LoadExtension(public_old.get());
  LoadExtension(public_new.get());

  base::span<const ExtensionId> candidates =
      registry()->GetHandlersForMimeType(kPdfMimeType);
  ASSERT_EQ(candidates.size(), 4u);
  EXPECT_EQ(candidates[0], public_new->id());
  EXPECT_EQ(candidates[1], public_old->id());
  EXPECT_EQ(candidates[2], qoffice->id());
  EXPECT_EQ(candidates[3], pdf->id());
}

TEST_F(MimeHandlerRegistryTest, MultipleMimeTypesWithOverlap) {
  constexpr char kDocMimeType[] = "application/msword";

  // Allowlisted QuickOffice extension registers for BOTH application/pdf
  // and application/msword via a dict-format manifest. Allowlisted
  // extensions bypass `GetPublicAllowedMIMETypeList()`, so msword is
  // accepted.
  auto allowlisted =
      ExtensionBuilder("AllowlistedMulti")
          .SetID(extension_misc::kQuickOfficeExtensionId)
          .AddJSON(R"("mime_types_handler": {)"
                   R"("application/pdf": {"handler_url": "q.html"},)"
                   R"("application/msword": {"handler_url": "q.html"}})")
          .Build();
  auto public_pdf =
      CreateMimeHandlerExtension("PublicPdf", kPdfMimeType, kViewerUrl);

  LoadExtension(allowlisted.get());
  LoadExtension(public_pdf.get());

  // PDF: public beats allowlisted.
  EXPECT_EQ(public_pdf->id(), registry()->GetHandlerForMimeType(kPdfMimeType));
  // msword: only the allowlisted extension registers for it.
  EXPECT_EQ(allowlisted->id(), registry()->GetHandlerForMimeType(kDocMimeType));

  base::span<const ExtensionId> pdf_candidates =
      registry()->GetHandlersForMimeType(kPdfMimeType);
  ASSERT_EQ(pdf_candidates.size(), 2u);
  EXPECT_EQ(pdf_candidates[0], public_pdf->id());
  EXPECT_EQ(pdf_candidates[1], allowlisted->id());

  base::span<const ExtensionId> doc_candidates =
      registry()->GetHandlersForMimeType(kDocMimeType);
  ASSERT_EQ(doc_candidates.size(), 1u);
  EXPECT_EQ(doc_candidates[0], allowlisted->id());
}

TEST_F(MimeHandlerRegistryTest, AllowlistedExtensionRegisteredWhenUsingDict) {
  auto ext = ExtensionBuilder("Allowlisted PDF")
                 .SetID(extension_misc::kPdfExtensionId)
                 .AddJSON(R"("mime_types_handler": {"application/pdf": )"
                          R"({"handler_url": "viewer.html"}})")
                 .Build();
  LoadExtension(ext.get());
  EXPECT_EQ(ext->id(), registry()->GetHandlerForMimeType(kPdfMimeType));
}

TEST_F(MimeHandlerRegistryTest, FlagDisabledRegistrationByChannel) {
  base::test::ScopedFeatureList disable_flag;
  disable_flag.InitAndDisableFeature(extensions_features::kApiMimeHandler);

  for (auto channel :
       {version_info::Channel::UNKNOWN, version_info::Channel::CANARY,
        version_info::Channel::DEV, version_info::Channel::BETA,
        version_info::Channel::STABLE}) {
    SCOPED_TRACE(version_info::GetChannelString(channel));
    ScopedCurrentChannel scoped_channel(channel);
    auto ext = CreateMimeHandlerExtension(
        std::string(version_info::GetChannelString(channel)), kPdfMimeType,
        kViewerUrl);
    LoadExtension(ext.get());
    EXPECT_TRUE(registry()->GetHandlerForMimeType(kPdfMimeType).empty());
    UnloadExtension(ext.get());
  }
}

TEST_F(MimeHandlerRegistryTest, FlagDefaultRegistrationByChannel) {
  // Without an explicit override, availability follows the flag's
  // default state with the dev-channel fallback.
  for (auto channel :
       {version_info::Channel::UNKNOWN, version_info::Channel::CANARY,
        version_info::Channel::DEV, version_info::Channel::BETA,
        version_info::Channel::STABLE}) {
    SCOPED_TRACE(testing::Message()
                 << "channel=" << version_info::GetChannelString(channel));
    ScopedCurrentChannel scoped_channel(channel);
    base::test::ScopedFeatureList no_override;
    no_override.InitWithEmptyFeatureAndFieldTrialLists();

    auto ext = CreateMimeHandlerExtension(
        std::string(version_info::GetChannelString(channel)), kPdfMimeType,
        kViewerUrl);
    LoadExtension(ext.get());

    if (ExpectHandlerRegistered(channel)) {
      EXPECT_EQ(ext->id(), registry()->GetHandlerForMimeType(kPdfMimeType));
    } else {
      EXPECT_TRUE(registry()->GetHandlerForMimeType(kPdfMimeType).empty());
    }

    UnloadExtension(ext.get());
  }
}

TEST_F(MimeHandlerRegistryTest, FlagEnabledRegistrationByChannel) {
  base::test::ScopedFeatureList enable_flag;
  enable_flag.InitAndEnableFeature(extensions_features::kApiMimeHandler);

  for (auto channel :
       {version_info::Channel::UNKNOWN, version_info::Channel::CANARY,
        version_info::Channel::DEV, version_info::Channel::BETA,
        version_info::Channel::STABLE}) {
    SCOPED_TRACE(version_info::GetChannelString(channel));
    ScopedCurrentChannel scoped_channel(channel);
    auto ext = CreateMimeHandlerExtension(
        std::string(version_info::GetChannelString(channel)), kPdfMimeType,
        kViewerUrl);
    LoadExtension(ext.get());
    EXPECT_EQ(ext->id(), registry()->GetHandlerForMimeType(kPdfMimeType));
    UnloadExtension(ext.get());
  }
}

TEST_F(MimeHandlerRegistryTest, EnabledByDefaultUntilDisabled) {
  auto ext =
      CreateMimeHandlerExtension("PDF Handler", kPdfMimeType, kViewerUrl);
  LoadExtension(ext.get());

  // Default: enabled, present in every accessor.
  EXPECT_TRUE(registry()->IsEnabledForMimeType(ext->id(), kPdfMimeType));
  EXPECT_EQ(ext->id(), registry()->GetHandlerForMimeType(kPdfMimeType));
  EXPECT_EQ(registry()->GetHandlersForMimeType(kPdfMimeType).size(), 1u);
  EXPECT_EQ(registry()->GetHandlersByMimeType().count(kPdfMimeType), 1u);

  // After disabling: gone from every accessor.
  registry()->SetEnabledForMimeType(ext->id(), kPdfMimeType, false);
  EXPECT_FALSE(registry()->IsEnabledForMimeType(ext->id(), kPdfMimeType));
  EXPECT_TRUE(registry()->GetHandlerForMimeType(kPdfMimeType).empty());
  EXPECT_TRUE(registry()->GetHandlersForMimeType(kPdfMimeType).empty());
  EXPECT_EQ(registry()->GetHandlersByMimeType().count(kPdfMimeType), 0u);

  // Re-enable: comes back.
  registry()->SetEnabledForMimeType(ext->id(), kPdfMimeType, true);
  EXPECT_EQ(ext->id(), registry()->GetHandlerForMimeType(kPdfMimeType));
}

TEST_F(MimeHandlerRegistryTest, DisableRollsBackToPreviouslyInstalledHandler) {
  // Precedence for public handlers is "newest GetFirstInstallTime
  // wins". Disabling the newer handler must roll the active handler
  // back to the previously-installed one; re-enabling must restore
  // the precedence order.
  auto ext_old = CreateMimeHandlerExtension("Older", kPdfMimeType, "a.html");
  auto ext_new = CreateMimeHandlerExtension("Newer", kPdfMimeType, "b.html");
  const base::Time t0 = base::Time::Now();
  SetFirstInstallTime(ext_old.get(), t0);
  SetFirstInstallTime(ext_new.get(), t0 + base::Hours(1));
  LoadExtension(ext_old.get());
  LoadExtension(ext_new.get());
  ASSERT_EQ(ext_new->id(), registry()->GetHandlerForMimeType(kPdfMimeType));

  registry()->SetEnabledForMimeType(ext_new->id(), kPdfMimeType, false);
  EXPECT_EQ(ext_old->id(), registry()->GetHandlerForMimeType(kPdfMimeType));

  // Re-enable: ext_new wins again (newer install time beats ext_old).
  registry()->SetEnabledForMimeType(ext_new->id(), kPdfMimeType, true);
  EXPECT_EQ(ext_new->id(), registry()->GetHandlerForMimeType(kPdfMimeType));
}

TEST_F(MimeHandlerRegistryTest, SetEnabledForUnclaimedMimeTypeDoesNotRegister) {
  constexpr char kPngMimeType[] = "image/png";
  // Anchor on the actual invariant: `kPngMimeType` is not in the public
  // allowed MIME types, so no non-allowlisted extension's manifest can
  // claim it (the parser drops unsupported entries with an install
  // warning). If this list ever grows to include `image/png`, pick a
  // different MIME type for this test.
  ASSERT_TRUE(
      std::ranges::find(MimeTypesHandler::GetPublicAllowedMIMETypeList(),
                        kPngMimeType) ==
      MimeTypesHandler::GetPublicAllowedMIMETypeList().end());

  auto ext =
      CreateMimeHandlerExtension("PDF Handler", kPdfMimeType, kViewerUrl);
  LoadExtension(ext.get());

  // Sanity: extension is the active handler for the MIME type it claims,
  // and `image/png` has no registered handler.
  ASSERT_EQ(ext->id(), registry()->GetHandlerForMimeType(kPdfMimeType));
  ASSERT_TRUE(registry()->GetHandlersForMimeType(kPngMimeType).empty());

  // Call for a MIME type the manifest does not claim. The pref write
  // happens unconditionally, but the in-memory registry must not gain
  // an entry for `image/png`.
  registry()->SetEnabledForMimeType(ext->id(), kPngMimeType, true);

  EXPECT_TRUE(registry()->IsEnabledForMimeType(ext->id(), kPngMimeType));
  EXPECT_TRUE(registry()->GetHandlersForMimeType(kPngMimeType).empty());
  EXPECT_EQ(registry()->GetHandlersByMimeType().count(kPngMimeType), 0u);
}

}  // namespace
}  // namespace extensions
