diff --git a/tf-psa-crypto/core/psa_crypto_random_impl.h b/tf-psa-crypto/core/psa_crypto_random_impl.h
index 533fb2e..9049822 100644
--- a/tf-psa-crypto/core/psa_crypto_random_impl.h
+++ b/tf-psa-crypto/core/psa_crypto_random_impl.h
@@ -32,7 +32,7 @@
 #elif defined(MBEDTLS_HMAC_DRBG_C)
 
 #include "mbedtls/hmac_drbg.h"
-#if defined(MBEDTLS_MD_CAN_SHA512) && defined(MBEDTLS_MD_CAN_SHA256)
+#if defined(MBEDTLS_MD_CAN_SHA512) && defined(PSA_WANT_ALG_SHA_256)
 #include <limits.h>
 #if SIZE_MAX > 0xffffffff
 /* Looks like a 64-bit system, so prefer SHA-512. */
@@ -43,7 +43,7 @@
 #endif
 #elif defined(MBEDTLS_MD_CAN_SHA512)
 #define MBEDTLS_PSA_HMAC_DRBG_MD_TYPE MBEDTLS_MD_SHA512
-#elif defined(MBEDTLS_MD_CAN_SHA256)
+#elif defined(PSA_WANT_ALG_SHA_256)
 #define MBEDTLS_PSA_HMAC_DRBG_MD_TYPE MBEDTLS_MD_SHA256
 #else
 #error "No hash algorithm available for HMAC_DBRG."
diff --git a/tf-psa-crypto/drivers/builtin/src/ecjpake.c b/tf-psa-crypto/drivers/builtin/src/ecjpake.c
index 628fb65..ebdae92 100644
--- a/tf-psa-crypto/drivers/builtin/src/ecjpake.c
+++ b/tf-psa-crypto/drivers/builtin/src/ecjpake.c
@@ -823,7 +823,7 @@
 #include "mbedtls/platform.h"
 
 #if !defined(MBEDTLS_ECP_DP_SECP256R1_ENABLED) || \
-    !defined(MBEDTLS_MD_CAN_SHA256)
+    !defined(PSA_WANT_ALG_SHA_256)
 int mbedtls_ecjpake_self_test(int verbose)
 {
     (void) verbose;
@@ -1199,7 +1199,7 @@
 
 #undef TEST_ASSERT
 
-#endif /* MBEDTLS_ECP_DP_SECP256R1_ENABLED && MBEDTLS_MD_CAN_SHA256 */
+#endif /* MBEDTLS_ECP_DP_SECP256R1_ENABLED && PSA_WANT_ALG_SHA_256 */
 
 #endif /* MBEDTLS_SELF_TEST */
 
diff --git a/tf-psa-crypto/drivers/builtin/src/md.c b/tf-psa-crypto/drivers/builtin/src/md.c
index 7e0465d..af17526 100644
--- a/tf-psa-crypto/drivers/builtin/src/md.c
+++ b/tf-psa-crypto/drivers/builtin/src/md.c
@@ -94,7 +94,7 @@
 };
 #endif
 
-#if defined(MBEDTLS_MD_CAN_SHA256)
+#if defined(PSA_WANT_ALG_SHA_256)
 static const mbedtls_md_info_t mbedtls_sha256_info = {
     MD_INFO(MBEDTLS_MD_SHA256, 32, 64)
 };
@@ -155,7 +155,7 @@
         case MBEDTLS_MD_SHA224:
             return &mbedtls_sha224_info;
 #endif
-#if defined(MBEDTLS_MD_CAN_SHA256)
+#if defined(PSA_WANT_ALG_SHA_256)
         case MBEDTLS_MD_SHA256:
             return &mbedtls_sha256_info;
 #endif
@@ -789,7 +789,7 @@
     MBEDTLS_MD_SHA384,
 #endif
 
-#if defined(MBEDTLS_MD_CAN_SHA256)
+#if defined(PSA_WANT_ALG_SHA_256)
     MBEDTLS_MD_SHA256,
 #endif
 #if defined(PSA_WANT_ALG_SHA_224)
@@ -851,7 +851,7 @@
 #if defined(PSA_WANT_ALG_SHA_224)
     { "SHA224", MBEDTLS_MD_SHA224 },
 #endif
-#if defined(MBEDTLS_MD_CAN_SHA256)
+#if defined(PSA_WANT_ALG_SHA_256)
     { "SHA256", MBEDTLS_MD_SHA256 },
 #endif
 #if defined(PSA_WANT_ALG_SHA_384)
diff --git a/tf-psa-crypto/drivers/builtin/src/oid.c b/tf-psa-crypto/drivers/builtin/src/oid.c
index 972a859..20a5630 100644
--- a/tf-psa-crypto/drivers/builtin/src/oid.c
+++ b/tf-psa-crypto/drivers/builtin/src/oid.c
@@ -398,13 +398,13 @@
         MBEDTLS_MD_SHA224,   MBEDTLS_PK_RSA,
     },
 #endif /* PSA_WANT_ALG_SHA_224 */
-#if defined(MBEDTLS_MD_CAN_SHA256)
+#if defined(PSA_WANT_ALG_SHA_256)
     {
         OID_DESCRIPTOR(MBEDTLS_OID_PKCS1_SHA256,     "sha256WithRSAEncryption",
                        "RSA with SHA-256"),
         MBEDTLS_MD_SHA256,   MBEDTLS_PK_RSA,
     },
-#endif /* MBEDTLS_MD_CAN_SHA256 */
+#endif /* PSA_WANT_ALG_SHA_256 */
 #if defined(PSA_WANT_ALG_SHA_384)
     {
         OID_DESCRIPTOR(MBEDTLS_OID_PKCS1_SHA384,     "sha384WithRSAEncryption",
@@ -439,12 +439,12 @@
         MBEDTLS_MD_SHA224,   MBEDTLS_PK_ECDSA,
     },
 #endif
-#if defined(MBEDTLS_MD_CAN_SHA256)
+#if defined(PSA_WANT_ALG_SHA_256)
     {
         OID_DESCRIPTOR(MBEDTLS_OID_ECDSA_SHA256,     "ecdsa-with-SHA256",    "ECDSA with SHA256"),
         MBEDTLS_MD_SHA256,   MBEDTLS_PK_ECDSA,
     },
-#endif /* MBEDTLS_MD_CAN_SHA256 */
+#endif /* PSA_WANT_ALG_SHA_256 */
 #if defined(PSA_WANT_ALG_SHA_384)
     {
         OID_DESCRIPTOR(MBEDTLS_OID_ECDSA_SHA384,     "ecdsa-with-SHA384",    "ECDSA with SHA384"),
@@ -737,7 +737,7 @@
         MBEDTLS_MD_SHA224,
     },
 #endif
-#if defined(MBEDTLS_MD_CAN_SHA256)
+#if defined(PSA_WANT_ALG_SHA_256)
     {
         OID_DESCRIPTOR(MBEDTLS_OID_DIGEST_ALG_SHA256,    "id-sha256",    "SHA-256"),
         MBEDTLS_MD_SHA256,
@@ -821,12 +821,12 @@
         MBEDTLS_MD_SHA224,
     },
 #endif /* PSA_WANT_ALG_SHA_224 */
-#if defined(MBEDTLS_MD_CAN_SHA256)
+#if defined(PSA_WANT_ALG_SHA_256)
     {
         OID_DESCRIPTOR(MBEDTLS_OID_HMAC_SHA256,    "hmacSHA256",    "HMAC-SHA-256"),
         MBEDTLS_MD_SHA256,
     },
-#endif /* MBEDTLS_MD_CAN_SHA256 */
+#endif /* PSA_WANT_ALG_SHA_256 */
 #if defined(PSA_WANT_ALG_SHA_384)
     {
         OID_DESCRIPTOR(MBEDTLS_OID_HMAC_SHA384,    "hmacSHA384",    "HMAC-SHA-384"),