s2n_drbg: de-allocate thread-local allocations at thread exit

tests/unit/s2n_drbg_init_test.c

@@ -0,0 +1,56 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License").
+ * You may not use this file except in compliance with the License.
+ * A copy of the License is located at
+ *
+ *  http://aws.amazon.com/apache2.0
+ *
+ * or in the "license" file accompanying this file. This file is distributed
+ * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ * express or implied. See the License for the specific language governing
+ * permissions and limitations under the License.
+ */
+
+#include "s2n_test.h"
+#include "api/s2n.h"
+#include "utils/s2n_random.h"
+
+#include <pthread.h>
+
+/**
+ * This is an ASAN test which ensures that if s2n_init_drbgs() (via s2n_rand_init()) and
+ * s2n_rand_cleanup_thread() are called from different threads, the thread-local memory
+ * allocated in s2n_init_drbgs() will leak if no thread exit handler is configured.
+ */
+static void * s2n_drbg_thread_initialization_cb(void *_unused_arg)
+{
+    (void)_unused_arg;
+
+    /* Called from a separate thread: */
+    EXPECT_OK(s2n_rand_init());
+    return NULL;
+}
+
+static S2N_RESULT s2n_drbg_thread_initialization(void) 
+{
+    pthread_t init_thread;
+
+    EXPECT_EQUAL(pthread_create(&init_thread, NULL, s2n_drbg_thread_initialization_cb, NULL), 0);
+    EXPECT_EQUAL(pthread_join(init_thread, NULL), 0);
+
+    return S2N_RESULT_OK;
+}
+
+int main(int argc, char **argv)
+{
+    BEGIN_TEST_NO_INIT();
+
+    EXPECT_OK(s2n_drbg_thread_initialization());
+    /* s2n_rand_cleanup is called from main only. It has no effect on the thread-local allocations
+     * made by s2n_init_drbgs() via s2n_rand_init() in the separate thread above.  */
+    EXPECT_OK(s2n_rand_cleanup());
+
+    END_TEST_NO_INIT();
+}

utils/s2n_random.c

@@ -70,6 +70,10 @@ static __thread struct s2n_rand_state s2n_per_thread_rand_state = {
     .drbgs_initialized = false
 };
 
+/* These 2 variables control the de-allocation of memory to fields inside of s2n_per_thread_rand_state: */
+static pthread_key_t s2n_per_thread_rand_state_alloc_key;
+static pthread_once_t s2n_per_thread_alloc_key_once = PTHREAD_ONCE_INIT;
+
 static int s2n_rand_init_impl(void);
 static int s2n_rand_cleanup_impl(void);
 static int s2n_rand_urandom_impl(void *ptr, uint32_t size);
@@ -130,18 +134,37 @@ S2N_RESULT s2n_get_mix_entropy(struct s2n_blob *blob)
     return S2N_RESULT_OK;
 }
 
+static void s2n_drbg_destructor(void *_unused_argument) {
+    (void)_unused_argument;
+    s2n_result_ignore(s2n_rand_cleanup_thread());
+}
+
+static void s2n_drbg_make_alloc_key(void)
+{
+    (void)pthread_key_create(&s2n_per_thread_rand_state_alloc_key, s2n_drbg_destructor);
+}
+
 static S2N_RESULT s2n_init_drbgs(void)
 {
     uint8_t s2n_public_drbg[] = "s2n public drbg";
     uint8_t s2n_private_drbg[] = "s2n private drbg";
     struct s2n_blob public = { .data = s2n_public_drbg, .size = sizeof(s2n_public_drbg) };
     struct s2n_blob private = { .data = s2n_private_drbg, .size = sizeof(s2n_private_drbg) };
 
+    RESULT_ENSURE(pthread_once(&s2n_per_thread_alloc_key_once, s2n_drbg_make_alloc_key) == 0, S2N_ERR_DRBG);
+
     RESULT_GUARD(s2n_drbg_instantiate(&s2n_per_thread_rand_state.public_drbg, &public, S2N_AES_128_CTR_NO_DF_PR));
     RESULT_GUARD(s2n_drbg_instantiate(&s2n_per_thread_rand_state.private_drbg, &private, S2N_AES_256_CTR_NO_DF_PR));
 
     s2n_per_thread_rand_state.drbgs_initialized = true;
 
+    /**
+     * Note: even though s2n_per_thread_rand_state is thread-local, s2n_drbg_instantiate() is allocating
+     * memory to fields of that struct. To avoid leaking these allocations, s2n_drbg_destructor() ensures
+     * that s2n_rand_cleanup_thread() is called when this thread exits.
+     */
+    RESULT_ENSURE(pthread_setspecific(s2n_per_thread_rand_state_alloc_key, &s2n_per_thread_rand_state) == 0, S2N_ERR_DRBG);
+
     return S2N_RESULT_OK;
 }
 
@@ -436,6 +459,9 @@ S2N_RESULT s2n_rand_cleanup_thread(void)
 {
     /* Currently, it is only safe for this function to mutate the drbg states
      * in the per thread rand state. See s2n_ensure_uniqueness().
+     *
+     * This function has no effect when called from any other thread than
+     * the one which allocated the fields that s2n_drbg_wipe de-allocates.
      */
     RESULT_GUARD(s2n_drbg_wipe(&s2n_per_thread_rand_state.private_drbg));
     RESULT_GUARD(s2n_drbg_wipe(&s2n_per_thread_rand_state.public_drbg));