diff --git a/library/pk.c b/library/pk.c
index e4a3ac0375..b5d76dbd91 100644
--- a/library/pk.c
+++ b/library/pk.c
@@ -1245,7 +1245,7 @@ int mbedtls_pk_verify_ext(mbedtls_pk_type_t type, const void *options,
 
 #if defined(MBEDTLS_USE_PSA_CRYPTO)
     if (pss_opts->mgf1_hash_id == md_alg) {
-        unsigned char buf[MBEDTLS_PK_RSA_PUB_DER_MAX_BYTES];
+        unsigned char buf[PSA_KEY_EXPORT_RSA_PUBLIC_KEY_MAX_SIZE(PSA_VENDOR_RSA_MAX_KEY_BITS)];
         unsigned char *p;
         int key_len;
         size_t signature_length;
diff --git a/library/pk_wrap.c b/library/pk_wrap.c
index 8c038e7db6..5265469349 100644
--- a/library/pk_wrap.c
+++ b/library/pk_wrap.c
@@ -72,7 +72,7 @@ static int rsa_verify_wrap(mbedtls_pk_context *pk, mbedtls_md_type_t md_alg,
     mbedtls_svc_key_id_t key_id = MBEDTLS_SVC_KEY_ID_INIT;
     psa_status_t status;
     int key_len;
-    unsigned char buf[MBEDTLS_PK_RSA_PUB_DER_MAX_BYTES];
+    unsigned char buf[PSA_KEY_EXPORT_RSA_PUBLIC_KEY_MAX_SIZE(PSA_VENDOR_RSA_MAX_KEY_BITS)];
     unsigned char *p = buf + sizeof(buf);
     psa_algorithm_t psa_alg_md;
     size_t rsa_len = mbedtls_rsa_get_len(rsa);
@@ -375,7 +375,7 @@ static int rsa_encrypt_wrap(mbedtls_pk_context *pk,
     psa_algorithm_t psa_md_alg, psa_encrypt_alg;
     psa_status_t status;
     int key_len;
-    unsigned char buf[MBEDTLS_PK_RSA_PUB_DER_MAX_BYTES];
+    unsigned char buf[PSA_KEY_EXPORT_RSA_PUBLIC_KEY_MAX_SIZE(PSA_VENDOR_RSA_MAX_KEY_BITS)];
     unsigned char *p = buf + sizeof(buf);
 
     ((void) f_rng);