From 0baf6e7b011f4723f37539dc184cc8fb1018bf1a Mon Sep 17 00:00:00 2001
From: William Yang <mscame@gmail.com>
Date: Fri, 20 Oct 2023 16:09:19 +0200
Subject: [PATCH] feat: thread safe r_ctx

---
 c_src/quicer_connection.c |  5 +++++
 c_src/quicer_ctx.c        | 13 +++++++++++++
 c_src/quicer_ctx.h        |  5 +++++
 c_src/quicer_listener.c   | 10 ++++++++++
 c_src/quicer_reg.c        | 11 +++++++----
 5 files changed, 40 insertions(+), 4 deletions(-)

diff --git a/c_src/quicer_connection.c b/c_src/quicer_connection.c
index 1d804eb4..de9ed2c4 100644
--- a/c_src/quicer_connection.c
+++ b/c_src/quicer_connection.c
@@ -554,6 +554,11 @@ open_connectionX(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[])
       r_ctx = G_r_ctx;
     }
 
+  if (!get_reg_handle(r_ctx))
+    {
+      return ERROR_TUPLE_2(ATOM_QUIC_REGISTRATION);
+    }
+
   QuicerConnCTX *c_ctx = init_c_ctx();
 
   if (!c_ctx)
diff --git a/c_src/quicer_ctx.c b/c_src/quicer_ctx.c
index d4e45921..e1ecdd0c 100644
--- a/c_src/quicer_ctx.c
+++ b/c_src/quicer_ctx.c
@@ -110,6 +110,7 @@ destroy_l_ctx(QuicerListenerCTX *l_ctx)
 
   if (r_ctx)
     {
+      put_reg_handle(r_ctx);
       enif_mutex_lock(r_ctx->lock);
       CxPlatListEntryRemove(&l_ctx->RegistrationLink);
       enif_mutex_unlock(r_ctx->lock);
@@ -361,3 +362,15 @@ get_listener_handle(QuicerListenerCTX *l_ctx)
 {
   return CxPlatRefIncrementNonZero(&l_ctx->ref_count, 1);
 }
+
+inline void
+put_reg_handle(QuicerRegistrationCTX *r_ctx)
+{
+  CxPlatRefDecrement(&r_ctx->ref_count);
+}
+
+inline BOOLEAN
+get_reg_handle(QuicerRegistrationCTX *r_ctx)
+{
+  return CxPlatRefIncrementNonZero(&r_ctx->ref_count, 1);
+}
diff --git a/c_src/quicer_ctx.h b/c_src/quicer_ctx.h
index 5a914116..080c7bb2 100644
--- a/c_src/quicer_ctx.h
+++ b/c_src/quicer_ctx.h
@@ -34,6 +34,8 @@ typedef struct QuicerRegistrationCTX
 {
   ErlNifEnv *env;
   HQUIC Registration;
+  // Tracking lifetime of Registration handle
+  CXPLAT_REF_COUNT ref_count;
   BOOLEAN is_released;
   char name[UINT8_MAX + 1];
   ErlNifMutex *lock;
@@ -182,4 +184,7 @@ BOOLEAN get_conn_handle(QuicerConnCTX *c_ctx);
 void put_listener_handle(QuicerListenerCTX *l_ctx);
 BOOLEAN get_listener_handle(QuicerListenerCTX *l_ctx);
 
+void put_reg_handle(QuicerRegistrationCTX *r_ctx);
+BOOLEAN get_reg_handle(QuicerRegistrationCTX *r_ctx);
+
 #endif // __QUICER_CTX_H_
diff --git a/c_src/quicer_listener.c b/c_src/quicer_listener.c
index f5b66802..7234dfc5 100644
--- a/c_src/quicer_listener.c
+++ b/c_src/quicer_listener.c
@@ -334,6 +334,11 @@ listen2(ErlNifEnv *env, __unused_parm__ int argc, const ERL_NIF_TERM argv[])
     {
       // quic_registration is set
       enif_keep_resource(l_ctx->r_ctx);
+      if (!get_reg_handle(l_ctx->r_ctx))
+        {
+          ret = ERROR_TUPLE_2(ATOM_QUIC_REGISTRATION);
+          goto exit;
+        }
       Registration = l_ctx->r_ctx->Registration;
       target_r_ctx = l_ctx->r_ctx;
     }
@@ -351,6 +356,11 @@ listen2(ErlNifEnv *env, __unused_parm__ int argc, const ERL_NIF_TERM argv[])
         }
 
       enif_keep_resource(G_r_ctx);
+      if (!get_reg_handle(G_r_ctx))
+        {
+          ret = ERROR_TUPLE_2(ATOM_QUIC_REGISTRATION);
+          goto exit;
+        }
       Registration = G_r_ctx->Registration;
       pthread_mutex_unlock(&GRegLock);
     }
diff --git a/c_src/quicer_reg.c b/c_src/quicer_reg.c
index 7377edda..9b87cc6f 100644
--- a/c_src/quicer_reg.c
+++ b/c_src/quicer_reg.c
@@ -23,7 +23,7 @@ QuicerRegistrationCTX *G_r_ctx = NULL;
 pthread_mutex_t GRegLock = PTHREAD_MUTEX_INITIALIZER;
 
 /*
-** For global registration only
+** Open global registration.
 */
 ERL_NIF_TERM
 registration(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[])
@@ -64,7 +64,7 @@ registration(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[])
       res = ERROR_TUPLE_2(ATOM_STATUS(status));
       goto exit;
     }
-
+  CxPlatRefInitialize(&r_ctx->ref_count);
   G_r_ctx = r_ctx;
   pthread_mutex_unlock(&GRegLock);
 
@@ -147,7 +147,7 @@ new_registration2(ErlNifEnv *env, int argc, const ERL_NIF_TERM argv[])
       res = ERROR_TUPLE_2(ATOM_STATUS(status));
       goto exit;
     }
-
+  CxPlatRefInitialize(&r_ctx->ref_count);
   return SUCCESS(enif_make_resource(env, r_ctx));
 
 exit:
@@ -231,7 +231,10 @@ get_registration_name1(ErlNifEnv *env,
       return ERROR_TUPLE_2(ATOM_BADARG);
     }
 
-  return SUCCESS(enif_make_string(env, r_ctx->name, ERL_NIF_LATIN1));
+  enif_mutex_lock(r_ctx->lock);
+  ERL_NIF_TERM name = enif_make_string(env, r_ctx->name, ERL_NIF_LATIN1);
+  enif_mutex_unlock(r_ctx->lock);
+  return SUCCESS(name);
 }
 
 BOOLEAN