diff --git a/library/psa_crypto_ecp.c b/library/psa_crypto_ecp.c
index 7edea81..e373ad9 100644
--- a/library/psa_crypto_ecp.c
+++ b/library/psa_crypto_ecp.c
@@ -281,20 +281,8 @@
 
         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 @@
     }
 
     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 */