PSA: implement key derivation for ECC keys

Signed-off-by: Przemyslaw Stekiel <przemyslaw.stekiel@mobica.com>
diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index 642fc13..9731932 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -4834,21 +4834,107 @@
     size_t storage_size = bytes;
     psa_status_t status;
 
-    if( ! key_type_is_raw_bytes( slot->attr.type ) )
-        return( PSA_ERROR_INVALID_ARGUMENT );
-    if( bits % 8 != 0 )
-        return( PSA_ERROR_INVALID_ARGUMENT );
-    data = mbedtls_calloc( 1, bytes );
-    if( data == NULL )
-        return( PSA_ERROR_INSUFFICIENT_MEMORY );
+    /*
+    * ECC key types require the generation of a private key which is an integer
+    * in the range [1, N - 1], where N is the boundary of the private key domain:
+    * N is the prime p for Diffie-Hellman, or the order of the
+    * curve’s base point for ECC.
+    *
+    * Let m be the bit size of N, such that 2^m > N >= 2^(m-1).
+    * This function generates the private key using the following process:
+    *
+    * 1. Draw a byte string of length ceiling(m/8) bytes.
+    * 2. If m is not a multiple of 8, set the most significant
+    *    (8 * ceiling(m/8) - m) bits of the first byte in the string to zero.
+    * 3. Convert the string to integer k by decoding it as a big-endian byte string.
+    * 4. If k > N - 2, discard the result and return to step 1.
+    * 5. Output k + 1 as the private key.
+    *
+    * This method allows compliance to NIST standards
+    */
+    if ( PSA_KEY_TYPE_IS_ECC( slot->attr.type ) )
+    {
+        int cmp_result;
+        do {
+            int ret;
+            psa_ecc_family_t curve = PSA_KEY_TYPE_ECC_GET_FAMILY(
+                                        slot->attr.type );
+            mbedtls_ecp_group_id grp_id =
+                mbedtls_ecc_group_of_psa( curve, bits, 0 );
 
-    status = psa_key_derivation_output_bytes( operation, data, bytes );
-    if( status != PSA_SUCCESS )
-        goto exit;
+            mbedtls_ecp_keypair ecp;
+            mbedtls_ecp_keypair_init( &ecp );
+
+            if( ( ret = mbedtls_ecp_group_load( &ecp.grp, grp_id ) ) != 0 )
+                return( ret );
+
+            /* N is the boundary of the private key domain */
+            mbedtls_mpi N = ecp.grp.N;
+            /* Let m be the bit size of N */
+            size_t m = ecp.grp.nbits;
+
+            size_t m_bytes = PSA_BITS_TO_BYTES( m );
+
+            /* Alloc buffer once */
+            if ( data == NULL )
+                data = mbedtls_calloc( 1, m_bytes );
+            if( data == NULL )
+                return( PSA_ERROR_INSUFFICIENT_MEMORY );
+
+            /* 1. Draw a byte string of length ceiling(m/8) bytes. */
+            status = psa_key_derivation_output_bytes( operation, data, m_bytes );
+            if( status != PSA_SUCCESS )
+                goto exit;
+
+            /* 2. If m is not a multiple of 8 */
+            if (m % 8)
+            {
+                /* set the most significant
+                 * (8 * ceiling(m/8) - m) bits of the first byte in
+                 * the string to zero.
+                 */
+                uint8_t clear_bit_count = ( 8 * m_bytes - m );
+                uint8_t clear_bit_mask = ( ( 1 << clear_bit_count ) - 1 );
+                clear_bit_mask = ~( clear_bit_mask << ( 8 - clear_bit_count ) );
+                data[0] = ( data[0] & clear_bit_mask );
+            }
+
+            /* 3. Convert the string to integer k by decoding it as a
+             *    big-endian byte string.
+             */
+            mbedtls_mpi k;
+            mbedtls_mpi_init( &k );
+            mbedtls_mpi_read_binary( &k, data, m_bytes);
+
+            /* 4. If k > N - 2, discard the result and return to step 1. */
+            mbedtls_mpi diff_N_2;
+            mbedtls_mpi_init( &diff_N_2 );
+            mbedtls_mpi_sub_int( &diff_N_2, &N, 2);
+            cmp_result = mbedtls_mpi_cmp_mpi( &k, &diff_N_2 );
+
+            /* 5. Output k + 1 as the private key. */
+            mbedtls_mpi sum_k_1;
+            mbedtls_mpi_init( &sum_k_1 );
+            mbedtls_mpi_add_int( &sum_k_1, &k, 1);
+            mbedtls_mpi_write_binary( &sum_k_1, data, m_bytes);
+        } while ( cmp_result == 1 );
+    } else {
+        if( ! key_type_is_raw_bytes( slot->attr.type ) )
+            return( PSA_ERROR_INVALID_ARGUMENT );
+        if( bits % 8 != 0 )
+            return( PSA_ERROR_INVALID_ARGUMENT );
+        data = mbedtls_calloc( 1, bytes );
+        if( data == NULL )
+            return( PSA_ERROR_INSUFFICIENT_MEMORY );
+
+        status = psa_key_derivation_output_bytes( operation, data, bytes );
+        if( status != PSA_SUCCESS )
+            goto exit;
 #if defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_DES)
-    if( slot->attr.type == PSA_KEY_TYPE_DES )
-        psa_des_set_key_parity( data, bytes );
+        if( slot->attr.type == PSA_KEY_TYPE_DES )
+            psa_des_set_key_parity( data, bytes );
 #endif /* MBEDTLS_PSA_BUILTIN_KEY_TYPE_DES */
+    }
 
     slot->attr.bits = (psa_key_bits_t) bits;
     psa_key_attributes_t attributes = {