diff --git a/library/pk.c b/library/pk.c
index 1ded4872f9..c647b459a2 100644
--- a/library/pk.c
+++ b/library/pk.c
@@ -1401,8 +1401,7 @@ int mbedtls_pk_wrap_as_opaque(mbedtls_pk_context *pk,
         mbedtls_ecp_keypair *ec = mbedtls_pk_ec_rw(*pk);
         int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
 
-        d_len = PSA_BITS_TO_BYTES(ec->grp.nbits);
-        if ((ret = mbedtls_ecp_write_key(ec, d, d_len)) != 0) {
+        if ((ret = mbedtls_ecp_write_key_ext(ec, &d_len, d, sizeof(d))) != 0) {
             return ret;
         }
 
diff --git a/library/pkwrite.c b/library/pkwrite.c
index b9ddcf1d8d..5e009c565e 100644
--- a/library/pkwrite.c
+++ b/library/pkwrite.c
@@ -202,7 +202,7 @@ static int pk_write_ec_private(unsigned char **p, unsigned char *start,
         mbedtls_ecp_keypair *ec = mbedtls_pk_ec_rw(*pk);
         byte_length = (ec->grp.pbits + 7) / 8;
 
-        ret = mbedtls_ecp_write_key(ec, tmp, byte_length);
+        ret = mbedtls_ecp_write_key_ext(ec, &byte_length, tmp, sizeof(tmp));
         if (ret != 0) {
             goto exit;
         }
diff --git a/library/psa_crypto_ecp.c b/library/psa_crypto_ecp.c
index 7edea81646..e373ad9f66 100644
--- a/library/psa_crypto_ecp.c
+++ b/library/psa_crypto_ecp.c
@@ -281,20 +281,8 @@ psa_status_t mbedtls_psa_ecp_export_key(psa_key_type_t type,
 
         return status;
     } else {
-        if (data_size < PSA_BITS_TO_BYTES(ecp->grp.nbits)) {
-            return PSA_ERROR_BUFFER_TOO_SMALL;
-        }
-
         status = mbedtls_to_psa_error(
-            mbedtls_ecp_write_key(ecp,
-                                  data,
-                                  PSA_BITS_TO_BYTES(ecp->grp.nbits)));
-        if (status == PSA_SUCCESS) {
-            *data_length = PSA_BITS_TO_BYTES(ecp->grp.nbits);
-        } else {
-            memset(data, 0, data_size);
-        }
-
+            mbedtls_ecp_write_key_ext(ecp, data_length, data, data_size));
         return status;
     }
 }
@@ -359,14 +347,11 @@ psa_status_t mbedtls_psa_ecp_generate_key(
     }
 
     status = mbedtls_to_psa_error(
-        mbedtls_ecp_write_key(&ecp, key_buffer, key_buffer_size));
+        mbedtls_ecp_write_key_ext(&ecp, key_buffer_length,
+                                  key_buffer, key_buffer_size));
 
     mbedtls_ecp_keypair_free(&ecp);
 
-    if (status == PSA_SUCCESS) {
-        *key_buffer_length = key_buffer_size;
-    }
-
     return status;
 }
 #endif /* MBEDTLS_PSA_BUILTIN_KEY_TYPE_ECC_KEY_PAIR_GENERATE */
diff --git a/library/ssl_tls12_server.c b/library/ssl_tls12_server.c
index 53a9ce206c..5bee188235 100644
--- a/library/ssl_tls12_server.c
+++ b/library/ssl_tls12_server.c
@@ -2703,8 +2703,7 @@ static int ssl_get_ecdh_params_from_cert(mbedtls_ssl_context *ssl)
                              PSA_KEY_TYPE_ECC_KEY_PAIR(ssl->handshake->xxdh_psa_type));
             psa_set_key_bits(&key_attributes, ssl->handshake->xxdh_psa_bits);
 
-            key_len = PSA_BITS_TO_BYTES(key->grp.pbits);
-            ret = mbedtls_ecp_write_key(key, buf, key_len);
+            ret = mbedtls_ecp_write_key_ext(key, &key_len, buf, sizeof(buf));
             if (ret != 0) {
                 mbedtls_platform_zeroize(buf, sizeof(buf));
                 break;
diff --git a/tests/suites/test_suite_ecp.function b/tests/suites/test_suite_ecp.function
index c0612755ff..ef0781bb77 100644
--- a/tests/suites/test_suite_ecp.function
+++ b/tests/suites/test_suite_ecp.function
@@ -1204,27 +1204,40 @@ void mbedtls_ecp_read_key(int grp_id, data_t *in_key, int expected, int canonica
         TEST_EQUAL(mbedtls_mpi_cmp_int(&key.Q.Y, 2), 0);
         TEST_EQUAL(mbedtls_mpi_cmp_int(&key.Q.Z, 3), 0);
 
-        if (canonical) {
+        if (canonical && in_key->len == (key.grp.nbits + 7) / 8) {
             unsigned char buf[MBEDTLS_ECP_MAX_BYTES];
+            size_t length = 0xdeadbeef;
 
-            ret = mbedtls_ecp_write_key(&key, buf, in_key->len);
-            TEST_ASSERT(ret == 0);
+            TEST_EQUAL(mbedtls_ecp_write_key_ext(&key,
+                                                 &length, buf, in_key->len), 0);
+            TEST_MEMORY_COMPARE(in_key->x, in_key->len,
+                                buf, length);
 
+            memset(buf, 0, sizeof(buf));
+            TEST_EQUAL(mbedtls_ecp_write_key(&key, buf, in_key->len), 0);
             TEST_MEMORY_COMPARE(in_key->x, in_key->len,
                                 buf, in_key->len);
         } else {
             unsigned char export1[MBEDTLS_ECP_MAX_BYTES];
             unsigned char export2[MBEDTLS_ECP_MAX_BYTES];
 
-            ret = mbedtls_ecp_write_key(&key, export1, in_key->len);
-            TEST_ASSERT(ret == 0);
-
-            ret = mbedtls_ecp_read_key(grp_id, &key2, export1, in_key->len);
-            TEST_ASSERT(ret == expected);
-
-            ret = mbedtls_ecp_write_key(&key2, export2, in_key->len);
-            TEST_ASSERT(ret == 0);
+            size_t length1 = 0xdeadbeef;
+            TEST_EQUAL(mbedtls_ecp_write_key_ext(&key, &length1,
+                                                 export1, sizeof(export1)), 0);
+            TEST_EQUAL(mbedtls_ecp_read_key(grp_id, &key2, export1, length1),
+                       expected);
+            size_t length2 = 0xdeadbeef;
+            TEST_EQUAL(mbedtls_ecp_write_key_ext(&key2, &length2,
+                                                 export2, sizeof(export2)), 0);
+            TEST_MEMORY_COMPARE(export1, length1,
+                                export2, length2);
 
+            memset(export1, 0, sizeof(export1));
+            memset(export2, 0, sizeof(export2));
+            TEST_EQUAL(mbedtls_ecp_write_key(&key, export1, in_key->len), 0);
+            TEST_EQUAL(mbedtls_ecp_read_key(grp_id, &key2, export1, in_key->len),
+                       expected);
+            TEST_EQUAL(mbedtls_ecp_write_key(&key2, export2, in_key->len), 0);
             TEST_MEMORY_COMPARE(export1, in_key->len,
                                 export2, in_key->len);
         }