modify TSI handshake to correctly handle the case of TSI_INCOMPLETE_DATA

2 files changed, 22 insertions, 17 deletions

diff --git a/test/core/tsi/transport_security_test_lib.c b/test/core/tsi/transport_security_test_lib.c
index 329b2371bf..1f32eab21c 100644
--- a/test/core/tsi/transport_security_test_lib.c
+++ b/test/core/tsi/transport_security_test_lib.c
@@ -350,25 +350,28 @@ static void do_handshaker_next(handshaker_args *args) {
   tsi_handshaker_result *handshaker_result = NULL;
   unsigned char *bytes_to_send = NULL;
   size_t bytes_to_send_size = 0;
+  tsi_result result = TSI_OK;
   /* Receive data from peer, if available. */
-  size_t buf_size = args->handshake_buffer_size;
-  receive_bytes_from_peer(args->fixture, &args->handshake_buffer, &buf_size,
-                          args->is_client);
-  if (buf_size > 0) {
-    args->transferred_data = true;
-  }
-  /* Peform handshaker next. */
-  tsi_result result = tsi_handshaker_next(
-      handshaker, args->handshake_buffer, buf_size,
-      (const unsigned char **)&bytes_to_send, &bytes_to_send_size,
-      &handshaker_result, &on_handshake_next_done_wrapper, args);
-  if (result != TSI_ASYNC) {
-    args->error = on_handshake_next_done(result, args, bytes_to_send,
-                                         bytes_to_send_size, handshaker_result);
-    if (args->error != GRPC_ERROR_NONE) {
-      return;
+  do {
+    size_t buf_size = args->handshake_buffer_size;
+    receive_bytes_from_peer(args->fixture, &args->handshake_buffer, &buf_size,
+                            args->is_client);
+    if (buf_size > 0) {
+      args->transferred_data = true;
     }
-  }
+    /* Peform handshaker next. */
+    result = tsi_handshaker_next(handshaker, args->handshake_buffer, buf_size,
+                                 (const unsigned char **)&bytes_to_send,
+                                 &bytes_to_send_size, &handshaker_result,
+                                 &on_handshake_next_done_wrapper, args);
+    if (result != TSI_ASYNC) {
+      args->error = on_handshake_next_done(
+          result, args, bytes_to_send, bytes_to_send_size, handshaker_result);
+      if (args->error != GRPC_ERROR_NONE) {
+        return;
+      }
+    }
+  } while (result == TSI_INCOMPLETE_DATA);
   notification_wait(fixture);
 }
 
diff --git a/test/core/tsi/transport_security_test_lib.h b/test/core/tsi/transport_security_test_lib.h
index ed8ff856df..74f4378bc1 100644
--- a/test/core/tsi/transport_security_test_lib.h
+++ b/test/core/tsi/transport_security_test_lib.h
@@ -21,6 +21,8 @@
 
 #include "src/core/tsi/transport_security_interface.h"
 
+#include <grpc/support/sync.h>
+
 #ifdef __cplusplus
 extern "C" {
 #endif