//
//
// Copyright 2025 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 <grpc/event_engine/event_engine.h>
#include <grpc/impl/channel_arg_names.h>
#include <grpcpp/channel.h>
#include <grpcpp/client_context.h>
#include <grpcpp/create_channel.h>
#include <grpcpp/security/credentials.h>
#include <grpcpp/security/server_credentials.h>
#include <grpcpp/security/tls_certificate_provider.h>
#include <grpcpp/security/tls_credentials_options.h>
#include <grpcpp/security/tls_private_key_signer.h>
#include <grpcpp/server.h>
#include <grpcpp/server_builder.h>
#include <grpcpp/support/channel_arguments.h>
#include <grpcpp/support/status.h>

#include <functional>
#include <memory>
#include <string>
#include <thread>
#include <utility>
#include <vector>

#include "src/proto/grpc/testing/echo_messages.pb.h"
#include "test/core/test_util/port.h"
#include "test/core/test_util/test_config.h"
#include "test/core/test_util/tls_utils.h"
#include "test/cpp/end2end/test_service_impl.h"
#include "gmock/gmock.h"
#include "gtest/gtest.h"
#include "absl/log/check.h"
#include "absl/status/statusor.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/string_view.h"
#include "absl/synchronization/notification.h"
#include "absl/time/time.h"

#if defined(OPENSSL_IS_BORINGSSL)

#include <openssl/digest.h>
#include <openssl/ec.h>
#include <openssl/evp.h>
#include <openssl/rsa.h>
#include <openssl/ssl.h>

namespace grpc {
namespace testing {
namespace {

constexpr absl::string_view kMessage = "Hello";
constexpr absl::string_view kCaPemPath = "src/core/tsi/test_creds/ca.pem";
constexpr absl::string_view kServerKeyPath =
    "src/core/tsi/test_creds/server1.key";
constexpr absl::string_view kServerCertPath =
    "src/core/tsi/test_creds/server1.pem";
constexpr absl::string_view kClientKeyPath =
    "src/core/tsi/test_creds/client.key";
constexpr absl::string_view kClientCertPath =
    "src/core/tsi/test_creds/client.pem";

bssl::UniquePtr<EVP_PKEY> LoadPrivateKeyFromString(
    absl::string_view private_pem) {
  bssl::UniquePtr<BIO> bio(
      BIO_new_mem_buf(private_pem.data(), private_pem.size()));
  return bssl::UniquePtr<EVP_PKEY>(
      PEM_read_bio_PrivateKey(bio.get(), nullptr, nullptr, nullptr));
}

uint16_t GetBoringSslAlgorithm(
    grpc::experimental::PrivateKeySigner::SignatureAlgorithm
        signature_algorithm) {
  switch (signature_algorithm) {
    case grpc::experimental::PrivateKeySigner::SignatureAlgorithm::
        kRsaPkcs1Sha256:
      return SSL_SIGN_RSA_PKCS1_SHA256;
    case grpc::experimental::PrivateKeySigner::SignatureAlgorithm::
        kRsaPkcs1Sha384:
      return SSL_SIGN_RSA_PKCS1_SHA384;
    case grpc::experimental::PrivateKeySigner::SignatureAlgorithm::
        kRsaPkcs1Sha512:
      return SSL_SIGN_RSA_PKCS1_SHA512;
    case grpc::experimental::PrivateKeySigner::SignatureAlgorithm::
        kEcdsaSecp256r1Sha256:
      return SSL_SIGN_ECDSA_SECP256R1_SHA256;
    case grpc::experimental::PrivateKeySigner::SignatureAlgorithm::
        kEcdsaSecp384r1Sha384:
      return SSL_SIGN_ECDSA_SECP384R1_SHA384;
    case grpc::experimental::PrivateKeySigner::SignatureAlgorithm::
        kEcdsaSecp521r1Sha512:
      return SSL_SIGN_ECDSA_SECP521R1_SHA512;
    case grpc::experimental::PrivateKeySigner::SignatureAlgorithm::
        kRsaPssRsaeSha256:
      return SSL_SIGN_RSA_PSS_RSAE_SHA256;
    case grpc::experimental::PrivateKeySigner::SignatureAlgorithm::
        kRsaPssRsaeSha384:
      return SSL_SIGN_RSA_PSS_RSAE_SHA384;
    case grpc::experimental::PrivateKeySigner::SignatureAlgorithm::
        kRsaPssRsaeSha512:
      return SSL_SIGN_RSA_PSS_RSAE_SHA512;
  }
  return -1;
}

class TlsPrivateKeyOffloadTest : public ::testing::Test {
 protected:
  void StartServer(std::shared_ptr<experimental::CertificateProviderInterface>
                       server_certificate_provider,
                   int handshake_timeout_ms = 0) {
    absl::Notification notification;
    server_thread_ = std::make_unique<std::thread>([this, &notification,
                                                    server_certificate_provider,
                                                    handshake_timeout_ms]() {
      RunServer(&notification, server_certificate_provider,
                handshake_timeout_ms);
    });
    notification.WaitForNotification();
  }

  void RunServer(absl::Notification* notification,
                 std::shared_ptr<experimental::CertificateProviderInterface>
                     server_certificate_provider,
                 int handshake_timeout_ms = 0) {
    grpc::experimental::TlsServerCredentialsOptions options(
        std::move(server_certificate_provider));
    options.watch_root_certs();
    options.set_root_cert_name("root");
    options.watch_identity_key_cert_pairs();
    options.set_identity_cert_name("identity");
    options.set_cert_request_type(
        GRPC_SSL_REQUEST_AND_REQUIRE_CLIENT_CERTIFICATE_AND_VERIFY);
    auto server_credentials = grpc::experimental::TlsServerCredentials(options);
    CHECK_NE(server_credentials.get(), nullptr);

    grpc::ServerBuilder builder;

    if (handshake_timeout_ms > 0) {
      builder.AddChannelArgument(GRPC_ARG_SERVER_HANDSHAKE_TIMEOUT_MS,
                                 handshake_timeout_ms);
    }

    builder.AddListeningPort(server_addr_, server_credentials);
    builder.RegisterService("foo.test.google.fr", &service_);
    server_ = builder.BuildAndStart();
    notification->Notify();
    server_->Wait();
  }

  void TearDown() override {
    if (server_ != nullptr) {
      server_->Shutdown();
    }
    if (server_thread_ != nullptr) {
      server_thread_->join();
    }
  }

  TestServiceImpl service_;
  std::unique_ptr<Server> server_ = nullptr;
  std::unique_ptr<std::thread> server_thread_;
  std::string server_addr_;
  std::shared_ptr<grpc::experimental::PrivateKeySigner> signer_;
};

void DoRpc(const std::string& server_addr,
           const experimental::TlsChannelCredentialsOptions& tls_options) {
  ChannelArguments channel_args;
  channel_args.SetSslTargetNameOverride("foo.test.google.fr");
  std::shared_ptr<Channel> channel = grpc::CreateCustomChannel(
      server_addr, grpc::experimental::TlsCredentials(tls_options),
      channel_args);

  auto stub = grpc::testing::EchoTestService::NewStub(channel);
  grpc::testing::EchoRequest request;
  grpc::testing::EchoResponse response;
  request.set_message(kMessage);
  ClientContext context;
  context.set_deadline(grpc_timeout_seconds_to_deadline(/*time_s=*/5));
  grpc::Status result = stub->Echo(&context, request, &response);
  EXPECT_TRUE(result.ok()) << result.error_message().c_str() << ", "
                           << result.error_details().c_str();
  EXPECT_EQ(response.message(), kMessage);
}

// Performs an RPC and expects it to fail.
// on_rpc_stalled is an optional callback that will be executed after a delay (1
// second) after the RPC is initiated. This is useful for cases where the RPC is
// expected to wait and an action is needed to unblock the test.
void DoRpcAndExpectFailure(
    const std::string& server_addr,
    const experimental::TlsChannelCredentialsOptions& tls_options,
    const std::function<void(ClientContext*)>& on_rpc_stalled = nullptr,
    absl::string_view expected_error_message = "",
    ChannelArguments channel_args = ChannelArguments()) {
  channel_args.SetSslTargetNameOverride("foo.test.google.fr");
  std::shared_ptr<Channel> channel = grpc::CreateCustomChannel(
      server_addr, grpc::experimental::TlsCredentials(tls_options),
      channel_args);

  auto stub = grpc::testing::EchoTestService::NewStub(channel);
  grpc::testing::EchoRequest request;
  grpc::testing::EchoResponse response;
  request.set_message(kMessage);
  ClientContext context;
  context.set_deadline(grpc_timeout_seconds_to_deadline(/*time_s=*/5));
  absl::Notification notification;
  bool wait_for_notification = false;
  if (on_rpc_stalled != nullptr) {
    wait_for_notification = true;
    grpc_core::ExecCtx exec_ctx;
    grpc_event_engine::experimental::GetDefaultEventEngine()->RunAfter(
        std::chrono::seconds(1),
        [on_rpc_stalled, &context, notification = &notification]() {
          on_rpc_stalled(&context);
          notification->Notify();
        });
  }
  grpc::Status result = stub->Echo(&context, request, &response);
  if (wait_for_notification) {
    notification.WaitForNotificationWithTimeout(absl::Seconds(10));
  }
  EXPECT_FALSE(result.ok());
  if (!expected_error_message.empty()) {
    EXPECT_THAT(result.error_message(),
                ::testing::HasSubstr(std::string(expected_error_message)));
  }
}

absl::StatusOr<std::string> SignWithBoringSSL(
    absl::string_view data_to_sign,
    grpc::experimental::PrivateKeySigner::SignatureAlgorithm
        signature_algorithm,
    EVP_PKEY* private_key) {
  const uint8_t* in = reinterpret_cast<const uint8_t*>(data_to_sign.data());
  const size_t in_len = data_to_sign.size();

  uint16_t boring_signature_algorithm =
      GetBoringSslAlgorithm(signature_algorithm);
  if (EVP_PKEY_id(private_key) !=
      SSL_get_signature_algorithm_key_type(boring_signature_algorithm)) {
    fprintf(stderr, "Key type does not match signature algorithm.\n");
  }

  // Determine the hash.
  const EVP_MD* md =
      SSL_get_signature_algorithm_digest(boring_signature_algorithm);
  bssl::ScopedEVP_MD_CTX ctx;
  EVP_PKEY_CTX* pctx;
  if (!EVP_DigestSignInit(ctx.get(), &pctx, md, nullptr, private_key)) {
    return absl::InternalError("EVP_DigestSignInit failed");
  }

  // Configure additional signature parameters.
  if (SSL_is_signature_algorithm_rsa_pss(boring_signature_algorithm)) {
    if (!EVP_PKEY_CTX_set_rsa_padding(pctx, RSA_PKCS1_PSS_PADDING) ||
        !EVP_PKEY_CTX_set_rsa_pss_saltlen(pctx, -1)) {
      return absl::InternalError("EVP_PKEY_CTX failed");
    }
  }

  size_t len = 0;
  if (!EVP_DigestSign(ctx.get(), nullptr, &len, in, in_len)) {
    return absl::InternalError("EVP_DigestSign failed");
  }
  std::vector<uint8_t> private_key_result;
  private_key_result.resize(len);
  if (!EVP_DigestSign(ctx.get(), private_key_result.data(), &len, in, in_len)) {
    return absl::InternalError("EVP_DigestSign failed");
  }
  private_key_result.resize(len);
  return std::string(private_key_result.begin(), private_key_result.end());
}

class SyncTestPrivateKeySigner final
    : public grpc::experimental::PrivateKeySigner {
 public:
  enum class Mode { kSuccess, kError };

  explicit SyncTestPrivateKeySigner(absl::string_view private_key,
                                    Mode mode = Mode::kSuccess)
      : pkey_(LoadPrivateKeyFromString(private_key)), mode_(mode) {}

  std::variant<absl::StatusOr<std::string>,
               std::shared_ptr<grpc_core::PrivateKeySigner::AsyncSigningHandle>>
  Sign(absl::string_view data_to_sign, SignatureAlgorithm signature_algorithm,
       OnSignComplete /*on_sign_complete*/) override {
    if (mode_ == Mode::kError) {
      return absl::InternalError("Test error sync");
    }
    return SignWithBoringSSL(data_to_sign, signature_algorithm, pkey_.get());
  }

  void Cancel(std::shared_ptr<grpc_core::PrivateKeySigner::AsyncSigningHandle>
              /*handle*/) override {}

 private:
  bssl::UniquePtr<EVP_PKEY> pkey_;
  Mode mode_;
};

class AsyncTestPrivateKeySigner final
    : public grpc::experimental::PrivateKeySigner,
      public std::enable_shared_from_this<AsyncTestPrivateKeySigner> {
 public:
  enum class Mode { kSuccess, kDelayed, kCancellation, kError };

  struct AsyncSigningHandleInternal
      : public grpc_core::PrivateKeySigner::AsyncSigningHandle {
    grpc_event_engine::experimental::EventEngine::TaskHandle task_handle;
  };

  explicit AsyncTestPrivateKeySigner(
      absl::string_view private_key, Mode mode = Mode::kSuccess,
      absl::Duration delay = absl::ZeroDuration())
      : pkey_(LoadPrivateKeyFromString(private_key)),
        mode_(mode),
        delay_(delay) {}

  std::variant<absl::StatusOr<std::string>,
               std::shared_ptr<grpc_core::PrivateKeySigner::AsyncSigningHandle>>
  Sign(absl::string_view data_to_sign, SignatureAlgorithm signature_algorithm,
       OnSignComplete on_sign_complete) override {
    grpc_core::ExecCtx exec_ctx;
    auto event_engine =
        grpc_event_engine::experimental::GetDefaultEventEngine();
    auto handle = std::make_shared<AsyncSigningHandleInternal>();

    if (mode_ == Mode::kCancellation) {
      event_engine->Run([self = shared_from_this()]() {
        while (!self->was_cancelled_.load()) {
          absl::SleepFor(absl::Milliseconds(10));
        }
      });
    } else {
      handle->task_handle = event_engine->RunAfter(
          std::chrono::nanoseconds(absl::ToInt64Nanoseconds(delay_)),
          [self = shared_from_this(), data_to_sign = std::string(data_to_sign),
           signature_algorithm,
           on_sign_complete = std::move(on_sign_complete)]() mutable {
            if (self->mode_ == Mode::kError) {
              on_sign_complete(absl::InternalError("Test error async"));
            } else {
              on_sign_complete(SignWithBoringSSL(
                  data_to_sign, signature_algorithm, self->pkey_.get()));
            }
          });
    }
    return handle;
  }

  void Cancel(std::shared_ptr<grpc_core::PrivateKeySigner::AsyncSigningHandle>
                  handle) override {
    if (mode_ == Mode::kDelayed) {
      grpc_core::ExecCtx exec_ctx;
      auto event_engine =
          grpc_event_engine::experimental::GetDefaultEventEngine();
      auto internal_handle =
          std::static_pointer_cast<AsyncSigningHandleInternal>(handle);
      event_engine->Cancel(internal_handle->task_handle);
    }
  }

  bool WasCancelled() { return was_cancelled_.load(); }

 private:
  bssl::UniquePtr<EVP_PKEY> pkey_;
  Mode mode_;
  absl::Duration delay_;
  std::atomic<bool> was_cancelled_{false};
};

// Verifies that the server can successfully offload signing to an asynchronous
// custom signer.
TEST_F(TlsPrivateKeyOffloadTest, OffloadWithCustomKeySignerAsync) {
  server_addr_ = absl::StrCat("localhost:", grpc_pick_unused_port_or_die());
  std::string server_key =
      grpc_core::testing::GetFileContents(std::string(kServerKeyPath));
  std::string server_cert =
      grpc_core::testing::GetFileContents(std::string(kServerCertPath));
  std::string ca_cert =
      grpc_core::testing::GetFileContents(std::string(kCaPemPath));

  std::vector<experimental::IdentityKeyOrSignerCertPair>
      server_identity_key_cert_pairs;
  signer_ = std::make_shared<AsyncTestPrivateKeySigner>(server_key);
  server_identity_key_cert_pairs.emplace_back(
      grpc::experimental::IdentityKeyOrSignerCertPair{signer_, server_cert});
  auto server_certificate_provider =
      std::make_shared<experimental::InMemoryCertificateProvider>();
  ASSERT_TRUE(server_certificate_provider
                  ->UpdateIdentityKeyCertPair(server_identity_key_cert_pairs)
                  .ok());
  ASSERT_TRUE(server_certificate_provider->UpdateRoot(ca_cert).ok());

  StartServer(server_certificate_provider);

  std::string client_key =
      grpc_core::testing::GetFileContents(std::string(kClientKeyPath));
  std::string client_cert =
      grpc_core::testing::GetFileContents(std::string(kClientCertPath));
  auto client_certificate_provider =
      std::make_shared<experimental::InMemoryCertificateProvider>();
  std::vector<experimental::IdentityKeyCertPair> identity_key_cert_pairs;
  identity_key_cert_pairs.emplace_back(
      experimental::IdentityKeyCertPair{client_key, client_cert});
  ASSERT_TRUE(client_certificate_provider
                  ->UpdateIdentityKeyCertPair(identity_key_cert_pairs)
                  .ok());
  ASSERT_TRUE(client_certificate_provider->UpdateRoot(ca_cert).ok());

  grpc::experimental::TlsChannelCredentialsOptions options;
  options.set_certificate_provider(client_certificate_provider);
  options.watch_root_certs();
  options.set_root_cert_name("root");
  options.watch_identity_key_cert_pairs();
  options.set_identity_cert_name("identity");
  options.set_check_call_host(false);

  DoRpc(server_addr_, options);
}

// Verifies that the server handshake succeeds even when the asynchronous signer
// has a significant delay.
TEST_F(TlsPrivateKeyOffloadTest, OffloadWithCustomKeySignerAsyncDelayed) {
  server_addr_ = absl::StrCat("localhost:", grpc_pick_unused_port_or_die());
  std::string server_key =
      grpc_core::testing::GetFileContents(std::string(kServerKeyPath));
  std::string server_cert =
      grpc_core::testing::GetFileContents(std::string(kServerCertPath));
  std::string ca_cert =
      grpc_core::testing::GetFileContents(std::string(kCaPemPath));

  std::vector<experimental::IdentityKeyOrSignerCertPair>
      server_identity_key_cert_pairs;
  signer_ = std::make_shared<AsyncTestPrivateKeySigner>(
      server_key, AsyncTestPrivateKeySigner::Mode::kDelayed, absl::Seconds(1));
  server_identity_key_cert_pairs.emplace_back(
      grpc::experimental::IdentityKeyOrSignerCertPair{signer_, server_cert});
  auto server_certificate_provider =
      std::make_shared<experimental::InMemoryCertificateProvider>();
  ASSERT_TRUE(server_certificate_provider
                  ->UpdateIdentityKeyCertPair(server_identity_key_cert_pairs)
                  .ok());
  ASSERT_TRUE(server_certificate_provider->UpdateRoot(ca_cert).ok());

  StartServer(server_certificate_provider);

  std::string client_key =
      grpc_core::testing::GetFileContents(std::string(kClientKeyPath));
  std::string client_cert =
      grpc_core::testing::GetFileContents(std::string(kClientCertPath));
  auto client_certificate_provider =
      std::make_shared<experimental::InMemoryCertificateProvider>();
  std::vector<experimental::IdentityKeyCertPair> identity_key_cert_pairs;
  identity_key_cert_pairs.emplace_back(
      experimental::IdentityKeyCertPair{client_key, client_cert});
  ASSERT_TRUE(client_certificate_provider
                  ->UpdateIdentityKeyCertPair(identity_key_cert_pairs)
                  .ok());
  ASSERT_TRUE(client_certificate_provider->UpdateRoot(ca_cert).ok());

  grpc::experimental::TlsChannelCredentialsOptions options;
  options.set_certificate_provider(client_certificate_provider);
  options.watch_root_certs();
  options.set_root_cert_name("root");
  options.watch_identity_key_cert_pairs();
  options.set_identity_cert_name("identity");
  options.set_check_call_host(false);

  DoRpc(server_addr_, options);
}

// Verifies that the client can successfully offload signing to an asynchronous
// custom signer during mTLS.
TEST_F(TlsPrivateKeyOffloadTest, OffloadWithCustomKeySignerClientAsync) {
  server_addr_ = absl::StrCat("localhost:", grpc_pick_unused_port_or_die());
  std::string server_key =
      grpc_core::testing::GetFileContents(std::string(kServerKeyPath));
  std::string server_cert =
      grpc_core::testing::GetFileContents(std::string(kServerCertPath));
  std::string ca_cert =
      grpc_core::testing::GetFileContents(std::string(kCaPemPath));

  auto server_certificate_provider =
      std::make_shared<experimental::InMemoryCertificateProvider>();
  std::vector<experimental::IdentityKeyCertPair> server_identity_key_cert_pairs;
  server_identity_key_cert_pairs.emplace_back(
      experimental::IdentityKeyCertPair{server_key, server_cert});
  ASSERT_TRUE(server_certificate_provider
                  ->UpdateIdentityKeyCertPair(server_identity_key_cert_pairs)
                  .ok());
  ASSERT_TRUE(server_certificate_provider->UpdateRoot(ca_cert).ok());

  StartServer(server_certificate_provider);

  std::string client_key =
      grpc_core::testing::GetFileContents(std::string(kClientKeyPath));
  std::string client_cert =
      grpc_core::testing::GetFileContents(std::string(kClientCertPath));
  auto client_certificate_provider =
      std::make_shared<experimental::InMemoryCertificateProvider>();
  signer_ = std::make_shared<AsyncTestPrivateKeySigner>(client_key);
  std::vector<grpc::experimental::IdentityKeyOrSignerCertPair> identity_pairs;
  identity_pairs.emplace_back(
      grpc::experimental::IdentityKeyOrSignerCertPair{signer_, client_cert});
  ASSERT_TRUE(
      client_certificate_provider->UpdateIdentityKeyCertPair(identity_pairs)
          .ok());
  ASSERT_TRUE(client_certificate_provider->UpdateRoot(ca_cert).ok());

  grpc::experimental::TlsChannelCredentialsOptions options;
  options.set_certificate_provider(client_certificate_provider);
  options.watch_root_certs();
  options.set_root_cert_name("root");
  options.watch_identity_key_cert_pairs();
  options.set_identity_cert_name("identity");
  options.set_check_call_host(false);

  DoRpc(server_addr_, options);
}

// Verifies that the server can successfully offload signing to a synchronous
// custom signer.
TEST_F(TlsPrivateKeyOffloadTest, OffloadWithCustomKeySignerSync) {
  server_addr_ = absl::StrCat("localhost:", grpc_pick_unused_port_or_die());
  std::string server_key =
      grpc_core::testing::GetFileContents(std::string(kServerKeyPath));
  std::string server_cert =
      grpc_core::testing::GetFileContents(std::string(kServerCertPath));
  std::string ca_cert =
      grpc_core::testing::GetFileContents(std::string(kCaPemPath));

  std::vector<experimental::IdentityKeyOrSignerCertPair>
      server_identity_key_cert_pairs;
  signer_ = std::make_shared<SyncTestPrivateKeySigner>(server_key);
  server_identity_key_cert_pairs.emplace_back(
      grpc::experimental::IdentityKeyOrSignerCertPair{signer_, server_cert});
  auto server_certificate_provider =
      std::make_shared<experimental::InMemoryCertificateProvider>();
  ASSERT_TRUE(server_certificate_provider
                  ->UpdateIdentityKeyCertPair(server_identity_key_cert_pairs)
                  .ok());
  ASSERT_TRUE(server_certificate_provider->UpdateRoot(ca_cert).ok());

  StartServer(server_certificate_provider);

  std::string client_key =
      grpc_core::testing::GetFileContents(std::string(kClientKeyPath));
  std::string client_cert =
      grpc_core::testing::GetFileContents(std::string(kClientCertPath));
  auto client_certificate_provider =
      std::make_shared<experimental::InMemoryCertificateProvider>();
  std::vector<experimental::IdentityKeyCertPair> identity_key_cert_pairs;
  identity_key_cert_pairs.emplace_back(
      experimental::IdentityKeyCertPair{client_key, client_cert});
  ASSERT_TRUE(client_certificate_provider
                  ->UpdateIdentityKeyCertPair(identity_key_cert_pairs)
                  .ok());
  ASSERT_TRUE(client_certificate_provider->UpdateRoot(ca_cert).ok());

  grpc::experimental::TlsChannelCredentialsOptions options;
  options.set_certificate_provider(client_certificate_provider);
  options.watch_root_certs();
  options.set_root_cert_name("root");
  options.watch_identity_key_cert_pairs();
  options.set_identity_cert_name("identity");
  options.set_check_call_host(false);

  DoRpc(server_addr_, options);
}

// Verifies that the client can successfully offload signing to a synchronous
// custom signer during mTLS.
TEST_F(TlsPrivateKeyOffloadTest, OffloadWithCustomKeySignerClientSync) {
  server_addr_ = absl::StrCat("localhost:", grpc_pick_unused_port_or_die());
  std::string server_key =
      grpc_core::testing::GetFileContents(std::string(kServerKeyPath));
  std::string server_cert =
      grpc_core::testing::GetFileContents(std::string(kServerCertPath));
  std::string ca_cert =
      grpc_core::testing::GetFileContents(std::string(kCaPemPath));

  auto server_certificate_provider =
      std::make_shared<experimental::InMemoryCertificateProvider>();
  std::vector<experimental::IdentityKeyCertPair> server_identity_key_cert_pairs;
  server_identity_key_cert_pairs.emplace_back(
      experimental::IdentityKeyCertPair{server_key, server_cert});
  ASSERT_TRUE(server_certificate_provider
                  ->UpdateIdentityKeyCertPair(server_identity_key_cert_pairs)
                  .ok());
  ASSERT_TRUE(server_certificate_provider->UpdateRoot(ca_cert).ok());

  StartServer(server_certificate_provider);

  std::string client_key =
      grpc_core::testing::GetFileContents(std::string(kClientKeyPath));
  std::string client_cert =
      grpc_core::testing::GetFileContents(std::string(kClientCertPath));
  auto client_certificate_provider =
      std::make_shared<experimental::InMemoryCertificateProvider>();
  signer_ = std::make_shared<SyncTestPrivateKeySigner>(client_key);
  std::vector<grpc::experimental::IdentityKeyOrSignerCertPair> identity_pairs;
  identity_pairs.emplace_back(
      grpc::experimental::IdentityKeyOrSignerCertPair{signer_, client_cert});
  ASSERT_TRUE(
      client_certificate_provider->UpdateIdentityKeyCertPair(identity_pairs)
          .ok());
  ASSERT_TRUE(client_certificate_provider->UpdateRoot(ca_cert).ok());

  grpc::experimental::TlsChannelCredentialsOptions options;
  options.set_certificate_provider(client_certificate_provider);
  options.watch_root_certs();
  options.set_root_cert_name("root");
  options.watch_identity_key_cert_pairs();
  options.set_identity_cert_name("identity");
  options.set_check_call_host(false);

  DoRpc(server_addr_, options);
}

// Verifies that an immediate error status returned by a synchronous signer
// correctly fails the RPC.
TEST_F(TlsPrivateKeyOffloadTest, OffloadWithCustomKeySignerErrorSync) {
  server_addr_ = absl::StrCat("localhost:", grpc_pick_unused_port_or_die());
  std::string server_key =
      grpc_core::testing::GetFileContents(std::string(kServerKeyPath));
  std::string server_cert =
      grpc_core::testing::GetFileContents(std::string(kServerCertPath));
  std::string ca_cert =
      grpc_core::testing::GetFileContents(std::string(kCaPemPath));

  auto server_certificate_provider =
      std::make_shared<experimental::InMemoryCertificateProvider>();
  std::vector<experimental::IdentityKeyCertPair> server_identity_key_cert_pairs;
  server_identity_key_cert_pairs.emplace_back(
      experimental::IdentityKeyCertPair{server_key, server_cert});
  ASSERT_TRUE(server_certificate_provider
                  ->UpdateIdentityKeyCertPair(server_identity_key_cert_pairs)
                  .ok());
  ASSERT_TRUE(server_certificate_provider->UpdateRoot(ca_cert).ok());

  StartServer(server_certificate_provider);

  std::string client_cert =
      grpc_core::testing::GetFileContents(std::string(kClientCertPath));
  auto client_certificate_provider =
      std::make_shared<experimental::InMemoryCertificateProvider>();
  auto signer = std::make_shared<SyncTestPrivateKeySigner>(
      "", SyncTestPrivateKeySigner::Mode::kError);
  std::vector<grpc::experimental::IdentityKeyOrSignerCertPair> identity_pairs;
  identity_pairs.emplace_back(
      grpc::experimental::IdentityKeyOrSignerCertPair{signer, client_cert});
  ASSERT_TRUE(
      client_certificate_provider->UpdateIdentityKeyCertPair(identity_pairs)
          .ok());
  ASSERT_TRUE(client_certificate_provider->UpdateRoot(ca_cert).ok());

  grpc::experimental::TlsChannelCredentialsOptions options;
  options.set_certificate_provider(client_certificate_provider);
  options.watch_root_certs();
  options.set_root_cert_name("root");
  options.watch_identity_key_cert_pairs();
  options.set_identity_cert_name("identity");
  options.set_check_call_host(false);

  DoRpcAndExpectFailure(server_addr_, options, nullptr, "Test error sync");
}

// Verifies that an error status returned via a callback from an asynchronous
// signer correctly fails the RPC.
TEST_F(TlsPrivateKeyOffloadTest, OffloadWithCustomKeySignerErrorAsync) {
  server_addr_ = absl::StrCat("localhost:", grpc_pick_unused_port_or_die());
  std::string server_key =
      grpc_core::testing::GetFileContents(std::string(kServerKeyPath));
  std::string server_cert =
      grpc_core::testing::GetFileContents(std::string(kServerCertPath));
  std::string ca_cert =
      grpc_core::testing::GetFileContents(std::string(kCaPemPath));

  auto server_certificate_provider =
      std::make_shared<experimental::InMemoryCertificateProvider>();
  std::vector<experimental::IdentityKeyCertPair> server_identity_key_cert_pairs;
  server_identity_key_cert_pairs.emplace_back(
      experimental::IdentityKeyCertPair{server_key, server_cert});
  ASSERT_TRUE(server_certificate_provider
                  ->UpdateIdentityKeyCertPair(server_identity_key_cert_pairs)
                  .ok());
  ASSERT_TRUE(server_certificate_provider->UpdateRoot(ca_cert).ok());

  StartServer(server_certificate_provider);

  std::string client_cert =
      grpc_core::testing::GetFileContents(std::string(kClientCertPath));
  auto client_certificate_provider =
      std::make_shared<experimental::InMemoryCertificateProvider>();
  auto signer = std::make_shared<AsyncTestPrivateKeySigner>(
      "", AsyncTestPrivateKeySigner::Mode::kError);
  std::vector<grpc::experimental::IdentityKeyOrSignerCertPair> identity_pairs;
  identity_pairs.emplace_back(
      grpc::experimental::IdentityKeyOrSignerCertPair{signer, client_cert});
  ASSERT_TRUE(
      client_certificate_provider->UpdateIdentityKeyCertPair(identity_pairs)
          .ok());
  ASSERT_TRUE(client_certificate_provider->UpdateRoot(ca_cert).ok());

  grpc::experimental::TlsChannelCredentialsOptions options;
  options.set_certificate_provider(client_certificate_provider);
  options.watch_root_certs();
  options.set_root_cert_name("root");
  options.watch_identity_key_cert_pairs();
  options.set_identity_cert_name("identity");
  options.set_check_call_host(false);

  DoRpcAndExpectFailure(server_addr_, options, nullptr, "Test error async");
}

// Verifies that the TLS handshake fails when a synchronous signer provides a
// signature from the wrong key.
TEST_F(TlsPrivateKeyOffloadTest,
       OffloadWithCustomKeySignerInvalidSignatureSync) {
  server_addr_ = absl::StrCat("localhost:", grpc_pick_unused_port_or_die());
  std::string server_key =
      grpc_core::testing::GetFileContents(std::string(kServerKeyPath));
  std::string server_cert =
      grpc_core::testing::GetFileContents(std::string(kServerCertPath));
  std::string ca_cert =
      grpc_core::testing::GetFileContents(std::string(kCaPemPath));

  auto server_certificate_provider =
      std::make_shared<experimental::InMemoryCertificateProvider>();
  std::vector<experimental::IdentityKeyCertPair> server_identity_key_cert_pairs;
  server_identity_key_cert_pairs.emplace_back(
      experimental::IdentityKeyCertPair{server_key, server_cert});
  ASSERT_TRUE(server_certificate_provider
                  ->UpdateIdentityKeyCertPair(server_identity_key_cert_pairs)
                  .ok());
  ASSERT_TRUE(server_certificate_provider->UpdateRoot(ca_cert).ok());

  StartServer(server_certificate_provider);

  // Load the wrong key for signing.
  std::string server_key_wrong =
      grpc_core::testing::GetFileContents(std::string(kServerKeyPath));
  std::string client_cert =
      grpc_core::testing::GetFileContents(std::string(kClientCertPath));
  auto client_certificate_provider =
      std::make_shared<experimental::InMemoryCertificateProvider>();
  auto signer = std::make_shared<SyncTestPrivateKeySigner>(server_key_wrong);
  std::vector<grpc::experimental::IdentityKeyOrSignerCertPair> identity_pairs;
  identity_pairs.emplace_back(
      grpc::experimental::IdentityKeyOrSignerCertPair{signer, client_cert});
  ASSERT_TRUE(
      client_certificate_provider->UpdateIdentityKeyCertPair(identity_pairs)
          .ok());
  ASSERT_TRUE(client_certificate_provider->UpdateRoot(ca_cert).ok());

  grpc::experimental::TlsChannelCredentialsOptions options;
  options.set_certificate_provider(client_certificate_provider);
  options.watch_root_certs();
  options.set_root_cert_name("root");
  options.watch_identity_key_cert_pairs();
  options.set_identity_cert_name("identity");
  options.set_check_call_host(false);

  DoRpcAndExpectFailure(server_addr_, options);
}

// Verifies that the TLS handshake fails when an asynchronous signer provides a
// signature from the wrong key.
TEST_F(TlsPrivateKeyOffloadTest,
       OffloadWithCustomKeySignerInvalidSignatureAsync) {
  server_addr_ = absl::StrCat("localhost:", grpc_pick_unused_port_or_die());
  std::string server_key =
      grpc_core::testing::GetFileContents(std::string(kServerKeyPath));
  std::string server_cert =
      grpc_core::testing::GetFileContents(std::string(kServerCertPath));
  std::string ca_cert =
      grpc_core::testing::GetFileContents(std::string(kCaPemPath));

  auto server_certificate_provider =
      std::make_shared<experimental::InMemoryCertificateProvider>();
  std::vector<experimental::IdentityKeyCertPair> server_identity_key_cert_pairs;
  server_identity_key_cert_pairs.emplace_back(
      experimental::IdentityKeyCertPair{server_key, server_cert});
  ASSERT_TRUE(server_certificate_provider
                  ->UpdateIdentityKeyCertPair(server_identity_key_cert_pairs)
                  .ok());
  ASSERT_TRUE(server_certificate_provider->UpdateRoot(ca_cert).ok());

  StartServer(server_certificate_provider);

  // Load the WRONG key for signing.
  std::string server_key_wrong =
      grpc_core::testing::GetFileContents(std::string(kServerKeyPath));
  std::string client_cert =
      grpc_core::testing::GetFileContents(std::string(kClientCertPath));
  auto client_certificate_provider =
      std::make_shared<experimental::InMemoryCertificateProvider>();
  auto signer = std::make_shared<AsyncTestPrivateKeySigner>(server_key_wrong);
  std::vector<grpc::experimental::IdentityKeyOrSignerCertPair> identity_pairs;
  identity_pairs.emplace_back(
      grpc::experimental::IdentityKeyOrSignerCertPair{signer, client_cert});
  ASSERT_TRUE(
      client_certificate_provider->UpdateIdentityKeyCertPair(identity_pairs)
          .ok());
  ASSERT_TRUE(client_certificate_provider->UpdateRoot(ca_cert).ok());

  grpc::experimental::TlsChannelCredentialsOptions options;
  options.set_certificate_provider(client_certificate_provider);
  options.watch_root_certs();
  options.set_root_cert_name("root");
  options.watch_identity_key_cert_pairs();
  options.set_identity_cert_name("identity");
  options.set_check_call_host(false);

  DoRpcAndExpectFailure(server_addr_, options);
}

// Verifies that the server correctly times out the handshake if the custom
// signer takes too long.
TEST_F(TlsPrivateKeyOffloadTest, OffloadWithCustomKeySignerHandshakeTimeout) {
  server_addr_ = absl::StrCat("localhost:", grpc_pick_unused_port_or_die());
  std::string server_key =
      grpc_core::testing::GetFileContents(std::string(kServerKeyPath));
  std::string server_cert =
      grpc_core::testing::GetFileContents(std::string(kServerCertPath));
  std::string ca_cert =
      grpc_core::testing::GetFileContents(std::string(kCaPemPath));

  std::vector<experimental::IdentityKeyOrSignerCertPair>
      server_identity_key_cert_pairs;
  // Signer with a long delay.
  signer_ = std::make_shared<AsyncTestPrivateKeySigner>(
      server_key, AsyncTestPrivateKeySigner::Mode::kDelayed, absl::Seconds(10));
  server_identity_key_cert_pairs.emplace_back(
      grpc::experimental::IdentityKeyOrSignerCertPair{signer_, server_cert});
  auto server_certificate_provider =
      std::make_shared<experimental::InMemoryCertificateProvider>();
  ASSERT_TRUE(server_certificate_provider
                  ->UpdateIdentityKeyCertPair(server_identity_key_cert_pairs)
                  .ok());
  ASSERT_TRUE(server_certificate_provider->UpdateRoot(ca_cert).ok());

  // Start server with a short handshake timeout.
  StartServer(server_certificate_provider, /*handshake_timeout_ms=*/500);

  std::string client_key =
      grpc_core::testing::GetFileContents(std::string(kClientKeyPath));
  std::string client_cert =
      grpc_core::testing::GetFileContents(std::string(kClientCertPath));
  auto client_certificate_provider =
      std::make_shared<experimental::InMemoryCertificateProvider>();
  std::vector<experimental::IdentityKeyCertPair> identity_key_cert_pairs;
  identity_key_cert_pairs.emplace_back(
      experimental::IdentityKeyCertPair{client_key, client_cert});
  ASSERT_TRUE(client_certificate_provider
                  ->UpdateIdentityKeyCertPair(identity_key_cert_pairs)
                  .ok());
  ASSERT_TRUE(client_certificate_provider->UpdateRoot(ca_cert).ok());

  grpc::experimental::TlsChannelCredentialsOptions options;
  options.set_certificate_provider(client_certificate_provider);
  options.watch_root_certs();
  options.set_root_cert_name("root");
  options.watch_identity_key_cert_pairs();
  options.set_identity_cert_name("identity");
  options.set_check_call_host(false);

  DoRpcAndExpectFailure(server_addr_, options);
}

// Verifies that the client correctly times out the handshake if the custom
// signer takes too long.
TEST_F(TlsPrivateKeyOffloadTest,
       OffloadWithCustomKeySignerClientHandshakeTimeout) {
  server_addr_ = absl::StrCat("localhost:", grpc_pick_unused_port_or_die());
  std::string server_key =
      grpc_core::testing::GetFileContents(std::string(kServerKeyPath));
  std::string server_cert =
      grpc_core::testing::GetFileContents(std::string(kServerCertPath));
  std::string ca_cert =
      grpc_core::testing::GetFileContents(std::string(kCaPemPath));

  auto server_certificate_provider =
      std::make_shared<experimental::InMemoryCertificateProvider>();
  std::vector<experimental::IdentityKeyCertPair> server_identity_key_cert_pairs;
  server_identity_key_cert_pairs.emplace_back(
      experimental::IdentityKeyCertPair{server_key, server_cert});
  ASSERT_TRUE(server_certificate_provider
                  ->UpdateIdentityKeyCertPair(server_identity_key_cert_pairs)
                  .ok());
  ASSERT_TRUE(server_certificate_provider->UpdateRoot(ca_cert).ok());

  StartServer(server_certificate_provider);

  std::string client_key =
      grpc_core::testing::GetFileContents(std::string(kClientKeyPath));
  std::string client_cert =
      grpc_core::testing::GetFileContents(std::string(kClientCertPath));
  auto client_certificate_provider =
      std::make_shared<experimental::InMemoryCertificateProvider>();

  // Signer with a long delay.
  auto signer = std::make_shared<AsyncTestPrivateKeySigner>(
      client_key, AsyncTestPrivateKeySigner::Mode::kDelayed, absl::Seconds(10));
  std::vector<grpc::experimental::IdentityKeyOrSignerCertPair> identity_pairs;
  identity_pairs.emplace_back(
      grpc::experimental::IdentityKeyOrSignerCertPair{signer, client_cert});
  ASSERT_TRUE(
      client_certificate_provider->UpdateIdentityKeyCertPair(identity_pairs)
          .ok());
  ASSERT_TRUE(client_certificate_provider->UpdateRoot(ca_cert).ok());

  grpc::experimental::TlsChannelCredentialsOptions options;
  options.set_certificate_provider(client_certificate_provider);
  options.watch_root_certs();
  options.set_root_cert_name("root");
  options.watch_identity_key_cert_pairs();
  options.set_identity_cert_name("identity");
  options.set_check_call_host(false);

  // Set a short handshake timeout on the client.
  ChannelArguments channel_args;
  channel_args.SetInt("grpc.testing.fixed_reconnect_backoff_ms", 500);

  DoRpcAndExpectFailure(server_addr_, options, nullptr, "Handshaker shutdown",
                        channel_args);
}

}  // namespace
}  // namespace testing
}  // namespace grpc

#endif

int main(int argc, char** argv) {
  grpc::testing::TestEnvironment env(&argc, argv);
  ::testing::InitGoogleTest(&argc, argv);
  int ret = RUN_ALL_TESTS();
  return ret;
}
