mbedtls_ecp_write_key_ext(): new function

Same as mbedtls_ecp_write_key(), but doesn't require the caller to figure out
the length of the output and possibly distinguish between Weierstrass and
Montgomery curves.

Signed-off-by: Gilles Peskine <Gilles.Peskine@arm.com>
diff --git a/ChangeLog.d/ecp_write_key.txt b/ChangeLog.d/ecp_write_key.txt
new file mode 100644
index 0000000..1961239
--- /dev/null
+++ b/ChangeLog.d/ecp_write_key.txt
@@ -0,0 +1,4 @@
+Features
+   * The new function mbedtls_ecp_write_key_ext() is similar to
+     mbedtls_ecp_write_key(), but can be used without separately calculating
+     the output length.
diff --git a/include/mbedtls/ecp.h b/include/mbedtls/ecp.h
index 0201963..b46d6c3 100644
--- a/include/mbedtls/ecp.h
+++ b/include/mbedtls/ecp.h
@@ -1371,6 +1371,26 @@
                           unsigned char *buf, size_t buflen);
 
 /**
+ * \brief           This function exports an elliptic curve private key.
+ *
+ * \param key       The private key.
+ * \param olen      On success, the length of the private key.
+ *                  This is always (`grp->nbits` + 7) / 8 bytes
+ *                  where `grp->nbits` is the private key size in bits.
+ * \param buf       The output buffer for containing the binary representation
+ *                  of the key.
+ * \param buflen    The total length of the buffer in bytes.
+ *                  #MBEDTLS_ECP_MAX_BYTES is always sufficient.
+ *
+ * \return          \c 0 on success.
+ * \return          #MBEDTLS_ERR_ECP_BUFFER_TOO_SMALL if the \p key
+ *                  representation is larger than the available space in \p buf.
+ * \return          Another negative error code on different kinds of failure.
+ */
+int mbedtls_ecp_write_key_ext(mbedtls_ecp_keypair *key,
+                              size_t *olen, unsigned char *buf, size_t buflen);
+
+/**
  * \brief           This function exports an elliptic curve public key.
  *
  * \note            If the public key was not set in \p key,
diff --git a/library/ecp.c b/library/ecp.c
index 66b3dc1..930102f 100644
--- a/library/ecp.c
+++ b/library/ecp.c
@@ -3333,6 +3333,38 @@
     return ret;
 }
 
+int mbedtls_ecp_write_key_ext(mbedtls_ecp_keypair *key,
+                              size_t *olen, unsigned char *buf, size_t buflen)
+{
+    size_t len = (key->grp.nbits + 7) / 8;
+    if (len > buflen) {
+        /* For robustness, ensure *olen <= buflen even on error. */
+        *olen = 0;
+        return MBEDTLS_ERR_ECP_BUFFER_TOO_SMALL;
+    }
+    *olen = len;
+
+    /* Private key not set */
+    if (key->d.n == 0) {
+        return MBEDTLS_ERR_ECP_BAD_INPUT_DATA;
+    }
+
+#if defined(MBEDTLS_ECP_MONTGOMERY_ENABLED)
+    if (mbedtls_ecp_get_type(&key->grp) == MBEDTLS_ECP_TYPE_MONTGOMERY) {
+        return mbedtls_mpi_write_binary_le(&key->d, buf, len);
+    }
+#endif
+
+#if defined(MBEDTLS_ECP_SHORT_WEIERSTRASS_ENABLED)
+    if (mbedtls_ecp_get_type(&key->grp) == MBEDTLS_ECP_TYPE_SHORT_WEIERSTRASS) {
+        return mbedtls_mpi_write_binary(&key->d, buf, len);
+    }
+#endif
+
+    /* Private key set but no recognized curve type? This shouldn't happen. */
+    return MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
+}
+
 /*
  * Write a public key.
  */
diff --git a/tests/suites/test_suite_ecp.data b/tests/suites/test_suite_ecp.data
index 1dd963a..fd63657 100644
--- a/tests/suites/test_suite_ecp.data
+++ b/tests/suites/test_suite_ecp.data
@@ -888,6 +888,109 @@
 depends_on:MBEDTLS_ECP_DP_CURVE448_ENABLED
 ecp_write_key:MBEDTLS_ECP_DP_CURVE448:"0000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000080":55:MBEDTLS_ERR_ECP_BUFFER_TOO_SMALL
 
+ECP write key ext: secp256r1, nominal
+depends_on:MBEDTLS_ECP_DP_SECP256R1_ENABLED
+ecp_write_key_ext:MBEDTLS_ECP_DP_SECP256R1:"f12a1320760270a83cbffd53f6031ef76a5d86c8a204f2c30ca9ebf51f0f0ea7":32:0
+
+ECP write key ext: secp256r1, output longer by 1
+depends_on:MBEDTLS_ECP_DP_SECP256R1_ENABLED
+ecp_write_key_ext:MBEDTLS_ECP_DP_SECP256R1:"f12a1320760270a83cbffd53f6031ef76a5d86c8a204f2c30ca9ebf51f0f0ea7":33:0
+
+ECP write key ext: secp256r1, output short by 1
+depends_on:MBEDTLS_ECP_DP_SECP256R1_ENABLED
+ecp_write_key_ext:MBEDTLS_ECP_DP_SECP256R1:"f12a1320760270a83cbffd53f6031ef76a5d86c8a204f2c30ca9ebf51f0f0ea7":31:MBEDTLS_ERR_ECP_BUFFER_TOO_SMALL
+
+ECP write key ext: secp256r1, output_size=0
+depends_on:MBEDTLS_ECP_DP_SECP256R1_ENABLED
+ecp_write_key_ext:MBEDTLS_ECP_DP_SECP256R1:"f12a1320760270a83cbffd53f6031ef76a5d86c8a204f2c30ca9ebf51f0f0ea7":0:MBEDTLS_ERR_ECP_BUFFER_TOO_SMALL
+
+ECP write key ext: secp256r1, top byte = 0, output_size=32
+depends_on:MBEDTLS_ECP_DP_SECP256R1_ENABLED
+ecp_write_key_ext:MBEDTLS_ECP_DP_SECP256R1:"00ffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff":32:0
+
+ECP write key ext: secp256r1, top byte = 0, output_size=31
+depends_on:MBEDTLS_ECP_DP_SECP256R1_ENABLED
+ecp_write_key_ext:MBEDTLS_ECP_DP_SECP256R1:"00ffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff":31:MBEDTLS_ERR_ECP_BUFFER_TOO_SMALL
+
+ECP write key ext: secp256r1, top byte = 0, output_size=30
+depends_on:MBEDTLS_ECP_DP_SECP256R1_ENABLED
+ecp_write_key_ext:MBEDTLS_ECP_DP_SECP256R1:"00ffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff":30:MBEDTLS_ERR_ECP_BUFFER_TOO_SMALL
+
+ECP write key ext: secp256r1, mostly-0 key, output_size=32
+depends_on:MBEDTLS_ECP_DP_SECP256R1_ENABLED
+ecp_write_key_ext:MBEDTLS_ECP_DP_SECP256R1:"0000000000000000000000000000000000000000000000000000000000000001":32:0
+
+ECP write key ext: secp256r1, mostly-0 key, output_size=1
+depends_on:MBEDTLS_ECP_DP_SECP256R1_ENABLED
+ecp_write_key_ext:MBEDTLS_ECP_DP_SECP256R1:"0000000000000000000000000000000000000000000000000000000000000001":1:MBEDTLS_ERR_ECP_BUFFER_TOO_SMALL
+
+ECP write key ext: secp256r1, private key not set
+depends_on:MBEDTLS_ECP_DP_SECP256R1_ENABLED
+ecp_write_key_ext:MBEDTLS_ECP_DP_SECP256R1:"":32:MBEDTLS_ERR_ECP_BAD_INPUT_DATA
+
+ECP write key ext: secp384r1, nominal
+depends_on:MBEDTLS_ECP_DP_SECP384R1_ENABLED
+ecp_write_key_ext:MBEDTLS_ECP_DP_SECP384R1:"d27335ea71664af244dd14e9fd1260715dfd8a7965571c48d709ee7a7962a156d706a90cbcb5df2986f05feadb9376f1":48:0
+
+ECP write key ext: secp384r1, output longer by 1
+depends_on:MBEDTLS_ECP_DP_SECP384R1_ENABLED
+ecp_write_key_ext:MBEDTLS_ECP_DP_SECP384R1:"d27335ea71664af244dd14e9fd1260715dfd8a7965571c48d709ee7a7962a156d706a90cbcb5df2986f05feadb9376f1":49:0
+
+ECP write key ext: secp384r1, output short by 1
+depends_on:MBEDTLS_ECP_DP_SECP384R1_ENABLED
+ecp_write_key_ext:MBEDTLS_ECP_DP_SECP384R1:"d27335ea71664af244dd14e9fd1260715dfd8a7965571c48d709ee7a7962a156d706a90cbcb5df2986f05feadb9376f1":47:MBEDTLS_ERR_ECP_BUFFER_TOO_SMALL
+
+ECP write key ext: Curve25519, nominal
+depends_on:MBEDTLS_ECP_DP_CURVE25519_ENABLED
+ecp_write_key_ext:MBEDTLS_ECP_DP_CURVE25519:"a046e36bf0527c9d3b16154b82465edd62144c0ac1fc5a18506a2244ba449a44":32:0
+
+ECP write key ext: Curve25519, output longer by 1
+depends_on:MBEDTLS_ECP_DP_CURVE25519_ENABLED
+ecp_write_key_ext:MBEDTLS_ECP_DP_CURVE25519:"a046e36bf0527c9d3b16154b82465edd62144c0ac1fc5a18506a2244ba449a44":33:0
+
+ECP write key ext: Curve25519, output short by 1
+depends_on:MBEDTLS_ECP_DP_CURVE25519_ENABLED
+ecp_write_key_ext:MBEDTLS_ECP_DP_CURVE25519:"a046e36bf0527c9d3b16154b82465edd62144c0ac1fc5a18506a2244ba449a44":31:MBEDTLS_ERR_ECP_BUFFER_TOO_SMALL
+
+ECP write key ext: Curve25519, output_size=0
+depends_on:MBEDTLS_ECP_DP_CURVE25519_ENABLED
+ecp_write_key_ext:MBEDTLS_ECP_DP_CURVE25519:"a046e36bf0527c9d3b16154b82465edd62144c0ac1fc5a18506a2244ba449a44":0:MBEDTLS_ERR_ECP_BUFFER_TOO_SMALL
+
+ECP write key ext: Curve25519, mostly-0 key, output_size=32
+depends_on:MBEDTLS_ECP_DP_CURVE25519_ENABLED
+ecp_write_key_ext:MBEDTLS_ECP_DP_CURVE25519:"0000000000000000000000000000000000000000000000000000000000000040":32:0
+
+ECP write key ext: Curve25519, mostly-0 key, output_size=31
+depends_on:MBEDTLS_ECP_DP_CURVE25519_ENABLED
+ecp_write_key_ext:MBEDTLS_ECP_DP_CURVE25519:"0000000000000000000000000000000000000000000000000000000000000040":31:MBEDTLS_ERR_ECP_BUFFER_TOO_SMALL
+
+ECP write key ext: Curve25519, private key not set
+depends_on:MBEDTLS_ECP_DP_CURVE25519_ENABLED
+ecp_write_key_ext:MBEDTLS_ECP_DP_CURVE25519:"":32:MBEDTLS_ERR_ECP_BAD_INPUT_DATA
+
+ECP write key ext: Curve448, nominal
+depends_on:MBEDTLS_ECP_DP_CURVE448_ENABLED
+ecp_write_key_ext:MBEDTLS_ECP_DP_CURVE448:"3c262fddf9ec8e88495266fea19a34d28882acef045104d0d1aae121700a779c984c24f8cdd78fbff44943eba368f54b29259a4f1c600ad3":56:0
+
+ECP write key ext: Curve448, output longer by 1
+depends_on:MBEDTLS_ECP_DP_CURVE448_ENABLED
+ecp_write_key_ext:MBEDTLS_ECP_DP_CURVE448:"3c262fddf9ec8e88495266fea19a34d28882acef045104d0d1aae121700a779c984c24f8cdd78fbff44943eba368f54b29259a4f1c600ad3":57:0
+
+ECP write key ext: Curve448, output short by 1
+depends_on:MBEDTLS_ECP_DP_CURVE448_ENABLED
+ecp_write_key_ext:MBEDTLS_ECP_DP_CURVE448:"3c262fddf9ec8e88495266fea19a34d28882acef045104d0d1aae121700a779c984c24f8cdd78fbff44943eba368f54b29259a4f1c600ad3":55:MBEDTLS_ERR_ECP_BUFFER_TOO_SMALL
+
+ECP write key ext: Curve448, mostly-0 key, output_size=56
+depends_on:MBEDTLS_ECP_DP_CURVE448_ENABLED
+ecp_write_key_ext:MBEDTLS_ECP_DP_CURVE448:"0000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000080":56:0
+
+ECP write key ext: Curve448, mostly-0 key, output_size=55
+depends_on:MBEDTLS_ECP_DP_CURVE448_ENABLED
+ecp_write_key_ext:MBEDTLS_ECP_DP_CURVE448:"0000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000080":55:MBEDTLS_ERR_ECP_BUFFER_TOO_SMALL
+
+ECP write key ext: group not set
+ecp_write_key_ext:MBEDTLS_ECP_DP_NONE:"":32:MBEDTLS_ERR_ECP_BAD_INPUT_DATA
+
 ECP mod p192 small (more than 192 bits, less limbs than 2 * 192 bits)
 depends_on:MBEDTLS_ECP_DP_SECP192R1_ENABLED:MBEDTLS_ECP_NIST_OPTIM
 ecp_fast_mod:MBEDTLS_ECP_DP_SECP192R1:"0100000000000103010000000000010201000000000001010100000000000100"
diff --git a/tests/suites/test_suite_ecp.function b/tests/suites/test_suite_ecp.function
index 9cf0ce1..c061275 100644
--- a/tests/suites/test_suite_ecp.function
+++ b/tests/suites/test_suite_ecp.function
@@ -1296,6 +1296,42 @@
 }
 /* END_CASE */
 
+/* BEGIN_CASE */
+void ecp_write_key_ext(int grp_id, data_t *in_key,
+                       int exported_size, int expected_ret)
+{
+    mbedtls_ecp_keypair key;
+    mbedtls_ecp_keypair_init(&key);
+    unsigned char *exported = NULL;
+
+    if (in_key->len != 0) {
+        TEST_EQUAL(mbedtls_ecp_read_key(grp_id, &key, in_key->x, in_key->len), 0);
+    } else if (grp_id != MBEDTLS_ECP_DP_NONE) {
+        TEST_EQUAL(mbedtls_ecp_group_load(&key.grp, grp_id), 0);
+    }
+
+    TEST_CALLOC(exported, exported_size);
+    size_t olen = 0xdeadbeef;
+    TEST_EQUAL(mbedtls_ecp_write_key_ext(&key, &olen, exported, exported_size),
+               expected_ret);
+
+    if (expected_ret == 0) {
+        TEST_EQUAL(olen, (key.grp.nbits + 7) / 8);
+        TEST_LE_U(olen, MBEDTLS_ECP_MAX_BYTES);
+        TEST_MEMORY_COMPARE(in_key->x, in_key->len,
+                            exported, olen);
+    } else {
+        /* Robustness check: even in the error case, insist that olen is less
+         * than the buffer size. */
+        TEST_LE_U(olen, exported_size);
+    }
+
+exit:
+    mbedtls_ecp_keypair_free(&key);
+    mbedtls_free(exported);
+}
+/* END_CASE */
+
 /* BEGIN_CASE depends_on:MBEDTLS_TEST_HOOKS:MBEDTLS_ECP_MONTGOMERY_ENABLED:MBEDTLS_ECP_LIGHT */
 void genkey_mx_known_answer(int bits, data_t *seed, data_t *expected)
 {