Fix initialization errors in unit tests (#4370)

aws · Jan 19, 2024 · 670cb43 · 670cb43
1 parent 5744711
commit 670cb43
Show file tree

Hide file tree

Showing 5 changed files with 79 additions and 22 deletions.
diff --git a/.github/s2n_osx.sh b/.github/s2n_osx.sh
@@ -28,3 +28,12 @@ cmake . -Bbuild -GNinja \
 
 cmake --build ./build -j $(nproc)
 time CTEST_PARALLEL_LEVEL=$(nproc) ninja -C build test
+
+# Build shared library
+cmake . -Bbuild -GNinja \
+-DCMAKE_BUILD_TYPE=Debug \
+-DCMAKE_PREFIX_PATH=${OPENSSL_1_1_1_INSTALL_DIR} .. \
+-DBUILD_SHARED_LIBS=ON
+
+cmake --build ./build -j $(nproc)
+time CTEST_PARALLEL_LEVEL=$(nproc) ninja -C build test
diff --git a/tests/testlib/s2n_mem_testlib.c b/tests/testlib/s2n_mem_testlib.c
@@ -16,9 +16,8 @@
 #include "testlib/s2n_mem_testlib.h"
 
 #include "stuffer/s2n_stuffer.h"
-
-/* Required to override memory callbacks at runtime */
-#include "utils/s2n_mem.c"
+#include "utils/s2n_mem.h"
+#include "utils/s2n_safety.h"
 
 struct s2n_mem_test_cb_ctx {
     struct s2n_stuffer mallocs;
@@ -32,25 +31,43 @@ static int s2n_mem_test_free_cb(void *ptr, uint32_t size);
 
 static S2N_RESULT s2n_mem_test_set_callbacks()
 {
-    if (s2n_mem_malloc_cb != s2n_mem_test_malloc_cb) {
-        s2n_mem_malloc_cb_backup = s2n_mem_malloc_cb;
-        s2n_mem_malloc_cb = s2n_mem_test_malloc_cb;
+    s2n_mem_init_callback mem_init_cb = NULL;
+    s2n_mem_cleanup_callback mem_cleanup_cb = NULL;
+    s2n_mem_malloc_callback mem_malloc_cb = NULL;
+    s2n_mem_free_callback mem_free_cb = NULL;
+    RESULT_GUARD(s2n_mem_get_callbacks(&mem_init_cb, &mem_cleanup_cb, &mem_malloc_cb, &mem_free_cb));
+
+    if (mem_malloc_cb != s2n_mem_test_malloc_cb) {
+        s2n_mem_malloc_cb_backup = mem_malloc_cb;
+        mem_malloc_cb = s2n_mem_test_malloc_cb;
     }
-    if (s2n_mem_free_cb != s2n_mem_test_free_cb) {
-        s2n_mem_free_cb_backup = s2n_mem_free_cb;
-        s2n_mem_free_cb = s2n_mem_test_free_cb;
+    if (mem_free_cb != s2n_mem_test_free_cb) {
+        s2n_mem_free_cb_backup = mem_free_cb;
+        mem_free_cb = s2n_mem_test_free_cb;
     }
+
+    RESULT_GUARD(s2n_mem_override_callbacks(mem_init_cb, mem_cleanup_cb, mem_malloc_cb, mem_free_cb));
+
     return S2N_RESULT_OK;
 }
 
 static S2N_RESULT s2n_mem_test_unset_callbacks()
 {
+    s2n_mem_init_callback mem_init_cb = NULL;
+    s2n_mem_cleanup_callback mem_cleanup_cb = NULL;
+    s2n_mem_malloc_callback mem_malloc_cb = NULL;
+    s2n_mem_free_callback mem_free_cb = NULL;
+    RESULT_GUARD(s2n_mem_get_callbacks(&mem_init_cb, &mem_cleanup_cb, &mem_malloc_cb, &mem_free_cb));
+
     if (s2n_mem_malloc_cb_backup != NULL) {
-        s2n_mem_malloc_cb = s2n_mem_malloc_cb_backup;
+        mem_malloc_cb = s2n_mem_malloc_cb_backup;
     }
     if (s2n_mem_free_cb_backup != NULL) {
-        s2n_mem_free_cb = s2n_mem_free_cb_backup;
+        mem_free_cb = s2n_mem_free_cb_backup;
     }
+
+    RESULT_GUARD(s2n_mem_override_callbacks(mem_init_cb, mem_cleanup_cb, mem_malloc_cb, mem_free_cb));
+
     return S2N_RESULT_OK;
 }
 

diff --git a/tests/unit/s2n_mem_test.c b/tests/unit/s2n_mem_test.c
@@ -13,12 +13,11 @@
  * permissions and limitations under the License.
  */
 
+#include "utils/s2n_mem.h"
+
 #include "s2n_test.h"
 #include "utils/s2n_safety.h"
 
-/* Required to override memory callbacks at runtime */
-#include "utils/s2n_mem.c"
-
 int s2n_strict_mem_free_cb(void *ptr, uint32_t size)
 {
     POSIX_ENSURE_REF(ptr);
@@ -69,10 +68,14 @@ int main(int argc, char **argv)
          */
         {
             /* Save real free callback */
-            s2n_mem_free_callback saved_free_cb = s2n_mem_free_cb;
+            s2n_mem_init_callback mem_init_cb = NULL;
+            s2n_mem_cleanup_callback mem_cleanup_cb = NULL;
+            s2n_mem_malloc_callback mem_malloc_cb = NULL;
+            s2n_mem_free_callback mem_free_cb = NULL;
+            EXPECT_OK(s2n_mem_get_callbacks(&mem_init_cb, &mem_cleanup_cb, &mem_malloc_cb, &mem_free_cb));
 
             /* Set callback that won't accepts NULLs / zeros */
-            s2n_mem_free_cb = s2n_strict_mem_free_cb;
+            EXPECT_OK(s2n_mem_override_callbacks(mem_init_cb, mem_cleanup_cb, mem_malloc_cb, s2n_strict_mem_free_cb));
 
             /* No-op for empty blob */
             struct s2n_blob blob = { 0 };
@@ -81,7 +84,7 @@ int main(int argc, char **argv)
             EXPECT_SUCCESS(s2n_free_or_wipe(&blob));
 
             /* Restore real free callback */
-            s2n_mem_free_cb = saved_free_cb;
+            EXPECT_OK(s2n_mem_override_callbacks(mem_init_cb, mem_cleanup_cb, mem_malloc_cb, mem_free_cb));
         };
     };
 

diff --git a/utils/s2n_mem.c b/utils/s2n_mem.c
@@ -135,18 +135,41 @@ int s2n_mem_set_callbacks(s2n_mem_init_callback mem_init_callback, s2n_mem_clean
         s2n_mem_malloc_callback mem_malloc_callback, s2n_mem_free_callback mem_free_callback)
 {
     POSIX_ENSURE(!initialized, S2N_ERR_INITIALIZED);
+    POSIX_GUARD_RESULT(s2n_mem_override_callbacks(mem_init_callback, mem_cleanup_callback,
+            mem_malloc_callback, mem_free_callback));
+    return S2N_SUCCESS;
+}
 
-    POSIX_ENSURE_REF(mem_init_callback);
-    POSIX_ENSURE_REF(mem_cleanup_callback);
-    POSIX_ENSURE_REF(mem_malloc_callback);
-    POSIX_ENSURE_REF(mem_free_callback);
+S2N_RESULT s2n_mem_override_callbacks(s2n_mem_init_callback mem_init_callback, s2n_mem_cleanup_callback mem_cleanup_callback,
+        s2n_mem_malloc_callback mem_malloc_callback, s2n_mem_free_callback mem_free_callback)
+{
+    RESULT_ENSURE_REF(mem_init_callback);
+    RESULT_ENSURE_REF(mem_cleanup_callback);
+    RESULT_ENSURE_REF(mem_malloc_callback);
+    RESULT_ENSURE_REF(mem_free_callback);
 
     s2n_mem_init_cb = mem_init_callback;
     s2n_mem_cleanup_cb = mem_cleanup_callback;
     s2n_mem_malloc_cb = mem_malloc_callback;
     s2n_mem_free_cb = mem_free_callback;
 
-    return S2N_SUCCESS;
+    return S2N_RESULT_OK;
+}
+
+S2N_RESULT s2n_mem_get_callbacks(s2n_mem_init_callback *mem_init_callback, s2n_mem_cleanup_callback *mem_cleanup_callback,
+        s2n_mem_malloc_callback *mem_malloc_callback, s2n_mem_free_callback *mem_free_callback)
+{
+    RESULT_ENSURE_REF(mem_init_callback);
+    RESULT_ENSURE_REF(mem_cleanup_callback);
+    RESULT_ENSURE_REF(mem_malloc_callback);
+    RESULT_ENSURE_REF(mem_free_callback);
+
+    *mem_init_callback = s2n_mem_init_cb;
+    *mem_cleanup_callback = s2n_mem_cleanup_cb;
+    *mem_malloc_callback = s2n_mem_malloc_cb;
+    *mem_free_callback = s2n_mem_free_cb;
+
+    return S2N_RESULT_OK;
 }
 
 int s2n_alloc(struct s2n_blob *b, uint32_t size)

diff --git a/utils/s2n_mem.h b/utils/s2n_mem.h
@@ -44,3 +44,8 @@ int s2n_dup(struct s2n_blob *from, struct s2n_blob *to);
  * Prefer s2n_free. Only use this method if completely necessary.
  */
 int s2n_free_or_wipe(struct s2n_blob *b);
+
+S2N_RESULT s2n_mem_override_callbacks(s2n_mem_init_callback mem_init_callback, s2n_mem_cleanup_callback mem_cleanup_callback,
+        s2n_mem_malloc_callback mem_malloc_callback, s2n_mem_free_callback mem_free_callback);
+S2N_RESULT s2n_mem_get_callbacks(s2n_mem_init_callback *mem_init_callback, s2n_mem_cleanup_callback *mem_cleanup_callback,
+        s2n_mem_malloc_callback *mem_malloc_callback, s2n_mem_free_callback *mem_free_callback);