New function mbedtls_ecp_set_public_key

Set the public key in a key pair. This complements mbedtls_ecp_read_key and
the functions can be used in either order.

Document the need to call check functions separately.

Signed-off-by: Gilles Peskine <Gilles.Peskine@arm.com>
diff --git a/include/mbedtls/ecp.h b/include/mbedtls/ecp.h
index f169008..96f030d 100644
--- a/include/mbedtls/ecp.h
+++ b/include/mbedtls/ecp.h
@@ -1259,6 +1259,38 @@
                         int (*f_rng)(void *, unsigned char *, size_t),
                         void *p_rng);
 
+/** \brief          Set the public key in a key pair object.
+ *
+ * \note            This function does not check that the point actually
+ *                  belongs to the given group. Call mbedtls_ecp_check_pubkey()
+ *                  on \p Q before calling this function to check that.
+ *
+ * \note            This function does not check that the public key matches
+ *                  the private key that is already in \p key, if any.
+ *                  To check the consistency of the resulting key pair object,
+ *                  call mbedtls_ecp_check_pub_priv() after setting both
+ *                  the public key and the private key.
+ *
+ * \param grp_id    The ECP group identifier.
+ * \param key       The key pair object. It must be initialized.
+ *                  If its group has already been set, it must match \p grp_id.
+ *                  If its group has not been set, it will be set to \p grp_id.
+ *                  If the public key has already been set, it is overwritten.
+ * \param Q         The public key to copy. This must be a point on the
+ *                  curve indicated by \p grp_id.
+ *
+ * \return          \c 0 on success.
+ * \return          #MBEDTLS_ERR_ECP_BAD_INPUT_DATA if \p key does not
+ *                  match \p grp_id.
+ * \return          #MBEDTLS_ERR_ECP_FEATURE_UNAVAILABLE if the operation for
+ *                  the group is not implemented.
+ * \return          #MBEDTLS_ERR_MPI_ALLOC_FAILED on memory-allocation failure.
+ * \return          Another negative error code on other kinds of failure.
+ */
+int mbedtls_ecp_set_public_key(mbedtls_ecp_group_id grp_id,
+                               mbedtls_ecp_keypair *key,
+                               const mbedtls_ecp_point *Q);
+
 /**
  * \brief           This function reads an elliptic curve private key.
  *
diff --git a/library/ecp.c b/library/ecp.c
index b4da3c5..bb0cf69 100644
--- a/library/ecp.c
+++ b/library/ecp.c
@@ -3198,6 +3198,25 @@
 }
 #endif /* MBEDTLS_ECP_C */
 
+int mbedtls_ecp_set_public_key(mbedtls_ecp_group_id grp_id,
+                               mbedtls_ecp_keypair *key,
+                               const mbedtls_ecp_point *Q)
+{
+    int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
+
+    if (key->grp.id == MBEDTLS_ECP_DP_NONE) {
+        /* Group not set yet */
+        if ((ret = mbedtls_ecp_group_load(&key->grp, grp_id)) != 0) {
+            return ret;
+        }
+    } else if (key->grp.id != grp_id) {
+        /* Group mismatch */
+        return MBEDTLS_ERR_ECP_BAD_INPUT_DATA;
+    }
+    return mbedtls_ecp_copy(&key->Q, Q);
+}
+
+
 #define ECP_CURVE25519_KEY_SIZE 32
 #define ECP_CURVE448_KEY_SIZE   56
 /*
diff --git a/tests/suites/test_suite_ecp.data b/tests/suites/test_suite_ecp.data
index 1002991..8bf288b 100644
--- a/tests/suites/test_suite_ecp.data
+++ b/tests/suites/test_suite_ecp.data
@@ -581,6 +581,48 @@
 ECP generate Montgomery key: Curve448, not enough entropy
 genkey_mx_known_answer:447:"4f0102030405060708090a0b0c0d0e0f101112131415161718191a1b1c1d1e1f202122232425262728292a2b2c2d2e2f30313233343536":""
 
+ECP set public key: invalid group (0)
+ecp_set_public_key_group_check:MBEDTLS_ECP_DP_NONE:MBEDTLS_ERR_ECP_FEATURE_UNAVAILABLE
+
+ECP set public key: valid group (secp256r1)
+depends_on:MBEDTLS_ECP_DP_SECP256R1_ENABLED
+ecp_set_public_key_group_check:MBEDTLS_ECP_DP_SECP256R1:0
+
+ECP set public key: group not supported (secp256r1)
+depends_on:!MBEDTLS_ECP_DP_SECP256R1_ENABLED
+ecp_set_public_key_group_check:MBEDTLS_ECP_DP_SECP256R1:MBEDTLS_ERR_ECP_FEATURE_UNAVAILABLE
+
+ECP set public key: bad group (not in enum)
+ecp_set_public_key_group_check:MBEDTLS_ECP_DP_MAX:MBEDTLS_ERR_ECP_FEATURE_UNAVAILABLE
+
+ECP set public key: good, secp256r1
+depends_on:MBEDTLS_ECP_DP_SECP256R1_ENABLED
+ecp_set_public_key_good:MBEDTLS_ECP_DP_SECP256R1:"04e1e1e1e1e1e1e1e1e1e1e1e1e1e1e1e1e1e1e1e1e1e0e1ff20e1ffe120e1e1e173287170a761308491683e345cacaebb500c96e1a7bbd37772968b2c951f0579"
+
+ECP set public key: good, Curve25519
+depends_on:MBEDTLS_ECP_DP_CURVE25519_ENABLED
+ecp_set_public_key_good:MBEDTLS_ECP_DP_CURVE25519:"8520f0098930a754748b7ddcb43ef75a0dbf3a0d26381af4eba4a98eaa9b4e6a"
+
+ECP set public key after private: good, secp256r1
+depends_on:MBEDTLS_ECP_DP_SECP256R1_ENABLED
+ecp_set_public_key_after_private:MBEDTLS_ECP_DP_SECP256R1:"70726976617465206b6579":MBEDTLS_ECP_DP_SECP256R1:"04e1e1e1e1e1e1e1e1e1e1e1e1e1e1e1e1e1e1e1e1e1e0e1ff20e1ffe120e1e1e173287170a761308491683e345cacaebb500c96e1a7bbd37772968b2c951f0579"
+
+ECP set public key after private: good, Curve25519
+depends_on:MBEDTLS_ECP_DP_CURVE25519_ENABLED
+ecp_set_public_key_after_private:MBEDTLS_ECP_DP_CURVE25519:"70076d0a7318a57d3c16c17251b26645df4c2f87ebc0992ab177fba51db92c6a":MBEDTLS_ECP_DP_CURVE25519:"8520f0098930a754748b7ddcb43ef75a0dbf3a0d26381af4eba4a98eaa9b4e6a"
+
+ECP set public key after private: secp256r1 then secp256k1
+depends_on:MBEDTLS_ECP_DP_SECP256R1_ENABLED:MBEDTLS_ECP_DP_SECP256K1_ENABLED
+ecp_set_public_key_after_private:MBEDTLS_ECP_DP_SECP256R1:"70726976617465206b6579":MBEDTLS_ECP_DP_SECP256K1:"04e1e1e1e1e1e1e1e1e1e1e1e1e1e1e1e1e1e1e1e1e1e0e1ff20e1ffe120e1e1e173287170a761308491683e345cacaebb500c96e1a7bbd37772968b2c951f0579"
+
+ECP set public key after private: secp256r1 then secp384r1
+depends_on:MBEDTLS_ECP_DP_SECP256R1_ENABLED:MBEDTLS_ECP_DP_SECP384R1_ENABLED
+ecp_set_public_key_after_private:MBEDTLS_ECP_DP_SECP256R1:"70726976617465206b6579":MBEDTLS_ECP_DP_SECP384R1:"04aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaae1e1e1e1e1e1e1e1e1e1e1e1e1e1e1e1e1e1e1e1e1e0e1ff20e1ffe120e1e1e173287170a761308491683e345cacaebb500c96e1a7bbd37772968b2c951f0579bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb"
+
+ECP set public key after private: secp384r1 then secp256r1
+depends_on:MBEDTLS_ECP_DP_SECP384R1_ENABLED:MBEDTLS_ECP_DP_SECP256R1_ENABLED
+ecp_set_public_key_after_private:MBEDTLS_ECP_DP_SECP384R1:"70726976617465206b6579":MBEDTLS_ECP_DP_SECP256R1:"04e1e1e1e1e1e1e1e1e1e1e1e1e1e1e1e1e1e1e1e1e1e0e1ff20e1ffe120e1e1e173287170a761308491683e345cacaebb500c96e1a7bbd37772968b2c951f0579"
+
 ECP read key #1 (short weierstrass, too small)
 depends_on:MBEDTLS_ECP_DP_SECP192R1_ENABLED
 mbedtls_ecp_read_key:MBEDTLS_ECP_DP_SECP192R1:"00":MBEDTLS_ERR_ECP_INVALID_KEY:0
diff --git a/tests/suites/test_suite_ecp.function b/tests/suites/test_suite_ecp.function
index aefb57a..53b78d9 100644
--- a/tests/suites/test_suite_ecp.function
+++ b/tests/suites/test_suite_ecp.function
@@ -1040,6 +1040,109 @@
 /* END_CASE */
 
 /* BEGIN_CASE */
+void ecp_set_public_key_group_check(int grp_id, int expected_ret)
+{
+    mbedtls_ecp_keypair key;
+    mbedtls_ecp_keypair_init(&key);
+    mbedtls_ecp_point Q;
+    mbedtls_ecp_point_init(&Q);
+
+    TEST_EQUAL(mbedtls_ecp_set_public_key(grp_id, &key, &Q),
+               expected_ret);
+
+exit:
+    mbedtls_ecp_keypair_free(&key);
+    mbedtls_ecp_point_free(&Q);
+}
+/* END_CASE */
+
+/* BEGIN_CASE */
+void ecp_set_public_key_good(int grp_id, data_t *public_data)
+{
+    mbedtls_ecp_keypair key;
+    mbedtls_ecp_keypair_init(&key);
+    mbedtls_ecp_group grp;
+    mbedtls_ecp_group_init(&grp);
+    mbedtls_ecp_point Q;
+    mbedtls_ecp_point_init(&Q);
+
+    TEST_EQUAL(mbedtls_ecp_group_load(&grp, grp_id), 0);
+    TEST_EQUAL(mbedtls_ecp_point_read_binary(&grp, &Q,
+                                             public_data->x, public_data->len),
+               0);
+
+    /* Freshly initialized key */
+    TEST_EQUAL(mbedtls_ecp_set_public_key(grp_id, &key, &Q), 0);
+    TEST_EQUAL(key.grp.id, grp_id);
+    TEST_EQUAL(mbedtls_ecp_point_cmp(&key.Q, &Q), 0);
+
+#if defined(MBEDTLS_BIGNUM_C)
+    /* Key with a public key already set to a different value */
+    TEST_EQUAL(mbedtls_mpi_add_int(&key.Q.X, &key.Q.X, 1), 0);
+    TEST_EQUAL(mbedtls_mpi_add_int(&key.Q.Y, &key.Q.Y, 1), 0);
+    TEST_EQUAL(mbedtls_mpi_add_int(&key.Q.Z, &key.Q.Z, 1), 0);
+    TEST_EQUAL(mbedtls_ecp_set_public_key(grp_id, &key, &Q), 0);
+    TEST_EQUAL(key.grp.id, grp_id);
+    TEST_EQUAL(mbedtls_ecp_point_cmp(&key.Q, &Q), 0);
+#endif
+
+exit:
+    mbedtls_ecp_keypair_free(&key);
+    mbedtls_ecp_group_free(&grp);
+    mbedtls_ecp_point_free(&Q);
+}
+/* END_CASE */
+
+/* BEGIN_CASE */
+void ecp_set_public_key_after_private(int private_grp_id, data_t *private_data,
+                                      int public_grp_id, data_t *public_data)
+{
+    mbedtls_ecp_keypair key;
+    mbedtls_ecp_keypair_init(&key);
+    mbedtls_ecp_group grp;
+    mbedtls_ecp_group_init(&grp);
+    mbedtls_ecp_point Q;
+    mbedtls_ecp_point_init(&Q);
+#if defined(MBEDTLS_BIGNUM_C)
+    mbedtls_mpi d;
+    mbedtls_mpi_init(&d);
+#endif
+
+    TEST_EQUAL(mbedtls_ecp_group_load(&grp, public_grp_id), 0);
+    TEST_EQUAL(mbedtls_ecp_point_read_binary(&grp, &Q,
+                                             public_data->x, public_data->len),
+               0);
+    TEST_EQUAL(mbedtls_ecp_read_key(private_grp_id, &key,
+                                    private_data->x, private_data->len),
+               0);
+#if defined(MBEDTLS_BIGNUM_C)
+    TEST_EQUAL(mbedtls_mpi_copy(&d, &key.d), 0);
+#endif
+
+    int ret = mbedtls_ecp_set_public_key(public_grp_id, &key, &Q);
+
+    if (private_grp_id == public_grp_id) {
+        TEST_EQUAL(ret, 0);
+        TEST_EQUAL(key.grp.id, public_grp_id);
+        TEST_EQUAL(mbedtls_ecp_point_cmp(&key.Q, &Q), 0);
+#if defined(MBEDTLS_BIGNUM_C)
+        TEST_EQUAL(mbedtls_mpi_cmp_mpi(&d, &key.d), 0);
+#endif
+    } else {
+        TEST_EQUAL(ret, MBEDTLS_ERR_ECP_BAD_INPUT_DATA);
+    }
+
+exit:
+    mbedtls_ecp_keypair_free(&key);
+    mbedtls_ecp_group_free(&grp);
+    mbedtls_ecp_point_free(&Q);
+#if defined(MBEDTLS_BIGNUM_C)
+    mbedtls_mpi_free(&d);
+#endif
+}
+/* END_CASE */
+
+/* BEGIN_CASE */
 void mbedtls_ecp_read_key(int grp_id, data_t *in_key, int expected, int canonical)
 {
     int ret = 0;