grpc 1.6.7 → 1.7.0.pre1

data/src/core/lib/transport/status_conversion.c

@@ -18,7 +18,7 @@
 
 #include "src/core/lib/transport/status_conversion.h"
 
-int grpc_status_to_http2_error(grpc_status_code status) {
+grpc_http2_error_code grpc_status_to_http2_error(grpc_status_code status) {
   switch (status) {
     case GRPC_STATUS_OK:
       return GRPC_HTTP2_NO_ERROR;

data/src/core/lib/transport/transport.c

@@ -72,7 +72,8 @@ void grpc_stream_unref(grpc_exec_ctx *exec_ctx,
          cope with.
          Throw this over to the executor (on a core-owned thread) and process it
          there. */
-      refcount->destroy.scheduler = grpc_executor_scheduler;
+      refcount->destroy.scheduler =
+          grpc_executor_scheduler(GRPC_EXECUTOR_SHORT);
     }
     GRPC_CLOSURE_SCHED(exec_ctx, &refcount->destroy, GRPC_ERROR_NONE);
   }
@@ -101,8 +102,11 @@ static void slice_stream_unref(grpc_exec_ctx *exec_ctx, void *p) {
 grpc_slice grpc_slice_from_stream_owned_buffer(grpc_stream_refcount *refcount,
                                                void *buffer, size_t length) {
   slice_stream_ref(&refcount->slice_refcount);
-  return (grpc_slice){.refcount = &refcount->slice_refcount,
-                      .data.refcounted = {.bytes = buffer, .length = length}};
+  grpc_slice res;
+  res.refcount = &refcount->slice_refcount,
+  res.data.refcounted.bytes = (uint8_t *)buffer;
+  res.data.refcounted.length = length;
+  return res;
 }
 
 static const grpc_slice_refcount_vtable stream_ref_slice_vtable = {
@@ -197,11 +201,6 @@ void grpc_transport_destroy_stream(grpc_exec_ctx *exec_ctx,
                                     then_schedule_closure);
 }
 
-char *grpc_transport_get_peer(grpc_exec_ctx *exec_ctx,
-                              grpc_transport *transport) {
-  return transport->vtable->get_peer(exec_ctx, transport);
-}
-
 grpc_endpoint *grpc_transport_get_endpoint(grpc_exec_ctx *exec_ctx,
                                            grpc_transport *transport) {
   return transport->vtable->get_endpoint(exec_ctx, transport);
@@ -214,24 +213,24 @@ grpc_endpoint *grpc_transport_get_endpoint(grpc_exec_ctx *exec_ctx,
 // is a function that must always unref cancel_error
 // though it lives in lib, it handles transport stream ops sure
 // it's grpc_transport_stream_op_batch_finish_with_failure
-
 void grpc_transport_stream_op_batch_finish_with_failure(
     grpc_exec_ctx *exec_ctx, grpc_transport_stream_op_batch *batch,
-    grpc_error *error) {
+    grpc_error *error, grpc_call_combiner *call_combiner) {
   if (batch->send_message) {
     grpc_byte_stream_destroy(exec_ctx,
                              batch->payload->send_message.send_message);
   }
   if (batch->recv_message) {
-    GRPC_CLOSURE_SCHED(exec_ctx,
-                       batch->payload->recv_message.recv_message_ready,
-                       GRPC_ERROR_REF(error));
+    GRPC_CALL_COMBINER_START(exec_ctx, call_combiner,
+                             batch->payload->recv_message.recv_message_ready,
+                             GRPC_ERROR_REF(error),
+                             "failing recv_message_ready");
   }
   if (batch->recv_initial_metadata) {
-    GRPC_CLOSURE_SCHED(
-        exec_ctx,
+    GRPC_CALL_COMBINER_START(
+        exec_ctx, call_combiner,
         batch->payload->recv_initial_metadata.recv_initial_metadata_ready,
-        GRPC_ERROR_REF(error));
+        GRPC_ERROR_REF(error), "failing recv_initial_metadata_ready");
   }
   GRPC_CLOSURE_SCHED(exec_ctx, batch->on_complete, error);
   if (batch->cancel_stream) {
@@ -247,13 +246,13 @@ typedef struct {
 
 static void destroy_made_transport_op(grpc_exec_ctx *exec_ctx, void *arg,
                                       grpc_error *error) {
-  made_transport_op *op = arg;
+  made_transport_op *op = (made_transport_op *)arg;
   GRPC_CLOSURE_SCHED(exec_ctx, op->inner_on_complete, GRPC_ERROR_REF(error));
   gpr_free(op);
 }
 
 grpc_transport_op *grpc_make_transport_op(grpc_closure *on_complete) {
-  made_transport_op *op = gpr_malloc(sizeof(*op));
+  made_transport_op *op = (made_transport_op *)gpr_malloc(sizeof(*op));
   GRPC_CLOSURE_INIT(&op->outer_on_complete, destroy_made_transport_op, op,
                     grpc_schedule_on_exec_ctx);
   op->inner_on_complete = on_complete;
@@ -271,7 +270,7 @@ typedef struct {
 
 static void destroy_made_transport_stream_op(grpc_exec_ctx *exec_ctx, void *arg,
                                              grpc_error *error) {
-  made_transport_stream_op *op = arg;
+  made_transport_stream_op *op = (made_transport_stream_op *)arg;
   grpc_closure *c = op->inner_on_complete;
   gpr_free(op);
   GRPC_CLOSURE_RUN(exec_ctx, c, GRPC_ERROR_REF(error));
@@ -279,7 +278,8 @@ static void destroy_made_transport_stream_op(grpc_exec_ctx *exec_ctx, void *arg,
 
 grpc_transport_stream_op_batch *grpc_make_transport_stream_op(
     grpc_closure *on_complete) {
-  made_transport_stream_op *op = gpr_zalloc(sizeof(*op));
+  made_transport_stream_op *op =
+      (made_transport_stream_op *)gpr_zalloc(sizeof(*op));
   op->op.payload = &op->payload;
   GRPC_CLOSURE_INIT(&op->outer_on_complete, destroy_made_transport_stream_op,
                     op, grpc_schedule_on_exec_ctx);

data/src/core/lib/transport/transport.h

@@ -22,6 +22,7 @@
 #include <stddef.h>
 
 #include "src/core/lib/channel/context.h"
+#include "src/core/lib/iomgr/call_combiner.h"
 #include "src/core/lib/iomgr/endpoint.h"
 #include "src/core/lib/iomgr/polling_entity.h"
 #include "src/core/lib/iomgr/pollset.h"
@@ -152,6 +153,9 @@ struct grpc_transport_stream_op_batch_payload {
     /** Iff send_initial_metadata != NULL, flags associated with
         send_initial_metadata: a bitfield of GRPC_INITIAL_METADATA_xxx */
     uint32_t send_initial_metadata_flags;
+    // If non-NULL, will be set by the transport to the peer string
+    // (a char*, which the caller takes ownership of).
+    gpr_atm *peer_string;
   } send_initial_metadata;
 
   struct {
@@ -176,6 +180,9 @@ struct grpc_transport_stream_op_batch_payload {
     // immediately available.  This may be a signal that we received a
     // Trailers-Only response.
     bool *trailing_metadata_available;
+    // If non-NULL, will be set by the transport to the peer string
+    // (a char*, which the caller takes ownership of).
+    gpr_atm *peer_string;
   } recv_initial_metadata;
 
   struct {
@@ -293,7 +300,7 @@ void grpc_transport_destroy_stream(grpc_exec_ctx *exec_ctx,
 
 void grpc_transport_stream_op_batch_finish_with_failure(
     grpc_exec_ctx *exec_ctx, grpc_transport_stream_op_batch *op,
-    grpc_error *error);
+    grpc_error *error, grpc_call_combiner *call_combiner);
 
 char *grpc_transport_stream_op_batch_string(grpc_transport_stream_op_batch *op);
 char *grpc_transport_op_string(grpc_transport_op *op);
@@ -332,10 +339,6 @@ void grpc_transport_close(grpc_transport *transport);
 /* Destroy the transport */
 void grpc_transport_destroy(grpc_exec_ctx *exec_ctx, grpc_transport *transport);
 
-/* Get the transports peer */
-char *grpc_transport_get_peer(grpc_exec_ctx *exec_ctx,
-                              grpc_transport *transport);
-
 /* Get the endpoint used by \a transport */
 grpc_endpoint *grpc_transport_get_endpoint(grpc_exec_ctx *exec_ctx,
                                            grpc_transport *transport);

data/src/core/lib/transport/transport_impl.h

@@ -59,9 +59,6 @@ typedef struct grpc_transport_vtable {
   /* implementation of grpc_transport_destroy */
   void (*destroy)(grpc_exec_ctx *exec_ctx, grpc_transport *self);
 
-  /* implementation of grpc_transport_get_peer */
-  char *(*get_peer)(grpc_exec_ctx *exec_ctx, grpc_transport *self);
-
   /* implementation of grpc_transport_get_endpoint */
   grpc_endpoint *(*get_endpoint)(grpc_exec_ctx *exec_ctx, grpc_transport *self);
 } grpc_transport_vtable;

data/src/core/lib/transport/transport_op_string.c

@@ -112,6 +112,13 @@ char *grpc_transport_stream_op_batch_string(
     gpr_strvec_add(&b, tmp);
   }
 
+  if (op->collect_stats) {
+    gpr_strvec_add(&b, gpr_strdup(" "));
+    gpr_asprintf(&tmp, "COLLECT_STATS:%p",
+                 op->payload->collect_stats.collect_stats);
+    gpr_strvec_add(&b, tmp);
+  }
+
   out = gpr_strvec_flatten(&b, NULL);
   gpr_strvec_destroy(&b);
 
@@ -190,7 +197,7 @@ char *grpc_transport_op_string(grpc_transport_op *op) {
   return out;
 }
 
-void grpc_call_log_op(char *file, int line, gpr_log_severity severity,
+void grpc_call_log_op(const char *file, int line, gpr_log_severity severity,
                       grpc_call_element *elem,
                       grpc_transport_stream_op_batch *op) {
   char *str = grpc_transport_stream_op_batch_string(op);

data/src/core/plugin_registry/grpc_plugin_registry.c

@@ -44,8 +44,8 @@ extern void grpc_resolver_dns_native_init(void);
 extern void grpc_resolver_dns_native_shutdown(void);
 extern void grpc_resolver_sockaddr_init(void);
 extern void grpc_resolver_sockaddr_shutdown(void);
-extern void grpc_load_reporting_plugin_init(void);
-extern void grpc_load_reporting_plugin_shutdown(void);
+extern void grpc_server_load_reporting_plugin_init(void);
+extern void grpc_server_load_reporting_plugin_shutdown(void);
 extern void census_grpc_plugin_init(void);
 extern void census_grpc_plugin_shutdown(void);
 extern void grpc_max_age_filter_init(void);
@@ -82,8 +82,8 @@ void grpc_register_built_in_plugins(void) {
                        grpc_resolver_dns_native_shutdown);
   grpc_register_plugin(grpc_resolver_sockaddr_init,
                        grpc_resolver_sockaddr_shutdown);
-  grpc_register_plugin(grpc_load_reporting_plugin_init,
-                       grpc_load_reporting_plugin_shutdown);
+  grpc_register_plugin(grpc_server_load_reporting_plugin_init,
+                       grpc_server_load_reporting_plugin_shutdown);
   grpc_register_plugin(census_grpc_plugin_init,
                        census_grpc_plugin_shutdown);
   grpc_register_plugin(grpc_max_age_filter_init,

data/src/core/tsi/fake_transport_security.c

@@ -25,7 +25,8 @@
 #include <grpc/support/log.h>
 #include <grpc/support/port_platform.h>
 #include <grpc/support/useful.h>
-#include "src/core/tsi/transport_security.h"
+#include "src/core/lib/slice/slice_internal.h"
+#include "src/core/tsi/transport_security_grpc.h"
 
 /* --- Constants. ---*/
 #define TSI_FAKE_FRAME_HEADER_SIZE 4
@@ -74,6 +75,14 @@ typedef struct {
   size_t max_frame_size;
 } tsi_fake_frame_protector;
 
+typedef struct {
+  tsi_zero_copy_grpc_protector base;
+  grpc_slice_buffer header_sb;
+  grpc_slice_buffer protected_sb;
+  size_t max_frame_size;
+  size_t parsed_frame_size;
+} tsi_fake_zero_copy_grpc_protector;
+
 /* --- Utils. ---*/
 
 static const char *tsi_fake_handshake_message_strings[] = {
@@ -113,6 +122,28 @@ static void store32_little_endian(uint32_t value, unsigned char *buf) {
   buf[0] = (unsigned char)((value)&0xFF);
 }
 
+static uint32_t read_frame_size(const grpc_slice_buffer *sb) {
+  GPR_ASSERT(sb != NULL && sb->length >= TSI_FAKE_FRAME_HEADER_SIZE);
+  uint8_t frame_size_buffer[TSI_FAKE_FRAME_HEADER_SIZE];
+  uint8_t *buf = frame_size_buffer;
+  /* Copies the first 4 bytes to a temporary buffer.  */
+  size_t remaining = TSI_FAKE_FRAME_HEADER_SIZE;
+  for (size_t i = 0; i < sb->count; i++) {
+    size_t slice_length = GRPC_SLICE_LENGTH(sb->slices[i]);
+    if (remaining <= slice_length) {
+      memcpy(buf, GRPC_SLICE_START_PTR(sb->slices[i]), remaining);
+      remaining = 0;
+      break;
+    } else {
+      memcpy(buf, GRPC_SLICE_START_PTR(sb->slices[i]), slice_length);
+      buf += slice_length;
+      remaining -= slice_length;
+    }
+  }
+  GPR_ASSERT(remaining == 0);
+  return load32_little_endian(frame_size_buffer);
+}
+
 static void tsi_fake_frame_reset(tsi_fake_frame *frame, int needs_draining) {
   frame->offset = 0;
   frame->needs_draining = needs_draining;
@@ -363,6 +394,84 @@ static const tsi_frame_protector_vtable frame_protector_vtable = {
     fake_protector_unprotect, fake_protector_destroy,
 };
 
+/* --- tsi_zero_copy_grpc_protector methods implementation. ---*/
+
+static tsi_result fake_zero_copy_grpc_protector_protect(
+    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;
+  }
+  tsi_fake_zero_copy_grpc_protector *impl =
+      (tsi_fake_zero_copy_grpc_protector *)self;
+  /* Protects each frame. */
+  while (unprotected_slices->length > 0) {
+    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);
+    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;
+    grpc_slice_buffer_move_first(unprotected_slices, data_length,
+                                 protected_slices);
+  }
+  return TSI_OK;
+}
+
+static tsi_result fake_zero_copy_grpc_protector_unprotect(
+    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;
+  }
+  tsi_fake_zero_copy_grpc_protector *impl =
+      (tsi_fake_zero_copy_grpc_protector *)self;
+  grpc_slice_buffer_move_into(protected_slices, &impl->protected_sb);
+  /* Unprotect each frame, if we get a full frame. */
+  while (impl->protected_sb.length >= TSI_FAKE_FRAME_HEADER_SIZE) {
+    if (impl->parsed_frame_size == 0) {
+      impl->parsed_frame_size = read_frame_size(&impl->protected_sb);
+      if (impl->parsed_frame_size <= 4) {
+        gpr_log(GPR_ERROR, "Invalid frame size.");
+        return TSI_DATA_CORRUPTED;
+      }
+    }
+    /* If we do not have a full frame, return with OK status. */
+    if (impl->protected_sb.length < impl->parsed_frame_size) break;
+    /* Strips header bytes. */
+    grpc_slice_buffer_move_first(&impl->protected_sb,
+                                 TSI_FAKE_FRAME_HEADER_SIZE, &impl->header_sb);
+    /* Moves data to unprotected slices. */
+    grpc_slice_buffer_move_first(
+        &impl->protected_sb,
+        impl->parsed_frame_size - TSI_FAKE_FRAME_HEADER_SIZE,
+        unprotected_slices);
+    impl->parsed_frame_size = 0;
+    grpc_slice_buffer_reset_and_unref_internal(exec_ctx, &impl->header_sb);
+  }
+  return TSI_OK;
+}
+
+static void fake_zero_copy_grpc_protector_destroy(
+    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_internal(exec_ctx, &impl->header_sb);
+  grpc_slice_buffer_destroy_internal(exec_ctx, &impl->protected_sb);
+  gpr_free(impl);
+}
+
+static const tsi_zero_copy_grpc_protector_vtable
+    zero_copy_grpc_protector_vtable = {
+        fake_zero_copy_grpc_protector_protect,
+        fake_zero_copy_grpc_protector_unprotect,
+        fake_zero_copy_grpc_protector_destroy,
+};
+
 /* --- tsi_handshaker_result methods implementation. ---*/
 
 typedef struct {
@@ -383,6 +492,15 @@ static tsi_result fake_handshaker_result_extract_peer(
   return result;
 }
 
+static tsi_result fake_handshaker_result_create_zero_copy_grpc_protector(
+    void *exec_ctx, const tsi_handshaker_result *self,
+    size_t *max_output_protected_frame_size,
+    tsi_zero_copy_grpc_protector **protector) {
+  *protector =
+      tsi_create_fake_zero_copy_grpc_protector(max_output_protected_frame_size);
+  return TSI_OK;
+}
+
 static tsi_result fake_handshaker_result_create_frame_protector(
     const tsi_handshaker_result *self, size_t *max_output_protected_frame_size,
     tsi_frame_protector **protector) {
@@ -407,7 +525,7 @@ static void fake_handshaker_result_destroy(tsi_handshaker_result *self) {
 
 static const tsi_handshaker_result_vtable handshaker_result_vtable = {
     fake_handshaker_result_extract_peer,
-    NULL, /* create_zero_copy_grpc_protector */
+    fake_handshaker_result_create_zero_copy_grpc_protector,
     fake_handshaker_result_create_frame_protector,
     fake_handshaker_result_get_unused_bytes,
     fake_handshaker_result_destroy,
@@ -631,3 +749,16 @@ tsi_frame_protector *tsi_create_fake_frame_protector(
   impl->base.vtable = &frame_protector_vtable;
   return &impl->base;
 }
+
+tsi_zero_copy_grpc_protector *tsi_create_fake_zero_copy_grpc_protector(
+    size_t *max_protected_frame_size) {
+  tsi_fake_zero_copy_grpc_protector *impl = gpr_zalloc(sizeof(*impl));
+  grpc_slice_buffer_init(&impl->header_sb);
+  grpc_slice_buffer_init(&impl->protected_sb);
+  impl->max_frame_size = (max_protected_frame_size == NULL)
+                             ? TSI_FAKE_DEFAULT_FRAME_SIZE
+                             : *max_protected_frame_size;
+  impl->parsed_frame_size = 0;
+  impl->base.vtable = &zero_copy_grpc_protector_vtable;
+  return &impl->base;
+}

data/src/core/tsi/fake_transport_security.h

@@ -39,6 +39,11 @@ tsi_handshaker *tsi_create_fake_handshaker(int is_client);
 tsi_frame_protector *tsi_create_fake_frame_protector(
     size_t *max_protected_frame_size);
 
+/* Creates a zero-copy protector directly without going through the handshake
+ * phase. */
+tsi_zero_copy_grpc_protector *tsi_create_fake_zero_copy_grpc_protector(
+    size_t *max_protected_frame_size);
+
 #ifdef __cplusplus
 }
 #endif

data/src/core/tsi/ssl_transport_security.c

@@ -67,7 +67,13 @@
 
 /* --- Structure definitions. ---*/
 
+struct tsi_ssl_handshaker_factory {
+  const tsi_ssl_handshaker_factory_vtable *vtable;
+  gpr_refcount refcount;
+};
+
 struct tsi_ssl_client_handshaker_factory {
+  tsi_ssl_handshaker_factory base;
   SSL_CTX *ssl_context;
   unsigned char *alpn_protocol_list;
   size_t alpn_protocol_list_length;
@@ -77,6 +83,7 @@ struct tsi_ssl_server_handshaker_factory {
   /* Several contexts to support SNI.
      The tsi_peer array contains the subject names of the server certificates
      associated with the contexts at the same index.  */
+  tsi_ssl_handshaker_factory base;
   SSL_CTX **ssl_contexts;
   tsi_peer *ssl_context_x509_subject_names;
   size_t ssl_context_count;
@@ -90,6 +97,7 @@ typedef struct {
   BIO *into_ssl;
   BIO *from_ssl;
   tsi_result result;
+  tsi_ssl_handshaker_factory *factory_ref;
 } tsi_ssl_handshaker;
 
 typedef struct {
@@ -846,6 +854,47 @@ static const tsi_frame_protector_vtable frame_protector_vtable = {
     ssl_protector_destroy,
 };
 
+/* --- tsi_server_handshaker_factory methods implementation. --- */
+
+static void tsi_ssl_handshaker_factory_destroy(
+    tsi_ssl_handshaker_factory *self) {
+  if (self == NULL) return;
+
+  if (self->vtable != NULL && self->vtable->destroy != NULL) {
+    self->vtable->destroy(self);
+  }
+  /* Note, we don't free(self) here because this object is always directly
+   * embedded in another object. If tsi_ssl_handshaker_factory_init allocates
+   * any memory, it should be free'd here. */
+}
+
+static tsi_ssl_handshaker_factory *tsi_ssl_handshaker_factory_ref(
+    tsi_ssl_handshaker_factory *self) {
+  if (self == NULL) return NULL;
+  gpr_refn(&self->refcount, 1);
+  return self;
+}
+
+static void tsi_ssl_handshaker_factory_unref(tsi_ssl_handshaker_factory *self) {
+  if (self == NULL) return;
+
+  if (gpr_unref(&self->refcount)) {
+    tsi_ssl_handshaker_factory_destroy(self);
+  }
+}
+
+static tsi_ssl_handshaker_factory_vtable handshaker_factory_vtable = {NULL};
+
+/* Initializes a tsi_ssl_handshaker_factory object. Caller is responsible for
+ * allocating memory for the factory. */
+static void tsi_ssl_handshaker_factory_init(
+    tsi_ssl_handshaker_factory *factory) {
+  GPR_ASSERT(factory != NULL);
+
+  factory->vtable = &handshaker_factory_vtable;
+  gpr_ref_init(&factory->refcount, 1);
+}
+
 /* --- tsi_handshaker methods implementation. ---*/
 
 static tsi_result ssl_handshaker_get_bytes_to_send_to_peer(tsi_handshaker *self,
@@ -1013,6 +1062,7 @@ static tsi_result ssl_handshaker_create_frame_protector(
 static void ssl_handshaker_destroy(tsi_handshaker *self) {
   tsi_ssl_handshaker *impl = (tsi_ssl_handshaker *)self;
   SSL_free(impl->ssl); /* The BIO objects are owned by ssl */
+  tsi_ssl_handshaker_factory_unref(impl->factory_ref);
   gpr_free(impl);
 }
 
@@ -1030,6 +1080,7 @@ static const tsi_handshaker_vtable handshaker_vtable = {
 
 static tsi_result create_tsi_ssl_handshaker(SSL_CTX *ctx, int is_client,
                                             const char *server_name_indication,
+                                            tsi_ssl_handshaker_factory *factory,
                                             tsi_handshaker **handshaker) {
   SSL *ssl = SSL_new(ctx);
   BIO *into_ssl = NULL;
@@ -1085,6 +1136,8 @@ static tsi_result create_tsi_ssl_handshaker(SSL_CTX *ctx, int is_client,
   impl->from_ssl = from_ssl;
   impl->result = TSI_HANDSHAKE_IN_PROGRESS;
   impl->base.vtable = &handshaker_vtable;
+  impl->factory_ref = tsi_ssl_handshaker_factory_ref(factory);
+
   *handshaker = &impl->base;
   return TSI_OK;
 }
@@ -1121,11 +1174,20 @@ tsi_result tsi_ssl_client_handshaker_factory_create_handshaker(
     tsi_ssl_client_handshaker_factory *self, const char *server_name_indication,
     tsi_handshaker **handshaker) {
   return create_tsi_ssl_handshaker(self->ssl_context, 1, server_name_indication,
-                                   handshaker);
+                                   &self->base, handshaker);
 }
 
-void tsi_ssl_client_handshaker_factory_destroy(
+void tsi_ssl_client_handshaker_factory_unref(
     tsi_ssl_client_handshaker_factory *self) {
+  if (self == NULL) return;
+  tsi_ssl_handshaker_factory_unref(&self->base);
+}
+
+static void tsi_ssl_client_handshaker_factory_destroy(
+    tsi_ssl_handshaker_factory *factory) {
+  if (factory == NULL) return;
+  tsi_ssl_client_handshaker_factory *self =
+      (tsi_ssl_client_handshaker_factory *)factory;
   if (self->ssl_context != NULL) SSL_CTX_free(self->ssl_context);
   if (self->alpn_protocol_list != NULL) gpr_free(self->alpn_protocol_list);
   gpr_free(self);
@@ -1150,11 +1212,21 @@ tsi_result tsi_ssl_server_handshaker_factory_create_handshaker(
   if (self->ssl_context_count == 0) return TSI_INVALID_ARGUMENT;
   /* Create the handshaker with the first context. We will switch if needed
      because of SNI in ssl_server_handshaker_factory_servername_callback.  */
-  return create_tsi_ssl_handshaker(self->ssl_contexts[0], 0, NULL, handshaker);
+  return create_tsi_ssl_handshaker(self->ssl_contexts[0], 0, NULL, &self->base,
+                                   handshaker);
 }
 
-void tsi_ssl_server_handshaker_factory_destroy(
+void tsi_ssl_server_handshaker_factory_unref(
     tsi_ssl_server_handshaker_factory *self) {
+  if (self == NULL) return;
+  tsi_ssl_handshaker_factory_unref(&self->base);
+}
+
+static void tsi_ssl_server_handshaker_factory_destroy(
+    tsi_ssl_handshaker_factory *factory) {
+  if (factory == NULL) return;
+  tsi_ssl_server_handshaker_factory *self =
+      (tsi_ssl_server_handshaker_factory *)factory;
   size_t i;
   for (i = 0; i < self->ssl_context_count; i++) {
     if (self->ssl_contexts[i] != NULL) {
@@ -1263,6 +1335,9 @@ static int server_handshaker_factory_npn_advertised_callback(
 
 /* --- tsi_ssl_handshaker_factory constructors. --- */
 
+static tsi_ssl_handshaker_factory_vtable client_handshaker_factory_vtable = {
+    tsi_ssl_client_handshaker_factory_destroy};
+
 tsi_result tsi_create_ssl_client_handshaker_factory(
     const tsi_ssl_pem_key_cert_pair *pem_key_cert_pair,
     const char *pem_root_certs, const char *cipher_suites,
@@ -1285,6 +1360,9 @@ tsi_result tsi_create_ssl_client_handshaker_factory(
   }
 
   impl = gpr_zalloc(sizeof(*impl));
+  tsi_ssl_handshaker_factory_init(&impl->base);
+  impl->base.vtable = &client_handshaker_factory_vtable;
+
   impl->ssl_context = ssl_context;
 
   do {
@@ -1322,7 +1400,7 @@ tsi_result tsi_create_ssl_client_handshaker_factory(
     }
   } while (0);
   if (result != TSI_OK) {
-    tsi_ssl_client_handshaker_factory_destroy(impl);
+    tsi_ssl_handshaker_factory_unref(&impl->base);
     return result;
   }
   SSL_CTX_set_verify(ssl_context, SSL_VERIFY_PEER, NULL);
@@ -1332,6 +1410,9 @@ tsi_result tsi_create_ssl_client_handshaker_factory(
   return TSI_OK;
 }
 
+static tsi_ssl_handshaker_factory_vtable server_handshaker_factory_vtable = {
+    tsi_ssl_server_handshaker_factory_destroy};
+
 tsi_result tsi_create_ssl_server_handshaker_factory(
     const tsi_ssl_pem_key_cert_pair *pem_key_cert_pairs,
     size_t num_key_cert_pairs, const char *pem_client_root_certs,
@@ -1364,12 +1445,15 @@ tsi_result tsi_create_ssl_server_handshaker_factory_ex(
   }
 
   impl = gpr_zalloc(sizeof(*impl));
+  tsi_ssl_handshaker_factory_init(&impl->base);
+  impl->base.vtable = &server_handshaker_factory_vtable;
+
   impl->ssl_contexts = gpr_zalloc(num_key_cert_pairs * sizeof(SSL_CTX *));
   impl->ssl_context_x509_subject_names =
       gpr_zalloc(num_key_cert_pairs * sizeof(tsi_peer));
   if (impl->ssl_contexts == NULL ||
       impl->ssl_context_x509_subject_names == NULL) {
-    tsi_ssl_server_handshaker_factory_destroy(impl);
+    tsi_ssl_handshaker_factory_unref(&impl->base);
     return TSI_OUT_OF_RESOURCES;
   }
   impl->ssl_context_count = num_key_cert_pairs;
@@ -1379,7 +1463,7 @@ tsi_result tsi_create_ssl_server_handshaker_factory_ex(
                                            &impl->alpn_protocol_list,
                                            &impl->alpn_protocol_list_length);
     if (result != TSI_OK) {
-      tsi_ssl_server_handshaker_factory_destroy(impl);
+      tsi_ssl_handshaker_factory_unref(&impl->base);
       return result;
     }
   }
@@ -1451,10 +1535,11 @@ tsi_result tsi_create_ssl_server_handshaker_factory_ex(
     } while (0);
 
     if (result != TSI_OK) {
-      tsi_ssl_server_handshaker_factory_destroy(impl);
+      tsi_ssl_handshaker_factory_unref(&impl->base);
       return result;
     }
   }
+
   *factory = impl;
   return TSI_OK;
 }
@@ -1501,3 +1586,15 @@ int tsi_ssl_peer_matches_name(const tsi_peer *peer, const char *name) {
 
   return 0; /* Not found. */
 }
+
+/* --- Testing support. --- */
+const tsi_ssl_handshaker_factory_vtable *tsi_ssl_handshaker_factory_swap_vtable(
+    tsi_ssl_handshaker_factory *factory,
+    tsi_ssl_handshaker_factory_vtable *new_vtable) {
+  GPR_ASSERT(factory != NULL);
+  GPR_ASSERT(factory->vtable != NULL);
+
+  const tsi_ssl_handshaker_factory_vtable *orig_vtable = factory->vtable;
+  factory->vtable = new_vtable;
+  return orig_vtable;
+}