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

#include "components/webapps/isolated_web_apps/identity/iwa_identity_validator.h"

#include <variant>

#include "base/base64.h"
#include "base/containers/map_util.h"
#include "base/metrics/histogram_functions.h"
#include "base/no_destructor.h"
#include "base/strings/stringprintf.h"
#include "components/web_package/signed_web_bundles/ecdsa_p256_public_key.h"
#include "components/web_package/signed_web_bundles/ed25519_public_key.h"
#include "components/webapps/isolated_web_apps/client.h"
#include "components/webapps/isolated_web_apps/public/iwa_runtime_data_provider.h"

namespace web_app {

namespace {

using KeyRotationInfo = IwaRuntimeDataProvider::KeyRotationInfo;

constexpr char kKeyRotationBundleIdentityCheckHistogram[] =
    "WebApp.Isolated.KeyDistributionComponent.KeyRotationBundleIdentityCheck";

// These values are persisted to logs. Entries should not be renumbered and
// numeric values should never be reused.
// LINT.IfChange(KeyRotationBundleIdentityCheckResult)
enum class KeyRotationBundleIdentityCheckResult {
  kTrustedRotatedKey = 0,
  kTrustedPreviousKey = 1,
  kFailed = 2,
  kMaxValue = kFailed,
};
// LINT.ThenChange(//tools/metrics/histograms/metadata/webapps/enums.xml:KeyRotationBundleIdentityCheckResult)

bool MatchesRotatedKey(const web_package::PublicKey& public_key,
                       const KeyRotationInfo& kr_info) {
  return std::visit(
      [&](const auto& public_key) {
        return std::ranges::equal(public_key.bytes(), kr_info.public_key);
      },
      public_key);
}

bool MatchesPreviousKey(const web_package::PublicKey& public_key,
                        const KeyRotationInfo& kr_info) {
  return kr_info.previous_key &&
         std::visit(
             [&](const auto& public_key) {
               return std::ranges::equal(public_key.bytes(),
                                         *kr_info.previous_key);
             },
             public_key);
}

base::expected<void, std::string>
ValidateWebBundleIdentityAgainstKeyRotationInfo(
    const std::string& web_bundle_id,
    const std::vector<web_package::PublicKey>& public_keys,
    const KeyRotationInfo& kr_info,
    bool allow_soft_key_rotation) {
  KeyRotationBundleIdentityCheckResult result =
      KeyRotationBundleIdentityCheckResult::kFailed;

  for (const auto& public_key : public_keys) {
    if (MatchesRotatedKey(public_key, kr_info)) {
      result = KeyRotationBundleIdentityCheckResult::kTrustedRotatedKey;
      break;
    }
    // When key is rotated, we can still accept the old key for a grace period.
    // See go/iwa-soft-key-rotation for more details.
    if (allow_soft_key_rotation && MatchesPreviousKey(public_key, kr_info)) {
      result = KeyRotationBundleIdentityCheckResult::kTrustedPreviousKey;
      // Keep looking in case another public key matches the rotated key.
    }
  }

  base::UmaHistogramEnumeration(kKeyRotationBundleIdentityCheckHistogram,
                                result);

  if (result == KeyRotationBundleIdentityCheckResult::kFailed) {
    return base::unexpected(
        base::StringPrintf("Rotated key for Web Bundle ID <%s> doesn't match "
                           "any public key in the signature list.",
                           web_bundle_id.c_str()));
  }
  return base::ok();
}

}  // namespace

void IwaIdentityValidator::CreateSingleton() {
  static base::NoDestructor<IwaIdentityValidator> instance;
  instance.get();
}

base::expected<void, std::string>
IwaIdentityValidator::ValidateWebBundleIdentity(
    const std::string& web_bundle_id,
    const std::vector<web_package::PublicKey>& public_keys) const {
  if (const auto* provider =
          IwaClient::GetInstance()->GetRuntimeDataProvider()) {
    if (const auto* kr_info = provider->GetKeyRotationInfo(web_bundle_id)) {
      return ValidateWebBundleIdentityAgainstKeyRotationInfo(
          web_bundle_id, public_keys, *kr_info,
          /*allow_soft_key_rotation=*/true);
    }
  }

  return IdentityValidator::ValidateWebBundleIdentity(web_bundle_id,
                                                      public_keys);
}

// static
base::expected<void, std::string>
IwaIdentityValidator::ValidateWebBundleIdentity(
    const std::string& web_bundle_id,
    const std::vector<web_package::PublicKey>& public_keys,
    bool allow_soft_key_rotation) {
  if (const auto* provider =
          IwaClient::GetInstance()->GetRuntimeDataProvider()) {
    if (const auto* kr_info = provider->GetKeyRotationInfo(web_bundle_id)) {
      return ValidateWebBundleIdentityAgainstKeyRotationInfo(
          web_bundle_id, public_keys, *kr_info, allow_soft_key_rotation);
    }
  }

  return web_package::IdentityValidator::GetInstance()
      ->ValidateWebBundleIdentity(web_bundle_id, public_keys);
}

}  // namespace web_app