Revise zero-copy protector interface

4 files changed, 40 insertions, 29 deletions

diff --git a/src/core/lib/security/transport/secure_endpoint.c b/src/core/lib/security/transport/secure_endpoint.c
index e0520eb72e..ae5633b82c 100644
--- a/src/core/lib/security/transport/secure_endpoint.c
+++ b/src/core/lib/security/transport/secure_endpoint.c
@@ -68,7 +68,7 @@ static void destroy(grpc_exec_ctx *exec_ctx, secure_endpoint *secure_ep) {
   secure_endpoint *ep = secure_ep;
   grpc_endpoint_destroy(exec_ctx, ep->wrapped_ep);
   tsi_frame_protector_destroy(ep->protector);
-  tsi_zero_copy_grpc_protector_destroy(ep->zero_copy_protector);
+  tsi_zero_copy_grpc_protector_destroy(exec_ctx, ep->zero_copy_protector);
   grpc_slice_buffer_destroy_internal(exec_ctx, &ep->leftover_bytes);
   grpc_slice_unref_internal(exec_ctx, ep->read_staging_buffer);
   grpc_slice_unref_internal(exec_ctx, ep->write_staging_buffer);
@@ -164,7 +164,7 @@ static void on_read(grpc_exec_ctx *exec_ctx, void *user_data,
   if (ep->zero_copy_protector != NULL) {
     // Use zero-copy grpc protector to unprotect.
     result = tsi_zero_copy_grpc_protector_unprotect(
-        ep->zero_copy_protector, &ep->source_buffer, ep->read_buffer);
+        exec_ctx, ep->zero_copy_protector, &ep->source_buffer, ep->read_buffer);
   } else {
     // Use frame protector to unprotect.
     /* TODO(yangg) check error, maybe bail out early */
@@ -281,8 +281,8 @@ static void endpoint_write(grpc_exec_ctx *exec_ctx, grpc_endpoint *secure_ep,
 
   if (ep->zero_copy_protector != NULL) {
     // Use zero-copy grpc protector to protect.
-    result = tsi_zero_copy_grpc_protector_protect(ep->zero_copy_protector,
-                                                  slices, &ep->output_buffer);
+    result = tsi_zero_copy_grpc_protector_protect(
+        exec_ctx, ep->zero_copy_protector, slices, &ep->output_buffer);
   } else {
     // Use frame protector to protect.
     for (i = 0; i < slices->count; i++) {
diff --git a/src/core/tsi/fake_transport_security.c b/src/core/tsi/fake_transport_security.c
index 17bbd74002..e7b3be3d86 100644
--- a/src/core/tsi/fake_transport_security.c
+++ b/src/core/tsi/fake_transport_security.c
@@ -25,6 +25,7 @@
 #include <grpc/support/log.h>
 #include <grpc/support/port_platform.h>
 #include <grpc/support/useful.h>
+#include "src/core/lib/slice/slice_internal.h"
 #include "src/core/tsi/transport_security_grpc.h"
 
 /* --- Constants. ---*/
@@ -396,7 +397,8 @@ static const tsi_frame_protector_vtable frame_protector_vtable = {
 /* --- tsi_zero_copy_grpc_protector methods implementation. ---*/
 
 static tsi_result fake_zero_copy_grpc_protector_protect(
-    tsi_zero_copy_grpc_protector *self, grpc_slice_buffer *unprotected_slices,
+    grpc_exec_ctx *exec_ctx, tsi_zero_copy_grpc_protector *self,
+    grpc_slice_buffer *unprotected_slices,
     grpc_slice_buffer *protected_slices) {
   if (self == NULL || unprotected_slices == NULL || protected_slices == NULL) {
     return TSI_INVALID_ARGUMENT;
@@ -408,7 +410,7 @@ static tsi_result fake_zero_copy_grpc_protector_protect(
     size_t frame_length =
         GPR_MIN(impl->max_frame_size,
                 unprotected_slices->length + TSI_FAKE_FRAME_HEADER_SIZE);
-    grpc_slice slice = grpc_slice_malloc(TSI_FAKE_FRAME_HEADER_SIZE);
+    grpc_slice slice = GRPC_SLICE_MALLOC(TSI_FAKE_FRAME_HEADER_SIZE);
     store32_little_endian((uint32_t)frame_length, GRPC_SLICE_START_PTR(slice));
     grpc_slice_buffer_add(protected_slices, slice);
     size_t data_length = frame_length - TSI_FAKE_FRAME_HEADER_SIZE;
@@ -419,7 +421,8 @@ static tsi_result fake_zero_copy_grpc_protector_protect(
 }
 
 static tsi_result fake_zero_copy_grpc_protector_unprotect(
-    tsi_zero_copy_grpc_protector *self, grpc_slice_buffer *protected_slices,
+    grpc_exec_ctx *exec_ctx, tsi_zero_copy_grpc_protector *self,
+    grpc_slice_buffer *protected_slices,
     grpc_slice_buffer *unprotected_slices) {
   if (self == NULL || unprotected_slices == NULL || protected_slices == NULL) {
     return TSI_INVALID_ARGUMENT;
@@ -447,18 +450,18 @@ static tsi_result fake_zero_copy_grpc_protector_unprotect(
         impl->parsed_frame_size - TSI_FAKE_FRAME_HEADER_SIZE,
         unprotected_slices);
     impl->parsed_frame_size = 0;
-    grpc_slice_buffer_reset_and_unref(&impl->header_sb);
+    grpc_slice_buffer_reset_and_unref_internal(exec_ctx, &impl->header_sb);
   }
   return TSI_OK;
 }
 
 static void fake_zero_copy_grpc_protector_destroy(
-    tsi_zero_copy_grpc_protector *self) {
+    grpc_exec_ctx *exec_ctx, tsi_zero_copy_grpc_protector *self) {
   if (self == NULL) return;
   tsi_fake_zero_copy_grpc_protector *impl =
       (tsi_fake_zero_copy_grpc_protector *)self;
-  grpc_slice_buffer_destroy(&impl->header_sb);
-  grpc_slice_buffer_destroy(&impl->protected_sb);
+  grpc_slice_buffer_destroy_internal(exec_ctx, &impl->header_sb);
+  grpc_slice_buffer_destroy_internal(exec_ctx, &impl->protected_sb);
   gpr_free(impl);
 }
 
diff --git a/src/core/tsi/transport_security_grpc.c b/src/core/tsi/transport_security_grpc.c
index 5bcfdfa61f..773b35e717 100644
--- a/src/core/tsi/transport_security_grpc.c
+++ b/src/core/tsi/transport_security_grpc.c
@@ -37,28 +37,33 @@ tsi_result tsi_handshaker_result_create_zero_copy_grpc_protector(
    Calls specific implementation after state/input validation. */
 
 tsi_result tsi_zero_copy_grpc_protector_protect(
-    tsi_zero_copy_grpc_protector *self, grpc_slice_buffer *unprotected_slices,
+    grpc_exec_ctx *exec_ctx, tsi_zero_copy_grpc_protector *self,
+    grpc_slice_buffer *unprotected_slices,
     grpc_slice_buffer *protected_slices) {
-  if (self == NULL || self->vtable == NULL || unprotected_slices == NULL ||
-      protected_slices == NULL) {
+  if (exec_ctx == NULL || self == NULL || self->vtable == NULL ||
+      unprotected_slices == NULL || protected_slices == NULL) {
     return TSI_INVALID_ARGUMENT;
   }
   if (self->vtable->protect == NULL) return TSI_UNIMPLEMENTED;
-  return self->vtable->protect(self, unprotected_slices, protected_slices);
+  return self->vtable->protect(exec_ctx, self, unprotected_slices,
+                               protected_slices);
 }
 
 tsi_result tsi_zero_copy_grpc_protector_unprotect(
-    tsi_zero_copy_grpc_protector *self, grpc_slice_buffer *protected_slices,
+    grpc_exec_ctx *exec_ctx, tsi_zero_copy_grpc_protector *self,
+    grpc_slice_buffer *protected_slices,
     grpc_slice_buffer *unprotected_slices) {
-  if (self == NULL || self->vtable == NULL || protected_slices == NULL ||
-      unprotected_slices == NULL) {
+  if (exec_ctx == NULL || self == NULL || self->vtable == NULL ||
+      protected_slices == NULL || unprotected_slices == NULL) {
     return TSI_INVALID_ARGUMENT;
   }
   if (self->vtable->unprotect == NULL) return TSI_UNIMPLEMENTED;
-  return self->vtable->unprotect(self, protected_slices, unprotected_slices);
+  return self->vtable->unprotect(exec_ctx, self, protected_slices,
+                                 unprotected_slices);
 }
 
-void tsi_zero_copy_grpc_protector_destroy(tsi_zero_copy_grpc_protector *self) {
+void tsi_zero_copy_grpc_protector_destroy(grpc_exec_ctx *exec_ctx,
+                                          tsi_zero_copy_grpc_protector *self) {
   if (self == NULL) return;
-  self->vtable->destroy(self);
+  self->vtable->destroy(exec_ctx, self);
 }
diff --git a/src/core/tsi/transport_security_grpc.h b/src/core/tsi/transport_security_grpc.h
index 5ab5297cc4..375a758888 100644
--- a/src/core/tsi/transport_security_grpc.h
+++ b/src/core/tsi/transport_security_grpc.h
@@ -42,8 +42,8 @@ tsi_result tsi_handshaker_result_create_zero_copy_grpc_protector(
    - This method returns TSI_OK in case of success or a specific error code in
      case of failure.  */
 tsi_result tsi_zero_copy_grpc_protector_protect(
-    tsi_zero_copy_grpc_protector *self, grpc_slice_buffer *unprotected_slices,
-    grpc_slice_buffer *protected_slices);
+    grpc_exec_ctx *exec_ctx, tsi_zero_copy_grpc_protector *self,
+    grpc_slice_buffer *unprotected_slices, grpc_slice_buffer *protected_slices);
 
 /* Outputs unprotected bytes.
    - protected_slices is the bytes of protected frames.
@@ -52,21 +52,24 @@ tsi_result tsi_zero_copy_grpc_protector_protect(
      there is not enough data to output in which case unprotected_slices has 0
      bytes.  */
 tsi_result tsi_zero_copy_grpc_protector_unprotect(
-    tsi_zero_copy_grpc_protector *self, grpc_slice_buffer *protected_slices,
-    grpc_slice_buffer *unprotected_slices);
+    grpc_exec_ctx *exec_ctx, tsi_zero_copy_grpc_protector *self,
+    grpc_slice_buffer *protected_slices, grpc_slice_buffer *unprotected_slices);
 
 /* Destroys the tsi_zero_copy_grpc_protector object.  */
-void tsi_zero_copy_grpc_protector_destroy(tsi_zero_copy_grpc_protector *self);
+void tsi_zero_copy_grpc_protector_destroy(grpc_exec_ctx *exec_ctx,
+                                          tsi_zero_copy_grpc_protector *self);
 
 /* Base for tsi_zero_copy_grpc_protector implementations.  */
 typedef struct {
-  tsi_result (*protect)(tsi_zero_copy_grpc_protector *self,
+  tsi_result (*protect)(grpc_exec_ctx *exec_ctx,
+                        tsi_zero_copy_grpc_protector *self,
                         grpc_slice_buffer *unprotected_slices,
                         grpc_slice_buffer *protected_slices);
-  tsi_result (*unprotect)(tsi_zero_copy_grpc_protector *self,
+  tsi_result (*unprotect)(grpc_exec_ctx *exec_ctx,
+                          tsi_zero_copy_grpc_protector *self,
                           grpc_slice_buffer *protected_slices,
                           grpc_slice_buffer *unprotected_slices);
-  void (*destroy)(tsi_zero_copy_grpc_protector *self);
+  void (*destroy)(grpc_exec_ctx *exec_ctx, tsi_zero_copy_grpc_protector *self);
 } tsi_zero_copy_grpc_protector_vtable;
 
 struct tsi_zero_copy_grpc_protector {