Add cancellation to asynchronous security APIs.

3 files changed, 210 insertions, 62 deletions

diff --git a/src/core/lib/security/transport/client_auth_filter.c b/src/core/lib/security/transport/client_auth_filter.c
index 50a51b31cd..531a88434f 100644
--- a/src/core/lib/security/transport/client_auth_filter.c
+++ b/src/core/lib/security/transport/client_auth_filter.c
@@ -51,8 +51,15 @@ typedef struct {
   grpc_polling_entity *pollent;
   gpr_atm security_context_set;
   gpr_mu security_context_mu;
+  grpc_credentials_mdelem_array md_array;
   grpc_linked_mdelem md_links[MAX_CREDENTIALS_METADATA_COUNT];
   grpc_auth_metadata_context auth_md_context;
+  grpc_closure closure;
+  // Either 0 (no cancellation and no async operation in flight),
+  // a grpc_closure* (if the lowest bit is 0),
+  // or a grpc_error* (if the lowest bit is 1).
+  gpr_atm cancellation_state;
+  grpc_closure cancel_closure;
 } call_data;
 
 /* We can have a per-channel credentials. */
@@ -61,6 +68,43 @@ typedef struct {
   grpc_auth_context *auth_context;
 } channel_data;
 
+static void decode_cancel_state(gpr_atm cancel_state, grpc_closure **func,
+                                grpc_error **error) {
+  // If the lowest bit is 1, the value is a grpc_error*.
+  // Otherwise, if non-zdero, the value is a grpc_closure*.
+  if (cancel_state & 1) {
+    *error = (grpc_error *)(cancel_state & ~(gpr_atm)1);
+  } else if (cancel_state != 0) {
+    *func = (grpc_closure *)cancel_state;
+  }
+}
+
+static gpr_atm encode_cancel_state_error(grpc_error *error) {
+  // Set the lowest bit to 1 to indicate that it's an error.
+  return (gpr_atm)1 | (gpr_atm)error;
+}
+
+// Returns an error if the call has been cancelled.  Otherwise, sets the
+// cancellation function to be called upon cancellation.
+static grpc_error *set_cancel_func(grpc_call_element *elem,
+                                   grpc_iomgr_cb_func func) {
+  call_data *calld = (call_data *)elem->call_data;
+  // Decode original state.
+  gpr_atm original_state = gpr_atm_acq_load(&calld->cancellation_state);
+  grpc_error *original_error = GRPC_ERROR_NONE;
+  grpc_closure *original_func = NULL;
+  decode_cancel_state(original_state, &original_func, &original_error);
+  // If error is set, return it.
+  if (original_error != GRPC_ERROR_NONE) return GRPC_ERROR_REF(original_error);
+  // Otherwise, store func.
+  GRPC_CLOSURE_INIT(&calld->cancel_closure, func, elem,
+                    grpc_schedule_on_exec_ctx);
+  GPR_ASSERT(((gpr_atm)&calld->cancel_closure & (gpr_atm)1) == 0);
+  gpr_atm_rel_store(&calld->cancellation_state,
+                    (gpr_atm)&calld->cancel_closure);
+  return GRPC_ERROR_NONE;
+}
+
 static void reset_auth_metadata_context(
     grpc_auth_metadata_context *auth_md_context) {
   if (auth_md_context->service_url != NULL) {
@@ -86,41 +130,29 @@ static void add_error(grpc_error **combined, grpc_error *error) {
   *combined = grpc_error_add_child(*combined, error);
 }
 
-static void on_credentials_metadata(grpc_exec_ctx *exec_ctx, void *user_data,
-                                    grpc_credentials_md *md_elems,
-                                    size_t num_md,
-                                    grpc_credentials_status status,
-                                    const char *error_details) {
-  grpc_transport_stream_op_batch *batch =
-      (grpc_transport_stream_op_batch *)user_data;
+static void on_credentials_metadata(grpc_exec_ctx *exec_ctx, void *arg,
+                                    grpc_error *input_error) {
+  grpc_transport_stream_op_batch *batch = (grpc_transport_stream_op_batch *)arg;
   grpc_call_element *elem = batch->handler_private.extra_arg;
   call_data *calld = elem->call_data;
   reset_auth_metadata_context(&calld->auth_md_context);
-  grpc_error *error = GRPC_ERROR_NONE;
-  if (status != GRPC_CREDENTIALS_OK) {
-    error = grpc_error_set_int(
-        GRPC_ERROR_CREATE_FROM_COPIED_STRING(
-            error_details != NULL && strlen(error_details) > 0
-                ? error_details
-                : "Credentials failed to get metadata."),
-        GRPC_ERROR_INT_GRPC_STATUS, GRPC_STATUS_UNAUTHENTICATED);
-  } else {
-    GPR_ASSERT(num_md <= MAX_CREDENTIALS_METADATA_COUNT);
+  grpc_error *error = GRPC_ERROR_REF(input_error);
+  if (error == GRPC_ERROR_NONE) {
+    GPR_ASSERT(calld->md_array.size <= MAX_CREDENTIALS_METADATA_COUNT);
     GPR_ASSERT(batch->send_initial_metadata);
     grpc_metadata_batch *mdb =
         batch->payload->send_initial_metadata.send_initial_metadata;
-    for (size_t i = 0; i < num_md; i++) {
-      add_error(&error,
-                grpc_metadata_batch_add_tail(
-                    exec_ctx, mdb, &calld->md_links[i],
-                    grpc_mdelem_from_slices(
-                        exec_ctx, grpc_slice_ref_internal(md_elems[i].key),
-                        grpc_slice_ref_internal(md_elems[i].value))));
+    for (size_t i = 0; i < calld->md_array.size; ++i) {
+      add_error(&error, grpc_metadata_batch_add_tail(
+                            exec_ctx, mdb, &calld->md_links[i],
+                            GRPC_MDELEM_REF(calld->md_array.md[i])));
     }
   }
   if (error == GRPC_ERROR_NONE) {
     grpc_call_next_op(exec_ctx, elem, batch);
   } else {
+    error = grpc_error_set_int(error, GRPC_ERROR_INT_GRPC_STATUS,
+                               GRPC_STATUS_UNAUTHENTICATED);
     grpc_transport_stream_op_batch_finish_with_failure(exec_ctx, batch, error);
   }
 }
@@ -155,6 +187,14 @@ void build_auth_metadata_context(grpc_security_connector *sc,
   gpr_free(host);
 }
 
+static void cancel_get_request_metadata(grpc_exec_ctx *exec_ctx, void *arg,
+                                        grpc_error *error) {
+  grpc_call_element *elem = (grpc_call_element *)arg;
+  call_data *calld = (call_data *)elem->call_data;
+  grpc_call_credentials_cancel_get_request_metadata(
+      exec_ctx, calld->creds, &calld->md_array, GRPC_ERROR_REF(error));
+}
+
 static void send_security_metadata(grpc_exec_ctx *exec_ctx,
                                    grpc_call_element *elem,
                                    grpc_transport_stream_op_batch *batch) {
@@ -193,20 +233,33 @@ static void send_security_metadata(grpc_exec_ctx *exec_ctx,
 
   build_auth_metadata_context(&chand->security_connector->base,
                               chand->auth_context, calld);
+
+  grpc_error *cancel_error = set_cancel_func(elem, cancel_get_request_metadata);
+  if (cancel_error != GRPC_ERROR_NONE) {
+    grpc_transport_stream_op_batch_finish_with_failure(exec_ctx, batch,
+                                                       cancel_error);
+    return;
+  }
   GPR_ASSERT(calld->pollent != NULL);
-  grpc_call_credentials_get_request_metadata(
-      exec_ctx, calld->creds, calld->pollent, calld->auth_md_context,
-      on_credentials_metadata, batch);
+  GRPC_CLOSURE_INIT(&calld->closure, on_credentials_metadata, batch,
+                    grpc_schedule_on_exec_ctx);
+  grpc_error *error = GRPC_ERROR_NONE;
+  if (grpc_call_credentials_get_request_metadata(
+          exec_ctx, calld->creds, calld->pollent, calld->auth_md_context,
+          &calld->md_array, &calld->closure, &error)) {
+    // Synchronous return; invoke on_credentials_metadata() directly.
+    on_credentials_metadata(exec_ctx, batch, error);
+    GRPC_ERROR_UNREF(error);
+  }
 }
 
-static void on_host_checked(grpc_exec_ctx *exec_ctx, void *user_data,
-                            grpc_security_status status) {
-  grpc_transport_stream_op_batch *batch =
-      (grpc_transport_stream_op_batch *)user_data;
+static void on_host_checked(grpc_exec_ctx *exec_ctx, void *arg,
+                            grpc_error *error) {
+  grpc_transport_stream_op_batch *batch = (grpc_transport_stream_op_batch *)arg;
   grpc_call_element *elem = batch->handler_private.extra_arg;
   call_data *calld = elem->call_data;
 
-  if (status == GRPC_SECURITY_OK) {
+  if (error == GRPC_ERROR_NONE) {
     send_security_metadata(exec_ctx, elem, batch);
   } else {
     char *error_msg;
@@ -223,6 +276,16 @@ static void on_host_checked(grpc_exec_ctx *exec_ctx, void *user_data,
   }
 }
 
+static void cancel_check_call_host(grpc_exec_ctx *exec_ctx, void *arg,
+                                   grpc_error *error) {
+  grpc_call_element *elem = (grpc_call_element *)arg;
+  call_data *calld = (call_data *)elem->call_data;
+  channel_data *chand = (channel_data *)elem->channel_data;
+  grpc_channel_security_connector_cancel_check_call_host(
+      exec_ctx, chand->security_connector, &calld->closure,
+      GRPC_ERROR_REF(error));
+}
+
 static void auth_start_transport_stream_op_batch(
     grpc_exec_ctx *exec_ctx, grpc_call_element *elem,
     grpc_transport_stream_op_batch *batch) {
@@ -232,7 +295,32 @@ static void auth_start_transport_stream_op_batch(
   call_data *calld = elem->call_data;
   channel_data *chand = elem->channel_data;
 
-  if (!batch->cancel_stream) {
+  if (batch->cancel_stream) {
+    while (true) {
+      // Decode the original cancellation state.
+      gpr_atm original_state = gpr_atm_acq_load(&calld->cancellation_state);
+      grpc_error *cancel_error = GRPC_ERROR_NONE;
+      grpc_closure *func = NULL;
+      decode_cancel_state(original_state, &func, &cancel_error);
+      // If we had already set a cancellation error, there's nothing
+      // more to do.
+      if (cancel_error != GRPC_ERROR_NONE) break;
+      // If there's a cancel func, call it.
+      // Note that even if the cancel func has been changed by some
+      // other thread between when we decoded it and now, it will just
+      // be a no-op.
+      cancel_error = GRPC_ERROR_REF(batch->payload->cancel_stream.cancel_error);
+      if (func != NULL) {
+        GRPC_CLOSURE_SCHED(exec_ctx, func, GRPC_ERROR_REF(cancel_error));
+      }
+      // Encode the new error into cancellation state.
+      if (gpr_atm_full_cas(&calld->cancellation_state, original_state,
+                           encode_cancel_state_error(cancel_error))) {
+        break;  // Success.
+      }
+      // The cas failed, so try again.
+    }
+  } else {
     /* double checked lock over security context to ensure it's set once */
     if (gpr_atm_acq_load(&calld->security_context_set) == 0) {
       gpr_mu_lock(&calld->security_context_mu);
@@ -277,12 +365,26 @@ static void auth_start_transport_stream_op_batch(
       }
     }
     if (calld->have_host) {
-      char *call_host = grpc_slice_to_c_string(calld->host);
-      batch->handler_private.extra_arg = elem;
-      grpc_channel_security_connector_check_call_host(
-          exec_ctx, chand->security_connector, call_host, chand->auth_context,
-          on_host_checked, batch);
-      gpr_free(call_host);
+      grpc_error *cancel_error = set_cancel_func(elem, cancel_check_call_host);
+      if (cancel_error != GRPC_ERROR_NONE) {
+        grpc_transport_stream_op_batch_finish_with_failure(exec_ctx, batch,
+                                                           cancel_error);
+      } else {
+        char *call_host = grpc_slice_to_c_string(calld->host);
+        batch->handler_private.extra_arg = elem;
+        grpc_error *error = GRPC_ERROR_NONE;
+        if (grpc_channel_security_connector_check_call_host(
+                exec_ctx, chand->security_connector, call_host,
+                chand->auth_context,
+                GRPC_CLOSURE_INIT(&calld->closure, on_host_checked, batch,
+                                  grpc_schedule_on_exec_ctx),
+                &error)) {
+          // Synchronous return; invoke on_host_checked() directly.
+          on_host_checked(exec_ctx, batch, error);
+          GRPC_ERROR_UNREF(error);
+        }
+        gpr_free(call_host);
+      }
       GPR_TIMER_END("auth_start_transport_stream_op_batch", 0);
       return; /* early exit */
     }
@@ -315,6 +417,7 @@ static void destroy_call_elem(grpc_exec_ctx *exec_ctx, grpc_call_element *elem,
                               const grpc_call_final_info *final_info,
                               grpc_closure *ignored) {
   call_data *calld = elem->call_data;
+  grpc_credentials_mdelem_array_destroy(exec_ctx, &calld->md_array);
   grpc_call_credentials_unref(exec_ctx, calld->creds);
   if (calld->have_host) {
     grpc_slice_unref_internal(exec_ctx, calld->host);
@@ -324,6 +427,11 @@ static void destroy_call_elem(grpc_exec_ctx *exec_ctx, grpc_call_element *elem,
   }
   reset_auth_metadata_context(&calld->auth_md_context);
   gpr_mu_destroy(&calld->security_context_mu);
+  gpr_atm cancel_state = gpr_atm_acq_load(&calld->cancellation_state);
+  grpc_error *cancel_error = GRPC_ERROR_NONE;
+  grpc_closure *cancel_func = NULL;
+  decode_cancel_state(cancel_state, &cancel_func, &cancel_error);
+  GRPC_ERROR_UNREF(cancel_error);
 }
 
 /* Constructor for channel_data */
diff --git a/src/core/lib/security/transport/security_connector.c b/src/core/lib/security/transport/security_connector.c
index 6788126769..a7568b995f 100644
--- a/src/core/lib/security/transport/security_connector.c
+++ b/src/core/lib/security/transport/security_connector.c
@@ -136,15 +136,27 @@ void grpc_security_connector_check_peer(grpc_exec_ctx *exec_ctx,
   }
 }
 
-void grpc_channel_security_connector_check_call_host(
+bool grpc_channel_security_connector_check_call_host(
     grpc_exec_ctx *exec_ctx, grpc_channel_security_connector *sc,
     const char *host, grpc_auth_context *auth_context,
-    grpc_security_call_host_check_cb cb, void *user_data) {
+    grpc_closure *on_call_host_checked, grpc_error **error) {
   if (sc == NULL || sc->check_call_host == NULL) {
-    cb(exec_ctx, user_data, GRPC_SECURITY_ERROR);
-  } else {
-    sc->check_call_host(exec_ctx, sc, host, auth_context, cb, user_data);
+    *error = GRPC_ERROR_CREATE_FROM_STATIC_STRING(
+        "cannot check call host -- no security connector");
+    return true;
   }
+  return sc->check_call_host(exec_ctx, sc, host, auth_context,
+                             on_call_host_checked, error);
+}
+
+void grpc_channel_security_connector_cancel_check_call_host(
+    grpc_exec_ctx *exec_ctx, grpc_channel_security_connector *sc,
+    grpc_closure *on_call_host_checked, grpc_error *error) {
+  if (sc == NULL || sc->cancel_check_call_host == NULL) {
+    GRPC_ERROR_UNREF(error);
+    return;
+  }
+  sc->cancel_check_call_host(exec_ctx, sc, on_call_host_checked, error);
 }
 
 #ifndef NDEBUG
@@ -368,13 +380,19 @@ static void fake_server_check_peer(grpc_exec_ctx *exec_ctx,
   fake_check_peer(exec_ctx, sc, peer, auth_context, on_peer_checked);
 }
 
-static void fake_channel_check_call_host(grpc_exec_ctx *exec_ctx,
+static bool fake_channel_check_call_host(grpc_exec_ctx *exec_ctx,
                                          grpc_channel_security_connector *sc,
                                          const char *host,
                                          grpc_auth_context *auth_context,
-                                         grpc_security_call_host_check_cb cb,
-                                         void *user_data) {
-  cb(exec_ctx, user_data, GRPC_SECURITY_OK);
+                                         grpc_closure *on_call_host_checked,
+                                         grpc_error **error) {
+  return true;
+}
+
+static void fake_channel_cancel_check_call_host(
+    grpc_exec_ctx *exec_ctx, grpc_channel_security_connector *sc,
+    grpc_closure *on_call_host_checked, grpc_error *error) {
+  GRPC_ERROR_UNREF(error);
 }
 
 static void fake_channel_add_handshakers(
@@ -413,6 +431,7 @@ grpc_channel_security_connector *grpc_fake_channel_security_connector_create(
   c->base.request_metadata_creds =
       grpc_call_credentials_ref(request_metadata_creds);
   c->base.check_call_host = fake_channel_check_call_host;
+  c->base.cancel_check_call_host = fake_channel_cancel_check_call_host;
   c->base.add_handshakers = fake_channel_add_handshakers;
   c->target = gpr_strdup(target);
   const char *expected_targets = grpc_fake_transport_get_expected_targets(args);
@@ -663,26 +682,35 @@ void tsi_shallow_peer_destruct(tsi_peer *peer) {
   if (peer->properties != NULL) gpr_free(peer->properties);
 }
 
-static void ssl_channel_check_call_host(grpc_exec_ctx *exec_ctx,
+static bool ssl_channel_check_call_host(grpc_exec_ctx *exec_ctx,
                                         grpc_channel_security_connector *sc,
                                         const char *host,
                                         grpc_auth_context *auth_context,
-                                        grpc_security_call_host_check_cb cb,
-                                        void *user_data) {
+                                        grpc_closure *on_call_host_checked,
+                                        grpc_error **error) {
   grpc_ssl_channel_security_connector *c =
       (grpc_ssl_channel_security_connector *)sc;
   grpc_security_status status = GRPC_SECURITY_ERROR;
   tsi_peer peer = tsi_shallow_peer_from_ssl_auth_context(auth_context);
   if (ssl_host_matches_name(&peer, host)) status = GRPC_SECURITY_OK;
-
   /* If the target name was overridden, then the original target_name was
      'checked' transitively during the previous peer check at the end of the
      handshake. */
   if (c->overridden_target_name != NULL && strcmp(host, c->target_name) == 0) {
     status = GRPC_SECURITY_OK;
   }
-  cb(exec_ctx, user_data, status);
+  if (status != GRPC_SECURITY_OK) {
+    *error = GRPC_ERROR_CREATE_FROM_STATIC_STRING(
+        "call host does not match SSL server name");
+  }
   tsi_shallow_peer_destruct(&peer);
+  return true;
+}
+
+static void ssl_channel_cancel_check_call_host(
+    grpc_exec_ctx *exec_ctx, grpc_channel_security_connector *sc,
+    grpc_closure *on_call_host_checked, grpc_error *error) {
+  GRPC_ERROR_UNREF(error);
 }
 
 static grpc_security_connector_vtable ssl_channel_vtable = {
@@ -811,6 +839,7 @@ grpc_security_status grpc_ssl_channel_security_connector_create(
   c->base.request_metadata_creds =
       grpc_call_credentials_ref(request_metadata_creds);
   c->base.check_call_host = ssl_channel_check_call_host;
+  c->base.cancel_check_call_host = ssl_channel_cancel_check_call_host;
   c->base.add_handshakers = ssl_channel_add_handshakers;
   gpr_split_host_port(target_name, &c->target_name, &port);
   gpr_free(port);
diff --git a/src/core/lib/security/transport/security_connector.h b/src/core/lib/security/transport/security_connector.h
index 1c0fe40045..4f9b63ad20 100644
--- a/src/core/lib/security/transport/security_connector.h
+++ b/src/core/lib/security/transport/security_connector.h
@@ -117,27 +117,38 @@ grpc_security_connector *grpc_security_connector_find_in_args(
 
 typedef struct grpc_channel_security_connector grpc_channel_security_connector;
 
-typedef void (*grpc_security_call_host_check_cb)(grpc_exec_ctx *exec_ctx,
-                                                 void *user_data,
-                                                 grpc_security_status status);
-
 struct grpc_channel_security_connector {
   grpc_security_connector base;
   grpc_call_credentials *request_metadata_creds;
-  void (*check_call_host)(grpc_exec_ctx *exec_ctx,
+  bool (*check_call_host)(grpc_exec_ctx *exec_ctx,
                           grpc_channel_security_connector *sc, const char *host,
                           grpc_auth_context *auth_context,
-                          grpc_security_call_host_check_cb cb, void *user_data);
+                          grpc_closure *on_call_host_checked,
+                          grpc_error **error);
+  void (*cancel_check_call_host)(grpc_exec_ctx *exec_ctx,
+                                 grpc_channel_security_connector *sc,
+                                 grpc_closure *on_call_host_checked,
+                                 grpc_error *error);
   void (*add_handshakers)(grpc_exec_ctx *exec_ctx,
                           grpc_channel_security_connector *sc,
                           grpc_handshake_manager *handshake_mgr);
 };
 
-/* Checks that the host that will be set for a call is acceptable. */
-void grpc_channel_security_connector_check_call_host(
+/// Checks that the host that will be set for a call is acceptable.
+/// Returns true if completed synchronously, in which case \a error will
+/// be set to indicate the result.  Otherwise, \a on_call_host_checked
+/// will be invoked when complete.
+bool grpc_channel_security_connector_check_call_host(
     grpc_exec_ctx *exec_ctx, grpc_channel_security_connector *sc,
     const char *host, grpc_auth_context *auth_context,
-    grpc_security_call_host_check_cb cb, void *user_data);
+    grpc_closure *on_call_host_checked, grpc_error **error);
+
+/// Cancels a pending asychronous call to
+/// grpc_channel_security_connector_check_call_host() with
+/// \a on_call_host_checked as its callback.
+void grpc_channel_security_connector_cancel_check_call_host(
+    grpc_exec_ctx *exec_ctx, grpc_channel_security_connector *sc,
+    grpc_closure *on_call_host_checked, grpc_error *error);
 
 /* Registers handshakers with \a handshake_mgr. */
 void grpc_channel_security_connector_add_handshakers(