//
// Copyright 2024 gRPC authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//

#include "src/core/ext/filters/gcp_authentication/gcp_authentication_filter.h"

#include <memory>
#include <string>
#include <utility>

#include "src/core/call/security_context.h"
#include "src/core/config/core_configuration.h"
#include "src/core/credentials/call/gcp_service_account_identity/gcp_service_account_identity_credentials.h"
#include "src/core/lib/channel/channel_stack.h"
#include "src/core/lib/promise/context.h"
#include "src/core/lib/resource_quota/arena.h"
#include "src/core/lib/transport/transport.h"
#include "src/core/resolver/xds/xds_resolver_attributes.h"
#include "src/core/service_config/service_config_call_data.h"
#include "src/core/util/grpc_check.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/str_format.h"

namespace grpc_core {

//
// GcpAuthenticationFilter::Config
//

bool GcpAuthenticationFilter::Config::Equals(const FilterConfig& other) const {
  const auto& o = DownCast<const Config&>(other);
  return instance_name == o.instance_name && cache_size == o.cache_size &&
         cache == o.cache;
}

std::string GcpAuthenticationFilter::Config::ToString() const {
  return absl::StrFormat("{instance_name=\"%s\", cache_size=%u, cache=%p}",
                         instance_name, cache_size, cache.get());
}

//
// GcpAuthenticationFilter::Call
//

absl::Status GcpAuthenticationFilter::Call::OnClientInitialMetadata(
    ClientMetadata& /*md*/, GcpAuthenticationFilter* filter) {
  // Get the cluster name chosen for this RPC.
  auto* service_config_call_data = GetContext<ServiceConfigCallData>();
  auto cluster_attribute =
      service_config_call_data->GetCallAttribute<XdsClusterAttribute>();
  if (cluster_attribute == nullptr) {
    // Can't happen, but be defensive.
    return absl::InternalError(
        "GCP authentication filter: call has no xDS cluster attribute");
  }
  absl::string_view cluster_name = cluster_attribute->cluster();
  if (!absl::ConsumePrefix(&cluster_name, "cluster:")) {
    return absl::OkStatus();  // Cluster specifier plugin.
  }
  // Look up the CDS resource for the cluster.
  auto it = filter->xds_config_->clusters.find(cluster_name);
  if (it == filter->xds_config_->clusters.end()) {
    // Can't happen, but be defensive.
    return absl::InternalError(
        absl::StrCat("GCP authentication filter: xDS cluster ", cluster_name,
                     " not found in XdsConfig"));
  }
  if (!it->second.ok()) {
    // Cluster resource had an error, so fail the call.
    // Note: For wait_for_ready calls, this does the wrong thing by
    // failing the call instead of queuing it, but there's no easy
    // way to queue the call here until we get a valid CDS resource,
    // because once that happens, a new instance of this filter will be
    // swapped in for subsequent calls, but *this* call is already tied
    // to this filter instance, which will never see the update.
    return absl::UnavailableError(
        absl::StrCat("GCP authentication filter: CDS resource unavailable for ",
                     cluster_name));
  }
  if (it->second->cluster == nullptr) {
    // Can't happen, but be defensive.
    return absl::InternalError(absl::StrCat(
        "GCP authentication filter: CDS resource not present for cluster ",
        cluster_name));
  }
  auto& metadata_map = it->second->cluster->metadata;
  const XdsMetadataValue* metadata_value =
      metadata_map.Find(filter->filter_config_->instance_name);
  // If no audience in the cluster, then no need to add call creds.
  if (metadata_value == nullptr) return absl::OkStatus();
  // If the entry is present but the wrong type, fail the RPC.
  if (metadata_value->type() != XdsGcpAuthnAudienceMetadataValue::Type()) {
    return absl::UnavailableError(absl::StrCat(
        "GCP authentication filter: audience metadata in wrong format for "
        "cluster ",
        cluster_name));
  }
  // Get the call creds instance.
  auto creds = filter->filter_config_->cache->Get(
      DownCast<const XdsGcpAuthnAudienceMetadataValue*>(metadata_value)->url());
  // Add the call creds instance to the call.
  auto* arena = GetContext<Arena>();
  auto* security_ctx = DownCast<grpc_client_security_context*>(
      arena->GetContext<SecurityContext>());
  if (security_ctx == nullptr) {
    security_ctx = arena->New<grpc_client_security_context>(std::move(creds));
    arena->SetContext<SecurityContext>(security_ctx);
  } else {
    security_ctx->creds = std::move(creds);
  }
  return absl::OkStatus();
}

//
// GcpAuthenticationFilter::CallCredentialsCache
//

UniqueTypeName GcpAuthenticationFilter::CallCredentialsCache::Type() {
  static UniqueTypeName::Factory factory("gcp_auth_call_creds_cache");
  return factory.Create();
}

void GcpAuthenticationFilter::CallCredentialsCache::SetMaxSize(
    size_t max_size) {
  MutexLock lock(&mu_);
  cache_.SetMaxSize(max_size);
}

RefCountedPtr<grpc_call_credentials>
GcpAuthenticationFilter::CallCredentialsCache::Get(
    const std::string& audience) {
  MutexLock lock(&mu_);
  return cache_.GetOrInsert(audience, [](const std::string& audience) {
    return MakeRefCounted<GcpServiceAccountIdentityCallCredentials>(audience);
  });
}

//
// GcpAuthenticationFilter
//

const grpc_channel_filter GcpAuthenticationFilter::kFilterVtable =
    MakePromiseBasedFilter<GcpAuthenticationFilter, FilterEndpoint::kClient,
                           0>();

absl::StatusOr<std::unique_ptr<GcpAuthenticationFilter>>
GcpAuthenticationFilter::Create(const ChannelArgs& args,
                                ChannelFilter::Args filter_args) {
  // Get filter config.
  if (filter_args.config() == nullptr) {
    return absl::InternalError("gcp_auth: filter config not set");
  }
  if (filter_args.config()->type() != Config::Type()) {
    return absl::InternalError(
        absl::StrCat("wrong config type passed to GCP authn filter: ",
                     filter_args.config()->type().name()));
  }
  auto config = filter_args.config().TakeAsSubclass<const Config>();
  // Make sure cache is present in config.
  if (config->cache == nullptr) {
    return absl::InternalError(
        "gcp_auth: cache object not found in filter config");
  }
  // Get XdsConfig so that we can look up CDS resources.
  auto xds_config = args.GetObjectRef<XdsConfig>();
  if (xds_config == nullptr) {
    return absl::InternalError(
        "gcp_auth: xds config not found in channel args");
  }
  // Instantiate filter.
  return std::unique_ptr<GcpAuthenticationFilter>(
      new GcpAuthenticationFilter(std::move(config), std::move(xds_config)));
}

GcpAuthenticationFilter::GcpAuthenticationFilter(
    RefCountedPtr<const Config> filter_config,
    RefCountedPtr<const XdsConfig> xds_config)
    : filter_config_(std::move(filter_config)),
      xds_config_(std::move(xds_config)) {}

}  // namespace grpc_core
