From 08d34b8693a03727e6731ae5adf44a8e081dfbf9 Mon Sep 17 00:00:00 2001
From: Andrzej Kurek <andrzej.kurek@arm.com>
Date: Fri, 29 Jul 2022 10:00:16 -0400
Subject: [PATCH 01/18] 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>
---
 include/mbedtls/config_psa.h            |   8 ++
 include/psa/crypto_config.h             |   2 +
 include/psa/crypto_sizes.h              |   9 +++
 include/psa/crypto_struct.h             |   9 +++
 include/psa/crypto_values.h             |   8 ++
 library/psa_crypto.c                    | 103 +++++++++++++++++++++---
 scripts/mbedtls_dev/crypto_knowledge.py |   1 +
 7 files changed, 129 insertions(+), 11 deletions(-)

diff --git a/include/mbedtls/config_psa.h b/include/mbedtls/config_psa.h
index b84a80a30c..88052d228b 100644
--- a/include/mbedtls/config_psa.h
+++ b/include/mbedtls/config_psa.h
@@ -228,6 +228,12 @@ extern "C" {
 #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 @@ extern "C" {
 #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 991be96ef4..8737e2911c 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 1024d6b918..231ea624a7 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 957b4c6113..afba325022 100644
--- a/include/psa/crypto_struct.h
+++ b/include/psa/crypto_struct.h
@@ -202,6 +202,12 @@ typedef struct
 #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)
@@ -266,6 +272,9 @@ struct psa_key_derivation_s
 #if defined(MBEDTLS_PSA_BUILTIN_ALG_TLS12_PRF) || \
     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 5e6e5e352c..03438196b6 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 b0116ddfb4..5c05f79282 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -4243,7 +4243,8 @@ psa_status_t psa_aead_abort( psa_aead_operation_t *operation )
 
 #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 @@ psa_status_t psa_key_derivation_abort( psa_key_derivation_operation_t *operation
     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 @@ static psa_status_t psa_key_derivation_tls12_prf_read(
 #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 @@ psa_status_t psa_key_derivation_output_bytes(
     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 );
@@ -5076,6 +5119,10 @@ static int is_kdf_alg_supported( psa_algorithm_t kdf_alg )
 #if defined(MBEDTLS_PSA_BUILTIN_ALG_TLS12_PSK_TO_MS)
     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 @@ static psa_status_t psa_key_derivation_setup_kdf(
     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 @@ static psa_status_t psa_tls12_prf_psk_to_ms_input(
 }
 #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 @@ static psa_status_t psa_key_derivation_input_internal(
     }
     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 592fc0afe2..f52ca9ac8c 100644
--- a/scripts/mbedtls_dev/crypto_knowledge.py
+++ b/scripts/mbedtls_dev/crypto_knowledge.py
@@ -357,6 +357,7 @@ class Algorithm:
         '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,

From d8705bc7b77288e4c18e9a011d8d0d5097d02657 Mon Sep 17 00:00:00 2001
From: Andrzej Kurek <andrzej.kurek@arm.com>
Date: Fri, 29 Jul 2022 10:02:05 -0400
Subject: [PATCH 02/18] Add tests for the newly created ad-hoc EC J-PAKE KDF

Signed-off-by: Andrzej Kurek <andrzej.kurek@arm.com>
---
 tests/suites/test_suite_psa_crypto.data     | 33 +++++++++++++++++
 tests/suites/test_suite_psa_crypto.function | 39 +++++++++++++++++++++
 2 files changed, 72 insertions(+)

diff --git a/tests/suites/test_suite_psa_crypto.data b/tests/suites/test_suite_psa_crypto.data
index 1182c00693..e552fc1f5a 100644
--- a/tests/suites/test_suite_psa_crypto.data
+++ b/tests/suites/test_suite_psa_crypto.data
@@ -4824,6 +4824,10 @@ PSA key derivation setup: TLS 1.2 PRF SHA-256, good case
 depends_on:PSA_WANT_ALG_SHA_256:PSA_WANT_ALG_TLS12_PRF
 derive_setup:PSA_ALG_TLS12_PRF(PSA_ALG_SHA_256):PSA_SUCCESS
 
+PSA key derivation setup: TLS 1.2 ECJPAKE to PMS
+depends_on:PSA_WANT_ALG_SHA_256:PSA_WANT_ALG_TLS12_ECJPAKE_TO_PMS
+derive_setup:PSA_ALG_TLS12_ECJPAKE_TO_PMS:PSA_SUCCESS
+
 PSA key derivation setup: not a key derivation algorithm (HMAC)
 depends_on:PSA_WANT_ALG_HMAC:PSA_WANT_ALG_SHA_256
 derive_setup:PSA_ALG_HMAC(PSA_ALG_SHA_256):PSA_ERROR_INVALID_ARGUMENT
@@ -5793,6 +5797,35 @@ PSA key derivation: over capacity 42: output 43+1
 depends_on:PSA_WANT_ALG_HKDF:PSA_WANT_ALG_SHA_256
 derive_output:PSA_ALG_HKDF(PSA_ALG_SHA_256):PSA_KEY_DERIVATION_INPUT_SALT:"000102030405060708090a0b0c":PSA_SUCCESS:PSA_KEY_DERIVATION_INPUT_SECRET:"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":PSA_SUCCESS:PSA_KEY_DERIVATION_INPUT_INFO:"f0f1f2f3f4f5f6f7f8f9":PSA_SUCCESS:0:"":PSA_SUCCESS:"":42:"3cb25f25faacd57a90434f64d0362f2a2d2d0a90cf1a5a4c5db02d56ecc4c5bf34007208d5b887185865ff":"ff":0:1:0
 
+PSA key derivation: ECJPAKE to PMS, no input
+depends_on:PSA_WANT_ALG_SHA_256
+derive_ecjpake_to_pms:"":PSA_ERROR_INVALID_ARGUMENT:"":PSA_ERROR_INVALID_ARGUMENT
+
+PSA key derivation: ECJPAKE to PMS, input too short
+depends_on:PSA_WANT_ALG_SHA_256
+derive_ecjpake_to_pms:"deadbeef":PSA_ERROR_INVALID_ARGUMENT:"":PSA_ERROR_INVALID_ARGUMENT
+
+PSA key derivation: ECJPAKE to PMS, input too long
+depends_on:PSA_WANT_ALG_SHA_256
+derive_ecjpake_to_pms:"0400000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000de":PSA_ERROR_INVALID_ARGUMENT:"":PSA_ERROR_INVALID_ARGUMENT
+
+PSA key derivation: ECJPAKE to PMS, bad input format
+depends_on:PSA_WANT_ALG_SHA_256
+derive_ecjpake_to_pms:"0200000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000":PSA_ERROR_INVALID_ARGUMENT:"":PSA_ERROR_INVALID_ARGUMENT
+
+#NIST CAVS 11.0 SHA-256 ShortMSG vector for L=256
+PSA key derivation: ECJPAKE to PMS, good case
+depends_on:PSA_WANT_ALG_SHA_256
+derive_ecjpake_to_pms:"0409fc1accc230a205e4a208e64a8f204291f581a12756392da4b8c0cf5ef02b950000000000000000000000000000000000000000000000000000000000000000":PSA_SUCCESS:"4f44c1c7fbebb6f9601829f3897bfd650c56fa07844be76489076356ac1886a4":PSA_SUCCESS
+
+PSA key derivation: ECJPAKE to PMS, output too short
+depends_on:PSA_WANT_ALG_SHA_256
+derive_ecjpake_to_pms:"0409fc1accc230a205e4a208e64a8f204291f581a12756392da4b8c0cf5ef02b950000000000000000000000000000000000000000000000000000000000000000":PSA_SUCCESS:"4f":PSA_ERROR_INVALID_ARGUMENT
+
+PSA key derivation: ECJPAKE to PMS, output too long
+depends_on:PSA_WANT_ALG_SHA_256
+derive_ecjpake_to_pms:"0409fc1accc230a205e4a208e64a8f204291f581a12756392da4b8c0cf5ef02b950000000000000000000000000000000000000000000000000000000000000000":PSA_SUCCESS:"4f44c1c7fbebb6f9601829f3897bfd650c56fa07844be76489076356ac1886a400":PSA_ERROR_INVALID_ARGUMENT
+
 PSA key derivation: HKDF SHA-256, read maximum capacity minus 1
 depends_on:PSA_WANT_ALG_HKDF:PSA_WANT_ALG_SHA_256
 derive_full:PSA_ALG_HKDF(PSA_ALG_SHA_256):"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":"000102030405060708090a0b0c":"f0f1f2f3f4f5f6f7f8f9":255 * PSA_HASH_LENGTH(PSA_ALG_SHA_256) - 1
diff --git a/tests/suites/test_suite_psa_crypto.function b/tests/suites/test_suite_psa_crypto.function
index 7d368cf162..b04adcca4f 100644
--- a/tests/suites/test_suite_psa_crypto.function
+++ b/tests/suites/test_suite_psa_crypto.function
@@ -7233,6 +7233,45 @@ exit:
 }
 /* END_CASE */
 
+/* BEGIN_CASE depends_on:MBEDTLS_PSA_BUILTIN_ALG_TLS12_PSK_TO_MS:MBEDTLS_SHA256_C */
+void derive_ecjpake_to_pms( data_t *input, int expected_input_status,
+                            data_t *expected_output, int expected_output_status )
+{
+    psa_algorithm_t alg = PSA_ALG_TLS12_ECJPAKE_TO_PMS;
+    psa_key_derivation_operation_t operation = PSA_KEY_DERIVATION_OPERATION_INIT;
+
+    uint8_t *output_buffer = NULL;
+    psa_status_t status;
+
+    ASSERT_ALLOC( output_buffer, expected_output->len );
+    PSA_ASSERT( psa_crypto_init() );
+
+    PSA_ASSERT( psa_key_derivation_setup( &operation, alg ) );
+    PSA_ASSERT( psa_key_derivation_set_capacity( &operation,
+                                                 expected_output->len ) );
+
+    TEST_EQUAL( psa_key_derivation_input_bytes( &operation,
+                PSA_KEY_DERIVATION_INPUT_SECRET, input->x, input->len ),
+        (psa_status_t) expected_input_status );
+
+    if( ( (psa_status_t) expected_input_status ) != PSA_SUCCESS )
+        goto exit;
+
+    status = psa_key_derivation_output_bytes( &operation, output_buffer,
+        expected_output->len );
+
+    TEST_EQUAL( status, (psa_status_t) expected_output_status );
+    if( expected_output->len != 0 && expected_output_status == PSA_SUCCESS )
+        ASSERT_COMPARE( output_buffer, expected_output->len, expected_output->x,
+                        expected_output->len );
+
+exit:
+    mbedtls_free( output_buffer );
+    psa_key_derivation_abort( &operation );
+    PSA_DONE();
+}
+/* END_CASE */
+
 /* BEGIN_CASE */
 void derive_key_exercise( int alg_arg,
                           data_t *key_data,

From 18f8e8d62c33005f3d42aa3fa492560a4978a94a Mon Sep 17 00:00:00 2001
From: Andrzej Kurek <andrzej.kurek@arm.com>
Date: Wed, 14 Sep 2022 08:44:34 -0400
Subject: [PATCH 03/18] Document the input size restriction for EC J-PAKE to
 PMS

Signed-off-by: Andrzej Kurek <andrzej.kurek@arm.com>
---
 include/psa/crypto_values.h | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/include/psa/crypto_values.h b/include/psa/crypto_values.h
index 03438196b6..573ff9a16c 100644
--- a/include/psa/crypto_values.h
+++ b/include/psa/crypto_values.h
@@ -2023,7 +2023,8 @@
 
 /* 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.
+ * use to derive the session secret. Uses PSA_ALG_SHA_256. Only P-256 is
+ * supported, so the input has to be exactly 65 bytes.
  */
 #define PSA_ALG_TLS12_ECJPAKE_TO_PMS            ((psa_algorithm_t)0x08000600)
 #define PSA_ALG_IS_TLS12_ECJPAKE_TO_PMS(alg)                               \

From d60907b85d469b66d1b3c0fafe9ed92cd1d097b0 Mon Sep 17 00:00:00 2001
From: Andrzej Kurek <andrzej.kurek@arm.com>
Date: Wed, 14 Sep 2022 10:02:30 -0400
Subject: [PATCH 04/18] Define ECJPAKE_TO_PMS in config_psa only if SHA_256 is
 available

Signed-off-by: Andrzej Kurek <andrzej.kurek@arm.com>
---
 include/mbedtls/config_psa.h  | 7 +++++--
 library/check_crypto_config.h | 5 +++++
 2 files changed, 10 insertions(+), 2 deletions(-)

diff --git a/include/mbedtls/config_psa.h b/include/mbedtls/config_psa.h
index 88052d228b..61950cd9c9 100644
--- a/include/mbedtls/config_psa.h
+++ b/include/mbedtls/config_psa.h
@@ -635,8 +635,6 @@ extern "C" {
 #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)
@@ -714,6 +712,11 @@ extern "C" {
 #define MBEDTLS_PSA_BUILTIN_KEY_TYPE_DES 1
 #endif
 
+#if defined(MBEDTLS_PSA_BUILTIN_ALG_SHA_256)
+#define MBEDTLS_PSA_BUILTIN_ALG_TLS12_ECJPAKE_TO_PMS 1
+#define PSA_WANT_ALG_TLS12_ECJPAKE_TO_PMS 1
+#endif
+
 #if defined(MBEDTLS_CHACHA20_C)
 #define PSA_WANT_KEY_TYPE_CHACHA20 1
 #define PSA_WANT_ALG_STREAM_CIPHER 1
diff --git a/library/check_crypto_config.h b/library/check_crypto_config.h
index c74437e7de..e60e666162 100644
--- a/library/check_crypto_config.h
+++ b/library/check_crypto_config.h
@@ -93,4 +93,9 @@
 #error "MBEDTLS_SSL_PROTO_TLS1_2 defined, but not all prerequisites"
 #endif
 
+#if defined(PSA_WANT_ALG_TLS12_ECJPAKE_TO_PMS) && \
+    !defined(PSA_WANT_ALG_SHA_256)
+#error "PSA_WANT_ALG_TLS12_ECJPAKE_TO_PMS defined, but not all prerequisites"
+#endif
+
 #endif /* MBEDTLS_CHECK_CRYPTO_CONFIG_H */

From 4ba0e45f8ebb81d3e408ddf802346ddd50c83223 Mon Sep 17 00:00:00 2001
From: Andrzej Kurek <andrzej.kurek@arm.com>
Date: Wed, 14 Sep 2022 12:47:26 -0400
Subject: [PATCH 05/18] all.sh: don't build with ECJPAKE_TO_PMS if SHA256 is
 not available

Signed-off-by: Andrzej Kurek <andrzej.kurek@arm.com>
---
 tests/scripts/all.sh | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/tests/scripts/all.sh b/tests/scripts/all.sh
index 7d9fe1f4a3..ea7ba41101 100755
--- a/tests/scripts/all.sh
+++ b/tests/scripts/all.sh
@@ -2067,6 +2067,7 @@ component_build_psa_accel_alg_md5() {
     scripts/config.py -f include/psa/crypto_config.h unset PSA_WANT_ALG_SHA_256
     scripts/config.py -f include/psa/crypto_config.h unset PSA_WANT_ALG_SHA_384
     scripts/config.py -f include/psa/crypto_config.h unset PSA_WANT_ALG_SHA_512
+    scripts/config.py -f include/psa/crypto_config.h unset PSA_WANT_ALG_TLS12_ECJPAKE_TO_PMS
     # Need to define the correct symbol and include the test driver header path in order to build with the test driver
     make CC=gcc CFLAGS="$ASAN_CFLAGS -DPSA_CRYPTO_DRIVER_TEST -DMBEDTLS_PSA_ACCEL_ALG_MD5 -I../tests/include -O2" LDFLAGS="$ASAN_CFLAGS"
 }
@@ -2086,6 +2087,7 @@ component_build_psa_accel_alg_ripemd160() {
     scripts/config.py -f include/psa/crypto_config.h unset PSA_WANT_ALG_SHA_256
     scripts/config.py -f include/psa/crypto_config.h unset PSA_WANT_ALG_SHA_384
     scripts/config.py -f include/psa/crypto_config.h unset PSA_WANT_ALG_SHA_512
+    scripts/config.py -f include/psa/crypto_config.h unset PSA_WANT_ALG_TLS12_ECJPAKE_TO_PMS
     # Need to define the correct symbol and include the test driver header path in order to build with the test driver
     make CC=gcc CFLAGS="$ASAN_CFLAGS -DPSA_CRYPTO_DRIVER_TEST -DMBEDTLS_PSA_ACCEL_ALG_RIPEMD160 -I../tests/include -O2" LDFLAGS="$ASAN_CFLAGS"
 }
@@ -2105,6 +2107,7 @@ component_build_psa_accel_alg_sha1() {
     scripts/config.py -f include/psa/crypto_config.h unset PSA_WANT_ALG_SHA_256
     scripts/config.py -f include/psa/crypto_config.h unset PSA_WANT_ALG_SHA_384
     scripts/config.py -f include/psa/crypto_config.h unset PSA_WANT_ALG_SHA_512
+    scripts/config.py -f include/psa/crypto_config.h unset PSA_WANT_ALG_TLS12_ECJPAKE_TO_PMS
     # Need to define the correct symbol and include the test driver header path in order to build with the test driver
     make CC=gcc CFLAGS="$ASAN_CFLAGS -DPSA_CRYPTO_DRIVER_TEST -DMBEDTLS_PSA_ACCEL_ALG_SHA_1 -I../tests/include -O2" LDFLAGS="$ASAN_CFLAGS"
 }
@@ -2123,6 +2126,7 @@ component_build_psa_accel_alg_sha224() {
     scripts/config.py -f include/psa/crypto_config.h unset PSA_WANT_ALG_SHA_1
     scripts/config.py -f include/psa/crypto_config.h unset PSA_WANT_ALG_SHA_384
     scripts/config.py -f include/psa/crypto_config.h unset PSA_WANT_ALG_SHA_512
+    scripts/config.py -f include/psa/crypto_config.h unset PSA_WANT_ALG_TLS12_ECJPAKE_TO_PMS
     # Need to define the correct symbol and include the test driver header path in order to build with the test driver
     make CC=gcc CFLAGS="$ASAN_CFLAGS -DPSA_CRYPTO_DRIVER_TEST -DMBEDTLS_PSA_ACCEL_ALG_SHA_224 -I../tests/include -O2" LDFLAGS="$ASAN_CFLAGS"
 }
@@ -2160,6 +2164,7 @@ component_build_psa_accel_alg_sha384() {
     scripts/config.py -f include/psa/crypto_config.h unset PSA_WANT_ALG_SHA_1
     scripts/config.py -f include/psa/crypto_config.h unset PSA_WANT_ALG_SHA_224
     scripts/config.py -f include/psa/crypto_config.h unset PSA_WANT_ALG_SHA_256
+    scripts/config.py -f include/psa/crypto_config.h unset PSA_WANT_ALG_TLS12_ECJPAKE_TO_PMS
     # Need to define the correct symbol and include the test driver header path in order to build with the test driver
     make CC=gcc CFLAGS="$ASAN_CFLAGS -DPSA_CRYPTO_DRIVER_TEST -DMBEDTLS_PSA_ACCEL_ALG_SHA_384 -I../tests/include -O2" LDFLAGS="$ASAN_CFLAGS"
 }
@@ -2179,6 +2184,7 @@ component_build_psa_accel_alg_sha512() {
     scripts/config.py -f include/psa/crypto_config.h unset PSA_WANT_ALG_SHA_224
     scripts/config.py -f include/psa/crypto_config.h unset PSA_WANT_ALG_SHA_256
     scripts/config.py -f include/psa/crypto_config.h unset PSA_WANT_ALG_SHA_384
+    scripts/config.py -f include/psa/crypto_config.h unset PSA_WANT_ALG_TLS12_ECJPAKE_TO_PMS
     # Need to define the correct symbol and include the test driver header path in order to build with the test driver
     make CC=gcc CFLAGS="$ASAN_CFLAGS -DPSA_CRYPTO_DRIVER_TEST -DMBEDTLS_PSA_ACCEL_ALG_SHA_512 -I../tests/include -O2" LDFLAGS="$ASAN_CFLAGS"
 }

From 702776f7cc69a13057f13291942c46e46273bccf Mon Sep 17 00:00:00 2001
From: Andrzej Kurek <andrzej.kurek@arm.com>
Date: Fri, 16 Sep 2022 06:22:44 -0400
Subject: [PATCH 06/18] Restrict the EC J-PAKE to PMS input type to secret

Signed-off-by: Andrzej Kurek <andrzej.kurek@arm.com>
---
 library/psa_crypto.c | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index 5c05f79282..cbdc912930 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -5148,7 +5148,7 @@ static psa_status_t psa_key_derivation_setup_kdf(
         return( PSA_ERROR_NOT_SUPPORTED );
 
     /* All currently supported key derivation algorithms (apart from
-     * ecjpake to pms are based on a hash algorithm. */
+     * 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( !PSA_ALG_IS_TLS12_ECJPAKE_TO_PMS( kdf_alg ) )
@@ -5570,10 +5570,12 @@ static psa_status_t psa_tls12_prf_psk_to_ms_input(
 #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,
+    psa_key_derivation_step_t step,
     const uint8_t *data,
     size_t data_length )
 {
-    if( data_length != PSA_TLS12_ECJPAKE_TO_PMS_INPUT_SIZE )
+    if( data_length != PSA_TLS12_ECJPAKE_TO_PMS_INPUT_SIZE ||
+        step != PSA_KEY_DERIVATION_INPUT_SECRET )
         return( PSA_ERROR_INVALID_ARGUMENT );
 
     /* Check if the passed point is in an uncompressed form */
@@ -5668,7 +5670,7 @@ static psa_status_t psa_key_derivation_input_internal(
     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 );
+            &operation->ctx.tls12_ecjpake_to_pms, step, data, data_length );
     }
     else
 #endif /* MBEDTLS_PSA_BUILTIN_ALG_TLS12_ECJPAKE_TO_PMS */

From d37850404adfb8064f6e590197983c7c852daa30 Mon Sep 17 00:00:00 2001
From: Andrzej Kurek <andrzej.kurek@arm.com>
Date: Fri, 16 Sep 2022 06:45:44 -0400
Subject: [PATCH 07/18] Add derivation step testing to EC J-PAKE to PMS tests

Signed-off-by: Andrzej Kurek <andrzej.kurek@arm.com>
---
 tests/suites/test_suite_psa_crypto.data     | 18 +++++++++++-------
 tests/suites/test_suite_psa_crypto.function |  7 ++++---
 2 files changed, 15 insertions(+), 10 deletions(-)

diff --git a/tests/suites/test_suite_psa_crypto.data b/tests/suites/test_suite_psa_crypto.data
index e552fc1f5a..4a052c4335 100644
--- a/tests/suites/test_suite_psa_crypto.data
+++ b/tests/suites/test_suite_psa_crypto.data
@@ -5799,32 +5799,36 @@ derive_output:PSA_ALG_HKDF(PSA_ALG_SHA_256):PSA_KEY_DERIVATION_INPUT_SALT:"00010
 
 PSA key derivation: ECJPAKE to PMS, no input
 depends_on:PSA_WANT_ALG_SHA_256
-derive_ecjpake_to_pms:"":PSA_ERROR_INVALID_ARGUMENT:"":PSA_ERROR_INVALID_ARGUMENT
+derive_ecjpake_to_pms:"":PSA_ERROR_INVALID_ARGUMENT:PSA_KEY_DERIVATION_INPUT_SECRET:"":PSA_ERROR_INVALID_ARGUMENT
 
 PSA key derivation: ECJPAKE to PMS, input too short
 depends_on:PSA_WANT_ALG_SHA_256
-derive_ecjpake_to_pms:"deadbeef":PSA_ERROR_INVALID_ARGUMENT:"":PSA_ERROR_INVALID_ARGUMENT
+derive_ecjpake_to_pms:"deadbeef":PSA_ERROR_INVALID_ARGUMENT:PSA_KEY_DERIVATION_INPUT_SECRET:"":PSA_ERROR_INVALID_ARGUMENT
 
 PSA key derivation: ECJPAKE to PMS, input too long
 depends_on:PSA_WANT_ALG_SHA_256
-derive_ecjpake_to_pms:"0400000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000de":PSA_ERROR_INVALID_ARGUMENT:"":PSA_ERROR_INVALID_ARGUMENT
+derive_ecjpake_to_pms:"0400000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000de":PSA_ERROR_INVALID_ARGUMENT:PSA_KEY_DERIVATION_INPUT_SECRET:"":PSA_ERROR_INVALID_ARGUMENT
 
 PSA key derivation: ECJPAKE to PMS, bad input format
 depends_on:PSA_WANT_ALG_SHA_256
-derive_ecjpake_to_pms:"0200000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000":PSA_ERROR_INVALID_ARGUMENT:"":PSA_ERROR_INVALID_ARGUMENT
+derive_ecjpake_to_pms:"0200000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000":PSA_ERROR_INVALID_ARGUMENT:PSA_KEY_DERIVATION_INPUT_SECRET:"":PSA_ERROR_INVALID_ARGUMENT
 
 #NIST CAVS 11.0 SHA-256 ShortMSG vector for L=256
 PSA key derivation: ECJPAKE to PMS, good case
 depends_on:PSA_WANT_ALG_SHA_256
-derive_ecjpake_to_pms:"0409fc1accc230a205e4a208e64a8f204291f581a12756392da4b8c0cf5ef02b950000000000000000000000000000000000000000000000000000000000000000":PSA_SUCCESS:"4f44c1c7fbebb6f9601829f3897bfd650c56fa07844be76489076356ac1886a4":PSA_SUCCESS
+derive_ecjpake_to_pms:"0409fc1accc230a205e4a208e64a8f204291f581a12756392da4b8c0cf5ef02b950000000000000000000000000000000000000000000000000000000000000000":PSA_SUCCESS:PSA_KEY_DERIVATION_INPUT_SECRET:"4f44c1c7fbebb6f9601829f3897bfd650c56fa07844be76489076356ac1886a4":PSA_SUCCESS
+
+PSA key derivation: ECJPAKE to PMS, bad derivation step
+depends_on:PSA_WANT_ALG_SHA_256
+derive_ecjpake_to_pms:"0409fc1accc230a205e4a208e64a8f204291f581a12756392da4b8c0cf5ef02b950000000000000000000000000000000000000000000000000000000000000000":PSA_ERROR_INVALID_ARGUMENT:PSA_KEY_DERIVATION_INPUT_SEED:"4f44c1c7fbebb6f9601829f3897bfd650c56fa07844be76489076356ac1886a4":PSA_SUCCESS
 
 PSA key derivation: ECJPAKE to PMS, output too short
 depends_on:PSA_WANT_ALG_SHA_256
-derive_ecjpake_to_pms:"0409fc1accc230a205e4a208e64a8f204291f581a12756392da4b8c0cf5ef02b950000000000000000000000000000000000000000000000000000000000000000":PSA_SUCCESS:"4f":PSA_ERROR_INVALID_ARGUMENT
+derive_ecjpake_to_pms:"0409fc1accc230a205e4a208e64a8f204291f581a12756392da4b8c0cf5ef02b950000000000000000000000000000000000000000000000000000000000000000":PSA_SUCCESS:PSA_KEY_DERIVATION_INPUT_SECRET:"4f":PSA_ERROR_INVALID_ARGUMENT
 
 PSA key derivation: ECJPAKE to PMS, output too long
 depends_on:PSA_WANT_ALG_SHA_256
-derive_ecjpake_to_pms:"0409fc1accc230a205e4a208e64a8f204291f581a12756392da4b8c0cf5ef02b950000000000000000000000000000000000000000000000000000000000000000":PSA_SUCCESS:"4f44c1c7fbebb6f9601829f3897bfd650c56fa07844be76489076356ac1886a400":PSA_ERROR_INVALID_ARGUMENT
+derive_ecjpake_to_pms:"0409fc1accc230a205e4a208e64a8f204291f581a12756392da4b8c0cf5ef02b950000000000000000000000000000000000000000000000000000000000000000":PSA_SUCCESS:PSA_KEY_DERIVATION_INPUT_SECRET:"4f44c1c7fbebb6f9601829f3897bfd650c56fa07844be76489076356ac1886a400":PSA_ERROR_INVALID_ARGUMENT
 
 PSA key derivation: HKDF SHA-256, read maximum capacity minus 1
 depends_on:PSA_WANT_ALG_HKDF:PSA_WANT_ALG_SHA_256
diff --git a/tests/suites/test_suite_psa_crypto.function b/tests/suites/test_suite_psa_crypto.function
index b04adcca4f..0162c07d19 100644
--- a/tests/suites/test_suite_psa_crypto.function
+++ b/tests/suites/test_suite_psa_crypto.function
@@ -7235,11 +7235,12 @@ exit:
 
 /* BEGIN_CASE depends_on:MBEDTLS_PSA_BUILTIN_ALG_TLS12_PSK_TO_MS:MBEDTLS_SHA256_C */
 void derive_ecjpake_to_pms( data_t *input, int expected_input_status,
-                            data_t *expected_output, int expected_output_status )
+                            int derivation_step, data_t *expected_output,
+                            int expected_output_status )
 {
     psa_algorithm_t alg = PSA_ALG_TLS12_ECJPAKE_TO_PMS;
     psa_key_derivation_operation_t operation = PSA_KEY_DERIVATION_OPERATION_INIT;
-
+    psa_key_derivation_step_t step = (psa_key_derivation_step_t) derivation_step;
     uint8_t *output_buffer = NULL;
     psa_status_t status;
 
@@ -7251,7 +7252,7 @@ void derive_ecjpake_to_pms( data_t *input, int expected_input_status,
                                                  expected_output->len ) );
 
     TEST_EQUAL( psa_key_derivation_input_bytes( &operation,
-                PSA_KEY_DERIVATION_INPUT_SECRET, input->x, input->len ),
+                step, input->x, input->len ),
         (psa_status_t) expected_input_status );
 
     if( ( (psa_status_t) expected_input_status ) != PSA_SUCCESS )

From b093650033a2ede358f17d599668fd3746b10ace Mon Sep 17 00:00:00 2001
From: Andrzej Kurek <andrzej.kurek@arm.com>
Date: Fri, 16 Sep 2022 07:13:00 -0400
Subject: [PATCH 08/18] Add proper capacity calculation for EC J-PAKE to PMS
 KDF

Signed-off-by: Andrzej Kurek <andrzej.kurek@arm.com>
---
 library/psa_crypto.c | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index cbdc912930..981b6f8ffb 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -5175,7 +5175,8 @@ static psa_status_t psa_key_derivation_setup_kdf(
         return( PSA_ERROR_NOT_SUPPORTED );
     }
 #if defined(MBEDTLS_PSA_BUILTIN_ALG_HKDF_EXTRACT)
-    if( PSA_ALG_IS_HKDF_EXTRACT( kdf_alg ) )
+    if( PSA_ALG_IS_HKDF_EXTRACT( kdf_alg ) ||
+        PSA_ALG_IS_TLS12_ECJPAKE_TO_PMS ( kdf_alg ))
         operation->capacity = hash_size;
     else
 #endif /* MBEDTLS_PSA_BUILTIN_ALG_HKDF_EXTRACT */

From 2be16895045c0f712afd70df5f70f3b9c17608fc Mon Sep 17 00:00:00 2001
From: Andrzej Kurek <andrzej.kurek@arm.com>
Date: Fri, 16 Sep 2022 07:14:04 -0400
Subject: [PATCH 09/18] Add capacity testing to EC J-PAKE to PMS tests

Let the caller restrict the capacity but limit it to 32 bytes.
Signed-off-by: Andrzej Kurek <andrzej.kurek@arm.com>
---
 tests/suites/test_suite_psa_crypto.data     | 24 ++++++++++++++-------
 tests/suites/test_suite_psa_crypto.function |  8 ++++---
 2 files changed, 21 insertions(+), 11 deletions(-)

diff --git a/tests/suites/test_suite_psa_crypto.data b/tests/suites/test_suite_psa_crypto.data
index 4a052c4335..fd33eba721 100644
--- a/tests/suites/test_suite_psa_crypto.data
+++ b/tests/suites/test_suite_psa_crypto.data
@@ -5799,36 +5799,44 @@ derive_output:PSA_ALG_HKDF(PSA_ALG_SHA_256):PSA_KEY_DERIVATION_INPUT_SALT:"00010
 
 PSA key derivation: ECJPAKE to PMS, no input
 depends_on:PSA_WANT_ALG_SHA_256
-derive_ecjpake_to_pms:"":PSA_ERROR_INVALID_ARGUMENT:PSA_KEY_DERIVATION_INPUT_SECRET:"":PSA_ERROR_INVALID_ARGUMENT
+derive_ecjpake_to_pms:"":PSA_ERROR_INVALID_ARGUMENT:PSA_KEY_DERIVATION_INPUT_SECRET:32:PSA_SUCCESS:"":PSA_ERROR_INVALID_ARGUMENT
 
 PSA key derivation: ECJPAKE to PMS, input too short
 depends_on:PSA_WANT_ALG_SHA_256
-derive_ecjpake_to_pms:"deadbeef":PSA_ERROR_INVALID_ARGUMENT:PSA_KEY_DERIVATION_INPUT_SECRET:"":PSA_ERROR_INVALID_ARGUMENT
+derive_ecjpake_to_pms:"deadbeef":PSA_ERROR_INVALID_ARGUMENT:PSA_KEY_DERIVATION_INPUT_SECRET:32:PSA_SUCCESS:"":PSA_ERROR_INVALID_ARGUMENT
 
 PSA key derivation: ECJPAKE to PMS, input too long
 depends_on:PSA_WANT_ALG_SHA_256
-derive_ecjpake_to_pms:"0400000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000de":PSA_ERROR_INVALID_ARGUMENT:PSA_KEY_DERIVATION_INPUT_SECRET:"":PSA_ERROR_INVALID_ARGUMENT
+derive_ecjpake_to_pms:"0400000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000de":PSA_ERROR_INVALID_ARGUMENT:PSA_KEY_DERIVATION_INPUT_SECRET:32:PSA_SUCCESS:"":PSA_ERROR_INVALID_ARGUMENT
 
 PSA key derivation: ECJPAKE to PMS, bad input format
 depends_on:PSA_WANT_ALG_SHA_256
-derive_ecjpake_to_pms:"0200000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000":PSA_ERROR_INVALID_ARGUMENT:PSA_KEY_DERIVATION_INPUT_SECRET:"":PSA_ERROR_INVALID_ARGUMENT
+derive_ecjpake_to_pms:"0200000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000":PSA_ERROR_INVALID_ARGUMENT:PSA_KEY_DERIVATION_INPUT_SECRET:32:PSA_SUCCESS:"":PSA_ERROR_INVALID_ARGUMENT
 
 #NIST CAVS 11.0 SHA-256 ShortMSG vector for L=256
 PSA key derivation: ECJPAKE to PMS, good case
 depends_on:PSA_WANT_ALG_SHA_256
-derive_ecjpake_to_pms:"0409fc1accc230a205e4a208e64a8f204291f581a12756392da4b8c0cf5ef02b950000000000000000000000000000000000000000000000000000000000000000":PSA_SUCCESS:PSA_KEY_DERIVATION_INPUT_SECRET:"4f44c1c7fbebb6f9601829f3897bfd650c56fa07844be76489076356ac1886a4":PSA_SUCCESS
+derive_ecjpake_to_pms:"0409fc1accc230a205e4a208e64a8f204291f581a12756392da4b8c0cf5ef02b950000000000000000000000000000000000000000000000000000000000000000":PSA_SUCCESS:PSA_KEY_DERIVATION_INPUT_SECRET:32:PSA_SUCCESS:"4f44c1c7fbebb6f9601829f3897bfd650c56fa07844be76489076356ac1886a4":PSA_SUCCESS
 
 PSA key derivation: ECJPAKE to PMS, bad derivation step
 depends_on:PSA_WANT_ALG_SHA_256
-derive_ecjpake_to_pms:"0409fc1accc230a205e4a208e64a8f204291f581a12756392da4b8c0cf5ef02b950000000000000000000000000000000000000000000000000000000000000000":PSA_ERROR_INVALID_ARGUMENT:PSA_KEY_DERIVATION_INPUT_SEED:"4f44c1c7fbebb6f9601829f3897bfd650c56fa07844be76489076356ac1886a4":PSA_SUCCESS
+derive_ecjpake_to_pms:"0409fc1accc230a205e4a208e64a8f204291f581a12756392da4b8c0cf5ef02b950000000000000000000000000000000000000000000000000000000000000000":PSA_ERROR_INVALID_ARGUMENT:PSA_KEY_DERIVATION_INPUT_SEED:32:PSA_SUCCESS:"4f44c1c7fbebb6f9601829f3897bfd650c56fa07844be76489076356ac1886a4":PSA_SUCCESS
+
+PSA key derivation: ECJPAKE to PMS, capacity 1 byte too big
+depends_on:PSA_WANT_ALG_SHA_256
+derive_ecjpake_to_pms:"0409fc1accc230a205e4a208e64a8f204291f581a12756392da4b8c0cf5ef02b950000000000000000000000000000000000000000000000000000000000000000":PSA_SUCCESS:PSA_KEY_DERIVATION_INPUT_SECRET:33:PSA_ERROR_INVALID_ARGUMENT:"4f44c1c7fbebb6f9601829f3897bfd650c56fa07844be76489076356ac1886a4":PSA_SUCCESS
+
+PSA key derivation: ECJPAKE to PMS, capacity 1 byte too small
+depends_on:PSA_WANT_ALG_SHA_256
+derive_ecjpake_to_pms:"0409fc1accc230a205e4a208e64a8f204291f581a12756392da4b8c0cf5ef02b950000000000000000000000000000000000000000000000000000000000000000":PSA_SUCCESS:PSA_KEY_DERIVATION_INPUT_SECRET:31:PSA_SUCCESS:"4f44c1c7fbebb6f9601829f3897bfd650c56fa07844be76489076356ac1886a4":PSA_ERROR_INSUFFICIENT_DATA
 
 PSA key derivation: ECJPAKE to PMS, output too short
 depends_on:PSA_WANT_ALG_SHA_256
-derive_ecjpake_to_pms:"0409fc1accc230a205e4a208e64a8f204291f581a12756392da4b8c0cf5ef02b950000000000000000000000000000000000000000000000000000000000000000":PSA_SUCCESS:PSA_KEY_DERIVATION_INPUT_SECRET:"4f":PSA_ERROR_INVALID_ARGUMENT
+derive_ecjpake_to_pms:"0409fc1accc230a205e4a208e64a8f204291f581a12756392da4b8c0cf5ef02b950000000000000000000000000000000000000000000000000000000000000000":PSA_SUCCESS:PSA_KEY_DERIVATION_INPUT_SECRET:32:PSA_SUCCESS:"4f":PSA_ERROR_INVALID_ARGUMENT
 
 PSA key derivation: ECJPAKE to PMS, output too long
 depends_on:PSA_WANT_ALG_SHA_256
-derive_ecjpake_to_pms:"0409fc1accc230a205e4a208e64a8f204291f581a12756392da4b8c0cf5ef02b950000000000000000000000000000000000000000000000000000000000000000":PSA_SUCCESS:PSA_KEY_DERIVATION_INPUT_SECRET:"4f44c1c7fbebb6f9601829f3897bfd650c56fa07844be76489076356ac1886a400":PSA_ERROR_INVALID_ARGUMENT
+derive_ecjpake_to_pms:"0409fc1accc230a205e4a208e64a8f204291f581a12756392da4b8c0cf5ef02b950000000000000000000000000000000000000000000000000000000000000000":PSA_SUCCESS:PSA_KEY_DERIVATION_INPUT_SECRET:32:PSA_SUCCESS:"4f44c1c7fbebb6f9601829f3897bfd650c56fa07844be76489076356ac1886a400":PSA_ERROR_INSUFFICIENT_DATA
 
 PSA key derivation: HKDF SHA-256, read maximum capacity minus 1
 depends_on:PSA_WANT_ALG_HKDF:PSA_WANT_ALG_SHA_256
diff --git a/tests/suites/test_suite_psa_crypto.function b/tests/suites/test_suite_psa_crypto.function
index 0162c07d19..355cba533b 100644
--- a/tests/suites/test_suite_psa_crypto.function
+++ b/tests/suites/test_suite_psa_crypto.function
@@ -7235,7 +7235,9 @@ exit:
 
 /* BEGIN_CASE depends_on:MBEDTLS_PSA_BUILTIN_ALG_TLS12_PSK_TO_MS:MBEDTLS_SHA256_C */
 void derive_ecjpake_to_pms( data_t *input, int expected_input_status,
-                            int derivation_step, data_t *expected_output,
+                            int derivation_step,
+                            int capacity, int expected_capacity_status,
+                            data_t *expected_output,
                             int expected_output_status )
 {
     psa_algorithm_t alg = PSA_ALG_TLS12_ECJPAKE_TO_PMS;
@@ -7248,8 +7250,8 @@ void derive_ecjpake_to_pms( data_t *input, int expected_input_status,
     PSA_ASSERT( psa_crypto_init() );
 
     PSA_ASSERT( psa_key_derivation_setup( &operation, alg ) );
-    PSA_ASSERT( psa_key_derivation_set_capacity( &operation,
-                                                 expected_output->len ) );
+    TEST_EQUAL( psa_key_derivation_set_capacity( &operation, capacity ),
+        (psa_status_t) expected_capacity_status );
 
     TEST_EQUAL( psa_key_derivation_input_bytes( &operation,
                 step, input->x, input->len ),

From 1fafb1f778c3eb98eddd3474c93d0ecdc647621d Mon Sep 17 00:00:00 2001
From: Andrzej Kurek <andrzej.kurek@arm.com>
Date: Fri, 16 Sep 2022 07:19:49 -0400
Subject: [PATCH 10/18] Documentation clarifications for ECJPAKE-to-PMS

Signed-off-by: Andrzej Kurek <andrzej.kurek@arm.com>
---
 include/psa/crypto_values.h | 13 +++++++++----
 1 file changed, 9 insertions(+), 4 deletions(-)

diff --git a/include/psa/crypto_values.h b/include/psa/crypto_values.h
index 573ff9a16c..5a954ee399 100644
--- a/include/psa/crypto_values.h
+++ b/include/psa/crypto_values.h
@@ -2021,10 +2021,15 @@
 #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. Only P-256 is
- * supported, so the input has to be exactly 65 bytes.
+/* The TLS 1.2 ECJPAKE-to-PMS KDF. It 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, as defined by step 2 of
+ * https://datatracker.ietf.org/doc/html/draft-cragie-tls-ecjpake-01#section-8.7.
+ * Uses PSA_ALG_SHA_256.
+ * This function takes a single input:
+ * #PSA_KEY_DERIVATION_INPUT_SECRET is the shared secret K from EC J-PAKE.
+ * The only supported curve is secp256r1 (the 256-bit curve in
+ * #PSA_ECC_FAMILY_SECP_R1), so the input must be exactly 65 bytes.
  */
 #define PSA_ALG_TLS12_ECJPAKE_TO_PMS            ((psa_algorithm_t)0x08000600)
 #define PSA_ALG_IS_TLS12_ECJPAKE_TO_PMS(alg)                               \

From 3c4c51430255cf6afc714ce13c4c29ba02f9115a Mon Sep 17 00:00:00 2001
From: Andrzej Kurek <andrzej.kurek@arm.com>
Date: Fri, 16 Sep 2022 07:24:14 -0400
Subject: [PATCH 11/18] Remove `PSA_ALG_IS_TLS12_ECJPAKE_TO_PMS`

Signed-off-by: Andrzej Kurek <andrzej.kurek@arm.com>
---
 include/psa/crypto_values.h |  2 --
 library/psa_crypto.c        | 12 ++++++------
 2 files changed, 6 insertions(+), 8 deletions(-)

diff --git a/include/psa/crypto_values.h b/include/psa/crypto_values.h
index 5a954ee399..455b6388b3 100644
--- a/include/psa/crypto_values.h
+++ b/include/psa/crypto_values.h
@@ -2032,8 +2032,6 @@
  * #PSA_ECC_FAMILY_SECP_R1), so the input must be exactly 65 bytes.
  */
 #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
diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index 981b6f8ffb..8c59cf68fd 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -4352,7 +4352,7 @@ psa_status_t psa_key_derivation_abort( psa_key_derivation_operation_t *operation
 #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 ) )
+    if( kdf_alg == PSA_ALG_TLS12_ECJPAKE_TO_PMS )
     {
         mbedtls_platform_zeroize( operation->ctx.tls12_ecjpake_to_pms.data,
             PSA_TLS12_ECJPAKE_TO_PMS_DATA_SIZE );
@@ -4720,7 +4720,7 @@ psa_status_t psa_key_derivation_output_bytes(
 #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 ) )
+    if( kdf_alg == PSA_ALG_TLS12_ECJPAKE_TO_PMS )
     {
         status = psa_key_derivation_tls12_ecjpake_to_pms_read(
             &operation->ctx.tls12_ecjpake_to_pms, output, output_length );
@@ -5121,7 +5121,7 @@ static int is_kdf_alg_supported( psa_algorithm_t 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 ) )
+    if( kdf_alg == PSA_ALG_TLS12_ECJPAKE_TO_PMS )
         return( 1 );
 #endif
     return( 0 );
@@ -5151,7 +5151,7 @@ static psa_status_t psa_key_derivation_setup_kdf(
      * 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( !PSA_ALG_IS_TLS12_ECJPAKE_TO_PMS( kdf_alg ) )
+    if( kdf_alg != PSA_ALG_TLS12_ECJPAKE_TO_PMS )
     {
         if( hash_size == 0 )
             return( PSA_ERROR_NOT_SUPPORTED );
@@ -5176,7 +5176,7 @@ static psa_status_t psa_key_derivation_setup_kdf(
     }
 #if defined(MBEDTLS_PSA_BUILTIN_ALG_HKDF_EXTRACT)
     if( PSA_ALG_IS_HKDF_EXTRACT( kdf_alg ) ||
-        PSA_ALG_IS_TLS12_ECJPAKE_TO_PMS ( kdf_alg ))
+        ( kdf_alg == PSA_ALG_TLS12_ECJPAKE_TO_PMS ) )
         operation->capacity = hash_size;
     else
 #endif /* MBEDTLS_PSA_BUILTIN_ALG_HKDF_EXTRACT */
@@ -5668,7 +5668,7 @@ static psa_status_t psa_key_derivation_input_internal(
     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 ) )
+    if( kdf_alg == PSA_ALG_TLS12_ECJPAKE_TO_PMS )
     {
         status = psa_tls12_ecjpake_to_pms_input(
             &operation->ctx.tls12_ecjpake_to_pms, step, data, data_length );

From 7763829c5cf736b6274593aaa1ef1d81ae10b05f Mon Sep 17 00:00:00 2001
From: Andrzej Kurek <andrzej.kurek@arm.com>
Date: Fri, 16 Sep 2022 12:24:52 -0400
Subject: [PATCH 12/18] Add missing ifdef when calculating operation capacity

Signed-off-by: Andrzej Kurek <andrzej.kurek@arm.com>
---
 library/psa_crypto.c | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index 8c59cf68fd..663b645bbb 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -5174,7 +5174,8 @@ static psa_status_t psa_key_derivation_setup_kdf(
     {
         return( PSA_ERROR_NOT_SUPPORTED );
     }
-#if defined(MBEDTLS_PSA_BUILTIN_ALG_HKDF_EXTRACT)
+#if defined(MBEDTLS_PSA_BUILTIN_ALG_HKDF_EXTRACT) || \
+    defined(MBEDTLS_PSA_BUILTIN_ALG_TLS12_PSK_TO_MS)
     if( PSA_ALG_IS_HKDF_EXTRACT( kdf_alg ) ||
         ( kdf_alg == PSA_ALG_TLS12_ECJPAKE_TO_PMS ) )
         operation->capacity = hash_size;

From 96b9f23853213356fb4e0b2b60077d58feee7eb8 Mon Sep 17 00:00:00 2001
From: Andrzej Kurek <andrzej.kurek@arm.com>
Date: Mon, 26 Sep 2022 10:30:46 -0400
Subject: [PATCH 13/18] Adjust ECJPAKE_TO_PMS macro value

This way the low 8 bits of the identifier indicate that this algorithm is used
with SHA-256.
Signed-off-by: Andrzej Kurek <andrzej.kurek@arm.com>
---
 include/psa/crypto_values.h | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/include/psa/crypto_values.h b/include/psa/crypto_values.h
index 455b6388b3..f050ecc42f 100644
--- a/include/psa/crypto_values.h
+++ b/include/psa/crypto_values.h
@@ -2031,7 +2031,7 @@
  * The only supported curve is secp256r1 (the 256-bit curve in
  * #PSA_ECC_FAMILY_SECP_R1), so the input must be exactly 65 bytes.
  */
-#define PSA_ALG_TLS12_ECJPAKE_TO_PMS            ((psa_algorithm_t)0x08000600)
+#define PSA_ALG_TLS12_ECJPAKE_TO_PMS            ((psa_algorithm_t)0x08000609)
 
 /* This flag indicates whether the key derivation algorithm is suitable for
  * use on low-entropy secrets such as password - these algorithms are also

From 5603efd525f43865a42f0653048f7a0aebb11115 Mon Sep 17 00:00:00 2001
From: Andrzej Kurek <andrzej.kurek@arm.com>
Date: Mon, 26 Sep 2022 10:49:16 -0400
Subject: [PATCH 14/18] Improve readability and formatting

Also use a sizeof instead of a constant for zeroization, as
requested in review.
Signed-off-by: Andrzej Kurek <andrzej.kurek@arm.com>
---
 library/psa_crypto.c | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index 663b645bbb..0b21eb083d 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -4355,7 +4355,7 @@ psa_status_t psa_key_derivation_abort( psa_key_derivation_operation_t *operation
     if( kdf_alg == PSA_ALG_TLS12_ECJPAKE_TO_PMS )
     {
         mbedtls_platform_zeroize( operation->ctx.tls12_ecjpake_to_pms.data,
-            PSA_TLS12_ECJPAKE_TO_PMS_DATA_SIZE );
+            sizeof( operation->ctx.tls12_ecjpake_to_pms.data ) );
     }
     else
 #endif /* defined(MBEDTLS_PSA_BUILTIN_ALG_TLS12_ECJPAKE_TO_PMS) */
@@ -4647,7 +4647,7 @@ static psa_status_t psa_key_derivation_tls12_ecjpake_to_pms_read(
     size_t output_length )
 {
     psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
-    size_t output_size;
+    size_t output_size = 0;
 
     if( output_length != 32 )
         return ( PSA_ERROR_INVALID_ARGUMENT );
@@ -5578,7 +5578,9 @@ static psa_status_t psa_tls12_ecjpake_to_pms_input(
 {
     if( data_length != PSA_TLS12_ECJPAKE_TO_PMS_INPUT_SIZE ||
         step != PSA_KEY_DERIVATION_INPUT_SECRET )
+    {
         return( PSA_ERROR_INVALID_ARGUMENT );
+    }
 
     /* Check if the passed point is in an uncompressed form */
     if( data[0] != 0x04 )

From b510cd2c507559ebc6f86e5787464e67d9617a0d Mon Sep 17 00:00:00 2001
From: Andrzej Kurek <andrzej.kurek@arm.com>
Date: Mon, 26 Sep 2022 10:50:22 -0400
Subject: [PATCH 15/18] Fix a copy-paste error - wrong macro used

Signed-off-by: Andrzej Kurek <andrzej.kurek@arm.com>
---
 library/psa_crypto.c | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index 0b21eb083d..553834f76e 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -5175,12 +5175,13 @@ static psa_status_t psa_key_derivation_setup_kdf(
         return( PSA_ERROR_NOT_SUPPORTED );
     }
 #if defined(MBEDTLS_PSA_BUILTIN_ALG_HKDF_EXTRACT) || \
-    defined(MBEDTLS_PSA_BUILTIN_ALG_TLS12_PSK_TO_MS)
+    defined(MBEDTLS_PSA_BUILTIN_ALG_TLS12_ECJPAKE_TO_PMS)
     if( PSA_ALG_IS_HKDF_EXTRACT( kdf_alg ) ||
         ( kdf_alg == PSA_ALG_TLS12_ECJPAKE_TO_PMS ) )
         operation->capacity = hash_size;
     else
-#endif /* MBEDTLS_PSA_BUILTIN_ALG_HKDF_EXTRACT */
+#endif /* MBEDTLS_PSA_BUILTIN_ALG_HKDF_EXTRACT ||
+          MBEDTLS_PSA_BUILTIN_ALG_TLS12_ECJPAKE_TO_PMS */
         operation->capacity = 255 * hash_size;
     return( PSA_SUCCESS );
 }

From 3539f2c90b9e133773a082752bcda07324744003 Mon Sep 17 00:00:00 2001
From: Andrzej Kurek <andrzej.kurek@arm.com>
Date: Mon, 26 Sep 2022 10:56:02 -0400
Subject: [PATCH 16/18] Improve readability in ecjpake tests

Signed-off-by: Andrzej Kurek <andrzej.kurek@arm.com>
---
 tests/suites/test_suite_psa_crypto.function | 17 ++++++++++-------
 1 file changed, 10 insertions(+), 7 deletions(-)

diff --git a/tests/suites/test_suite_psa_crypto.function b/tests/suites/test_suite_psa_crypto.function
index 355cba533b..b4a83ab8da 100644
--- a/tests/suites/test_suite_psa_crypto.function
+++ b/tests/suites/test_suite_psa_crypto.function
@@ -7234,28 +7234,31 @@ exit:
 /* END_CASE */
 
 /* BEGIN_CASE depends_on:MBEDTLS_PSA_BUILTIN_ALG_TLS12_PSK_TO_MS:MBEDTLS_SHA256_C */
-void derive_ecjpake_to_pms( data_t *input, int expected_input_status,
+void derive_ecjpake_to_pms( data_t *input, int expected_input_status_arg,
                             int derivation_step,
-                            int capacity, int expected_capacity_status,
+                            int capacity, int expected_capacity_status_arg,
                             data_t *expected_output,
-                            int expected_output_status )
+                            int expected_output_status_arg )
 {
     psa_algorithm_t alg = PSA_ALG_TLS12_ECJPAKE_TO_PMS;
     psa_key_derivation_operation_t operation = PSA_KEY_DERIVATION_OPERATION_INIT;
     psa_key_derivation_step_t step = (psa_key_derivation_step_t) derivation_step;
     uint8_t *output_buffer = NULL;
     psa_status_t status;
+    psa_status_t expected_input_status = (psa_status_t) expected_input_status_arg;
+    psa_status_t expected_capacity_status = (psa_status_t) expected_capacity_status_arg;
+    psa_status_t expected_output_status = (psa_status_t) expected_output_status_arg;
 
     ASSERT_ALLOC( output_buffer, expected_output->len );
     PSA_ASSERT( psa_crypto_init() );
 
     PSA_ASSERT( psa_key_derivation_setup( &operation, alg ) );
     TEST_EQUAL( psa_key_derivation_set_capacity( &operation, capacity ),
-        (psa_status_t) expected_capacity_status );
+                expected_capacity_status);
 
     TEST_EQUAL( psa_key_derivation_input_bytes( &operation,
-                step, input->x, input->len ),
-        (psa_status_t) expected_input_status );
+                    step, input->x, input->len ),
+                expected_input_status );
 
     if( ( (psa_status_t) expected_input_status ) != PSA_SUCCESS )
         goto exit;
@@ -7263,7 +7266,7 @@ void derive_ecjpake_to_pms( data_t *input, int expected_input_status,
     status = psa_key_derivation_output_bytes( &operation, output_buffer,
         expected_output->len );
 
-    TEST_EQUAL( status, (psa_status_t) expected_output_status );
+    TEST_EQUAL( status, expected_output_status );
     if( expected_output->len != 0 && expected_output_status == PSA_SUCCESS )
         ASSERT_COMPARE( output_buffer, expected_output->len, expected_output->x,
                         expected_output->len );

From e09aff8f5a46cba41a280739d23fa44ec0940f15 Mon Sep 17 00:00:00 2001
From: Andrzej Kurek <andrzej.kurek@arm.com>
Date: Mon, 26 Sep 2022 10:59:31 -0400
Subject: [PATCH 17/18] Add information about ECJPAKE_TO_PMS output size
 expectations

Signed-off-by: Andrzej Kurek <andrzej.kurek@arm.com>
---
 include/psa/crypto_values.h | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/include/psa/crypto_values.h b/include/psa/crypto_values.h
index f050ecc42f..b465ddb14c 100644
--- a/include/psa/crypto_values.h
+++ b/include/psa/crypto_values.h
@@ -2030,6 +2030,8 @@
  * #PSA_KEY_DERIVATION_INPUT_SECRET is the shared secret K from EC J-PAKE.
  * The only supported curve is secp256r1 (the 256-bit curve in
  * #PSA_ECC_FAMILY_SECP_R1), so the input must be exactly 65 bytes.
+ * The output has to be read as a single chunk of 32 bytes, defined as
+ * PSA_TLS12_ECJPAKE_TO_PMS_DATA_SIZE.
  */
 #define PSA_ALG_TLS12_ECJPAKE_TO_PMS            ((psa_algorithm_t)0x08000609)
 

From f13925c02203c7cffe95c37846eeda24b1a109b3 Mon Sep 17 00:00:00 2001
From: Andrzej Kurek <andrzej.kurek@arm.com>
Date: Tue, 27 Sep 2022 05:16:10 -0400
Subject: [PATCH 18/18] Add a changelog entry for ECJPAKE to PMS KDF

Signed-off-by: Andrzej Kurek <andrzej.kurek@arm.com>
---
 ChangeLog.d/ecjpake_to_pms.txt | 5 +++++
 1 file changed, 5 insertions(+)
 create mode 100644 ChangeLog.d/ecjpake_to_pms.txt

diff --git a/ChangeLog.d/ecjpake_to_pms.txt b/ChangeLog.d/ecjpake_to_pms.txt
new file mode 100644
index 0000000000..4dd2075de0
--- /dev/null
+++ b/ChangeLog.d/ecjpake_to_pms.txt
@@ -0,0 +1,5 @@
+API changes
+   * Add an ad-hoc key derivation function handling ECJPAKE to PMS
+     calculation that can be used to derive the session secret in TLS 1.2,
+     as described in draft-cragie-tls-ecjpake-01. This can be achieved by
+     using PSA_ALG_TLS12_ECJPAKE_TO_PMS as the key derivation algorithm.