Add cancellation to asynchronous security APIs.

1 files changed, 147 insertions, 39 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 */