diff --git a/tf-psa-crypto/drivers/builtin/src/pkwrite.c b/tf-psa-crypto/drivers/builtin/src/pkwrite.c
index 8c01b44..0b57995 100644
--- a/tf-psa-crypto/drivers/builtin/src/pkwrite.c
+++ b/tf-psa-crypto/drivers/builtin/src/pkwrite.c
@@ -65,17 +65,21 @@
 #if defined(MBEDTLS_USE_PSA_CRYPTO)
     if (mbedtls_pk_get_type(pk) == MBEDTLS_PK_OPAQUE) {
         uint8_t tmp[PSA_EXPORT_KEY_PAIR_MAX_SIZE];
-        size_t len = 0, tmp_len = 0;
+        size_t tmp_len = 0;
 
         if (psa_export_key(pk->priv_id, tmp, sizeof(tmp), &tmp_len) != PSA_SUCCESS) {
             return MBEDTLS_ERR_PK_BAD_INPUT_DATA;
         }
+        /* Ensure there's enough space in the provided buffer before copying data into it. */
+        if (tmp_len > (size_t) (*p - buf)) {
+            mbedtls_platform_zeroize(tmp, sizeof(tmp));
+            return MBEDTLS_ERR_ASN1_BUF_TOO_SMALL;
+        }
         *p -= tmp_len;
         memcpy(*p, tmp, tmp_len);
-        len += tmp_len;
         mbedtls_platform_zeroize(tmp, sizeof(tmp));
 
-        return (int) len;
+        return (int) tmp_len;
     }
 #endif /* MBEDTLS_USE_PSA_CRYPTO */
     return mbedtls_rsa_write_key(mbedtls_pk_rsa(*pk), buf, p);
@@ -125,6 +129,10 @@
         if (psa_export_public_key(pk->priv_id, buf, sizeof(buf), &len) != PSA_SUCCESS) {
             return MBEDTLS_ERR_PK_BAD_INPUT_DATA;
         }
+        /* Ensure there's enough space in the provided buffer before copying data into it. */
+        if (len > (size_t) (*p - start)) {
+            return MBEDTLS_ERR_ASN1_BUF_TOO_SMALL;
+        }
         *p -= len;
         memcpy(*p, buf, len);
         return (int) len;
diff --git a/tf-psa-crypto/tests/suites/test_suite_pkwrite.function b/tf-psa-crypto/tests/suites/test_suite_pkwrite.function
index 735c125..3392528 100644
--- a/tf-psa-crypto/tests/suites/test_suite_pkwrite.function
+++ b/tf-psa-crypto/tests/suites/test_suite_pkwrite.function
@@ -2,6 +2,7 @@
 #include "pk_internal.h"
 #include "mbedtls/pem.h"
 #include "mbedtls/oid.h"
+#include "mbedtls/base64.h"
 #include "psa/crypto_sizes.h"
 
 typedef enum {
@@ -72,7 +73,8 @@
     unsigned char *buf = NULL;
     unsigned char *check_buf = NULL;
     unsigned char *start_buf;
-    size_t buf_len, check_buf_len;
+    size_t buf_len, check_buf_len, wrong_buf_len = 1;
+    int expected_result;
 #if defined(MBEDTLS_USE_PSA_CRYPTO)
     mbedtls_svc_key_id_t opaque_id = MBEDTLS_SVC_KEY_ID_INIT;
     psa_key_attributes_t key_attr = PSA_KEY_ATTRIBUTES_INIT;
@@ -109,6 +111,17 @@
 
     start_buf = buf;
     buf_len = check_buf_len;
+    if (is_der) {
+        expected_result = MBEDTLS_ERR_ASN1_BUF_TOO_SMALL;
+    } else {
+        expected_result = MBEDTLS_ERR_BASE64_BUFFER_TOO_SMALL;
+    }
+    /* Intentionally pass a wrong size for the provided output buffer and check
+     * that the writing functions fails as expected. */
+    TEST_EQUAL(pk_write_any_key(&key, &start_buf, &wrong_buf_len, is_public_key,
+                                is_der), expected_result);
+    TEST_EQUAL(pk_write_any_key(&key, &start_buf, &buf_len, is_public_key,
+                                is_der), 0);
     TEST_EQUAL(pk_write_any_key(&key, &start_buf, &buf_len, is_public_key,
                                 is_der), 0);
 
@@ -127,6 +140,10 @@
         TEST_EQUAL(mbedtls_pk_setup_opaque(&key, opaque_id), 0);
         start_buf = buf;
         buf_len = check_buf_len;
+        /* Intentionally pass a wrong size for the provided output buffer and check
+         * that the writing functions fails as expected. */
+        TEST_EQUAL(pk_write_any_key(&key, &start_buf, &wrong_buf_len, is_public_key,
+                                    is_der), expected_result);
         TEST_EQUAL(pk_write_any_key(&key, &start_buf, &buf_len, is_public_key,
                                     is_der), 0);