diff --git a/library/bignum_core_invasive.h b/library/bignum_core_invasive.h
new file mode 100644
index 0000000..167099d
--- /dev/null
+++ b/library/bignum_core_invasive.h
@@ -0,0 +1,23 @@
+/**
+ * \file bignum_core_invasive.h
+ *
+ * \brief Function declarations for invasive functions of bignum core.
+ */
+/**
+ *  Copyright The Mbed TLS Contributors
+ *  SPDX-License-Identifier: Apache-2.0 OR GPL-2.0-or-later
+ */
+
+#ifndef MBEDTLS_BIGNUM_CORE_INVASIVE_H
+#define MBEDTLS_BIGNUM_CORE_INVASIVE_H
+
+#include "bignum_core.h"
+
+#if defined(MBEDTLS_TEST_HOOKS) && !defined(MBEDTLS_THREADING_C)
+
+extern void (*mbedtls_safe_codepath_hook)(void);
+extern void (*mbedtls_unsafe_codepath_hook)(void);
+
+#endif /* MBEDTLS_TEST_HOOKS && !MBEDTLS_THREADING_C */
+
+#endif /* MBEDTLS_BIGNUM_CORE_INVASIVE_H */
diff --git a/tests/include/test/bignum_codepath_check.h b/tests/include/test/bignum_codepath_check.h
new file mode 100644
index 0000000..6ab68bb
--- /dev/null
+++ b/tests/include/test/bignum_codepath_check.h
@@ -0,0 +1,48 @@
+/** Support for path tracking in optionally safe bignum functions
+ *
+ * The functions are called when an optionally safe path is taken and logs it with a single
+ * variable. This variable is at any time in one of three states:
+ *      - MBEDTLS_MPI_IS_TEST: No optionally safe path has been taken since the last reset
+ *      - MBEDTLS_MPI_IS_SECRET: Only safe paths were teken since the last reset
+ *      - MBEDTLS_MPI_IS_PUBLIC: At least one unsafe path has been taken since the last reset
+ *
+ * Using a simple global variable to track execution path. Making it work with multithreading
+ * doesn't worth the effort as multithreaded tests add little to no value here.
+ */
+/*
+ *  Copyright The Mbed TLS Contributors
+ *  SPDX-License-Identifier: Apache-2.0 OR GPL-2.0-or-later
+ */
+
+#ifndef BIGNUM_CODEPATH_CHECK_H
+#define BIGNUM_CODEPATH_CHECK_H
+
+#include "bignum_core.h"
+
+#if defined(MBEDTLS_TEST_HOOKS) && !defined(MBEDTLS_THREADING_C)
+
+extern int mbedtls_codepath_check;
+
+/**
+ * \brief         Setup the codepath test hooks used by optionally safe bignum functions to signal
+ *                the path taken.
+ */
+void mbedtls_codepath_test_hooks_setup(void);
+
+/**
+ * \brief         Teardown the codepath test hooks used by optionally safe bignum functions to
+ *                signal the path taken.
+ */
+void mbedtls_codepath_test_hooks_teardown(void);
+
+/**
+ * \brief         Reset the state of the codepath to the initial state.
+ */
+static inline void mbedtls_codepath_reset(void)
+{
+    mbedtls_codepath_check = MBEDTLS_MPI_IS_TEST;
+}
+
+#endif /* MBEDTLS_TEST_HOOKS && !MBEDTLS_THREADING_C */
+
+#endif /* BIGNUM_CODEPATH_CHECK_H */
diff --git a/tests/src/bignum_codepath_check.c b/tests/src/bignum_codepath_check.c
new file mode 100644
index 0000000..b6b85d9
--- /dev/null
+++ b/tests/src/bignum_codepath_check.c
@@ -0,0 +1,39 @@
+/** Support for path tracking in optionally safe bignum functions
+ */
+/*
+ *  Copyright The Mbed TLS Contributors
+ *  SPDX-License-Identifier: Apache-2.0 OR GPL-2.0-or-later
+ */
+
+#include "test/bignum_codepath_check.h"
+#include "bignum_core_invasive.h"
+
+#if defined(MBEDTLS_TEST_HOOKS) && !defined(MBEDTLS_THREADING_C)
+int mbedtls_codepath_check = MBEDTLS_MPI_IS_TEST;
+
+void mbedtls_codepath_take_safe(void)
+{
+    if(mbedtls_codepath_check == MBEDTLS_MPI_IS_TEST) {
+        mbedtls_codepath_check = MBEDTLS_MPI_IS_SECRET;
+    }
+}
+
+void mbedtls_codepath_take_unsafe(void)
+{
+    mbedtls_codepath_check = MBEDTLS_MPI_IS_PUBLIC;
+}
+
+void mbedtls_codepath_test_hooks_setup(void)
+{
+    mbedtls_safe_codepath_hook = mbedtls_codepath_take_safe;
+    mbedtls_unsafe_codepath_hook = mbedtls_codepath_take_unsafe;
+}
+
+void mbedtls_codepath_test_hooks_teardown(void)
+{
+    mbedtls_safe_codepath_hook = NULL;
+    mbedtls_unsafe_codepath_hook = NULL;
+}
+
+#endif /* MBEDTLS_TEST_HOOKS && !MBEDTLS_THREADING_C */
+
diff --git a/tests/src/helpers.c b/tests/src/helpers.c
index 065d17d..db50296 100644
--- a/tests/src/helpers.c
+++ b/tests/src/helpers.c
@@ -16,6 +16,9 @@
 #if defined(MBEDTLS_TEST_HOOKS) && defined(MBEDTLS_PSA_CRYPTO_C)
 #include <test/psa_memory_poisoning_wrappers.h>
 #endif
+#if defined(MBEDTLS_TEST_HOOKS) && !defined(MBEDTLS_THREADING_C)
+#include <test/bignum_codepath_check.h>
+#endif
 #if defined(MBEDTLS_THREADING_C)
 #include "mbedtls/threading.h"
 #endif
@@ -342,6 +345,11 @@
     mbedtls_mutex_init(&mbedtls_test_info_mutex);
 #endif /* MBEDTLS_THREADING_C */
 
+
+#if defined(MBEDTLS_TEST_HOOKS) && !defined(MBEDTLS_THREADING_C)
+    mbedtls_codepath_test_hooks_setup();
+#endif /* MBEDTLS_TEST_HOOKS && !MBEDTLS_THREADING_C */
+
     return ret;
 }
 
@@ -359,6 +367,10 @@
 #if defined(MBEDTLS_PLATFORM_C)
     mbedtls_platform_teardown(&platform_ctx);
 #endif /* MBEDTLS_PLATFORM_C */
+
+#if defined(MBEDTLS_TEST_HOOKS) && !defined(MBEDTLS_THREADING_C)
+    mbedtls_codepath_test_hooks_teardown();
+#endif /* MBEDTLS_TEST_HOOKS && !MBEDTLS_THREADING_C */
 }
 
 int mbedtls_test_ascii2uc(const char c, unsigned char *uc)
diff --git a/tf-psa-crypto/drivers/builtin/src/bignum_core.c b/tf-psa-crypto/drivers/builtin/src/bignum_core.c
index 97e212d..da3b6f4 100644
--- a/tf-psa-crypto/drivers/builtin/src/bignum_core.c
+++ b/tf-psa-crypto/drivers/builtin/src/bignum_core.c
@@ -748,7 +748,8 @@
 }
 
 #if defined(MBEDTLS_TEST_HOOKS) && !defined(MBEDTLS_THREADING_C)
-int mbedtls_mpi_optionally_safe_codepath = MBEDTLS_MPI_IS_TEST;
+void (*mbedtls_safe_codepath_hook)(void) = NULL;
+void (*mbedtls_unsafe_codepath_hook)(void) = NULL;
 #endif
 
 /*
@@ -781,7 +782,8 @@
         *E_bit_index = E_bits % biL;
 
 #if defined(MBEDTLS_TEST_HOOKS) && !defined(MBEDTLS_THREADING_C)
-        mbedtls_mpi_optionally_safe_codepath = MBEDTLS_MPI_IS_PUBLIC;
+        if(mbedtls_unsafe_codepath_hook != NULL)
+            mbedtls_unsafe_codepath_hook();
 #endif
     } else {
         /*
@@ -791,10 +793,8 @@
         *E_limb_index = E_limbs;
         *E_bit_index = 0;
 #if defined(MBEDTLS_TEST_HOOKS) && !defined(MBEDTLS_THREADING_C)
-        // Only mark the codepath safe if there wasn't an unsafe codepath before
-        if (mbedtls_mpi_optionally_safe_codepath != MBEDTLS_MPI_IS_PUBLIC) {
-            mbedtls_mpi_optionally_safe_codepath = MBEDTLS_MPI_IS_SECRET;
-        }
+        if(mbedtls_safe_codepath_hook != NULL)
+            mbedtls_safe_codepath_hook();
 #endif
     }
 }
@@ -813,7 +813,8 @@
     if (window_public == MBEDTLS_MPI_IS_PUBLIC) {
         memcpy(Wselect, Wtable + window * AN_limbs, AN_limbs * ciL);
 #if defined(MBEDTLS_TEST_HOOKS) && !defined(MBEDTLS_THREADING_C)
-        mbedtls_mpi_optionally_safe_codepath = MBEDTLS_MPI_IS_PUBLIC;
+        if(mbedtls_unsafe_codepath_hook != NULL)
+            mbedtls_unsafe_codepath_hook();
 #endif
     } else {
         /* Select Wtable[window] without leaking window through
@@ -821,10 +822,8 @@
         mbedtls_mpi_core_ct_uint_table_lookup(Wselect, Wtable,
                                               AN_limbs, welem, window);
 #if defined(MBEDTLS_TEST_HOOKS) && !defined(MBEDTLS_THREADING_C)
-        // Only mark the codepath safe if there wasn't an unsafe codepath before
-        if (mbedtls_mpi_optionally_safe_codepath != MBEDTLS_MPI_IS_PUBLIC) {
-            mbedtls_mpi_optionally_safe_codepath = MBEDTLS_MPI_IS_SECRET;
-        }
+        if(mbedtls_safe_codepath_hook != NULL)
+            mbedtls_safe_codepath_hook();
 #endif
     }
 }
diff --git a/tf-psa-crypto/drivers/builtin/src/bignum_core.h b/tf-psa-crypto/drivers/builtin/src/bignum_core.h
index c7db590..3c6b105 100644
--- a/tf-psa-crypto/drivers/builtin/src/bignum_core.h
+++ b/tf-psa-crypto/drivers/builtin/src/bignum_core.h
@@ -837,16 +837,4 @@
                                     mbedtls_mpi_uint mm,
                                     mbedtls_mpi_uint *T);
 
-/*
- * Can't define thread local variables with our abstraction layer: do nothing if threading is on.
- */
-#if defined(MBEDTLS_TEST_HOOKS) && !defined(MBEDTLS_THREADING_C)
-extern int mbedtls_mpi_optionally_safe_codepath;
-
-static inline void mbedtls_mpi_optionally_safe_codepath_reset(void)
-{
-    mbedtls_mpi_optionally_safe_codepath = MBEDTLS_MPI_IS_TEST;
-}
-#endif
-
 #endif /* MBEDTLS_BIGNUM_CORE_H */
diff --git a/tf-psa-crypto/tests/suites/test_suite_bignum.function b/tf-psa-crypto/tests/suites/test_suite_bignum.function
index 5e18441..c71f3c8 100644
--- a/tf-psa-crypto/tests/suites/test_suite_bignum.function
+++ b/tf-psa-crypto/tests/suites/test_suite_bignum.function
@@ -5,6 +5,7 @@
 #include "bignum_core.h"
 #include "bignum_internal.h"
 #include "test/constant_flow.h"
+#include "test/bignum_codepath_check.h"
 
 #if MBEDTLS_MPI_MAX_BITS > 792
 #define MPI_MAX_BITS_LARGER_THAN_792
@@ -990,11 +991,11 @@
     TEST_LE_U(RR.n, N.n - 1);
 
 #if defined(MBEDTLS_TEST_HOOKS) && !defined(MBEDTLS_THREADING_C)
-    mbedtls_mpi_optionally_safe_codepath_reset();
+    mbedtls_codepath_reset();
 #endif
     res = mbedtls_mpi_exp_mod(&Z, &A, &E, &N, &RR);
 #if defined(MBEDTLS_TEST_HOOKS) && !defined(MBEDTLS_THREADING_C)
-    TEST_EQUAL(mbedtls_mpi_optionally_safe_codepath, MBEDTLS_MPI_IS_SECRET);
+    TEST_EQUAL(mbedtls_codepath_check, MBEDTLS_MPI_IS_SECRET);
 #endif
     /* We know that exp_mod internally needs RR to be as large as N.
      * Validate that it is the case now, otherwise there was probably
@@ -1029,11 +1030,11 @@
     TEST_ASSERT(mbedtls_test_read_mpi(&X, input_X) == 0);
 
 #if defined(MBEDTLS_TEST_HOOKS) && !defined(MBEDTLS_THREADING_C)
-    mbedtls_mpi_optionally_safe_codepath_reset();
+    mbedtls_codepath_reset();
 #endif
     res = mbedtls_mpi_exp_mod(&Z, &A, &E, &N, NULL);
 #if defined(MBEDTLS_TEST_HOOKS) && !defined(MBEDTLS_THREADING_C)
-    TEST_EQUAL(mbedtls_mpi_optionally_safe_codepath, MBEDTLS_MPI_IS_SECRET);
+    TEST_EQUAL(mbedtls_codepath_check, MBEDTLS_MPI_IS_SECRET);
 #endif
     TEST_ASSERT(res == exp_result);
     if (res == 0) {
@@ -1042,11 +1043,11 @@
     }
 
 #if defined(MBEDTLS_TEST_HOOKS) && !defined(MBEDTLS_THREADING_C)
-    mbedtls_mpi_optionally_safe_codepath_reset();
+    mbedtls_codepath_reset();
 #endif
     res = mbedtls_mpi_exp_mod_unsafe(&Z, &A, &E, &N, NULL);
 #if defined(MBEDTLS_TEST_HOOKS) && !defined(MBEDTLS_THREADING_C)
-    TEST_EQUAL(mbedtls_mpi_optionally_safe_codepath, MBEDTLS_MPI_IS_PUBLIC);
+    TEST_EQUAL(mbedtls_codepath_check, MBEDTLS_MPI_IS_PUBLIC);
 #endif
     TEST_ASSERT(res == exp_result);
     if (res == 0) {
@@ -1056,11 +1057,11 @@
 
     /* Now test again with the speed-up parameter supplied as an output. */
 #if defined(MBEDTLS_TEST_HOOKS) && !defined(MBEDTLS_THREADING_C)
-    mbedtls_mpi_optionally_safe_codepath_reset();
+    mbedtls_codepath_reset();
 #endif
     res = mbedtls_mpi_exp_mod(&Z, &A, &E, &N, &RR);
 #if defined(MBEDTLS_TEST_HOOKS) && !defined(MBEDTLS_THREADING_C)
-    TEST_EQUAL(mbedtls_mpi_optionally_safe_codepath, MBEDTLS_MPI_IS_SECRET);
+    TEST_EQUAL(mbedtls_codepath_check, MBEDTLS_MPI_IS_SECRET);
 #endif
     TEST_ASSERT(res == exp_result);
     if (res == 0) {
@@ -1070,11 +1071,11 @@
 
     /* Now test again with the speed-up parameter supplied in calculated form. */
 #if defined(MBEDTLS_TEST_HOOKS) && !defined(MBEDTLS_THREADING_C)
-    mbedtls_mpi_optionally_safe_codepath_reset();
+    mbedtls_codepath_reset();
 #endif
     res = mbedtls_mpi_exp_mod(&Z, &A, &E, &N, &RR);
 #if defined(MBEDTLS_TEST_HOOKS) && !defined(MBEDTLS_THREADING_C)
-    TEST_EQUAL(mbedtls_mpi_optionally_safe_codepath, MBEDTLS_MPI_IS_SECRET);
+    TEST_EQUAL(mbedtls_codepath_check, MBEDTLS_MPI_IS_SECRET);
 #endif
     TEST_ASSERT(res == exp_result);
     if (res == 0) {
@@ -1116,19 +1117,19 @@
     }
 
 #if defined(MBEDTLS_TEST_HOOKS) && !defined(MBEDTLS_THREADING_C)
-    mbedtls_mpi_optionally_safe_codepath_reset();
+    mbedtls_codepath_reset();
 #endif
     TEST_ASSERT(mbedtls_mpi_exp_mod(&Z, &A, &E, &N, &RR) == exp_result);
 #if defined(MBEDTLS_TEST_HOOKS) && !defined(MBEDTLS_THREADING_C)
-    TEST_EQUAL(mbedtls_mpi_optionally_safe_codepath, MBEDTLS_MPI_IS_SECRET);
+    TEST_EQUAL(mbedtls_codepath_check, MBEDTLS_MPI_IS_SECRET);
 #endif
 
 #if defined(MBEDTLS_TEST_HOOKS) && !defined(MBEDTLS_THREADING_C)
-    mbedtls_mpi_optionally_safe_codepath_reset();
+    mbedtls_codepath_reset();
 #endif
     TEST_ASSERT(mbedtls_mpi_exp_mod_unsafe(&Z, &A, &E, &N, &RR) == exp_result);
 #if defined(MBEDTLS_TEST_HOOKS) && !defined(MBEDTLS_THREADING_C)
-    TEST_EQUAL(mbedtls_mpi_optionally_safe_codepath, MBEDTLS_MPI_IS_SECRET);
+    TEST_EQUAL(mbedtls_codepath_check, MBEDTLS_MPI_IS_SECRET);
 #endif
 
 exit:
diff --git a/tf-psa-crypto/tests/suites/test_suite_bignum_core.function b/tf-psa-crypto/tests/suites/test_suite_bignum_core.function
index 3128307..d5cc08e 100644
--- a/tf-psa-crypto/tests/suites/test_suite_bignum_core.function
+++ b/tf-psa-crypto/tests/suites/test_suite_bignum_core.function
@@ -4,6 +4,7 @@
 #include "bignum_core.h"
 #include "constant_time_internal.h"
 #include "test/constant_flow.h"
+#include "test/bignum_codepath_check.h"
 
 /** Verifies mbedtls_mpi_core_add().
  *
@@ -1303,11 +1304,11 @@
     TEST_CF_SECRET(N, N_limbs * sizeof(mbedtls_mpi_uint));
     TEST_CF_SECRET(E, E_limbs * sizeof(mbedtls_mpi_uint));
 #if defined(MBEDTLS_TEST_HOOKS) && !defined(MBEDTLS_THREADING_C)
-    mbedtls_mpi_optionally_safe_codepath_reset();
+    mbedtls_codepath_reset();
 #endif
     mbedtls_mpi_core_exp_mod(Y, A, N, N_limbs, E, E_limbs, R2, T);
 #if defined(MBEDTLS_TEST_HOOKS) && !defined(MBEDTLS_THREADING_C)
-    TEST_EQUAL(mbedtls_mpi_optionally_safe_codepath, MBEDTLS_MPI_IS_SECRET);
+    TEST_EQUAL(mbedtls_codepath_check, MBEDTLS_MPI_IS_SECRET);
 #endif
     TEST_EQUAL(0, memcmp(X, Y, N_limbs * sizeof(mbedtls_mpi_uint)));
 
@@ -1318,11 +1319,11 @@
     /* Test the unsafe variant */
 
 #if defined(MBEDTLS_TEST_HOOKS) && !defined(MBEDTLS_THREADING_C)
-    mbedtls_mpi_optionally_safe_codepath_reset();
+    mbedtls_codepath_reset();
 #endif
     mbedtls_mpi_core_exp_mod_unsafe(Y, A, N, N_limbs, E, E_limbs, R2, T);
 #if defined(MBEDTLS_TEST_HOOKS) && !defined(MBEDTLS_THREADING_C)
-    TEST_EQUAL(mbedtls_mpi_optionally_safe_codepath, MBEDTLS_MPI_IS_PUBLIC);
+    TEST_EQUAL(mbedtls_codepath_check, MBEDTLS_MPI_IS_PUBLIC);
 #endif
     TEST_EQUAL(0, memcmp(X, Y, N_limbs * sizeof(mbedtls_mpi_uint)));
 
@@ -1335,22 +1336,22 @@
     TEST_CF_SECRET(N, N_limbs * sizeof(mbedtls_mpi_uint));
     TEST_CF_SECRET(E, E_limbs * sizeof(mbedtls_mpi_uint));
 #if defined(MBEDTLS_TEST_HOOKS) && !defined(MBEDTLS_THREADING_C)
-    mbedtls_mpi_optionally_safe_codepath_reset();
+    mbedtls_codepath_reset();
 #endif
     mbedtls_mpi_core_exp_mod(A, A, N, N_limbs, E, E_limbs, R2, T);
 #if defined(MBEDTLS_TEST_HOOKS) && !defined(MBEDTLS_THREADING_C)
-    TEST_EQUAL(mbedtls_mpi_optionally_safe_codepath, MBEDTLS_MPI_IS_SECRET);
+    TEST_EQUAL(mbedtls_codepath_check, MBEDTLS_MPI_IS_SECRET);
 #endif
     TEST_EQUAL(0, memcmp(X, A, N_limbs * sizeof(mbedtls_mpi_uint)));
 
     TEST_CF_PUBLIC(A, A_limbs * sizeof(mbedtls_mpi_uint));
     memcpy(A, A_copy, sizeof(*A) * A_limbs);
 #if defined(MBEDTLS_TEST_HOOKS) && !defined(MBEDTLS_THREADING_C)
-    mbedtls_mpi_optionally_safe_codepath_reset();
+    mbedtls_codepath_reset();
 #endif
     mbedtls_mpi_core_exp_mod_unsafe(A, A, N, N_limbs, E, E_limbs, R2, T);
 #if defined(MBEDTLS_TEST_HOOKS) && !defined(MBEDTLS_THREADING_C)
-    TEST_EQUAL(mbedtls_mpi_optionally_safe_codepath, MBEDTLS_MPI_IS_PUBLIC);
+    TEST_EQUAL(mbedtls_codepath_check, MBEDTLS_MPI_IS_PUBLIC);
 #endif
     TEST_EQUAL(0, memcmp(X, A, N_limbs * sizeof(mbedtls_mpi_uint)));
 
diff --git a/tf-psa-crypto/tests/suites/test_suite_rsa.function b/tf-psa-crypto/tests/suites/test_suite_rsa.function
index e0206ec..75f3f42 100644
--- a/tf-psa-crypto/tests/suites/test_suite_rsa.function
+++ b/tf-psa-crypto/tests/suites/test_suite_rsa.function
@@ -3,6 +3,7 @@
 #include "bignum_core.h"
 #include "rsa_alt_helpers.h"
 #include "rsa_internal.h"
+#include "test/bignum_codepath_check.h"
 /* END_HEADER */
 
 /* BEGIN_DEPENDENCIES
@@ -491,11 +492,11 @@
     TEST_ASSERT(mbedtls_rsa_check_pubkey(&ctx) == 0);
 
 #if defined(MBEDTLS_TEST_HOOKS) && !defined(MBEDTLS_THREADING_C)
-    mbedtls_mpi_optionally_safe_codepath_reset();
+    mbedtls_codepath_reset();
 #endif
     TEST_ASSERT(mbedtls_rsa_public(&ctx, message_str->x, output) == result);
 #if defined(MBEDTLS_TEST_HOOKS) && !defined(MBEDTLS_THREADING_C)
-    TEST_EQUAL(mbedtls_mpi_optionally_safe_codepath, MBEDTLS_MPI_IS_PUBLIC);
+    TEST_EQUAL(mbedtls_codepath_check, MBEDTLS_MPI_IS_PUBLIC);
 #endif
     if (result == 0) {
 
@@ -562,13 +563,13 @@
     for (i = 0; i < 3; i++) {
         memset(output, 0x00, sizeof(output));
 #if defined(MBEDTLS_TEST_HOOKS) && !defined(MBEDTLS_THREADING_C)
-        mbedtls_mpi_optionally_safe_codepath_reset();
+        mbedtls_codepath_reset();
 #endif
         TEST_ASSERT(mbedtls_rsa_private(&ctx, mbedtls_test_rnd_pseudo_rand,
                                         &rnd_info, message_str->x,
                                         output) == result);
 #if defined(MBEDTLS_TEST_HOOKS) && !defined(MBEDTLS_THREADING_C)
-        TEST_EQUAL(mbedtls_mpi_optionally_safe_codepath, MBEDTLS_MPI_IS_SECRET);
+        TEST_EQUAL(mbedtls_codepath_check, MBEDTLS_MPI_IS_SECRET);
 #endif
         if (result == 0) {