diff --git a/library/psa_crypto_ffdh.c b/library/psa_crypto_ffdh.c
index ae38f6d7c6..ec169b8a98 100644
--- a/library/psa_crypto_ffdh.c
+++ b/library/psa_crypto_ffdh.c
@@ -168,6 +168,10 @@ psa_status_t mbedtls_psa_ffdh_export_public_key(
     mbedtls_mpi_init(&X); mbedtls_mpi_init(&P);
 
     size_t key_len = PSA_BITS_TO_BYTES(attributes->bits);
+    if (key_len > data_size) {
+        status = PSA_ERROR_BUFFER_TOO_SMALL;
+        goto cleanup;
+    }
 
     status = mbedtls_psa_ffdh_set_prime_generator(key_len, &P, &G);