diff --git a/tf-psa-crypto/drivers/builtin/src/bignum_core.c b/tf-psa-crypto/drivers/builtin/src/bignum_core.c
index f8e97fbf7b..9eae74cf03 100644
--- a/tf-psa-crypto/drivers/builtin/src/bignum_core.c
+++ b/tf-psa-crypto/drivers/builtin/src/bignum_core.c
@@ -766,6 +766,9 @@ static inline void exp_mod_calc_first_bit_optionally_safe(const mbedtls_mpi_uint
             *E_limb_index = E_bits / biL;
             *E_bit_index = E_bits % biL;
         }
+#if defined(MBEDTLS_TEST_HOOKS)
+        mbedtls_mpi_optionally_safe_codepath = MBEDTLS_MPI_IS_PUBLIC;
+#endif
     } else {
         /*
          * Here we need to be constant time with respect to E and can't do anything better than
@@ -773,6 +776,12 @@ static inline void exp_mod_calc_first_bit_optionally_safe(const mbedtls_mpi_uint
          */
         *E_limb_index = E_limbs;
         *E_bit_index = 0;
+#if defined(MBEDTLS_TEST_HOOKS)
+        // 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;
+        }
+#endif
     }
 }
 
@@ -789,11 +798,20 @@ static inline void exp_mod_table_lookup_optionally_safe(mbedtls_mpi_uint *Wselec
 {
     if (window_public == MBEDTLS_MPI_IS_PUBLIC) {
         memcpy(Wselect, Wtable + window * AN_limbs, AN_limbs * ciL);
+#if defined(MBEDTLS_TEST_HOOKS)
+        mbedtls_mpi_optionally_safe_codepath = MBEDTLS_MPI_IS_PUBLIC;
+#endif
     } else {
         /* Select Wtable[window] without leaking window through
          * memory access patterns. */
         mbedtls_mpi_core_ct_uint_table_lookup(Wselect, Wtable,
                                               AN_limbs, welem, window);
+#if defined(MBEDTLS_TEST_HOOKS)
+        // 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;
+        }
+#endif
     }
 }
 
diff --git a/tf-psa-crypto/drivers/builtin/src/bignum_core.h b/tf-psa-crypto/drivers/builtin/src/bignum_core.h
index 69fa3063d2..eb84161279 100644
--- a/tf-psa-crypto/drivers/builtin/src/bignum_core.h
+++ b/tf-psa-crypto/drivers/builtin/src/bignum_core.h
@@ -830,4 +830,14 @@ void mbedtls_mpi_core_from_mont_rep(mbedtls_mpi_uint *X,
                                     mbedtls_mpi_uint mm,
                                     mbedtls_mpi_uint *T);
 
+#if defined(MBEDTLS_TEST_HOOKS)
+int mbedtls_mpi_optionally_safe_codepath;
+
+static inline void mbedtls_mpi_optionally_safe_codepath_reset()
+{
+    // Set to a default that is neither MBEDTLS_MPI_IS_PUBLIC nor MBEDTLS_MPI_IS_SECRET
+    mbedtls_mpi_optionally_safe_codepath = MBEDTLS_MPI_IS_PUBLIC + MBEDTLS_MPI_IS_SECRET + 1;
+}
+#endif
+
 #endif /* MBEDTLS_BIGNUM_CORE_H */
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 6c0bd1e0c4..fe47cc69a6 100644
--- a/tf-psa-crypto/tests/suites/test_suite_bignum_core.function
+++ b/tf-psa-crypto/tests/suites/test_suite_bignum_core.function
@@ -1301,7 +1301,13 @@ void mpi_core_exp_mod(char *input_N, char *input_A,
     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)
+    mbedtls_mpi_optionally_safe_codepath_reset();
+#endif
     mbedtls_mpi_core_exp_mod(Y, A, N, N_limbs, E, E_limbs, R2, T);
+#if defined(MBEDTLS_TEST_HOOKS)
+    TEST_EQUAL(mbedtls_mpi_optionally_safe_codepath, MBEDTLS_MPI_IS_SECRET);
+#endif
 
     TEST_CF_PUBLIC(Y, N_limbs * sizeof(mbedtls_mpi_uint));
 
@@ -1312,7 +1318,13 @@ void mpi_core_exp_mod(char *input_N, char *input_A,
     TEST_CF_SECRET(E, E_limbs * sizeof(mbedtls_mpi_uint));
 
     /* Check when output aliased to input */
+#if defined(MBEDTLS_TEST_HOOKS)
+    mbedtls_mpi_optionally_safe_codepath_reset();
+#endif
     mbedtls_mpi_core_exp_mod(A, A, N, N_limbs, E, E_limbs, R2, T);
+#if defined(MBEDTLS_TEST_HOOKS)
+    TEST_EQUAL(mbedtls_mpi_optionally_safe_codepath, MBEDTLS_MPI_IS_SECRET);
+#endif
 
     TEST_CF_PUBLIC(A, A_limbs * sizeof(mbedtls_mpi_uint));
     TEST_EQUAL(0, memcmp(X, A, N_limbs * sizeof(mbedtls_mpi_uint)));