use gRPC thread model

2 files changed, 262 insertions, 268 deletions

diff --git a/test/core/tsi/alts/handshaker/alts_handshaker_client_test.cc b/test/core/tsi/alts/handshaker/alts_handshaker_client_test.cc
index c8d88aa72c..9a4541bb7d 100644
--- a/test/core/tsi/alts/handshaker/alts_handshaker_client_test.cc
+++ b/test/core/tsi/alts/handshaker/alts_handshaker_client_test.cc
@@ -19,14 +19,14 @@
 #include <grpc/grpc.h>
 
 #include "src/core/tsi/alts/handshaker/alts_handshaker_client.h"
-#include "src/core/tsi/alts/handshaker/alts_tsi_event.h"
+#include "src/core/tsi/alts/handshaker/alts_shared_resource.h"
 #include "src/core/tsi/alts/handshaker/alts_tsi_handshaker.h"
+#include "src/core/tsi/alts/handshaker/alts_tsi_handshaker_private.h"
 #include "src/core/tsi/transport_security.h"
 #include "src/core/tsi/transport_security_interface.h"
 #include "test/core/tsi/alts/handshaker/alts_handshaker_service_api_test_lib.h"
 
 #define ALTS_HANDSHAKER_CLIENT_TEST_OUT_FRAME "Hello Google"
-#define ALTS_HANDSHAKER_CLIENT_TEST_HANDSHAKER_SERVICE_URL "lame"
 #define ALTS_HANDSHAKER_CLIENT_TEST_TARGET_NAME "bigtable.google.api.com"
 #define ALTS_HANDSHAKER_CLIENT_TEST_TARGET_SERVICE_ACCOUNT1 "A@google.com"
 #define ALTS_HANDSHAKER_CLIENT_TEST_TARGET_SERVICE_ACCOUNT2 "B@google.com"
@@ -37,37 +37,22 @@ const size_t kMaxRpcVersionMinor = 2;
 const size_t kMinRpcVersionMajor = 2;
 const size_t kMinRpcVersionMinor = 1;
 
+using grpc_core::internal::alts_handshaker_client_get_closure_for_testing;
+using grpc_core::internal::
+    alts_handshaker_client_get_initial_metadata_for_testing;
+using grpc_core::internal::
+    alts_handshaker_client_get_recv_buffer_addr_for_testing;
+using grpc_core::internal::alts_handshaker_client_get_send_buffer_for_testing;
 using grpc_core::internal::alts_handshaker_client_set_grpc_caller_for_testing;
 
 typedef struct alts_handshaker_client_test_config {
   grpc_channel* channel;
   grpc_completion_queue* cq;
   alts_handshaker_client* client;
+  alts_handshaker_client* server;
   grpc_slice out_frame;
 } alts_handshaker_client_test_config;
 
-static alts_tsi_event* alts_tsi_event_create_for_testing(bool is_client) {
-  alts_tsi_event* e = static_cast<alts_tsi_event*>(gpr_zalloc(sizeof(*e)));
-  grpc_metadata_array_init(&e->initial_metadata);
-  grpc_metadata_array_init(&e->trailing_metadata);
-  e->options = is_client ? grpc_alts_credentials_client_options_create()
-                         : grpc_alts_credentials_server_options_create();
-  if (is_client) {
-    grpc_alts_credentials_client_options_add_target_service_account(
-        e->options, ALTS_HANDSHAKER_CLIENT_TEST_TARGET_SERVICE_ACCOUNT1);
-    grpc_alts_credentials_client_options_add_target_service_account(
-        e->options, ALTS_HANDSHAKER_CLIENT_TEST_TARGET_SERVICE_ACCOUNT2);
-  }
-  grpc_gcp_rpc_protocol_versions* versions = &e->options->rpc_versions;
-  GPR_ASSERT(grpc_gcp_rpc_protocol_versions_set_max(
-      versions, kMaxRpcVersionMajor, kMaxRpcVersionMinor));
-  GPR_ASSERT(grpc_gcp_rpc_protocol_versions_set_min(
-      versions, kMinRpcVersionMajor, kMinRpcVersionMinor));
-  e->target_name =
-      grpc_slice_from_static_string(ALTS_HANDSHAKER_CLIENT_TEST_TARGET_NAME);
-  return e;
-}
-
 static void validate_rpc_protocol_versions(
     grpc_gcp_rpc_protocol_versions* versions) {
   GPR_ASSERT(versions != nullptr);
@@ -101,11 +86,11 @@ static void validate_target_identities(
 
 /**
  * Validate if grpc operation data is correctly populated with the fields of
- * ALTS TSI event.
+ * ALTS handshaker client.
  */
-static bool validate_op(alts_tsi_event* event, const grpc_op* op, size_t nops,
-                        bool is_start) {
-  GPR_ASSERT(event != nullptr && op != nullptr && nops != 0);
+static bool validate_op(alts_handshaker_client* c, const grpc_op* op,
+                        size_t nops, bool is_start) {
+  GPR_ASSERT(c != nullptr && op != nullptr && nops != 0);
   bool ok = true;
   grpc_op* start_op = const_cast<grpc_op*>(op);
   if (is_start) {
@@ -113,23 +98,22 @@ static bool validate_op(alts_tsi_event* event, const grpc_op* op, size_t nops,
     ok &= (op->data.send_initial_metadata.count == 0);
     op++;
     GPR_ASSERT((size_t)(op - start_op) <= kHandshakerClientOpNum);
-
     ok &= (op->op == GRPC_OP_RECV_INITIAL_METADATA);
     ok &= (op->data.recv_initial_metadata.recv_initial_metadata ==
-           &event->initial_metadata);
+           alts_handshaker_client_get_initial_metadata_for_testing(c));
     op++;
     GPR_ASSERT((size_t)(op - start_op) <= kHandshakerClientOpNum);
   }
   ok &= (op->op == GRPC_OP_SEND_MESSAGE);
-  ok &= (op->data.send_message.send_message == event->send_buffer);
+  ok &= (op->data.send_message.send_message ==
+         alts_handshaker_client_get_send_buffer_for_testing(c));
   op++;
   GPR_ASSERT((size_t)(op - start_op) <= kHandshakerClientOpNum);
-
   ok &= (op->op == GRPC_OP_RECV_MESSAGE);
-  ok &= (op->data.recv_message.recv_message == &event->recv_buffer);
+  ok &= (op->data.recv_message.recv_message ==
+         alts_handshaker_client_get_recv_buffer_addr_for_testing(c));
   op++;
   GPR_ASSERT((size_t)(op - start_op) <= kHandshakerClientOpNum);
-
   return ok;
 }
 
@@ -152,7 +136,7 @@ static grpc_gcp_handshaker_req* deserialize_handshaker_req(
  */
 static grpc_call_error check_must_not_be_called(grpc_call* call,
                                                 const grpc_op* ops, size_t nops,
-                                                void* tag) {
+                                                grpc_closure* tag) {
   GPR_ASSERT(0);
 }
 
@@ -164,10 +148,14 @@ static grpc_call_error check_must_not_be_called(grpc_call* call,
  */
 static grpc_call_error check_client_start_success(grpc_call* call,
                                                   const grpc_op* op,
-                                                  size_t nops, void* tag) {
-  alts_tsi_event* event = static_cast<alts_tsi_event*>(tag);
-  grpc_gcp_handshaker_req* req =
-      deserialize_handshaker_req(CLIENT_START_REQ, event->send_buffer);
+                                                  size_t nops,
+                                                  grpc_closure* closure) {
+  alts_handshaker_client* client =
+      static_cast<alts_handshaker_client*>(closure->cb_arg);
+  GPR_ASSERT(alts_handshaker_client_get_closure_for_testing(client) == closure);
+  grpc_gcp_handshaker_req* req = deserialize_handshaker_req(
+      CLIENT_START_REQ,
+      alts_handshaker_client_get_send_buffer_for_testing(client));
   GPR_ASSERT(req->client_start.handshake_security_protocol ==
              grpc_gcp_HandshakeProtocol_ALTS);
   const void* data = (static_cast<repeated_field*>(
@@ -194,7 +182,7 @@ static grpc_call_error check_client_start_success(grpc_call* call,
                     GRPC_SLICE_LENGTH(*target_name)) == 0);
   GPR_ASSERT(GRPC_SLICE_LENGTH(*target_name) ==
              strlen(ALTS_HANDSHAKER_CLIENT_TEST_TARGET_NAME));
-  GPR_ASSERT(validate_op(event, op, nops, true /* is_start */));
+  GPR_ASSERT(validate_op(client, op, nops, true /* is_start */));
   grpc_gcp_handshaker_req_destroy(req);
   return GRPC_CALL_OK;
 }
@@ -207,10 +195,14 @@ static grpc_call_error check_client_start_success(grpc_call* call,
  */
 static grpc_call_error check_server_start_success(grpc_call* call,
                                                   const grpc_op* op,
-                                                  size_t nops, void* tag) {
-  alts_tsi_event* event = static_cast<alts_tsi_event*>(tag);
-  grpc_gcp_handshaker_req* req =
-      deserialize_handshaker_req(SERVER_START_REQ, event->send_buffer);
+                                                  size_t nops,
+                                                  grpc_closure* closure) {
+  alts_handshaker_client* client =
+      static_cast<alts_handshaker_client*>(closure->cb_arg);
+  GPR_ASSERT(alts_handshaker_client_get_closure_for_testing(client) == closure);
+  grpc_gcp_handshaker_req* req = deserialize_handshaker_req(
+      SERVER_START_REQ,
+      alts_handshaker_client_get_send_buffer_for_testing(client));
   const void* data = (static_cast<repeated_field*>(
                           req->server_start.application_protocols.arg))
                          ->data;
@@ -231,7 +223,7 @@ static grpc_call_error check_server_start_success(grpc_call* call,
                     ALTS_RECORD_PROTOCOL,
                     GRPC_SLICE_LENGTH(*record_protocol)) == 0);
   validate_rpc_protocol_versions(&req->server_start.rpc_versions);
-  GPR_ASSERT(validate_op(event, op, nops, true /* is_start */));
+  GPR_ASSERT(validate_op(client, op, nops, true /* is_start */));
   grpc_gcp_handshaker_req_destroy(req);
   return GRPC_CALL_OK;
 }
@@ -242,16 +234,18 @@ static grpc_call_error check_server_start_success(grpc_call* call,
  * and op is correctly populated.
  */
 static grpc_call_error check_next_success(grpc_call* call, const grpc_op* op,
-                                          size_t nops, void* tag) {
-  alts_tsi_event* event = static_cast<alts_tsi_event*>(tag);
-  grpc_gcp_handshaker_req* req =
-      deserialize_handshaker_req(NEXT_REQ, event->send_buffer);
+                                          size_t nops, grpc_closure* closure) {
+  alts_handshaker_client* client =
+      static_cast<alts_handshaker_client*>(closure->cb_arg);
+  GPR_ASSERT(alts_handshaker_client_get_closure_for_testing(client) == closure);
+  grpc_gcp_handshaker_req* req = deserialize_handshaker_req(
+      NEXT_REQ, alts_handshaker_client_get_send_buffer_for_testing(client));
   grpc_slice* in_bytes = static_cast<grpc_slice*>(req->next.in_bytes.arg);
   GPR_ASSERT(in_bytes != nullptr);
   GPR_ASSERT(memcmp(GRPC_SLICE_START_PTR(*in_bytes),
                     ALTS_HANDSHAKER_CLIENT_TEST_OUT_FRAME,
                     GRPC_SLICE_LENGTH(*in_bytes)) == 0);
-  GPR_ASSERT(validate_op(event, op, nops, false /* is_start */));
+  GPR_ASSERT(validate_op(client, op, nops, false /* is_start */));
   grpc_gcp_handshaker_req_destroy(req);
   return GRPC_CALL_OK;
 }
@@ -262,21 +256,54 @@ static grpc_call_error check_next_success(grpc_call* call, const grpc_op* op,
  */
 static grpc_call_error check_grpc_call_failure(grpc_call* call,
                                                const grpc_op* op, size_t nops,
-                                               void* tag) {
+                                               grpc_closure* tag) {
   return GRPC_CALL_ERROR;
 }
 
+static grpc_alts_credentials_options* create_credentials_options(
+    bool is_client) {
+  grpc_alts_credentials_options* options =
+      is_client ? grpc_alts_credentials_client_options_create()
+                : grpc_alts_credentials_server_options_create();
+  if (is_client) {
+    grpc_alts_credentials_client_options_add_target_service_account(
+        options, ALTS_HANDSHAKER_CLIENT_TEST_TARGET_SERVICE_ACCOUNT1);
+    grpc_alts_credentials_client_options_add_target_service_account(
+        options, ALTS_HANDSHAKER_CLIENT_TEST_TARGET_SERVICE_ACCOUNT2);
+  }
+  grpc_gcp_rpc_protocol_versions* versions = &options->rpc_versions;
+  GPR_ASSERT(grpc_gcp_rpc_protocol_versions_set_max(
+      versions, kMaxRpcVersionMajor, kMaxRpcVersionMinor));
+  GPR_ASSERT(grpc_gcp_rpc_protocol_versions_set_min(
+      versions, kMinRpcVersionMajor, kMinRpcVersionMinor));
+  return options;
+}
+
 static alts_handshaker_client_test_config* create_config() {
   alts_handshaker_client_test_config* config =
       static_cast<alts_handshaker_client_test_config*>(
           gpr_zalloc(sizeof(*config)));
   config->channel = grpc_insecure_channel_create(
-      ALTS_HANDSHAKER_CLIENT_TEST_HANDSHAKER_SERVICE_URL, nullptr, nullptr);
+      ALTS_HANDSHAKER_SERVICE_URL_FOR_TESTING, nullptr, nullptr);
   config->cq = grpc_completion_queue_create_for_next(nullptr);
+  grpc_alts_credentials_options* client_options =
+      create_credentials_options(true /* is_client */);
+  grpc_alts_credentials_options* server_options =
+      create_credentials_options(false /*  is_client */);
+  config->server = alts_grpc_handshaker_client_create(
+      nullptr, config->channel, ALTS_HANDSHAKER_SERVICE_URL_FOR_TESTING,
+      nullptr, server_options,
+      grpc_slice_from_static_string(ALTS_HANDSHAKER_CLIENT_TEST_TARGET_NAME),
+      nullptr, nullptr, nullptr, nullptr, false);
   config->client = alts_grpc_handshaker_client_create(
-      config->channel, config->cq,
-      ALTS_HANDSHAKER_CLIENT_TEST_HANDSHAKER_SERVICE_URL);
+      nullptr, config->channel, ALTS_HANDSHAKER_SERVICE_URL_FOR_TESTING,
+      nullptr, client_options,
+      grpc_slice_from_static_string(ALTS_HANDSHAKER_CLIENT_TEST_TARGET_NAME),
+      nullptr, nullptr, nullptr, nullptr, true);
   GPR_ASSERT(config->client != nullptr);
+  GPR_ASSERT(config->server != nullptr);
+  grpc_alts_credentials_options_destroy(client_options);
+  grpc_alts_credentials_options_destroy(server_options);
   config->out_frame =
       grpc_slice_from_static_string(ALTS_HANDSHAKER_CLIENT_TEST_OUT_FRAME);
   return config;
@@ -289,6 +316,7 @@ static void destroy_config(alts_handshaker_client_test_config* config) {
   grpc_completion_queue_destroy(config->cq);
   grpc_channel_destroy(config->channel);
   alts_handshaker_client_destroy(config->client);
+  alts_handshaker_client_destroy(config->server);
   grpc_slice_unref(config->out_frame);
   gpr_free(config);
 }
@@ -296,73 +324,50 @@ static void destroy_config(alts_handshaker_client_test_config* config) {
 static void schedule_request_invalid_arg_test() {
   /* Initialization. */
   alts_handshaker_client_test_config* config = create_config();
-  alts_tsi_event* event = nullptr;
-
   /* Tests. */
   alts_handshaker_client_set_grpc_caller_for_testing(config->client,
                                                      check_must_not_be_called);
-  event = alts_tsi_event_create_for_testing(true /* is_client */);
   /* Check client_start. */
-  GPR_ASSERT(alts_handshaker_client_start_client(nullptr, event) ==
-             TSI_INVALID_ARGUMENT);
-  GPR_ASSERT(alts_handshaker_client_start_client(config->client, nullptr) ==
+  GPR_ASSERT(alts_handshaker_client_start_client(nullptr) ==
              TSI_INVALID_ARGUMENT);
-
   /* Check server_start. */
-  GPR_ASSERT(alts_handshaker_client_start_server(
-                 config->client, event, nullptr) == TSI_INVALID_ARGUMENT);
-  GPR_ASSERT(alts_handshaker_client_start_server(config->client, nullptr,
-                                                 &config->out_frame) ==
+  GPR_ASSERT(alts_handshaker_client_start_server(config->server, nullptr) ==
              TSI_INVALID_ARGUMENT);
-  GPR_ASSERT(alts_handshaker_client_start_server(
-                 nullptr, event, &config->out_frame) == TSI_INVALID_ARGUMENT);
-
-  /* Check next. */
-  GPR_ASSERT(alts_handshaker_client_next(config->client, event, nullptr) ==
+  GPR_ASSERT(alts_handshaker_client_start_server(nullptr, &config->out_frame) ==
              TSI_INVALID_ARGUMENT);
-  GPR_ASSERT(alts_handshaker_client_next(config->client, nullptr,
-                                         &config->out_frame) ==
+  /* Check next. */
+  GPR_ASSERT(alts_handshaker_client_next(config->client, nullptr) ==
              TSI_INVALID_ARGUMENT);
-  GPR_ASSERT(alts_handshaker_client_next(nullptr, event, &config->out_frame) ==
+  GPR_ASSERT(alts_handshaker_client_next(nullptr, &config->out_frame) ==
              TSI_INVALID_ARGUMENT);
-
   /* Check shutdown. */
   alts_handshaker_client_shutdown(nullptr);
-
   /* Cleanup. */
-  alts_tsi_event_destroy(event);
   destroy_config(config);
 }
 
 static void schedule_request_success_test() {
   /* Initialization. */
   alts_handshaker_client_test_config* config = create_config();
-  alts_tsi_event* event = nullptr;
-
   /* Check client_start success. */
   alts_handshaker_client_set_grpc_caller_for_testing(
       config->client, check_client_start_success);
-  event = alts_tsi_event_create_for_testing(true /* is_client. */);
-  GPR_ASSERT(alts_handshaker_client_start_client(config->client, event) ==
-             TSI_OK);
-  alts_tsi_event_destroy(event);
-
+  GPR_ASSERT(alts_handshaker_client_start_client(config->client) == TSI_OK);
   /* Check server_start success. */
   alts_handshaker_client_set_grpc_caller_for_testing(
-      config->client, check_server_start_success);
-  event = alts_tsi_event_create_for_testing(false /* is_client. */);
-  GPR_ASSERT(alts_handshaker_client_start_server(config->client, event,
+      config->server, check_server_start_success);
+  GPR_ASSERT(alts_handshaker_client_start_server(config->server,
                                                  &config->out_frame) == TSI_OK);
-  alts_tsi_event_destroy(event);
-
-  /* Check next success. */
+  /* Check client next success. */
   alts_handshaker_client_set_grpc_caller_for_testing(config->client,
                                                      check_next_success);
-  event = alts_tsi_event_create_for_testing(true /* is_client. */);
-  GPR_ASSERT(alts_handshaker_client_next(config->client, event,
-                                         &config->out_frame) == TSI_OK);
-  alts_tsi_event_destroy(event);
-
+  GPR_ASSERT(alts_handshaker_client_next(config->client, &config->out_frame) ==
+             TSI_OK);
+  /* Check server next success. */
+  alts_handshaker_client_set_grpc_caller_for_testing(config->server,
+                                                     check_next_success);
+  GPR_ASSERT(alts_handshaker_client_next(config->server, &config->out_frame) ==
+             TSI_OK);
   /* Cleanup. */
   destroy_config(config);
 }
@@ -370,30 +375,22 @@ static void schedule_request_success_test() {
 static void schedule_request_grpc_call_failure_test() {
   /* Initialization. */
   alts_handshaker_client_test_config* config = create_config();
-  alts_tsi_event* event = nullptr;
-
   /* Check client_start failure. */
   alts_handshaker_client_set_grpc_caller_for_testing(config->client,
                                                      check_grpc_call_failure);
-  event = alts_tsi_event_create_for_testing(true /* is_client. */);
-  GPR_ASSERT(alts_handshaker_client_start_client(config->client, event) ==
+  GPR_ASSERT(alts_handshaker_client_start_client(config->client) ==
              TSI_INTERNAL_ERROR);
-  alts_tsi_event_destroy(event);
-
   /* Check server_start failure. */
-  event = alts_tsi_event_create_for_testing(false /* is_client. */);
-  GPR_ASSERT(alts_handshaker_client_start_server(config->client, event,
-                                                 &config->out_frame) ==
+  alts_handshaker_client_set_grpc_caller_for_testing(config->server,
+                                                     check_grpc_call_failure);
+  GPR_ASSERT(alts_handshaker_client_start_server(
+                 config->server, &config->out_frame) == TSI_INTERNAL_ERROR);
+  /* Check client next failure. */
+  GPR_ASSERT(alts_handshaker_client_next(config->client, &config->out_frame) ==
+             TSI_INTERNAL_ERROR);
+  /* Check server next failure. */
+  GPR_ASSERT(alts_handshaker_client_next(config->server, &config->out_frame) ==
              TSI_INTERNAL_ERROR);
-  alts_tsi_event_destroy(event);
-
-  /* Check next failure. */
-  event = alts_tsi_event_create_for_testing(true /* is_cleint. */);
-  GPR_ASSERT(
-      alts_handshaker_client_next(config->client, event, &config->out_frame) ==
-      TSI_INTERNAL_ERROR);
-  alts_tsi_event_destroy(event);
-
   /* Cleanup. */
   destroy_config(config);
 }
@@ -401,13 +398,13 @@ static void schedule_request_grpc_call_failure_test() {
 int main(int argc, char** argv) {
   /* Initialization. */
   grpc_init();
-
+  grpc_alts_shared_resource_dedicated_init();
   /* Tests. */
   schedule_request_invalid_arg_test();
   schedule_request_success_test();
   schedule_request_grpc_call_failure_test();
-
   /* Cleanup. */
+  grpc_alts_shared_resource_dedicated_shutdown();
   grpc_shutdown();
   return 0;
 }
diff --git a/test/core/tsi/alts/handshaker/alts_tsi_handshaker_test.cc b/test/core/tsi/alts/handshaker/alts_tsi_handshaker_test.cc
index e9eb7e175f..316ff13816 100644
--- a/test/core/tsi/alts/handshaker/alts_tsi_handshaker_test.cc
+++ b/test/core/tsi/alts/handshaker/alts_tsi_handshaker_test.cc
@@ -24,7 +24,7 @@
 
 #include "src/core/lib/gprpp/thd.h"
 #include "src/core/tsi/alts/handshaker/alts_handshaker_client.h"
-#include "src/core/tsi/alts/handshaker/alts_tsi_event.h"
+#include "src/core/tsi/alts/handshaker/alts_shared_resource.h"
 #include "src/core/tsi/alts/handshaker/alts_tsi_handshaker.h"
 #include "src/core/tsi/alts/handshaker/alts_tsi_handshaker_private.h"
 #include "test/core/tsi/alts/handshaker/alts_handshaker_service_api_test_lib.h"
@@ -43,12 +43,18 @@
 #define ALTS_TSI_HANDSHAKER_TEST_MIN_RPC_VERSION_MAJOR 2
 #define ALTS_TSI_HANDSHAKER_TEST_MIN_RPC_VERSION_MINOR 1
 
+using grpc_core::internal::alts_handshaker_client_check_fields_for_testing;
+using grpc_core::internal::alts_handshaker_client_get_handshaker_for_testing;
 using grpc_core::internal::
-    alts_tsi_handshaker_get_has_sent_start_message_for_testing;
+    alts_handshaker_client_get_recv_buffer_addr_for_testing;
+using grpc_core::internal::alts_handshaker_client_set_cb_for_testing;
+using grpc_core::internal::alts_handshaker_client_set_fields_for_testing;
+using grpc_core::internal::alts_handshaker_client_set_recv_bytes_for_testing;
+using grpc_core::internal::alts_handshaker_client_set_vtable_for_testing;
+using grpc_core::internal::alts_tsi_handshaker_get_client_for_testing;
 using grpc_core::internal::alts_tsi_handshaker_get_is_client_for_testing;
-using grpc_core::internal::alts_tsi_handshaker_get_recv_bytes_for_testing;
-using grpc_core::internal::alts_tsi_handshaker_set_client_for_testing;
-using grpc_core::internal::alts_tsi_handshaker_set_recv_bytes_for_testing;
+using grpc_core::internal::alts_tsi_handshaker_set_client_vtable_for_testing;
+static bool should_handshaker_client_api_succeed = true;
 
 /* ALTS mock notification. */
 typedef struct notification {
@@ -57,12 +63,6 @@ typedef struct notification {
   bool notified;
 } notification;
 
-/* ALTS mock handshaker client. */
-typedef struct alts_mock_handshaker_client {
-  alts_handshaker_client base;
-  bool used_for_success_test;
-} alts_mock_handshaker_client;
-
 /* Type of ALTS handshaker response. */
 typedef enum {
   INVALID,
@@ -73,10 +73,7 @@ typedef enum {
   SERVER_NEXT,
 } alts_handshaker_response_type;
 
-static alts_tsi_event* client_start_event;
-static alts_tsi_event* client_next_event;
-static alts_tsi_event* server_start_event;
-static alts_tsi_event* server_next_event;
+static alts_handshaker_client* cb_event = nullptr;
 static notification caller_to_tsi_notification;
 static notification tsi_to_caller_notification;
 
@@ -311,89 +308,69 @@ static void on_server_next_success_cb(tsi_result status, void* user_data,
   signal(&tsi_to_caller_notification);
 }
 
-static tsi_result mock_client_start(alts_handshaker_client* self,
-                                    alts_tsi_event* event) {
-  alts_mock_handshaker_client* client =
-      reinterpret_cast<alts_mock_handshaker_client*>(self);
-  if (!client->used_for_success_test) {
-    alts_tsi_event_destroy(event);
+static tsi_result mock_client_start(alts_handshaker_client* client) {
+  if (!should_handshaker_client_api_succeed) {
     return TSI_INTERNAL_ERROR;
   }
-  GPR_ASSERT(event->cb == on_client_start_success_cb);
-  GPR_ASSERT(event->user_data == nullptr);
-  GPR_ASSERT(!alts_tsi_handshaker_get_has_sent_start_message_for_testing(
-      event->handshaker));
+  alts_handshaker_client_check_fields_for_testing(
+      client, on_client_start_success_cb, nullptr, false, nullptr);
   /* Populate handshaker response for client_start request. */
-  event->recv_buffer = generate_handshaker_response(CLIENT_START);
-  client_start_event = event;
+  grpc_byte_buffer** recv_buffer_ptr =
+      alts_handshaker_client_get_recv_buffer_addr_for_testing(client);
+  *recv_buffer_ptr = generate_handshaker_response(CLIENT_START);
+  cb_event = client;
   signal(&caller_to_tsi_notification);
   return TSI_OK;
 }
 
 static void mock_shutdown(alts_handshaker_client* self) {}
 
-static tsi_result mock_server_start(alts_handshaker_client* self,
-                                    alts_tsi_event* event,
+static tsi_result mock_server_start(alts_handshaker_client* client,
                                     grpc_slice* bytes_received) {
-  alts_mock_handshaker_client* client =
-      reinterpret_cast<alts_mock_handshaker_client*>(self);
-  if (!client->used_for_success_test) {
-    alts_tsi_event_destroy(event);
+  if (!should_handshaker_client_api_succeed) {
     return TSI_INTERNAL_ERROR;
   }
-  GPR_ASSERT(event->cb == on_server_start_success_cb);
-  GPR_ASSERT(event->user_data == nullptr);
+  alts_handshaker_client_check_fields_for_testing(
+      client, on_server_start_success_cb, nullptr, false, nullptr);
   grpc_slice slice = grpc_empty_slice();
   GPR_ASSERT(grpc_slice_cmp(*bytes_received, slice) == 0);
-  GPR_ASSERT(!alts_tsi_handshaker_get_has_sent_start_message_for_testing(
-      event->handshaker));
   /* Populate handshaker response for server_start request. */
-  event->recv_buffer = generate_handshaker_response(SERVER_START);
-  server_start_event = event;
+  grpc_byte_buffer** recv_buffer_ptr =
+      alts_handshaker_client_get_recv_buffer_addr_for_testing(client);
+  *recv_buffer_ptr = generate_handshaker_response(SERVER_START);
+  cb_event = client;
   grpc_slice_unref(slice);
   signal(&caller_to_tsi_notification);
   return TSI_OK;
 }
 
-static tsi_result mock_next(alts_handshaker_client* self, alts_tsi_event* event,
+static tsi_result mock_next(alts_handshaker_client* client,
                             grpc_slice* bytes_received) {
-  alts_mock_handshaker_client* client =
-      reinterpret_cast<alts_mock_handshaker_client*>(self);
-  if (!client->used_for_success_test) {
-    alts_tsi_event_destroy(event);
+  if (!should_handshaker_client_api_succeed) {
     return TSI_INTERNAL_ERROR;
   }
-  bool is_client =
-      alts_tsi_handshaker_get_is_client_for_testing(event->handshaker);
-  if (is_client) {
-    GPR_ASSERT(event->cb == on_client_next_success_cb);
-  } else {
-    GPR_ASSERT(event->cb == on_server_next_success_cb);
-  }
-  GPR_ASSERT(event->user_data == nullptr);
+  alts_tsi_handshaker* handshaker =
+      alts_handshaker_client_get_handshaker_for_testing(client);
+  bool is_client = alts_tsi_handshaker_get_is_client_for_testing(handshaker);
+  tsi_handshaker_on_next_done_cb cb =
+      is_client ? on_client_next_success_cb : on_server_next_success_cb;
+  alts_handshaker_client_set_cb_for_testing(client, cb);
+  alts_handshaker_client_set_recv_bytes_for_testing(client, bytes_received);
+  alts_handshaker_client_check_fields_for_testing(client, cb, nullptr, true,
+                                                  bytes_received);
   GPR_ASSERT(bytes_received != nullptr);
   GPR_ASSERT(memcmp(GRPC_SLICE_START_PTR(*bytes_received),
                     ALTS_TSI_HANDSHAKER_TEST_RECV_BYTES,
                     GRPC_SLICE_LENGTH(*bytes_received)) == 0);
-  GPR_ASSERT(grpc_slice_cmp(alts_tsi_handshaker_get_recv_bytes_for_testing(
-                                event->handshaker),
-                            *bytes_received) == 0);
-  GPR_ASSERT(alts_tsi_handshaker_get_has_sent_start_message_for_testing(
-      event->handshaker));
   /* Populate handshaker response for next request. */
   grpc_slice out_frame =
       grpc_slice_from_static_string(ALTS_TSI_HANDSHAKER_TEST_OUT_FRAME);
-  if (is_client) {
-    event->recv_buffer = generate_handshaker_response(CLIENT_NEXT);
-  } else {
-    event->recv_buffer = generate_handshaker_response(SERVER_NEXT);
-  }
-  alts_tsi_handshaker_set_recv_bytes_for_testing(event->handshaker, &out_frame);
-  if (is_client) {
-    client_next_event = event;
-  } else {
-    server_next_event = event;
-  }
+  grpc_byte_buffer** recv_buffer_ptr =
+      alts_handshaker_client_get_recv_buffer_addr_for_testing(client);
+  *recv_buffer_ptr = is_client ? generate_handshaker_response(CLIENT_NEXT)
+                               : generate_handshaker_response(SERVER_NEXT);
+  alts_handshaker_client_set_recv_bytes_for_testing(client, &out_frame);
+  cb_event = client;
   signal(&caller_to_tsi_notification);
   grpc_slice_unref(out_frame);
   return TSI_OK;
@@ -401,38 +378,27 @@ static tsi_result mock_next(alts_handshaker_client* self, alts_tsi_event* event,
 
 static void mock_destruct(alts_handshaker_client* client) {}
 
-static const alts_handshaker_client_vtable vtable = {
-    mock_client_start, mock_server_start, mock_next, mock_shutdown,
-    mock_destruct};
-
-static alts_handshaker_client* alts_mock_handshaker_client_create(
-    bool used_for_success_test) {
-  alts_mock_handshaker_client* client =
-      static_cast<alts_mock_handshaker_client*>(gpr_zalloc(sizeof(*client)));
-  client->base.vtable = &vtable;
-  client->used_for_success_test = used_for_success_test;
-  return &client->base;
-}
+static alts_handshaker_client_vtable vtable = {mock_client_start,
+                                               mock_server_start, mock_next,
+                                               mock_shutdown, mock_destruct};
 
-static tsi_handshaker* create_test_handshaker(bool used_for_success_test,
-                                              bool is_client) {
+static tsi_handshaker* create_test_handshaker(bool is_client) {
   tsi_handshaker* handshaker = nullptr;
-  alts_handshaker_client* client =
-      alts_mock_handshaker_client_create(used_for_success_test);
   grpc_alts_credentials_options* options =
       grpc_alts_credentials_client_options_create();
-  alts_tsi_handshaker_create(options, "target_name", "lame", is_client, nullptr,
-                             &handshaker);
+  alts_tsi_handshaker_create(options, "target_name",
+                             ALTS_HANDSHAKER_SERVICE_URL_FOR_TESTING, is_client,
+                             nullptr, &handshaker);
   alts_tsi_handshaker* alts_handshaker =
       reinterpret_cast<alts_tsi_handshaker*>(handshaker);
-  alts_tsi_handshaker_set_client_for_testing(alts_handshaker, client);
+  alts_tsi_handshaker_set_client_vtable_for_testing(alts_handshaker, &vtable);
   grpc_alts_credentials_options_destroy(options);
   return handshaker;
 }
 
 static void check_handshaker_next_invalid_input() {
   /* Initialization. */
-  tsi_handshaker* handshaker = create_test_handshaker(true, true);
+  tsi_handshaker* handshaker = create_test_handshaker(true);
   /* Check nullptr handshaker. */
   GPR_ASSERT(tsi_handshaker_next(nullptr, nullptr, 0, nullptr, nullptr, nullptr,
                                  check_must_not_be_called,
@@ -447,8 +413,7 @@ static void check_handshaker_next_invalid_input() {
 
 static void check_handshaker_shutdown_invalid_input() {
   /* Initialization. */
-  tsi_handshaker* handshaker = create_test_handshaker(
-      false /* used_for_success_test */, true /* is_client */);
+  tsi_handshaker* handshaker = create_test_handshaker(true /* is_client */);
   /* Check nullptr handshaker. */
   tsi_handshaker_shutdown(nullptr);
   /* Cleanup. */
@@ -460,10 +425,10 @@ static void check_handshaker_next_success() {
    * Create handshakers for which internal mock client is going to do
    * correctness check.
    */
-  tsi_handshaker* client_handshaker = create_test_handshaker(
-      true /* used_for_success_test */, true /* is_client */);
-  tsi_handshaker* server_handshaker = create_test_handshaker(
-      true /* used_for_success_test */, false /* is_client */);
+  tsi_handshaker* client_handshaker =
+      create_test_handshaker(true /* is_client */);
+  tsi_handshaker* server_handshaker =
+      create_test_handshaker(false /* is_client */);
   /* Client start. */
   GPR_ASSERT(tsi_handshaker_next(client_handshaker, nullptr, 0, nullptr,
                                  nullptr, nullptr, on_client_start_success_cb,
@@ -494,9 +459,7 @@ static void check_handshaker_next_success() {
 }
 
 static void check_handshaker_next_with_shutdown() {
-  /* Initialization. */
-  tsi_handshaker* handshaker = create_test_handshaker(
-      true /* used_for_success_test */, true /* is_client*/);
+  tsi_handshaker* handshaker = create_test_handshaker(true /* is_client*/);
   /* next(success) -- shutdown(success) -- next (fail) */
   GPR_ASSERT(tsi_handshaker_next(handshaker, nullptr, 0, nullptr, nullptr,
                                  nullptr, on_client_start_success_cb,
@@ -514,20 +477,18 @@ static void check_handshaker_next_with_shutdown() {
 }
 
 static void check_handle_response_with_shutdown(void* unused) {
-  /* Client start. */
   wait(&caller_to_tsi_notification);
-  alts_tsi_event_dispatch_to_handshaker(client_start_event, true /* is_ok */);
-  alts_tsi_event_destroy(client_start_event);
+  alts_handshaker_client_handle_response(cb_event, true /* is_ok */);
 }
 
 static void check_handshaker_next_failure() {
   /**
    * Create handshakers for which internal mock client is always going to fail.
    */
-  tsi_handshaker* client_handshaker = create_test_handshaker(
-      false /* used_for_success_test */, true /* is_client */);
-  tsi_handshaker* server_handshaker = create_test_handshaker(
-      false /* used_for_success_test */, false /* is_client */);
+  tsi_handshaker* client_handshaker =
+      create_test_handshaker(true /* is_client */);
+  tsi_handshaker* server_handshaker =
+      create_test_handshaker(false /* is_client */);
   /* Client start. */
   GPR_ASSERT(tsi_handshaker_next(client_handshaker, nullptr, 0, nullptr,
                                  nullptr, nullptr, check_must_not_be_called,
@@ -578,34 +539,46 @@ static void on_failed_grpc_call_cb(tsi_result status, void* user_data,
 }
 
 static void check_handle_response_invalid_input() {
+  /* Initialization. */
+  notification_init(&caller_to_tsi_notification);
+  notification_init(&tsi_to_caller_notification);
   /**
    * Create a handshaker at the client side, for which internal mock client is
    * always going to fail.
    */
-  tsi_handshaker* handshaker = create_test_handshaker(
-      false /* used_for_success_test */, true /* is_client */);
+  tsi_handshaker* handshaker = create_test_handshaker(true /* is_client */);
+  tsi_handshaker_next(handshaker, nullptr, 0, nullptr, nullptr, nullptr,
+                      on_client_start_success_cb, nullptr);
   alts_tsi_handshaker* alts_handshaker =
       reinterpret_cast<alts_tsi_handshaker*>(handshaker);
-  grpc_byte_buffer recv_buffer;
+  grpc_slice slice = grpc_empty_slice();
+  grpc_byte_buffer* recv_buffer = grpc_raw_byte_buffer_create(&slice, 1);
+  alts_handshaker_client* client =
+      alts_tsi_handshaker_get_client_for_testing(alts_handshaker);
   /* Check nullptr handshaker. */
-  alts_tsi_handshaker_handle_response(nullptr, &recv_buffer, GRPC_STATUS_OK,
-                                      nullptr, on_invalid_input_cb, nullptr,
-                                      true);
+  alts_handshaker_client_set_fields_for_testing(client, nullptr,
+                                                on_invalid_input_cb, nullptr,
+                                                recv_buffer, GRPC_STATUS_OK);
+  alts_handshaker_client_handle_response(client, true);
   /* Check nullptr recv_bytes. */
-  alts_tsi_handshaker_handle_response(alts_handshaker, nullptr, GRPC_STATUS_OK,
-                                      nullptr, on_invalid_input_cb, nullptr,
-                                      true);
+  alts_handshaker_client_set_fields_for_testing(client, alts_handshaker,
+                                                on_invalid_input_cb, nullptr,
+                                                nullptr, GRPC_STATUS_OK);
+  alts_handshaker_client_handle_response(client, true);
   /* Check failed grpc call made to handshaker service. */
-  alts_tsi_handshaker_handle_response(alts_handshaker, &recv_buffer,
-                                      GRPC_STATUS_UNKNOWN, nullptr,
-                                      on_failed_grpc_call_cb, nullptr, true);
-
-  alts_tsi_handshaker_handle_response(alts_handshaker, &recv_buffer,
-                                      GRPC_STATUS_OK, nullptr,
-                                      on_failed_grpc_call_cb, nullptr, false);
-
+  alts_handshaker_client_set_fields_for_testing(
+      client, alts_handshaker, on_failed_grpc_call_cb, nullptr, recv_buffer,
+      GRPC_STATUS_UNKNOWN);
+  alts_handshaker_client_handle_response(client, true);
+  alts_handshaker_client_set_fields_for_testing(client, alts_handshaker,
+                                                on_failed_grpc_call_cb, nullptr,
+                                                recv_buffer, GRPC_STATUS_OK);
+  alts_handshaker_client_handle_response(client, false);
   /* Cleanup. */
+  grpc_slice_unref(slice);
   tsi_handshaker_destroy(handshaker);
+  notification_destroy(&caller_to_tsi_notification);
+  notification_destroy(&tsi_to_caller_notification);
 }
 
 static void on_invalid_resp_cb(tsi_result status, void* user_data,
@@ -620,41 +593,45 @@ static void on_invalid_resp_cb(tsi_result status, void* user_data,
 }
 
 static void check_handle_response_invalid_resp() {
+  /* Initialization. */
+  notification_init(&caller_to_tsi_notification);
+  notification_init(&tsi_to_caller_notification);
   /**
    * Create a handshaker at the client side, for which internal mock client is
    * always going to fail.
    */
-  tsi_handshaker* handshaker = create_test_handshaker(
-      false /* used_for_success_test */, true /* is_client */);
+  tsi_handshaker* handshaker = create_test_handshaker(true /* is_client */);
+  tsi_handshaker_next(handshaker, nullptr, 0, nullptr, nullptr, nullptr,
+                      on_client_start_success_cb, nullptr);
   alts_tsi_handshaker* alts_handshaker =
       reinterpret_cast<alts_tsi_handshaker*>(handshaker);
+  alts_handshaker_client* client =
+      alts_tsi_handshaker_get_client_for_testing(alts_handshaker);
   /* Tests. */
   grpc_byte_buffer* recv_buffer = generate_handshaker_response(INVALID);
-  alts_tsi_handshaker_handle_response(alts_handshaker, recv_buffer,
-                                      GRPC_STATUS_OK, nullptr,
-                                      on_invalid_resp_cb, nullptr, true);
+  alts_handshaker_client_set_fields_for_testing(client, alts_handshaker,
+                                                on_invalid_resp_cb, nullptr,
+                                                recv_buffer, GRPC_STATUS_OK);
+  alts_handshaker_client_handle_response(client, true);
   /* Cleanup. */
-  grpc_byte_buffer_destroy(recv_buffer);
   tsi_handshaker_destroy(handshaker);
+  notification_destroy(&caller_to_tsi_notification);
+  notification_destroy(&tsi_to_caller_notification);
 }
 
 static void check_handle_response_success(void* unused) {
   /* Client start. */
   wait(&caller_to_tsi_notification);
-  alts_tsi_event_dispatch_to_handshaker(client_start_event, true /* is_ok */);
-  alts_tsi_event_destroy(client_start_event);
+  alts_handshaker_client_handle_response(cb_event, true /* is_ok */);
   /* Client next. */
   wait(&caller_to_tsi_notification);
-  alts_tsi_event_dispatch_to_handshaker(client_next_event, true /* is_ok */);
-  alts_tsi_event_destroy(client_next_event);
+  alts_handshaker_client_handle_response(cb_event, true /* is_ok */);
   /* Server start. */
   wait(&caller_to_tsi_notification);
-  alts_tsi_event_dispatch_to_handshaker(server_start_event, true /* is_ok */);
-  alts_tsi_event_destroy(server_start_event);
+  alts_handshaker_client_handle_response(cb_event, true /* is_ok */);
   /* Server next. */
   wait(&caller_to_tsi_notification);
-  alts_tsi_event_dispatch_to_handshaker(server_next_event, true /* is_ok */);
-  alts_tsi_event_destroy(server_next_event);
+  alts_handshaker_client_handle_response(cb_event, true /* is_ok */);
 }
 
 static void on_failed_resp_cb(tsi_result status, void* user_data,
@@ -669,22 +646,30 @@ static void on_failed_resp_cb(tsi_result status, void* user_data,
 }
 
 static void check_handle_response_failure() {
+  /* Initialization. */
+  notification_init(&caller_to_tsi_notification);
+  notification_init(&tsi_to_caller_notification);
   /**
    * Create a handshaker at the client side, for which internal mock client is
    * always going to fail.
    */
-  tsi_handshaker* handshaker = create_test_handshaker(
-      false /* used_for_success_test */, true /* is_client */);
+  tsi_handshaker* handshaker = create_test_handshaker(true /* is_client */);
+  tsi_handshaker_next(handshaker, nullptr, 0, nullptr, nullptr, nullptr,
+                      on_client_start_success_cb, nullptr);
   alts_tsi_handshaker* alts_handshaker =
       reinterpret_cast<alts_tsi_handshaker*>(handshaker);
+  alts_handshaker_client* client =
+      alts_tsi_handshaker_get_client_for_testing(alts_handshaker);
   /* Tests. */
   grpc_byte_buffer* recv_buffer = generate_handshaker_response(FAILED);
-  alts_tsi_handshaker_handle_response(alts_handshaker, recv_buffer,
-                                      GRPC_STATUS_OK, nullptr,
-                                      on_failed_resp_cb, nullptr, true);
-  grpc_byte_buffer_destroy(recv_buffer);
+  alts_handshaker_client_set_fields_for_testing(client, alts_handshaker,
+                                                on_failed_resp_cb, nullptr,
+                                                recv_buffer, GRPC_STATUS_OK);
+  alts_handshaker_client_handle_response(client, true /* is_ok*/);
   /* Cleanup. */
   tsi_handshaker_destroy(handshaker);
+  notification_destroy(&caller_to_tsi_notification);
+  notification_destroy(&tsi_to_caller_notification);
 }
 
 static void on_shutdown_resp_cb(tsi_result status, void* user_data,
@@ -699,26 +684,38 @@ static void on_shutdown_resp_cb(tsi_result status, void* user_data,
 }
 
 static void check_handle_response_after_shutdown() {
-  tsi_handshaker* handshaker = create_test_handshaker(
-      true /* used_for_success_test */, true /* is_client */);
+  /* Initialization. */
+  notification_init(&caller_to_tsi_notification);
+  notification_init(&tsi_to_caller_notification);
+  tsi_handshaker* handshaker = create_test_handshaker(true /* is_client */);
+  tsi_handshaker_next(handshaker, nullptr, 0, nullptr, nullptr, nullptr,
+                      on_client_start_success_cb, nullptr);
   alts_tsi_handshaker* alts_handshaker =
       reinterpret_cast<alts_tsi_handshaker*>(handshaker);
+  alts_handshaker_client* client =
+      alts_tsi_handshaker_get_client_for_testing(alts_handshaker);
+  grpc_byte_buffer** recv_buffer_ptr =
+      alts_handshaker_client_get_recv_buffer_addr_for_testing(client);
+  grpc_byte_buffer_destroy(*recv_buffer_ptr);
+
   /* Tests. */
   tsi_handshaker_shutdown(handshaker);
   grpc_byte_buffer* recv_buffer = generate_handshaker_response(CLIENT_START);
-  alts_tsi_handshaker_handle_response(alts_handshaker, recv_buffer,
-                                      GRPC_STATUS_OK, nullptr,
-                                      on_shutdown_resp_cb, nullptr, true);
-  grpc_byte_buffer_destroy(recv_buffer);
+  alts_handshaker_client_set_fields_for_testing(client, alts_handshaker,
+                                                on_shutdown_resp_cb, nullptr,
+                                                recv_buffer, GRPC_STATUS_OK);
+  alts_handshaker_client_handle_response(client, true);
   /* Cleanup. */
   tsi_handshaker_destroy(handshaker);
+  notification_destroy(&caller_to_tsi_notification);
+  notification_destroy(&tsi_to_caller_notification);
 }
 
 void check_handshaker_next_fails_after_shutdown() {
   /* Initialization. */
   notification_init(&caller_to_tsi_notification);
   notification_init(&tsi_to_caller_notification);
-  client_start_event = nullptr;
+  cb_event = nullptr;
   /* Tests. */
   grpc_core::Thread thd("alts_tsi_handshaker_test",
                         &check_handle_response_with_shutdown, nullptr);
@@ -734,10 +731,6 @@ void check_handshaker_success() {
   /* Initialization. */
   notification_init(&caller_to_tsi_notification);
   notification_init(&tsi_to_caller_notification);
-  client_start_event = nullptr;
-  client_next_event = nullptr;
-  server_start_event = nullptr;
-  server_next_event = nullptr;
   /* Tests. */
   grpc_core::Thread thd("alts_tsi_handshaker_test",
                         &check_handle_response_success, nullptr);
@@ -752,17 +745,21 @@ void check_handshaker_success() {
 int main(int argc, char** argv) {
   /* Initialization. */
   grpc_init();
+  grpc_alts_shared_resource_dedicated_init();
   /* Tests. */
+  should_handshaker_client_api_succeed = true;
   check_handshaker_success();
   check_handshaker_next_invalid_input();
-  check_handshaker_shutdown_invalid_input();
   check_handshaker_next_fails_after_shutdown();
+  check_handle_response_after_shutdown();
+  should_handshaker_client_api_succeed = false;
+  check_handshaker_shutdown_invalid_input();
   check_handshaker_next_failure();
   check_handle_response_invalid_input();
   check_handle_response_invalid_resp();
   check_handle_response_failure();
-  check_handle_response_after_shutdown();
   /* Cleanup. */
+  grpc_alts_shared_resource_dedicated_shutdown();
   grpc_shutdown();
   return 0;
 }