Fix use-after-free

1 files changed, 21 insertions, 11 deletions

diff --git a/src/core/lib/security/transport/handshake.c b/src/core/lib/security/transport/handshake.c
index ff338f8084..b374ca7633 100644
--- a/src/core/lib/security/transport/handshake.c
+++ b/src/core/lib/security/transport/handshake.c
@@ -63,6 +63,7 @@ typedef struct {
   grpc_closure on_handshake_data_received_from_peer;
   grpc_auth_context *auth_context;
   grpc_timer timer;
+  gpr_refcount refs;
 } grpc_security_handshake;
 
 static void on_handshake_data_received_from_peer(grpc_exec_ctx *exec_ctx,
@@ -99,6 +100,19 @@ static void security_connector_remove_handshake(grpc_security_handshake *h) {
   gpr_mu_unlock(&sc->mu);
 }
 
+static void unref_handshake(grpc_security_handshake *h) {
+  if (gpr_unref(&h->refs)) {
+    if (h->handshaker != NULL) tsi_handshaker_destroy(h->handshaker);
+    if (h->handshake_buffer != NULL) gpr_free(h->handshake_buffer);
+    gpr_slice_buffer_destroy(&h->left_overs);
+    gpr_slice_buffer_destroy(&h->outgoing);
+    gpr_slice_buffer_destroy(&h->incoming);
+    GRPC_AUTH_CONTEXT_UNREF(h->auth_context, "handshake");
+    GRPC_SECURITY_CONNECTOR_UNREF(h->connector, "handshake");
+    gpr_free(h);
+  }
+}
+
 static void security_handshake_done(grpc_exec_ctx *exec_ctx,
                                     grpc_security_handshake *h,
                                     grpc_error *error) {
@@ -122,14 +136,7 @@ static void security_handshake_done(grpc_exec_ctx *exec_ctx,
     }
     h->cb(exec_ctx, h->user_data, GRPC_SECURITY_ERROR, NULL, NULL);
   }
-  if (h->handshaker != NULL) tsi_handshaker_destroy(h->handshaker);
-  if (h->handshake_buffer != NULL) gpr_free(h->handshake_buffer);
-  gpr_slice_buffer_destroy(&h->left_overs);
-  gpr_slice_buffer_destroy(&h->outgoing);
-  gpr_slice_buffer_destroy(&h->incoming);
-  GRPC_AUTH_CONTEXT_UNREF(h->auth_context, "handshake");
-  GRPC_SECURITY_CONNECTOR_UNREF(h->connector, "handshake");
-  gpr_free(h);
+  unref_handshake(h);
   GRPC_ERROR_UNREF(error);
 }
 
@@ -308,9 +315,11 @@ static void on_handshake_data_sent_to_peer(grpc_exec_ctx *exec_ctx,
 }
 
 static void on_timeout(grpc_exec_ctx *exec_ctx, void *arg, grpc_error *error) {
+  grpc_security_handshake *h = arg;
   if (error == GRPC_ERROR_NONE) {
-    grpc_endpoint_shutdown(exec_ctx, arg);
+    grpc_endpoint_shutdown(exec_ctx, h->wrapped_endpoint);
   }
+  unref_handshake(h);
 }
 
 void grpc_do_security_handshake(
@@ -329,6 +338,7 @@ void grpc_do_security_handshake(
   h->wrapped_endpoint = nonsecure_endpoint;
   h->user_data = user_data;
   h->cb = cb;
+  gpr_ref_init(&h->refs, 2); /* timer and handshake proper each get a ref */
   grpc_closure_init(&h->on_handshake_data_sent_to_peer,
                     on_handshake_data_sent_to_peer, h);
   grpc_closure_init(&h->on_handshake_data_received_from_peer,
@@ -347,8 +357,8 @@ void grpc_do_security_handshake(
     gpr_mu_unlock(&server_connector->mu);
   }
   send_handshake_bytes_to_peer(exec_ctx, h);
-  grpc_timer_init(exec_ctx, &h->timer, deadline, on_timeout,
-                  h->wrapped_endpoint, gpr_now(deadline.clock_type));
+  grpc_timer_init(exec_ctx, &h->timer, deadline, on_timeout, h,
+                  gpr_now(deadline.clock_type));
 }
 
 void grpc_security_handshake_shutdown(grpc_exec_ctx *exec_ctx,