2 files changed, 142 insertions, 128 deletions

diff --git a/src/core/lib/channel/handshaker.c b/src/core/lib/channel/handshaker.c
index a45a39981c..89130e9951 100644
--- a/src/core/lib/channel/handshaker.c
+++ b/src/core/lib/channel/handshaker.c
@@ -38,67 +38,63 @@
 
 #include "src/core/lib/channel/channel_args.h"
 #include "src/core/lib/channel/handshaker.h"
+#include "src/core/lib/iomgr/timer.h"
 
 //
 // grpc_handshaker
 //
 
-void grpc_handshaker_init(const struct grpc_handshaker_vtable* vtable,
+void grpc_handshaker_init(const grpc_handshaker_vtable* vtable,
                           grpc_handshaker* handshaker) {
   handshaker->vtable = vtable;
 }
 
-void grpc_handshaker_destroy(grpc_exec_ctx* exec_ctx,
-                             grpc_handshaker* handshaker) {
+static void grpc_handshaker_destroy(grpc_exec_ctx* exec_ctx,
+                                    grpc_handshaker* handshaker) {
   handshaker->vtable->destroy(exec_ctx, handshaker);
 }
 
-void grpc_handshaker_shutdown(grpc_exec_ctx* exec_ctx,
-                              grpc_handshaker* handshaker) {
+static void grpc_handshaker_shutdown(grpc_exec_ctx* exec_ctx,
+                                     grpc_handshaker* handshaker) {
   handshaker->vtable->shutdown(exec_ctx, handshaker);
 }
 
-void grpc_handshaker_do_handshake(grpc_exec_ctx* exec_ctx,
-                                  grpc_handshaker* handshaker,
-                                  grpc_endpoint* endpoint,
-                                  grpc_channel_args* args,
-                                  grpc_slice_buffer* read_buffer,
-                                  gpr_timespec deadline,
-                                  grpc_tcp_server_acceptor* acceptor,
-                                  grpc_handshaker_done_cb cb, void* user_data) {
-  handshaker->vtable->do_handshake(exec_ctx, handshaker, endpoint, args,
-                                   read_buffer, deadline, acceptor, cb,
-                                   user_data);
+static void grpc_handshaker_do_handshake(grpc_exec_ctx* exec_ctx,
+                                         grpc_handshaker* handshaker,
+                                         grpc_tcp_server_acceptor* acceptor,
+                                         grpc_closure* on_handshake_done,
+                                         grpc_handshaker_args* args) {
+  handshaker->vtable->do_handshake(exec_ctx, handshaker, acceptor,
+                                   on_handshake_done, args);
 }
 
 //
 // grpc_handshake_manager
 //
 
-// State used while chaining handshakers.
-struct grpc_handshaker_state {
-  // The index of the handshaker to invoke next.
-  size_t index;
-  // The deadline for all handshakers.
-  gpr_timespec deadline;
-  // The acceptor to call the handshakers with.
-  grpc_tcp_server_acceptor* acceptor;
-  // The final callback and user_data to invoke after the last handshaker.
-  grpc_handshaker_done_cb final_cb;
-  void* final_user_data;
-};
-
 struct grpc_handshake_manager {
+  gpr_mu mu;
+  gpr_refcount refs;
   // An array of handshakers added via grpc_handshake_manager_add().
   size_t count;
   grpc_handshaker** handshakers;
-  // State used while chaining handshakers.
-  struct grpc_handshaker_state* state;
+  // The index of the handshaker to invoke next and closure to invoke it.
+  size_t index;
+  grpc_closure call_next_handshaker;
+  // The acceptor to call the handshakers with.
+  grpc_tcp_server_acceptor* acceptor;
+  // Deadline timer across all handshakers.
+  grpc_timer deadline_timer;
+  // The final callback and user_data to invoke after the last handshaker.
+  grpc_closure on_handshake_done;
+  void* user_data;
 };
 
 grpc_handshake_manager* grpc_handshake_manager_create() {
   grpc_handshake_manager* mgr = gpr_malloc(sizeof(grpc_handshake_manager));
   memset(mgr, 0, sizeof(*mgr));
+  gpr_mu_init(&mgr->mu);
+  gpr_ref_init(&mgr->refs, 1);
   return mgr;
 }
 
@@ -106,6 +102,7 @@ static bool is_power_of_2(size_t n) { return (n & (n - 1)) == 0; }
 
 void grpc_handshake_manager_add(grpc_handshake_manager* mgr,
                                 grpc_handshaker* handshaker) {
+  gpr_mu_lock(&mgr->mu);
   // To avoid allocating memory for each handshaker we add, we double
   // the number of elements every time we need more.
   size_t realloc_count = 0;
@@ -119,85 +116,119 @@ void grpc_handshake_manager_add(grpc_handshake_manager* mgr,
         gpr_realloc(mgr->handshakers, realloc_count * sizeof(grpc_handshaker*));
   }
   mgr->handshakers[mgr->count++] = handshaker;
+  gpr_mu_unlock(&mgr->mu);
+}
+
+static void grpc_handshake_manager_unref(grpc_exec_ctx* exec_ctx,
+                                         grpc_handshake_manager* mgr) {
+  if (gpr_unref(&mgr->refs)) {
+    for (size_t i = 0; i < mgr->count; ++i) {
+      grpc_handshaker_destroy(exec_ctx, mgr->handshakers[i]);
+    }
+    gpr_free(mgr->handshakers);
+    gpr_mu_destroy(&mgr->mu);
+    gpr_free(mgr);
+  }
 }
 
 void grpc_handshake_manager_destroy(grpc_exec_ctx* exec_ctx,
                                     grpc_handshake_manager* mgr) {
-  for (size_t i = 0; i < mgr->count; ++i) {
-    grpc_handshaker_destroy(exec_ctx, mgr->handshakers[i]);
-  }
-  gpr_free(mgr->handshakers);
-  gpr_free(mgr);
+  grpc_handshake_manager_unref(exec_ctx, mgr);
 }
 
 void grpc_handshake_manager_shutdown(grpc_exec_ctx* exec_ctx,
                                      grpc_handshake_manager* mgr) {
+  gpr_mu_lock(&mgr->mu);
   for (size_t i = 0; i < mgr->count; ++i) {
     grpc_handshaker_shutdown(exec_ctx, mgr->handshakers[i]);
   }
-  if (mgr->state != NULL) {
-    gpr_free(mgr->state);
-    mgr->state = NULL;
-  }
+  gpr_mu_unlock(&mgr->mu);
 }
 
-// A function used as the handshaker-done callback when chaining
-// handshakers together.
-static void call_next_handshaker(grpc_exec_ctx* exec_ctx,
-                                 grpc_endpoint* endpoint,
-                                 grpc_channel_args* args,
-                                 grpc_slice_buffer* read_buffer,
-                                 void* user_data, grpc_error* error) {
-  grpc_handshake_manager* mgr = user_data;
-  GPR_ASSERT(mgr->state != NULL);
-  GPR_ASSERT(mgr->state->index < mgr->count);
+static void call_next_handshaker(grpc_exec_ctx* exec_ctx, void* arg,
+                                 grpc_error* error);
+
+// Helper function to call either the next handshaker or the
+// on_handshake_done callback.
+static void call_next_handshaker_locked(grpc_exec_ctx* exec_ctx,
+                                        grpc_handshake_manager* mgr,
+                                        grpc_handshaker_args* args,
+                                        grpc_error* error) {
+  GPR_ASSERT(mgr->index <= mgr->count);
   // If we got an error, skip all remaining handshakers and invoke the
   // caller-supplied callback immediately.
-  if (error != GRPC_ERROR_NONE) {
-    mgr->state->final_cb(exec_ctx, endpoint, args, read_buffer,
-                         mgr->state->final_user_data, error);
+  // Otherwise, if this is the last handshaker, then call the on_handshake_done
+  // callback instead of chaining back to this function again.
+  if (error != GRPC_ERROR_NONE || mgr->index == mgr->count) {
+    // Cancel deadline timer, since we're invoking the on_handshake_done
+    // callback now.
+    grpc_timer_cancel(exec_ctx, &mgr->deadline_timer);
+    args->user_data = mgr->user_data;
+    grpc_exec_ctx_sched(exec_ctx, &mgr->on_handshake_done,
+                        GRPC_ERROR_REF(error), NULL);
+    // Since we're invoking the final callback, we won't be coming back
+    // to this function, so we can release our reference to the
+    // handshake manager.
+    grpc_handshake_manager_unref(exec_ctx, mgr);
     return;
   }
-  grpc_handshaker_done_cb cb = call_next_handshaker;
-  // If this is the last handshaker, use the caller-supplied callback
-  // and user_data instead of chaining back to this function again.
-  if (mgr->state->index == mgr->count - 1) {
-    cb = mgr->state->final_cb;
-    user_data = mgr->state->final_user_data;
-  }
-  // Invoke handshaker.
-  grpc_handshaker_do_handshake(
-      exec_ctx, mgr->handshakers[mgr->state->index], endpoint, args,
-      read_buffer, mgr->state->deadline, mgr->state->acceptor, cb, user_data);
-  ++mgr->state->index;
-  // If this is the last handshaker, clean up state.
-  if (mgr->state->index == mgr->count) {
-    gpr_free(mgr->state);
-    mgr->state = NULL;
+  // Call the next handshaker.
+  grpc_handshaker_do_handshake(exec_ctx, mgr->handshakers[mgr->index],
+                               mgr->acceptor, &mgr->call_next_handshaker, args);
+  ++mgr->index;
+}
+
+// A function used as the handshaker-done callback when chaining
+// handshakers together.
+static void call_next_handshaker(grpc_exec_ctx* exec_ctx, void* arg,
+                                 grpc_error* error) {
+  grpc_handshaker_args* args = arg;
+  grpc_handshake_manager* mgr = args->user_data;
+  gpr_mu_lock(&mgr->mu);
+  call_next_handshaker_locked(exec_ctx, mgr, args, error);
+  gpr_mu_unlock(&mgr->mu);
+}
+
+// Callback invoked when deadline is exceeded.
+static void on_timeout(grpc_exec_ctx* exec_ctx, void* arg, grpc_error* error) {
+  grpc_handshake_manager* mgr = arg;
+  if (error == GRPC_ERROR_NONE) {  // Timer fired, rather than being cancelled.
+    grpc_handshake_manager_shutdown(exec_ctx, mgr);
   }
+  grpc_handshake_manager_unref(exec_ctx, mgr);
 }
 
 void grpc_handshake_manager_do_handshake(
     grpc_exec_ctx* exec_ctx, grpc_handshake_manager* mgr,
-    grpc_endpoint* endpoint, const grpc_channel_args* args,
+    grpc_endpoint* endpoint, const grpc_channel_args* channel_args,
     gpr_timespec deadline, grpc_tcp_server_acceptor* acceptor,
-    grpc_handshaker_done_cb cb, void* user_data) {
-  grpc_channel_args* args_copy = grpc_channel_args_copy(args);
-  grpc_slice_buffer* read_buffer = gpr_malloc(sizeof(*read_buffer));
-  grpc_slice_buffer_init(read_buffer);
-  if (mgr->count == 0) {
-    // No handshakers registered, so we just immediately call the done
-    // callback with the passed-in endpoint.
-    cb(exec_ctx, endpoint, args_copy, read_buffer, user_data, GRPC_ERROR_NONE);
-  } else {
-    GPR_ASSERT(mgr->state == NULL);
-    mgr->state = gpr_malloc(sizeof(struct grpc_handshaker_state));
-    memset(mgr->state, 0, sizeof(*mgr->state));
-    mgr->state->deadline = deadline;
-    mgr->state->acceptor = acceptor;
-    mgr->state->final_cb = cb;
-    mgr->state->final_user_data = user_data;
-    call_next_handshaker(exec_ctx, endpoint, args_copy, read_buffer, mgr,
-                         GRPC_ERROR_NONE);
-  }
+    grpc_iomgr_cb_func on_handshake_done, void* user_data) {
+  // Construct handshaker args.  These will be passed through all
+  // handshakers and eventually be freed by the on_handshake_done callback.
+  grpc_handshaker_args* args = gpr_malloc(sizeof(*args));
+  args->endpoint = endpoint;
+  args->args = grpc_channel_args_copy(channel_args);
+  args->read_buffer = gpr_malloc(sizeof(*args->read_buffer));
+  grpc_slice_buffer_init(args->read_buffer);
+  // Initialize state needed for calling handshakers.
+  gpr_mu_lock(&mgr->mu);
+  GPR_ASSERT(mgr->index == 0);
+  mgr->acceptor = acceptor;
+  grpc_closure_init(&mgr->call_next_handshaker, call_next_handshaker, args);
+  grpc_closure_init(&mgr->on_handshake_done, on_handshake_done, args);
+  // While chaining between handshakers, we use args->user_data to
+  // store a pointer to the handshake manager.  This will be
+  // changed to point to the caller-supplied user_data before calling
+  // the on_handshake_done callback.
+  args->user_data = mgr;
+  mgr->user_data = user_data;
+  // Start deadline timer, which owns a ref.
+  gpr_ref(&mgr->refs);
+  grpc_timer_init(exec_ctx, &mgr->deadline_timer,
+                  gpr_convert_clock_type(deadline, GPR_CLOCK_MONOTONIC),
+                  on_timeout, mgr, gpr_now(GPR_CLOCK_MONOTONIC));
+  // Start first handshaker, which also owns a ref.
+  gpr_ref(&mgr->refs);
+  call_next_handshaker_locked(exec_ctx, mgr, args, GRPC_ERROR_NONE);
+  gpr_mu_unlock(&mgr->mu);
 }
diff --git a/src/core/lib/channel/handshaker.h b/src/core/lib/channel/handshaker.h
index f8a36c6473..8cad81c444 100644
--- a/src/core/lib/channel/handshaker.h
+++ b/src/core/lib/channel/handshaker.h
@@ -54,15 +54,16 @@
 
 typedef struct grpc_handshaker grpc_handshaker;
 
-/// Callback type invoked when a handshaker is done.
-/// Takes ownership of \a args and \a read_buffer.
-typedef void (*grpc_handshaker_done_cb)(grpc_exec_ctx* exec_ctx,
-                                        grpc_endpoint* endpoint,
-                                        grpc_channel_args* args,
-                                        grpc_slice_buffer* read_buffer,
-                                        void* user_data, grpc_error* error);
-
-struct grpc_handshaker_vtable {
+/// Arguments passed through handshakers and to the on_handshake_done callback.
+/// All data members are owned by the struct.
+typedef struct {
+  grpc_endpoint* endpoint;
+  grpc_channel_args* args;
+  void* user_data;
+  grpc_slice_buffer* read_buffer;
+} grpc_handshaker_args;
+
+typedef struct {
   /// Destroys the handshaker.
   void (*destroy)(grpc_exec_ctx* exec_ctx, grpc_handshaker* handshaker);
 
@@ -70,44 +71,26 @@ struct grpc_handshaker_vtable {
   /// aborted in the middle).
   void (*shutdown)(grpc_exec_ctx* exec_ctx, grpc_handshaker* handshaker);
 
-  /// Performs handshaking.  When finished, calls \a cb with \a user_data.
-  /// Takes ownership of \a args.
-  /// Takes ownership of \a read_buffer, which contains leftover bytes read
-  /// from the endpoint by the previous handshaker.
+  /// Performs handshaking, modifying \a args as needed (e.g., to
+  /// replace \a endpoint with a wrapped endpoint).
+  /// When finished, invokes \a on_handshake_done.
   /// \a acceptor will be NULL for client-side handshakers.
   void (*do_handshake)(grpc_exec_ctx* exec_ctx, grpc_handshaker* handshaker,
-                       grpc_endpoint* endpoint, grpc_channel_args* args,
-                       grpc_slice_buffer* read_buffer, gpr_timespec deadline,
                        grpc_tcp_server_acceptor* acceptor,
-                       grpc_handshaker_done_cb cb, void* user_data);
-};
+                       grpc_closure* on_handshake_done,
+                       grpc_handshaker_args* args);
+} grpc_handshaker_vtable;
 
 /// Base struct.  To subclass, make this the first member of the
 /// implementation struct.
 struct grpc_handshaker {
-  const struct grpc_handshaker_vtable* vtable;
+  const grpc_handshaker_vtable* vtable;
 };
 
 /// Called by concrete implementations to initialize the base struct.
-void grpc_handshaker_init(const struct grpc_handshaker_vtable* vtable,
+void grpc_handshaker_init(const grpc_handshaker_vtable* vtable,
                           grpc_handshaker* handshaker);
 
-/// Convenient wrappers for invoking methods via the vtable.
-/// These probably do not need to be called from anywhere but
-/// grpc_handshake_manager.
-void grpc_handshaker_destroy(grpc_exec_ctx* exec_ctx,
-                             grpc_handshaker* handshaker);
-void grpc_handshaker_shutdown(grpc_exec_ctx* exec_ctx,
-                              grpc_handshaker* handshaker);
-void grpc_handshaker_do_handshake(grpc_exec_ctx* exec_ctx,
-                                  grpc_handshaker* handshaker,
-                                  grpc_endpoint* endpoint,
-                                  grpc_channel_args* args,
-                                  grpc_slice_buffer* read_buffer,
-                                  gpr_timespec deadline,
-                                  grpc_tcp_server_acceptor* acceptor,
-                                  grpc_handshaker_done_cb cb, void* user_data);
-
 ///
 /// grpc_handshake_manager
 ///
@@ -137,12 +120,12 @@ void grpc_handshake_manager_shutdown(grpc_exec_ctx* exec_ctx,
 /// Does NOT take ownership of \a args.  Instead, makes a copy before
 /// invoking the first handshaker.
 /// \a acceptor will be NULL for client-side handshakers.
-/// Invokes \a cb with \a user_data after either a handshaker fails or
-/// all handshakers have completed successfully.
+/// When done, invokes \a on_handshake_done with an argument of a
+/// grpc_handshaker_args object, which the callback takes ownership of.
 void grpc_handshake_manager_do_handshake(
     grpc_exec_ctx* exec_ctx, grpc_handshake_manager* mgr,
-    grpc_endpoint* endpoint, const grpc_channel_args* args,
+    grpc_endpoint* endpoint, const grpc_channel_args* channel_args,
     gpr_timespec deadline, grpc_tcp_server_acceptor* acceptor,
-    grpc_handshaker_done_cb cb, void* user_data);
+    grpc_iomgr_cb_func on_handshake_done, void* user_data);
 
 #endif /* GRPC_CORE_LIB_CHANNEL_HANDSHAKER_H */