Add tsi_handshaker_shutdown to TSI

2 files changed, 90 insertions, 1 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 b9dd52a64a..c8d88aa72c 100644
--- a/test/core/tsi/alts/handshaker/alts_handshaker_client_test.cc
+++ b/test/core/tsi/alts/handshaker/alts_handshaker_client_test.cc
@@ -326,6 +326,9 @@ static void schedule_request_invalid_arg_test() {
   GPR_ASSERT(alts_handshaker_client_next(nullptr, event, &config->out_frame) ==
              TSI_INVALID_ARGUMENT);
 
+  /* Check shutdown. */
+  alts_handshaker_client_shutdown(nullptr);
+
   /* Cleanup. */
   alts_tsi_event_destroy(event);
   destroy_config(config);
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 95724f84f4..85a58114ba 100644
--- a/test/core/tsi/alts/handshaker/alts_tsi_handshaker_test.cc
+++ b/test/core/tsi/alts/handshaker/alts_tsi_handshaker_test.cc
@@ -330,6 +330,8 @@ static tsi_result mock_client_start(alts_handshaker_client* self,
   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,
                                     grpc_slice* bytes_received) {
@@ -400,7 +402,8 @@ 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_destruct};
+    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) {
@@ -442,6 +445,16 @@ static void check_handshaker_next_invalid_input() {
   tsi_handshaker_destroy(handshaker);
 }
 
+static void check_handshaker_shutdown_invalid_input() {
+  /* Initialization. */
+  tsi_handshaker* handshaker = create_test_handshaker(
+      false /* used_for_success_test */, true /* is_client */);
+  /* Check nullptr handshaker. */
+  tsi_handshaker_shutdown(nullptr);
+  /* Cleanup. */
+  tsi_handshaker_destroy(handshaker);
+}
+
 static void check_handshaker_next_success() {
   /**
    * Create handshakers for which internal mock client is going to do
@@ -480,6 +493,33 @@ static void check_handshaker_next_success() {
   tsi_handshaker_destroy(client_handshaker);
 }
 
+static void check_handshaker_next_with_shutdown() {
+  /* Initialization. */
+  tsi_handshaker* handshaker = create_test_handshaker(
+      true /* used_for_success_test */, 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,
+                                 nullptr) == TSI_ASYNC);
+  wait(&tsi_to_caller_notification);
+  tsi_handshaker_shutdown(handshaker);
+  GPR_ASSERT(tsi_handshaker_next(
+                 handshaker,
+                 (const unsigned char*)ALTS_TSI_HANDSHAKER_TEST_RECV_BYTES,
+                 strlen(ALTS_TSI_HANDSHAKER_TEST_RECV_BYTES), nullptr, nullptr,
+                 nullptr, on_client_next_success_cb,
+                 nullptr) == TSI_HANDSHAKE_SHUTDOWN);
+  /* Cleanup. */
+  tsi_handshaker_destroy(handshaker);
+}
+
+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);
+}
+
 static void check_handshaker_next_failure() {
   /**
    * Create handshakers for which internal mock client is always going to fail.
@@ -647,6 +687,49 @@ static void check_handle_response_failure() {
   tsi_handshaker_destroy(handshaker);
 }
 
+static void on_shutdown_resp_cb(tsi_result status, void* user_data,
+                                const unsigned char* bytes_to_send,
+                                size_t bytes_to_send_size,
+                                tsi_handshaker_result* result) {
+  GPR_ASSERT(status == TSI_HANDSHAKE_SHUTDOWN);
+  GPR_ASSERT(user_data == nullptr);
+  GPR_ASSERT(bytes_to_send == nullptr);
+  GPR_ASSERT(bytes_to_send_size == 0);
+  GPR_ASSERT(result == nullptr);
+}
+
+static void check_handle_response_after_shutdown() {
+  tsi_handshaker* handshaker = create_test_handshaker(
+      true /* used_for_success_test */, true /* is_client */);
+  alts_tsi_handshaker* alts_handshaker =
+      reinterpret_cast<alts_tsi_handshaker*>(handshaker);
+  /* 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);
+  /* Cleanup. */
+  tsi_handshaker_destroy(handshaker);
+}
+
+void check_handshaker_next_fails_after_shutdown() {
+  /* Initialization. */
+  notification_init(&caller_to_tsi_notification);
+  notification_init(&tsi_to_caller_notification);
+  client_start_event = nullptr;
+  /* Tests. */
+  grpc_core::Thread thd("alts_tsi_handshaker_test",
+                        &check_handle_response_with_shutdown, nullptr);
+  thd.Start();
+  check_handshaker_next_with_shutdown();
+  thd.Join();
+  /* Cleanup. */
+  notification_destroy(&caller_to_tsi_notification);
+  notification_destroy(&tsi_to_caller_notification);
+}
+
 void check_handshaker_success() {
   /* Initialization. */
   notification_init(&caller_to_tsi_notification);
@@ -672,10 +755,13 @@ int main(int argc, char** argv) {
   /* Tests. */
   check_handshaker_success();
   check_handshaker_next_invalid_input();
+  check_handshaker_shutdown_invalid_input();
+  check_handshaker_next_fails_after_shutdown();
   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_shutdown();
   return 0;