pkwrite: fix buffer overrun

This commit fixes potential buffer overrun in:
- pk_write_rsa_der
- pk_write_ec_pubkey

In both functions, when dealing with opaque keys, there was no
check that the provided buffer was large enough to contain the
key being exported. This commit fixes this problem and it also
adds some testing in test_suite_pkwrite to trigger these checks.

Signed-off-by: Valerio Setti <valerio.setti@nordicsemi.no>
diff --git a/library/pkwrite.c b/library/pkwrite.c
index 5e009c5..2a69844 100644
--- a/library/pkwrite.c
+++ b/library/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/tests/suites/test_suite_pkwrite.function b/tests/suites/test_suite_pkwrite.function
index 735c125..3392528 100644
--- a/tests/suites/test_suite_pkwrite.function
+++ b/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);