Add an EC J-PAKE KDF to transform K -> SHA256(K.X) for TLS 1.2

TLS uses it to derive the session secret. The algorithm takes a serialized
point in an uncompressed form, extracts the X coordinate and computes
SHA256 of it. It is only expected to work with P-256.
Fixes #5978.
Signed-off-by: Andrzej Kurek <andrzej.kurek@arm.com>
diff --git a/include/mbedtls/config_psa.h b/include/mbedtls/config_psa.h
index b84a80a..88052d2 100644
--- a/include/mbedtls/config_psa.h
+++ b/include/mbedtls/config_psa.h
@@ -228,6 +228,12 @@
 #endif /* !MBEDTLS_PSA_ACCEL_ALG_TLS12_PSK_TO_MS */
 #endif /* PSA_WANT_ALG_TLS12_PSK_TO_MS */
 
+#if defined(PSA_WANT_ALG_TLS12_ECJPAKE_TO_PMS)
+#if !defined(MBEDTLS_PSA_ACCEL_ALG_TLS12_ECJPAKE_TO_PMS)
+#define MBEDTLS_PSA_BUILTIN_ALG_TLS12_ECJPAKE_TO_PMS 1
+#endif /* !MBEDTLS_PSA_ACCEL_ALG_TLS12_ECJPAKE_TO_PMS */
+#endif /* PSA_WANT_ALG_TLS12_ECJPAKE_TO_PMS */
+
 #if defined(PSA_WANT_KEY_TYPE_ECC_KEY_PAIR)
 #if !defined(MBEDTLS_PSA_ACCEL_KEY_TYPE_ECC_KEY_PAIR)
 #define MBEDTLS_PSA_BUILTIN_KEY_TYPE_ECC_KEY_PAIR 1
@@ -629,6 +635,8 @@
 #define PSA_WANT_ALG_TLS12_PRF 1
 #define MBEDTLS_PSA_BUILTIN_ALG_TLS12_PSK_TO_MS 1
 #define PSA_WANT_ALG_TLS12_PSK_TO_MS 1
+#define MBEDTLS_PSA_BUILTIN_ALG_TLS12_ECJPAKE_TO_PMS 1
+#define PSA_WANT_ALG_TLS12_ECJPAKE_TO_PMS 1
 #endif /* MBEDTLS_MD_C */
 
 #if defined(MBEDTLS_MD5_C)
diff --git a/include/psa/crypto_config.h b/include/psa/crypto_config.h
index 991be96..8737e29 100644
--- a/include/psa/crypto_config.h
+++ b/include/psa/crypto_config.h
@@ -88,6 +88,8 @@
 #define PSA_WANT_ALG_STREAM_CIPHER              1
 #define PSA_WANT_ALG_TLS12_PRF                  1
 #define PSA_WANT_ALG_TLS12_PSK_TO_MS            1
+#define PSA_WANT_ALG_TLS12_ECJPAKE_TO_PMS       1
+
 /* PBKDF2-HMAC is not yet supported via the PSA API in Mbed TLS.
  * Note: when adding support, also adjust include/mbedtls/config_psa.h */
 //#define PSA_WANT_ALG_XTS                        1
diff --git a/include/psa/crypto_sizes.h b/include/psa/crypto_sizes.h
index 1024d6b..231ea62 100644
--- a/include/psa/crypto_sizes.h
+++ b/include/psa/crypto_sizes.h
@@ -239,6 +239,15 @@
  */
 #define PSA_TLS12_PSK_TO_MS_PSK_MAX_SIZE 128
 
+/* The expected size of input passed to psa_tls12_ecjpake_to_pms_input,
+ * which is expected to work with P-256 curve only. */
+#define PSA_TLS12_ECJPAKE_TO_PMS_INPUT_SIZE 65
+
+/* The size of a serialized K.X coordinate to be used in
+ * psa_tls12_ecjpake_to_pms_input. This function only accepts the P-256
+ * curve. */
+#define PSA_TLS12_ECJPAKE_TO_PMS_DATA_SIZE 32
+
 /** The maximum size of a block cipher. */
 #define PSA_BLOCK_CIPHER_BLOCK_MAX_SIZE 16
 
diff --git a/include/psa/crypto_struct.h b/include/psa/crypto_struct.h
index 957b4c6..afba325 100644
--- a/include/psa/crypto_struct.h
+++ b/include/psa/crypto_struct.h
@@ -202,6 +202,12 @@
 #endif /* MBEDTLS_PSA_BUILTIN_ALG_HKDF ||
           MBEDTLS_PSA_BUILTIN_ALG_HKDF_EXTRACT ||
           MBEDTLS_PSA_BUILTIN_ALG_HKDF_EXPAND */
+#if defined(MBEDTLS_PSA_BUILTIN_ALG_TLS12_ECJPAKE_TO_PMS)
+typedef struct
+{
+    uint8_t MBEDTLS_PRIVATE(data)[PSA_TLS12_ECJPAKE_TO_PMS_DATA_SIZE];
+} psa_tls12_ecjpake_to_pms_t;
+#endif /* MBEDTLS_PSA_BUILTIN_ALG_TLS12_ECJPAKE_TO_PMS */
 
 #if defined(MBEDTLS_PSA_BUILTIN_ALG_TLS12_PRF) || \
     defined(MBEDTLS_PSA_BUILTIN_ALG_TLS12_PSK_TO_MS)
@@ -267,6 +273,9 @@
     defined(MBEDTLS_PSA_BUILTIN_ALG_TLS12_PSK_TO_MS)
         psa_tls12_prf_key_derivation_t MBEDTLS_PRIVATE(tls12_prf);
 #endif
+#if defined(MBEDTLS_PSA_BUILTIN_ALG_TLS12_ECJPAKE_TO_PMS)
+        psa_tls12_ecjpake_to_pms_t MBEDTLS_PRIVATE(tls12_ecjpake_to_pms);
+#endif
     } MBEDTLS_PRIVATE(ctx);
 };
 
diff --git a/include/psa/crypto_values.h b/include/psa/crypto_values.h
index 5e6e5e3..0343819 100644
--- a/include/psa/crypto_values.h
+++ b/include/psa/crypto_values.h
@@ -2021,6 +2021,14 @@
 #define PSA_ALG_TLS12_PSK_TO_MS_GET_HASH(hkdf_alg)                         \
     (PSA_ALG_CATEGORY_HASH | ((hkdf_alg) & PSA_ALG_HASH_MASK))
 
+/* Macro to build a KDF that takes the shared secret K (an EC point in case
+ * of EC J-PAKE) and calculates SHA256(K.X) that the rest of TLS 1.2 will
+ * use to derive the session secret. Uses PSA_ALG_SHA_256.
+ */
+#define PSA_ALG_TLS12_ECJPAKE_TO_PMS            ((psa_algorithm_t)0x08000600)
+#define PSA_ALG_IS_TLS12_ECJPAKE_TO_PMS(alg)                               \
+    (alg == PSA_ALG_TLS12_ECJPAKE_TO_PMS)
+
 /* This flag indicates whether the key derivation algorithm is suitable for
  * use on low-entropy secrets such as password - these algorithms are also
  * known as key stretching or password hashing schemes. These are also the
diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index b0116dd..5c05f79 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -4243,7 +4243,8 @@
 
 #if defined(BUILTIN_ALG_ANY_HKDF) || \
     defined(MBEDTLS_PSA_BUILTIN_ALG_TLS12_PRF) || \
-    defined(MBEDTLS_PSA_BUILTIN_ALG_TLS12_PSK_TO_MS)
+    defined(MBEDTLS_PSA_BUILTIN_ALG_TLS12_PSK_TO_MS) || \
+    defined(MBEDTLS_PSA_BUILTIN_ALG_TLS12_ECJPAKE_TO_PMS)
 #define AT_LEAST_ONE_BUILTIN_KDF
 #endif /* At least one builtin KDF */
 
@@ -4350,6 +4351,14 @@
     else
 #endif /* defined(MBEDTLS_PSA_BUILTIN_ALG_TLS12_PRF) ||
         * defined(MBEDTLS_PSA_BUILTIN_ALG_TLS12_PSK_TO_MS) */
+#if defined(MBEDTLS_PSA_BUILTIN_ALG_TLS12_ECJPAKE_TO_PMS)
+    if( PSA_ALG_IS_TLS12_ECJPAKE_TO_PMS( kdf_alg ) )
+    {
+        mbedtls_platform_zeroize( operation->ctx.tls12_ecjpake_to_pms.data,
+            PSA_TLS12_ECJPAKE_TO_PMS_DATA_SIZE );
+    }
+    else
+#endif /* defined(MBEDTLS_PSA_BUILTIN_ALG_TLS12_ECJPAKE_TO_PMS) */
     {
         status = PSA_ERROR_BAD_STATE;
     }
@@ -4631,6 +4640,31 @@
 #endif /* MBEDTLS_PSA_BUILTIN_ALG_TLS12_PRF ||
         * MBEDTLS_PSA_BUILTIN_ALG_TLS12_PSK_TO_MS */
 
+#if defined(MBEDTLS_PSA_BUILTIN_ALG_TLS12_ECJPAKE_TO_PMS)
+static psa_status_t psa_key_derivation_tls12_ecjpake_to_pms_read(
+    psa_tls12_ecjpake_to_pms_t *ecjpake,
+    uint8_t *output,
+    size_t output_length )
+{
+    psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
+    size_t output_size;
+
+    if( output_length != 32 )
+        return ( PSA_ERROR_INVALID_ARGUMENT );
+
+    status = psa_hash_compute( PSA_ALG_SHA_256, ecjpake->data,
+        PSA_TLS12_ECJPAKE_TO_PMS_DATA_SIZE, output, output_length,
+        &output_size );
+    if( status != PSA_SUCCESS )
+        return ( status );
+
+    if( output_size != output_length )
+        return ( PSA_ERROR_GENERIC_ERROR );
+
+    return ( PSA_SUCCESS );
+}
+#endif
+
 psa_status_t psa_key_derivation_output_bytes(
     psa_key_derivation_operation_t *operation,
     uint8_t *output,
@@ -4685,6 +4719,15 @@
     else
 #endif /* MBEDTLS_PSA_BUILTIN_ALG_TLS12_PRF ||
         * MBEDTLS_PSA_BUILTIN_ALG_TLS12_PSK_TO_MS */
+#if defined(MBEDTLS_PSA_BUILTIN_ALG_TLS12_ECJPAKE_TO_PMS)
+    if( PSA_ALG_IS_TLS12_ECJPAKE_TO_PMS( kdf_alg ) )
+    {
+        status = psa_key_derivation_tls12_ecjpake_to_pms_read(
+            &operation->ctx.tls12_ecjpake_to_pms, output, output_length );
+    }
+    else
+#endif /* MBEDTLS_PSA_BUILTIN_ALG_TLS12_ECJPAKE_TO_PMS */
+
     {
         (void) kdf_alg;
         return( PSA_ERROR_BAD_STATE );
@@ -5077,6 +5120,10 @@
     if( PSA_ALG_IS_TLS12_PSK_TO_MS( kdf_alg ) )
         return( 1 );
 #endif
+#if defined(MBEDTLS_PSA_BUILTIN_ALG_TLS12_ECJPAKE_TO_PMS)
+    if( PSA_ALG_IS_TLS12_ECJPAKE_TO_PMS( kdf_alg ) )
+        return( 1 );
+#endif
     return( 0 );
 }
 
@@ -5100,19 +5147,26 @@
     if( ! is_kdf_alg_supported( kdf_alg ) )
         return( PSA_ERROR_NOT_SUPPORTED );
 
-    /* All currently supported key derivation algorithms are based on a
-     * hash algorithm. */
+    /* All currently supported key derivation algorithms (apart from
+     * ecjpake to pms are based on a hash algorithm. */
     psa_algorithm_t hash_alg = PSA_ALG_HKDF_GET_HASH( kdf_alg );
     size_t hash_size = PSA_HASH_LENGTH( hash_alg );
-    if( hash_size == 0 )
-        return( PSA_ERROR_NOT_SUPPORTED );
+    if( !PSA_ALG_IS_TLS12_ECJPAKE_TO_PMS( kdf_alg ) )
+    {
+        if( hash_size == 0 )
+            return( PSA_ERROR_NOT_SUPPORTED );
 
-    /* Make sure that hash_alg is a supported hash algorithm. Otherwise
-     * we might fail later, which is somewhat unfriendly and potentially
-     * risk-prone. */
-    psa_status_t status = psa_hash_try_support( hash_alg );
-    if( status != PSA_SUCCESS )
-        return( status );
+        /* Make sure that hash_alg is a supported hash algorithm. Otherwise
+         * we might fail later, which is somewhat unfriendly and potentially
+         * risk-prone. */
+        psa_status_t status = psa_hash_try_support( hash_alg );
+        if( status != PSA_SUCCESS )
+            return( status );
+    }
+    else
+    {
+        hash_size = PSA_HASH_LENGTH( PSA_ALG_SHA_256 );
+    }
 
     if( ( PSA_ALG_IS_TLS12_PRF( kdf_alg ) ||
           PSA_ALG_IS_TLS12_PSK_TO_MS( kdf_alg ) ) &&
@@ -5513,6 +5567,25 @@
 }
 #endif /* MBEDTLS_PSA_BUILTIN_ALG_TLS12_PSK_TO_MS */
 
+#if defined(MBEDTLS_PSA_BUILTIN_ALG_TLS12_ECJPAKE_TO_PMS)
+static psa_status_t psa_tls12_ecjpake_to_pms_input(
+    psa_tls12_ecjpake_to_pms_t *ecjpake,
+    const uint8_t *data,
+    size_t data_length )
+{
+    if( data_length != PSA_TLS12_ECJPAKE_TO_PMS_INPUT_SIZE )
+        return( PSA_ERROR_INVALID_ARGUMENT );
+
+    /* Check if the passed point is in an uncompressed form */
+    if( data[0] != 0x04 )
+        return( PSA_ERROR_INVALID_ARGUMENT );
+
+    /* Only K.X has to be extracted - bytes 1 to 32 inclusive. */
+    memcpy( ecjpake->data, data + 1, PSA_TLS12_ECJPAKE_TO_PMS_DATA_SIZE );
+
+    return( PSA_SUCCESS );
+}
+#endif /* MBEDTLS_PSA_BUILTIN_ALG_TLS12_ECJPAKE_TO_PMS */
 /** Check whether the given key type is acceptable for the given
  * input step of a key derivation.
  *
@@ -5591,6 +5664,14 @@
     }
     else
 #endif /* MBEDTLS_PSA_BUILTIN_ALG_TLS12_PSK_TO_MS */
+#if defined(MBEDTLS_PSA_BUILTIN_ALG_TLS12_ECJPAKE_TO_PMS)
+    if( PSA_ALG_IS_TLS12_ECJPAKE_TO_PMS( kdf_alg ) )
+    {
+        status = psa_tls12_ecjpake_to_pms_input(
+            &operation->ctx.tls12_ecjpake_to_pms, data, data_length );
+    }
+    else
+#endif /* MBEDTLS_PSA_BUILTIN_ALG_TLS12_ECJPAKE_TO_PMS */
     {
         /* This can't happen unless the operation object was not initialized */
         (void) data;
diff --git a/scripts/mbedtls_dev/crypto_knowledge.py b/scripts/mbedtls_dev/crypto_knowledge.py
index 592fc0a..f52ca9a 100644
--- a/scripts/mbedtls_dev/crypto_knowledge.py
+++ b/scripts/mbedtls_dev/crypto_knowledge.py
@@ -357,6 +357,7 @@
         'HKDF': AlgorithmCategory.KEY_DERIVATION,
         'TLS12_PRF': AlgorithmCategory.KEY_DERIVATION,
         'TLS12_PSK_TO_MS': AlgorithmCategory.KEY_DERIVATION,
+        'TLS12_ECJPAKE_TO_PMS': AlgorithmCategory.KEY_DERIVATION,
         'PBKDF': AlgorithmCategory.KEY_DERIVATION,
         'ECDH': AlgorithmCategory.KEY_AGREEMENT,
         'FFDH': AlgorithmCategory.KEY_AGREEMENT,