Fix comments and add client timeout test

Author: anniefrchz

src/core/tsi/ssl_transport_security.cc

@@ -189,7 +189,8 @@ struct tsi_ssl_handshaker : public tsi_handshaker,
   // Will be set if there is a pending call to tsi_handshaker_next(),
   // or nullopt if not.
   std::optional<HandshakerNextArgs> handshaker_next_args ABSL_GUARDED_BY(mu);
-  void MaybeSetError(std::string error) ABSL_EXCLUSIVE_LOCKS_REQUIRED(mu) {
+  const void MaybeSetError(std::string error)
+      ABSL_EXCLUSIVE_LOCKS_REQUIRED(mu) {
     if (!handshaker_next_args.has_value()) return;
     if (handshaker_next_args->error_ptr == nullptr) return;
     *handshaker_next_args->error_ptr = std::move(error);
@@ -455,11 +456,6 @@ enum ssl_private_key_result_t TlsPrivateKeySignWrapper(
       [&](absl::StatusOr<std::string>* status_or_string)
           ABSL_NO_THREAD_SAFETY_ANALYSIS {
             if (status_or_string->ok()) {
-              if ((*status_or_string)->size() > max_out) {
-                handshaker->MaybeSetError(
-                    "Private Key Signature exceeds maximum allowed size.");
-                return ssl_private_key_failure;
-              }
               handshaker->signed_bytes = std::move(*status_or_string);
               return TlsPrivateKeyOffloadComplete(ssl, out, out_len, max_out);
             } else {
@@ -2233,7 +2229,6 @@ static tsi_result ssl_handshaker_write_output_buffer(tsi_handshaker* self,
 
 static tsi_result ssl_handshaker_next_impl(tsi_ssl_handshaker* self)
     ABSL_EXCLUSIVE_LOCKS_REQUIRED(&self->mu) {
-  std::string* error = self->handshaker_next_args->error_ptr;
   // If there are received bytes, process them first.
   tsi_result status = TSI_OK;
   size_t bytes_written = 0;
@@ -2254,17 +2249,18 @@ static tsi_result ssl_handshaker_next_impl(tsi_ssl_handshaker* self)
           remaining_bytes_to_write_to_openssl_size;
       status = ssl_handshaker_process_bytes_from_peer(
           self, remaining_bytes_to_write_to_openssl, &bytes_written_to_openssl,
-          error);
+          self->handshaker_next_args->error_ptr);
       // As long as the BIO is full, drive the SSL handshake to consume bytes
       // from the BIO. If the SSL handshake returns any bytes, write them to
       // the peer.
       while (status == TSI_DRAIN_BUFFER) {
-        status =
-            ssl_handshaker_write_output_buffer(self, &bytes_written, error);
+        status = ssl_handshaker_write_output_buffer(
+            self, &bytes_written, self->handshaker_next_args->error_ptr);
         if (status != TSI_OK) {
           return status;
         }
-        status = ssl_handshaker_do_handshake(self, error);
+        status = ssl_handshaker_do_handshake(
+            self, self->handshaker_next_args->error_ptr);
       }
       // Move the pointer to the first byte not yet successfully written to
       // the BIO.
@@ -2287,16 +2283,17 @@ static tsi_result ssl_handshaker_next_impl(tsi_ssl_handshaker* self)
     // During the PrivateKeyOffload signature, an empty call to
     // ssl_handshaker_do_handshake needs to be forced  after the async offload
     // has completed.
-    status = ssl_handshaker_do_handshake(self, error);
-    self->signed_bytes = "";
+    status = ssl_handshaker_do_handshake(self,
+                                         self->handshaker_next_args->error_ptr);
 #endif
   }
 
   if (status != TSI_OK) {
     return status;
   }
   // Get bytes to send to the peer, if available.
-  status = ssl_handshaker_write_output_buffer(self, &bytes_written, error);
+  status = ssl_handshaker_write_output_buffer(
+      self, &bytes_written, self->handshaker_next_args->error_ptr);
   if (status != TSI_OK) {
     return status;
   }
@@ -2312,21 +2309,22 @@ static tsi_result ssl_handshaker_next_impl(tsi_ssl_handshaker* self)
     // bytes from the peer that must be processed.
     unsigned char* unused_bytes = nullptr;
     size_t unused_bytes_size = 0;
-    status =
-        ssl_bytes_remaining(self, &unused_bytes, &unused_bytes_size, error);
+    status = ssl_bytes_remaining(self, &unused_bytes, &unused_bytes_size,
+                                 self->handshaker_next_args->error_ptr);
     if (status != TSI_OK) {
       return status;
     }
     if (unused_bytes_size >
         self->handshaker_next_args->original_received_bytes_size) {
       LOG(ERROR) << "More unused bytes than received bytes.";
       gpr_free(unused_bytes);
-      if (error != nullptr) *error = "More unused bytes than received bytes.";
+      self->MaybeSetError("More unused bytes than received bytes.");
       return TSI_INTERNAL_ERROR;
     }
     status = ssl_handshaker_result_create(
         self, unused_bytes, unused_bytes_size,
-        &self->handshaker_next_args->handshaker_result, error);
+        &self->handshaker_next_args->handshaker_result,
+        self->handshaker_next_args->error_ptr);
     if (status == TSI_OK) {
       // Indicates that the handshake has completed and that a
       // handshaker_result has been created.

src/cpp/common/tls_certificate_provider.cc

@@ -25,6 +25,7 @@
 #include "src/core/credentials/transport/tls/grpc_tls_certificate_provider.h"
 #include "src/core/util/grpc_check.h"
 #include "src/core/util/match.h"
+#include "absl/status/statusor.h"
 
 namespace grpc {
 namespace experimental {
@@ -41,7 +42,7 @@ grpc_tls_identity_pairs* CreatePairsCore(
   return pairs_core;
 }
 
-grpc_tls_identity_pairs* CreatePairsCore(
+absl::StatusOr<grpc_tls_identity_pairs*> CreatePairsCore(
     std::vector<IdentityKeyOrSignerCertPair>
         identity_key_or_signer_cert_pairs) {
   grpc_tls_identity_pairs* pairs_core = grpc_tls_identity_pairs_create();
@@ -60,7 +61,7 @@ grpc_tls_identity_pairs* CreatePairsCore(
         });
     if (!status.ok()) {
       grpc_tls_identity_pairs_destroy(pairs_core);
-      return nullptr;
+      return status;
     }
   }
   return pairs_core;
@@ -147,13 +148,11 @@ absl::Status InMemoryCertificateProvider::UpdateIdentityKeyCertPair(
     std::vector<IdentityKeyOrSignerCertPair>
         identity_key_or_signer_cert_pairs) {
   GRPC_CHECK(!identity_key_or_signer_cert_pairs.empty());
-  auto* pairs_core =
+  auto pairs_core_or =
       CreatePairsCore(std::move(identity_key_or_signer_cert_pairs));
-  if (pairs_core == nullptr) {
-    return absl::InternalError("Unable to update identity certificate");
-  }
+  if (!pairs_core_or.ok()) return pairs_core_or.status();
   return grpc_tls_certificate_provider_in_memory_set_identity_certificate(
-             c_provider_, pairs_core)
+             c_provider_, *pairs_core_or)
              ? absl::OkStatus()
              : absl::InternalError("Unable to update identity certificate");
 }

test/core/test_util/tls_utils.cc

@@ -73,7 +73,8 @@ PemKeyCertPairList MakeCertKeyPairs(absl::string_view private_key,
   if (private_key.empty() && certs.empty()) {
     return {};
   }
-  return PemKeyCertPairList{PemKeyCertPair(private_key.data(), certs)};
+  return PemKeyCertPairList{PemKeyCertPair(
+      std::string(private_key.data(), private_key.length()), certs)};
 }
 
 std::string GetFileContents(const std::string& path) {

test/core/tsi/private_key_offload_test.cc

@@ -183,12 +183,7 @@ class AsyncTestPrivateKeySigner final
     return std::make_shared<AsyncSigningHandle>();
   }
 
-  void Cancel(std::shared_ptr<AsyncSigningHandle> /*handle*/) override {
-    LOG(ERROR) << "Here";
-    if (!was_cancelled_.exchange(true)) {
-      notification_.Notify();
-    }
-  }
+  void Cancel(std::shared_ptr<AsyncSigningHandle> /*handle*/) override {}
 
   bool WasCancelled() { return was_cancelled_.load(); }
 
@@ -197,7 +192,6 @@ class AsyncTestPrivateKeySigner final
       event_engine_;
   bssl::UniquePtr<EVP_PKEY> pkey_;
   Mode mode_;
-  absl::Notification notification_;
   std::atomic<bool> was_cancelled_{false};
 };
 
@@ -361,7 +355,6 @@ class SslOffloadTsiTestFixture {
   tsi_tls_version tls_version_;
   bool expect_success_ = false;
   bool expect_success_on_client_ = false;
-  Mutex mu_;
 };
 
 struct tsi_test_fixture_vtable SslOffloadTsiTestFixture::kVtable = {
@@ -389,7 +382,6 @@ class PrivateKeyOffloadTest : public ::testing::TestWithParam<tsi_tls_version> {
     event_engine_->UnsetGlobalHooks();
     WaitForSingleOwner(std::move(event_engine_));
     grpc_shutdown_blocking();
-    event_engine_.reset();
   }
 
   std::shared_ptr<grpc_event_engine::experimental::FuzzingEventEngine>

test/cpp/end2end/tls_private_key_signer_end2end_test.cc

@@ -203,8 +203,8 @@ 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;
+    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),
@@ -882,6 +882,64 @@ TEST_F(TlsPrivateKeyOffloadTest, OffloadWithCustomKeySignerHandshakeTimeout) {
   DoRpcAndExpectFailure(server_addr_, options, nullptr, "Socket closed");
 }
 
+// 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