/* 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 http://mozilla.org/MPL/2.0/. */

#include "gtest/gtest.h"
#include "mozilla/BasePrincipal.h"
#include "mozilla/ExpandedPrincipal.h"
#include "mozilla/ExtensionPolicyService.h"
#include "mozilla/GenericFactory.h"
#include "mozilla/NullPrincipal.h"
#include "mozilla/StaticPrefs_browser.h"
#include "mozilla/SystemPrincipal.h"
#include "mozilla/dom/ProcessIsolation.h"
#include "mozilla/dom/WorkerPrivate.h"
#include "mozilla/gtest/MozAssertions.h"
#include "mozilla/gtest/MozHelpers.h"
#include "nsComponentManager.h"
#include "nsIEnterprisePolicies.h"

using namespace mozilla;
using namespace mozilla::dom;

static nsCOMPtr<nsIPrincipal> MakeTestPrincipal(const char* aURI) {
  nsCOMPtr<nsIURI> uri;
  MOZ_ALWAYS_SUCCEEDS(NS_NewURI(getter_AddRefs(uri), aURI));
  return BasePrincipal::CreateContentPrincipal(uri, {});
}

namespace {

static bool gJitDisabled = false;

struct RemoteTypes {
  RemoteType mIsolated;
  RemoteType mUnisolated;
};

struct WorkerExpectation {
  nsCOMPtr<nsIPrincipal> mPrincipal;
  WorkerKind mWorkerKind = WorkerKindShared;
  bool mJitDisabled = false;
  Result<RemoteTypes, nsresult> mExpected = Err(NS_ERROR_FAILURE);
  RemoteType mCurrentRemoteType = RemoteType(RemoteType::Kind::WebContent);

  void Check(bool aUseRemoteSubframes) {
    nsAutoCString origin;
    ASSERT_NS_SUCCEEDED(mPrincipal->GetOrigin(origin));

    nsPrintfCString describe(
        "origin: %s, workerKind: %s, currentRemoteType: %s, "
        "useRemoteSubframes: %d",
        origin.get(), mWorkerKind == WorkerKindShared ? "shared" : "service",
        mCurrentRemoteType.Stringify().get(), aUseRemoteSubframes);

    gJitDisabled = mJitDisabled;
    auto result = IsolationOptionsForWorker(
        mPrincipal, mWorkerKind, mCurrentRemoteType, aUseRemoteSubframes);
    ASSERT_EQ(result.isOk(), mExpected.isOk())
        << "Unexpected status (expected " << (mExpected.isOk() ? "ok" : "err")
        << ") for " << describe;
    if (mExpected.isOk()) {
      const RemoteType& expected = aUseRemoteSubframes
                                       ? mExpected.inspect().mIsolated
                                       : mExpected.inspect().mUnisolated;
      EXPECT_TRUE(expected.IsKnown())
          << "invalid remote type expectation in test";
      ASSERT_EQ(result.inspect().mRemoteType, expected)
          << "Unexpected remote type (expected " << expected.Stringify()
          << ") for " << describe;
    }
  }
};

#define MOCK_ENTERPRISE_POLICIES_CID \
  {0xaabbc001, 0xdd00, 0x1234, {0xab, 0xcd, 0xef, 0x01, 0x23, 0x45, 0x67, 0x89}}
NS_DEFINE_NAMED_CID(MOCK_ENTERPRISE_POLICIES_CID);

class MockEnterprisePoliciesService final : public nsIEnterprisePolicies {
  ~MockEnterprisePoliciesService() = default;

 public:
  NS_DECL_ISUPPORTS

  NS_IMETHOD GetStatus(int16_t* aStatus) override {
    return NS_ERROR_NOT_IMPLEMENTED;
  }
  NS_IMETHOD GetIsEnterprise(bool* aIsEnterprise) override {
    return NS_ERROR_NOT_IMPLEMENTED;
  }
  NS_IMETHOD IsAllowed(const nsACString&, bool* aRetVal) override {
    return NS_ERROR_NOT_IMPLEMENTED;
  }
  NS_IMETHOD IsAllowedForURI(const nsACString&, nsIURI*,
                             bool* aRetVal) override {
    *aRetVal = !gJitDisabled;
    return NS_OK;
  }
  NS_IMETHOD GetActivePolicies(JS::MutableHandle<JS::Value>) override {
    return NS_ERROR_NOT_IMPLEMENTED;
  }
  NS_IMETHOD GetSupportMenu(JS::MutableHandle<JS::Value>) override {
    return NS_ERROR_NOT_IMPLEMENTED;
  }
  NS_IMETHOD GetExtensionPolicy(const nsACString&,
                                JS::MutableHandle<JS::Value>) override {
    return NS_ERROR_NOT_IMPLEMENTED;
  }
  NS_IMETHOD GetExtensionSettings(const nsACString&,
                                  JS::MutableHandle<JS::Value>) override {
    return NS_ERROR_NOT_IMPLEMENTED;
  }
  NS_IMETHOD MayInstallAddon(JS::Handle<JS::Value>, bool*) override {
    return NS_ERROR_NOT_IMPLEMENTED;
  }
  NS_IMETHOD IsAddonRequiredByPolicy(const nsACString&, bool*) override {
    return NS_ERROR_NOT_IMPLEMENTED;
  }
  NS_IMETHOD AllowedInstallSource(nsIURI*, bool*) override {
    return NS_ERROR_NOT_IMPLEMENTED;
  }
  NS_IMETHOD IsExemptExecutableExtension(const nsACString&, const nsACString&,
                                         bool*) override {
    return NS_ERROR_NOT_IMPLEMENTED;
  }
  NS_IMETHOD GetContainerForURI(nsIURI*, uint32_t*) override {
    return NS_ERROR_NOT_IMPLEMENTED;
  }
};

NS_IMPL_ISUPPORTS(MockEnterprisePoliciesService, nsIEnterprisePolicies)

static nsresult ConstructMockEnterprisePolicies(const nsIID& aIID,
                                                void** aResult) {
  RefPtr<MockEnterprisePoliciesService> service =
      new MockEnterprisePoliciesService();
  return service->QueryInterface(aIID, aResult);
}

StaticRefPtr<nsIFactory> gMockPolicyFactory;

static void RegisterMockPolicyService() {
  MOZ_ASSERT(!gMockPolicyFactory);
  gMockPolicyFactory =
      new mozilla::GenericFactory(ConstructMockEnterprisePolicies);
  MOZ_ALWAYS_SUCCEEDS(
      nsComponentManagerImpl::gComponentManager->RegisterFactory(
          kMOCK_ENTERPRISE_POLICIES_CID, "MockEnterprisePolicies",
          "@mozilla.org/enterprisepolicies;1", gMockPolicyFactory));
}

static void UnregisterMockPolicyService() {
  MOZ_ASSERT(gMockPolicyFactory);
  MOZ_ALWAYS_SUCCEEDS(
      nsComponentManagerImpl::gComponentManager->UnregisterFactory(
          kMOCK_ENTERPRISE_POLICIES_CID, gMockPolicyFactory));
  gMockPolicyFactory = nullptr;
}

}  // namespace

// When file URI process separation is disabled (as is the default on
// Android), a file: shared worker is allowed to load in any remote type,
// rather than being rejected when it isn't already in a file: process.
static Result<RemoteTypes, nsresult> FileWorkerOutsideFileProcessExpected(
    const RemoteType& aFileRemoteType) {
  if (StaticPrefs::browser_tabs_remote_separateFileUriProcess()) {
    return Err(NS_ERROR_UNEXPECTED);
  }
  return RemoteTypes{aFileRemoteType, aFileRemoteType};
}

TEST(ProcessIsolationTest, WorkerOptions)
{
  // Forcibly enable the privileged mozilla content process for the duration of
  // the test.
  MOZ_ALWAYS_SUCCEEDS(Preferences::SetCString(
      "browser.tabs.remote.separatedMozillaDomains", "addons.mozilla.org"));
  MOZ_ALWAYS_SUCCEEDS(Preferences::SetBool(
      "browser.tabs.remote.separatePrivilegedMozillaWebContentProcess", true));
  auto cleanup = MakeScopeExit([&] {
    MOZ_ALWAYS_SUCCEEDS(
        Preferences::ClearUser("browser.tabs.remote.separatedMozillaDomains"));
    MOZ_ALWAYS_SUCCEEDS(Preferences::ClearUser(
        "browser.tabs.remote.separatePrivilegedMozillaWebContentProcess"));
  });

  OriginAttributes containerOA;
  containerOA.mUserContextId = 1;

  nsCOMPtr<nsIPrincipal> systemPrincipal = SystemPrincipal::Get();
  nsCOMPtr<nsIPrincipal> nullPrincipal =
      NullPrincipal::CreateWithoutOriginAttributes();
  nsCOMPtr<nsIPrincipal> nullContainerPrincipal =
      NullPrincipal::Create(containerOA);
  nsCOMPtr<nsIPrincipal> secureComPrincipal =
      MakeTestPrincipal("https://example.com");
  nsCOMPtr<nsIPrincipal> secureOrgPrincipal =
      MakeTestPrincipal("https://example.org");
  nsCOMPtr<nsIPrincipal> insecureOrgPrincipal =
      MakeTestPrincipal("http://example.org");
  nsCOMPtr<nsIPrincipal> filePrincipal =
      MakeTestPrincipal("file:///path/to/dir");
  nsCOMPtr<nsIPrincipal> extensionPrincipal =
      MakeTestPrincipal("moz-extension://fake-uuid");
  nsCOMPtr<nsIPrincipal> privilegedMozillaPrincipal =
      MakeTestPrincipal("https://addons.mozilla.org");
  nsCOMPtr<nsIPrincipal> expandedPrincipal = ExpandedPrincipal::Create(
      nsTArray{secureComPrincipal, extensionPrincipal}, {});
  nsCOMPtr<nsIPrincipal> nullSecureComPrecursorPrincipal =
      NullPrincipal::CreateWithInheritedAttributes(secureComPrincipal);

  RemoteType extensionRemoteType =
      ExtensionPolicyService::GetSingleton().UseRemoteExtensions()
          ? RemoteType(RemoteType::Kind::Extension)
          : RemoteType::NotRemote();
  RemoteType fileRemoteType =
      StaticPrefs::browser_tabs_remote_separateFileUriProcess()
          ? RemoteType(RemoteType::Kind::File)
          : RemoteType(RemoteType::Kind::WebContent);

  WorkerExpectation expectations[] = {
      // Neither service not shared workers can have expanded principals
      {.mPrincipal = expandedPrincipal,
       .mWorkerKind = WorkerKindService,
       .mExpected = Err(NS_ERROR_UNEXPECTED)},
      {.mPrincipal = expandedPrincipal,
       .mWorkerKind = WorkerKindShared,
       .mExpected = Err(NS_ERROR_UNEXPECTED)},

      // Service workers cannot have system or null principals
      {.mPrincipal = systemPrincipal,
       .mWorkerKind = WorkerKindService,
       .mExpected = Err(NS_ERROR_UNEXPECTED)},
      {.mPrincipal = nullPrincipal,
       .mWorkerKind = WorkerKindService,
       .mExpected = Err(NS_ERROR_UNEXPECTED)},
      {.mPrincipal = nullContainerPrincipal,
       .mWorkerKind = WorkerKindService,
       .mExpected = Err(NS_ERROR_UNEXPECTED)},
      {.mPrincipal = nullSecureComPrecursorPrincipal,
       .mWorkerKind = WorkerKindService,
       .mExpected = Err(NS_ERROR_UNEXPECTED)},

      // Service workers with various content principals
      {.mPrincipal = secureComPrincipal,
       .mWorkerKind = WorkerKindService,
       .mExpected = RemoteTypes{RemoteType::Parse(
                                    "webServiceWorker=https://example.com"_ns),
                                RemoteType(RemoteType::Kind::WebContent)}},
      {.mPrincipal = secureOrgPrincipal,
       .mWorkerKind = WorkerKindService,
       .mExpected = RemoteTypes{RemoteType::Parse(
                                    "webServiceWorker=https://example.org"_ns),
                                RemoteType(RemoteType::Kind::WebContent)}},
      {.mPrincipal = extensionPrincipal,
       .mWorkerKind = WorkerKindService,
       .mExpected = RemoteTypes{extensionRemoteType, extensionRemoteType},
       .mCurrentRemoteType = RemoteType(RemoteType::Kind::Extension)},
      {.mPrincipal = privilegedMozillaPrincipal,
       .mWorkerKind = WorkerKindService,
       .mExpected = Err(NS_ERROR_UNEXPECTED)},
      {.mPrincipal = privilegedMozillaPrincipal,
       .mWorkerKind = WorkerKindService,
       .mExpected =
           RemoteTypes{RemoteType(RemoteType::Kind::PrivilegedMozilla),
                       RemoteType(RemoteType::Kind::PrivilegedMozilla)},
       .mCurrentRemoteType = RemoteType(RemoteType::Kind::PrivilegedMozilla)},

      // Shared Worker loaded from within a webCOOP+COEP remote type process,
      // should load elsewhere.
      {.mPrincipal = secureComPrincipal,
       .mWorkerKind = WorkerKindShared,
       .mExpected =
           RemoteTypes{RemoteType::Parse("webIsolated=https://example.com"_ns),
                       RemoteType(RemoteType::Kind::WebContent)},
       .mCurrentRemoteType =
           RemoteType::Parse("webCOOP+COEP=https://example.com"_ns)},

      // Even precursorless null principal should load elsewhere.
      {.mPrincipal = nullPrincipal,
       .mWorkerKind = WorkerKindShared,
       .mExpected = RemoteTypes{RemoteType(RemoteType::Kind::WebContent),
                                RemoteType(RemoteType::Kind::WebContent)},
       .mCurrentRemoteType =
           RemoteType::Parse("webCOOP+COEP=https://example.com"_ns)},
      {.mPrincipal = nullContainerPrincipal,
       .mWorkerKind = WorkerKindShared,
       .mExpected = RemoteTypes{RemoteType::Parse("web=^userContextId=1"_ns),
                                RemoteType::Parse("web=^userContextId=1"_ns)},
       .mCurrentRemoteType =
           RemoteType::Parse("webCOOP+COEP=https://example.com"_ns)},

      // System principal shared workers can only load in the parent process.
      {.mPrincipal = systemPrincipal,
       .mWorkerKind = WorkerKindShared,
       .mExpected =
           RemoteTypes{RemoteType::NotRemote(), RemoteType::NotRemote()},
       .mCurrentRemoteType = RemoteType::NotRemote()},
      {.mPrincipal = systemPrincipal,
       .mWorkerKind = WorkerKindShared,
       .mExpected = Err(NS_ERROR_UNEXPECTED),
       .mCurrentRemoteType = RemoteType(RemoteType::Kind::PrivilegedAbout)},
      {.mPrincipal = systemPrincipal,
       .mWorkerKind = WorkerKindShared,
       .mExpected = Err(NS_ERROR_UNEXPECTED)},
      {.mPrincipal = systemPrincipal,
       .mWorkerKind = WorkerKindShared,
       .mExpected = Err(NS_ERROR_UNEXPECTED)},

      // Content principals should load in the appropriate remote types.
      {.mPrincipal = secureComPrincipal,
       .mWorkerKind = WorkerKindShared,
       .mExpected =
           RemoteTypes{RemoteType::Parse("webIsolated=https://example.com"_ns),
                       RemoteType(RemoteType::Kind::WebContent)}},
      {.mPrincipal = secureOrgPrincipal,
       .mWorkerKind = WorkerKindShared,
       .mExpected =
           RemoteTypes{RemoteType::Parse("webIsolated=https://example.org"_ns),
                       RemoteType(RemoteType::Kind::WebContent)}},
      {.mPrincipal = insecureOrgPrincipal,
       .mWorkerKind = WorkerKindShared,
       .mExpected =
           RemoteTypes{RemoteType::Parse("webIsolated=http://example.org"_ns),
                       RemoteType(RemoteType::Kind::WebContent)}},
      {.mPrincipal = filePrincipal,
       .mWorkerKind = WorkerKindShared,
       .mExpected = FileWorkerOutsideFileProcessExpected(fileRemoteType)},
      {.mPrincipal = filePrincipal,
       .mWorkerKind = WorkerKindShared,
       .mExpected = RemoteTypes{fileRemoteType, fileRemoteType},
       .mCurrentRemoteType = RemoteType(RemoteType::Kind::File)},
      {.mPrincipal = extensionPrincipal,
       .mWorkerKind = WorkerKindShared,
       .mExpected = RemoteTypes{extensionRemoteType, extensionRemoteType},
       .mCurrentRemoteType = RemoteType(RemoteType::Kind::Extension)},
      {.mPrincipal = privilegedMozillaPrincipal,
       .mWorkerKind = WorkerKindShared,
       .mExpected = Err(NS_ERROR_UNEXPECTED)},
      {.mPrincipal = privilegedMozillaPrincipal,
       .mWorkerKind = WorkerKindShared,
       .mExpected =
           RemoteTypes{RemoteType(RemoteType::Kind::PrivilegedMozilla),
                       RemoteType(RemoteType::Kind::PrivilegedMozilla)},
       .mCurrentRemoteType = RemoteType(RemoteType::Kind::PrivilegedMozilla)},
      {.mPrincipal = nullSecureComPrecursorPrincipal,
       .mWorkerKind = WorkerKindShared,
       .mExpected =
           RemoteTypes{RemoteType::Parse("webIsolated=https://example.com"_ns),
                       RemoteType(RemoteType::Kind::WebContent)}},

      // When the policy service calls for the JIT to be disabled the remote
      // type should reflect that.
      {.mPrincipal = secureComPrincipal,
       .mWorkerKind = WorkerKindShared,
       .mJitDisabled = true,
       .mExpected =
           RemoteTypes{RemoteType::Parse(
                           "webIsolated=https://example.com^disableJit=1"_ns),
                       RemoteType::Parse("web=^disableJit=1"_ns)}},
      {.mPrincipal = secureComPrincipal,
       .mWorkerKind = WorkerKindService,
       .mJitDisabled = true,
       .mExpected =
           RemoteTypes{
               RemoteType::Parse(
                   "webServiceWorker=https://example.com^disableJit=1"_ns),
               RemoteType::Parse("web=^disableJit=1"_ns)}},
  };

  RegisterMockPolicyService();
  for (auto& expectation : expectations) {
    expectation.Check(true);
    expectation.Check(false);
  }
  UnregisterMockPolicyService();
}

// The file:// URI allowlist is populated from the `capability.policy.*` prefs,
// which in practice are written by the LocalFileLinks enterprise policy so that
// an intranet origin may link to files on a network share. Historically that
// allowlist only made documents load in the file: content process; workers were
// explicitly excluded (see the `!aIsWorker` guard that used to live in
// E10SUtils). Since bug 1850589 dropped that exclusion, a service worker on an
// allowlisted origin resolves to IsolationBehavior::File and is then rejected
// outright by ValidateBehaviorForWorker, because ServiceWorkerPrivate passes
// the shared "web" remote type rather than the file remote type.
TEST(ProcessIsolationTest, FileURIAllowlistedWorkerOptions)
{
  MOZ_ALWAYS_SUCCEEDS(Preferences::SetCString("capability.policy.policynames",
                                              "localfilelinks_policy"));
  MOZ_ALWAYS_SUCCEEDS(Preferences::SetCString(
      "capability.policy.localfilelinks_policy.checkloaduri.enabled",
      "allAccess"));
  MOZ_ALWAYS_SUCCEEDS(Preferences::SetCString(
      "capability.policy.localfilelinks_policy.sites", "https://example.com"));
  auto cleanup = MakeScopeExit([&] {
    MOZ_ALWAYS_SUCCEEDS(
        Preferences::ClearUser("capability.policy.policynames"));
    MOZ_ALWAYS_SUCCEEDS(Preferences::ClearUser(
        "capability.policy.localfilelinks_policy.checkloaduri.enabled"));
    MOZ_ALWAYS_SUCCEEDS(Preferences::ClearUser(
        "capability.policy.localfilelinks_policy.sites"));
  });

  nsCOMPtr<nsIPrincipal> allowlistedPrincipal =
      MakeTestPrincipal("https://example.com");
  nsCOMPtr<nsIPrincipal> filePrincipal =
      MakeTestPrincipal("file:///path/to/dir");

  RemoteType fileRemoteType(
      StaticPrefs::browser_tabs_remote_separateFileUriProcess()
          ? RemoteType::Kind::File
          : RemoteType::Kind::WebContent);

  WorkerExpectation expectations[] = {
      // Being in the file:// URI allowlist must not change worker process
      // selection: these are the same expectations as for a principal which
      // isn't allowlisted at all. The current remote type mirrors what
      // ServiceWorkerPrivate::Initialize() passes for a service worker.
      {.mPrincipal = allowlistedPrincipal,
       .mWorkerKind = WorkerKindService,
       .mExpected = RemoteTypes{RemoteType::Parse(
                                    "webServiceWorker=https://example.com"_ns),
                                RemoteType(RemoteType::Kind::WebContent)},
       .mCurrentRemoteType = RemoteType(RemoteType::Kind::WebContent)},
      {.mPrincipal = allowlistedPrincipal,
       .mWorkerKind = WorkerKindShared,
       .mExpected =
           RemoteTypes{RemoteType::Parse("webIsolated=https://example.com"_ns),
                       RemoteType(RemoteType::Kind::WebContent)},
       .mCurrentRemoteType = RemoteType(RemoteType::Kind::WebContent)},

      // An actual file: principal is still confined to the file process,
      // regardless of the allowlist.
      {.mPrincipal = filePrincipal,
       .mWorkerKind = WorkerKindShared,
       .mExpected = FileWorkerOutsideFileProcessExpected(fileRemoteType)},
      {.mPrincipal = filePrincipal,
       .mWorkerKind = WorkerKindShared,
       .mExpected = RemoteTypes{fileRemoteType, fileRemoteType},
       .mCurrentRemoteType = RemoteType(RemoteType::Kind::File)},
  };

  for (auto& expectation : expectations) {
    expectation.Check(true);
    expectation.Check(false);
  }
}