feat: add s2n_cleanup_thread (#4584)

api/unstable/cleanup.h

@@ -0,0 +1,29 @@
+/*
+ * 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.
+ */
+
+#pragma once
+
+#include <s2n.h>
+
+/**
+ * Cleans up any internal thread-local resources used by s2n-tls. This function
+ * is called by `s2n_cleanup`, but depending on your thread management model, 
+ * it may be called directly instead. 
+ * 
+ * See [Initialization](https://github.com/aws/s2n-tls/blob/main/docs/usage-guide/topics/ch02-initializing.md) for details.
+ *
+ * @returns S2N_SUCCESS on success. S2N_FAILURE on failure
+ */
+S2N_API extern int s2n_cleanup_thread(void);

docs/DEVELOPMENT-GUIDE.md

@@ -162,7 +162,9 @@ s2n states publicly that every `s2n_init()` call should be paired with an `s2n_c
 For every thread that s2n functions are called in, a small amount of thread-local memory also gets initialized. This is to ensure that our random number generator will output different numbers in different threads. This memory needs to be cleaned up per thread and users can do this themselves if they call `s2n_cleanup()` per thread. But if they forget, we utilize a pthread key that calls a destructor function that cleans up our thread-local memory when the thread closes.
 
 An important thing to note is that a call to `s2n_cleanup()` usually does not fully clean up s2n. It only cleans up the thread-local memory. This is because we have an atexit handler that does fully clean up s2n at process-exit.
-The behavior is different if the atexit handler is disabled by calling `s2n_disable_atexit()`. Then s2n is actually fully cleaned up if `s2n_cleanup()` is called on the thread that called `s2n_init()`.
+The behavior is different if the atexit handler is disabled by calling `s2n_disable_atexit()`. Then s2n is actually fully cleaned up if `s2n_cleanup()` is called on the thread that called `s2n_init()`. `s2n_cleanup()` attempts to track the thread that had originally called `s2n_init()` and only performs the full cleanup when that "main" thread exits. If this behavior does not work for your use case (for example, if the "main" thread does not outlive child threads using s2n), you can instead call `s2n_cleanup_thread()` on each thread exit.
+
+> Note: `s2n_cleanup_thread()` is currently considered unstable, meaning the API is subject to change in a future release. To access this API, include `api/unstable/cleanup.h`.
 
 ### Control flow and the state machine
 

docs/usage-guide/topics/ch02-initialization.md

@@ -1,5 +1,7 @@
 # Initialization and Teardown
-The s2n-tls library must be initialized with `s2n_init()` before calling most library functions. `s2n_init()` MUST NOT be called more than once, even when an application uses multiple threads or processes. s2n attempts to clean up its thread-local memory at thread-exit and all other memory at process-exit. However, this may not work if you are using a thread library other than pthreads. In that case you should call `s2n_cleanup()` from every thread or process created after `s2n_init()`.
+The s2n-tls library must be initialized with `s2n_init()` before calling most library functions. `s2n_init()` MUST NOT be called more than once, even when an application uses multiple threads or processes. s2n attempts to clean up its thread-local memory at thread-exit and all other memory at process-exit. However, this may not work if you are using a thread library other than pthreads or other threads using s2n outlive the thread that called `s2n_init()`. In that case you should call `s2n_cleanup_thread()` from every thread or process created after `s2n_init()`.
+
+> Note: `s2n_cleanup_thread()` is currently considered unstable, meaning the API is subject to change in a future release. To access this API, include `api/unstable/cleanup.h`.
 
 Initialization can be modified by calling `s2n_crypto_disable_init()` or `s2n_disable_atexit()` before `s2n_init()`.
 

tests/unit/s2n_init_test.c

@@ -15,8 +15,12 @@
 
 #include <pthread.h>
 
+#include "api/unstable/cleanup.h"
 #include "s2n_test.h"
 
+bool s2n_is_initialized(void);
+int s2n_enable_atexit(void);
+
 static void *s2n_init_fail_cb(void *_unused_arg)
 {
     (void) _unused_arg;
@@ -33,7 +37,13 @@ static void *s2n_init_success_cb(void *_unused_arg)
     return NULL;
 }
 
-int s2n_enable_atexit(void);
+static void *s2n_cleanup_thread_cb(void *_unused_arg)
+{
+    (void) _unused_arg;
+
+    EXPECT_SUCCESS(s2n_cleanup_thread());
+    return NULL;
+}
 
 int main(int argc, char **argv)
 {
@@ -77,6 +87,23 @@ int main(int argc, char **argv)
     EXPECT_EQUAL(pthread_join(init_thread, NULL), 0);
     EXPECT_SUCCESS(s2n_cleanup());
 
+    /* Calling s2n_init/s2n_cleanup in a different thread than s2n_cleanup_thread is called cleans up properly */
+    {
+        EXPECT_SUCCESS(s2n_init());
+        EXPECT_TRUE(s2n_is_initialized());
+        pthread_t s2n_cleanup_th = { 0 };
+        EXPECT_EQUAL(pthread_create(&s2n_cleanup_th, NULL, s2n_cleanup_thread_cb, NULL), 0);
+        EXPECT_EQUAL(pthread_join(s2n_cleanup_th, NULL), 0);
+        EXPECT_TRUE(s2n_is_initialized());
+        /* Calling s2n_cleanup_thread in the main thread leaves s2n initialized. */
+        EXPECT_SUCCESS(s2n_cleanup_thread());
+        EXPECT_TRUE(s2n_is_initialized());
+        EXPECT_SUCCESS(s2n_cleanup());
+        EXPECT_FALSE(s2n_is_initialized());
+        /* Second call to s2n_cleanup will fail, since the full cleanup is not idempotent. */
+        EXPECT_FAILURE_WITH_ERRNO(s2n_cleanup(), S2N_ERR_NOT_INITIALIZED);
+    }
+
     /* The following test requires atexit to be enabled. */
     EXPECT_SUCCESS(s2n_enable_atexit());
 

utils/s2n_init.c

@@ -12,8 +12,12 @@
  * express or implied. See the License for the specific language governing
  * permissions and limitations under the License.
  */
+
+#include "utils/s2n_init.h"
+
 #include <pthread.h>
 
+#include "api/unstable/cleanup.h"
 #include "crypto/s2n_fips.h"
 #include "crypto/s2n_libcrypto.h"
 #include "crypto/s2n_locking.h"
@@ -114,18 +118,31 @@ static bool s2n_cleanup_atexit_impl(void)
     return cleaned_up;
 }
 
-int s2n_cleanup(void)
+int s2n_cleanup_thread(void)
 {
-    /* s2n_cleanup is supposed to be called from each thread before exiting,
+    /* s2n_cleanup_thread is supposed to be called from each thread before exiting,
      * so ensure that whatever clean ups we have here are thread safe */
     POSIX_GUARD_RESULT(s2n_rand_cleanup_thread());
+    return S2N_SUCCESS;
+}
+
+int s2n_cleanup_final(void)
+{
+    /* some cleanups are not idempotent (rand_cleanup, mem_cleanup) so protect */
+    POSIX_ENSURE(initialized, S2N_ERR_NOT_INITIALIZED);
+    POSIX_ENSURE(s2n_cleanup_atexit_impl(), S2N_ERR_ATEXIT);
+
+    return S2N_SUCCESS;
+}
+
+int s2n_cleanup(void)
+{
+    POSIX_GUARD(s2n_cleanup_thread());
 
     /* If this is the main thread and atexit cleanup is disabled,
      * perform final cleanup now */
     if (pthread_equal(pthread_self(), main_thread) && !atexit_cleanup) {
-        /* some cleanups are not idempotent (rand_cleanup, mem_cleanup) so protect */
-        POSIX_ENSURE(initialized, S2N_ERR_NOT_INITIALIZED);
-        POSIX_ENSURE(s2n_cleanup_atexit_impl(), S2N_ERR_ATEXIT);
+        POSIX_GUARD(s2n_cleanup_final());
     }
 
     return 0;

utils/s2n_init.h

@@ -15,6 +15,8 @@
 
 #pragma once
 
+#include <stdbool.h>
+
 int s2n_init(void);
 int s2n_cleanup(void);
 bool s2n_is_initialized(void);