Merge branch 'security_handshaker1' into security_handshaker2

2 files changed, 55 insertions, 49 deletions

diff --git a/src/core/ext/client_channel/http_connect_handshaker.c b/src/core/ext/client_channel/http_connect_handshaker.c
index 61fec5cba9..6a34e390bc 100644
--- a/src/core/ext/client_channel/http_connect_handshaker.c
+++ b/src/core/ext/client_channel/http_connect_handshaker.c
@@ -106,6 +106,31 @@ static void cleanup_args_for_failure_locked(
   handshaker->args->args = NULL;
 }
 
+// If the handshake failed or we're shutting down, clean up and invoke the
+// callback with the error.
+static void handshake_failed_locked(grpc_exec_ctx* exec_ctx,
+                                    http_connect_handshaker* handshaker,
+                                    grpc_error* error) {
+  if (error == GRPC_ERROR_NONE) {
+    // If we were shut down after an endpoint operation succeeded but
+    // before the endpoint callback was invoked, we need to generate our
+    // own error.
+    error = GRPC_ERROR_CREATE("Handshaker shutdown");
+  }
+  if (!handshaker->shutdown) {
+    // TODO(ctiller): It is currently necessary to shutdown endpoints
+    // before destroying them, even if we know that there are no
+    // pending read/write callbacks.  This should be fixed, at which
+    // point this can be removed.
+    grpc_endpoint_shutdown(exec_ctx, handshaker->args->endpoint);
+    // Not shutting down, so the handshake failed.  Clean up before
+    // invoking the callback.
+    cleanup_args_for_failure_locked(handshaker);
+  }
+  // Invoke callback.
+  grpc_exec_ctx_sched(exec_ctx, handshaker->on_handshake_done, error, NULL);
+}
+
 // Callback invoked when finished writing HTTP CONNECT request.
 static void on_write_done(grpc_exec_ctx* exec_ctx, void* arg,
                           grpc_error* error) {
@@ -114,25 +139,7 @@ static void on_write_done(grpc_exec_ctx* exec_ctx, void* arg,
   if (error != GRPC_ERROR_NONE || handshaker->shutdown) {
     // If the write failed or we're shutting down, clean up and invoke the
     // callback with the error.
-    if (error == GRPC_ERROR_NONE) {
-      // If we were shut down after the write succeeded but before this
-      // callback was invoked, we need to generate our own error.
-      error = GRPC_ERROR_CREATE("Handshaker shutdown");
-    } else {
-      GRPC_ERROR_REF(error);  // Take ref for the handshake-done callback.
-    }
-    if (!handshaker->shutdown) {
-      // TODO(ctiller): It is currently necessary to shutdown endpoints
-      // before destroying them, even if we know that there are no
-      // pending read/write callbacks.  This should be fixed, at which
-      // point this can be removed.
-      grpc_endpoint_shutdown(exec_ctx, handshaker->args->endpoint);
-      // Not shutting down, so the write failed.  Clean up before
-      // invoking the callback.
-      cleanup_args_for_failure_locked(handshaker);
-    }
-    // Invoke callback.
-    grpc_exec_ctx_sched(exec_ctx, handshaker->on_handshake_done, error, NULL);
+    handshake_failed_locked(exec_ctx, handshaker, GRPC_ERROR_REF(error));
     gpr_mu_unlock(&handshaker->mu);
     http_connect_handshaker_unref(exec_ctx, handshaker);
   } else {
@@ -151,25 +158,9 @@ static void on_read_done(grpc_exec_ctx* exec_ctx, void* arg,
   http_connect_handshaker* handshaker = arg;
   gpr_mu_lock(&handshaker->mu);
   if (error != GRPC_ERROR_NONE || handshaker->shutdown) {
-    // If the write failed or we're shutting down, clean up and invoke the
+    // If the read failed or we're shutting down, clean up and invoke the
     // callback with the error.
-    if (error == GRPC_ERROR_NONE) {
-      // If we were shut down after the write succeeded but before this
-      // callback was invoked, we need to generate our own error.
-      error = GRPC_ERROR_CREATE("Handshaker shutdown");
-    } else {
-      GRPC_ERROR_REF(error);  // Take ref for the handshake-done callback.
-    }
-    if (!handshaker->shutdown) {
-      // TODO(ctiller): It is currently necessary to shutdown endpoints
-      // before destroying them, even if we know that there are no
-      // pending read/write callbacks.  This should be fixed, at which
-      // point this can be removed.
-      grpc_endpoint_shutdown(exec_ctx, handshaker->args->endpoint);
-      // Not shutting down, so the write failed.  Clean up before
-      // invoking the callback.
-      cleanup_args_for_failure_locked(handshaker);
-    }
+    handshake_failed_locked(exec_ctx, handshaker, GRPC_ERROR_REF(error));
     goto done;
   }
   // Add buffer to parser.
@@ -179,7 +170,10 @@ static void on_read_done(grpc_exec_ctx* exec_ctx, void* arg,
       error = grpc_http_parser_parse(&handshaker->http_parser,
                                      handshaker->args->read_buffer->slices[i],
                                      &body_start_offset);
-      if (error != GRPC_ERROR_NONE) goto done;
+      if (error != GRPC_ERROR_NONE) {
+        handshake_failed_locked(exec_ctx, handshaker, error);
+        goto done;
+      }
       if (handshaker->http_parser.state == GRPC_HTTP_BODY) {
         // Remove the data we've already read from the read buffer,
         // leaving only the leftover bytes (if any).
@@ -228,10 +222,12 @@ static void on_read_done(grpc_exec_ctx* exec_ctx, void* arg,
                  handshaker->http_response.status);
     error = GRPC_ERROR_CREATE(msg);
     gpr_free(msg);
+    handshake_failed_locked(exec_ctx, handshaker, error);
+    goto done;
   }
-done:
-  // Invoke handshake-done callback.
+  // Success.  Invoke handshake-done callback.
   grpc_exec_ctx_sched(exec_ctx, handshaker->on_handshake_done, error, NULL);
+done:
   gpr_mu_unlock(&handshaker->mu);
   http_connect_handshaker_unref(exec_ctx, handshaker);
 }
diff --git a/src/core/lib/channel/handshaker.c b/src/core/lib/channel/handshaker.c
index f3bd91284e..992cd420de 100644
--- a/src/core/lib/channel/handshaker.c
+++ b/src/core/lib/channel/handshaker.c
@@ -141,6 +141,7 @@ void grpc_handshake_manager_destroy(grpc_exec_ctx* exec_ctx,
 void grpc_handshake_manager_shutdown(grpc_exec_ctx* exec_ctx,
                                      grpc_handshake_manager* mgr) {
   gpr_mu_lock(&mgr->mu);
+  // Shutdown the handshaker that's currently in progress, if any.
   if (mgr->index > 0) {
     grpc_handshaker_shutdown(exec_ctx, mgr->handshakers[mgr->index - 1]);
   }
@@ -149,7 +150,8 @@ void grpc_handshake_manager_shutdown(grpc_exec_ctx* exec_ctx,
 
 // 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,
+// Returns true if we've scheduled the on_handshake_done callback.
+static bool call_next_handshaker_locked(grpc_exec_ctx* exec_ctx,
                                         grpc_handshake_manager* mgr,
                                         grpc_error* error) {
   GPR_ASSERT(mgr->index <= mgr->count);
@@ -160,17 +162,16 @@ static void call_next_handshaker_locked(grpc_exec_ctx* exec_ctx,
     // callback now.
     grpc_timer_cancel(exec_ctx, &mgr->deadline_timer);
     grpc_exec_ctx_sched(exec_ctx, &mgr->on_handshake_done, 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;
+    // Reset index to 0 so that we can start over if we re-attempt the
+    // connection.
+    mgr->index = 0;
+    return true;
   }
-  // Call the next handshaker.
   grpc_handshaker_do_handshake(exec_ctx, mgr->handshakers[mgr->index],
                                mgr->acceptor, &mgr->call_next_handshaker,
                                &mgr->args);
   ++mgr->index;
+  return false;
 }
 
 // A function used as the handshaker-done callback when chaining
@@ -179,8 +180,14 @@ static void call_next_handshaker(grpc_exec_ctx* exec_ctx, void* arg,
                                  grpc_error* error) {
   grpc_handshake_manager* mgr = arg;
   gpr_mu_lock(&mgr->mu);
-  call_next_handshaker_locked(exec_ctx, mgr, GRPC_ERROR_REF(error));
+  bool done = call_next_handshaker_locked(exec_ctx, mgr, GRPC_ERROR_REF(error));
   gpr_mu_unlock(&mgr->mu);
+  // If we're invoked the final callback, we won't be coming back
+  // to this function, so we can release our reference to the
+  // handshake manager.
+  if (done) {
+    grpc_handshake_manager_unref(exec_ctx, mgr);
+  }
 }
 
 // Callback invoked when deadline is exceeded.
@@ -217,6 +224,9 @@ void grpc_handshake_manager_do_handshake(
                   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, GRPC_ERROR_NONE);
+  bool done = call_next_handshaker_locked(exec_ctx, mgr, GRPC_ERROR_NONE);
   gpr_mu_unlock(&mgr->mu);
+  if (done) {
+    grpc_handshake_manager_unref(exec_ctx, mgr);
+  }
 }