Key derivation by small input steps: proof-of-concept

Document the new API. Keep the old one.

Implement for HKDF. Use it in a few test cases.

Key agreement is still unchanged.
diff --git a/include/psa/crypto.h b/include/psa/crypto.h
index 683feb8..6005269 100644
--- a/include/psa/crypto.h
+++ b/include/psa/crypto.h
@@ -1963,6 +1963,22 @@
 psa_status_t psa_get_generator_capacity(const psa_crypto_generator_t *generator,
                                         size_t *capacity);
 
+/** Set the maximum capacity of a generator.
+ *
+ * \param[in,out] generator The generator object to modify.
+ * \param capacity          The new capacity of the generator.
+ *                          It must be less or equal to the generator's
+ *                          current capacity.
+ *
+ * \retval #PSA_SUCCESS
+ * \retval #PSA_ERROR_INVALID_ARGUMENT
+ *         \p capacity is larger than the generator's current capacity.
+ * \retval #PSA_ERROR_BAD_STATE
+ * \retval #PSA_ERROR_COMMUNICATION_FAILURE
+ */
+psa_status_t psa_set_generator_capacity(psa_crypto_generator_t *generator,
+                                        size_t capacity);
+
 /** Read some data from a generator.
  *
  * This function reads and returns a sequence of bytes from a generator.
@@ -2090,6 +2106,131 @@
 
 /** Set up a key derivation operation.
  *
+ * A key derivation algorithm takes some inputs and uses them to create
+ * a byte generator which can be used to produce keys and other
+ * cryptographic material.
+ *
+ * To use a generator for key derivation:
+ * - Start with an initialized object of type #psa_crypto_generator_t.
+ * - Call psa_key_derivation_setup() to select the algorithm.
+ * - Provide the inputs for the key derivation by calling
+ *   psa_key_derivation_input_bytes() or psa_key_derivation_input_key()
+ *   as appropriate. Which inputs are needed, in what order, and whether
+ *   they may be keys and if so of what type depends on the algorithm.
+ * - Optionally set the generator's maximum capacity with
+ *   psa_set_generator_capacity(). You may do this before, in the middle of
+ *   or after providing inputs. For some algorithms, this step is mandatory
+ *   because the output depends on the maximum capacity.
+ * - Generate output with psa_generator_read() or
+ *   psa_generator_import_key(). Successive calls to these functions
+ *   use successive output bytes from the generator.
+ * - Clean up the generator object with psa_generator_abort().
+ *
+ * \param[in,out] generator       The generator object to set up. It must
+ *                                have been initialized but not set up yet.
+ * \param alg                     The key derivation algorithm to compute
+ *                                (\c PSA_ALG_XXX value such that
+ *                                #PSA_ALG_IS_KEY_DERIVATION(\p alg) is true).
+ *
+ * \retval #PSA_SUCCESS
+ *         Success.
+ * \retval #PSA_ERROR_INVALID_ARGUMENT
+ *         \c alg is not a key derivation algorithm.
+ * \retval #PSA_ERROR_NOT_SUPPORTED
+ *         \c alg is not supported or is not a key derivation algorithm.
+ * \retval #PSA_ERROR_INSUFFICIENT_MEMORY
+ * \retval #PSA_ERROR_COMMUNICATION_FAILURE
+ * \retval #PSA_ERROR_HARDWARE_FAILURE
+ * \retval #PSA_ERROR_TAMPERING_DETECTED
+ * \retval #PSA_ERROR_BAD_STATE
+ */
+psa_status_t psa_key_derivation_setup(psa_crypto_generator_t *generator,
+                                      psa_algorithm_t alg);
+
+/** Provide an input for key derivation.
+ *
+ * Which inputs are required and in what order depends on the type of
+ * key derivation algorithm.
+ *
+ * - For HKDF (#PSA_ALG_HKDF), the following inputs are supported:
+ *     - #PSA_KDF_STEP_SALT is the salt used in the "extract" step.
+ *       It is optional; if omitted, the derivation uses an empty salt.
+ *     - #PSA_KDF_STEP_SECRET is the secret key used in the "extract" step.
+ *       It may be a key of type #PSA_KEY_TYPE_DERIVE with the
+ *       usage flag #PSA_KEY_USAGE_DERIVE.
+ *     - #PSA_KDF_STEP_INFO is the info string used in the "expand" step.
+ *   You must pass #PSA_KDF_STEP_SALT before #PSA_KDF_STEP_SECRET.
+ *   #PSA_KDF_STEP_INFO may be passed at any time before starting to
+ *   generate output.
+ *
+ * \param[in,out] generator       The generator object to use. It must
+ *                                have been set up with
+ *                                psa_key_derivation_setup() and must not
+ *                                have produced any output yet.
+ * \param step                    Which step the input data is for.
+ *                                See above for the permitted values
+ *                                depending on the algorithm.
+ * \param[in] data                Input data to use.
+ * \param data_length             Size of the \p data buffer in bytes.
+ *
+ * \retval #PSA_SUCCESS
+ *         Success.
+ * \retval #PSA_ERROR_INVALID_ARGUMENT
+ *         \c step is not compatible with the generator's algorithm.
+ * \retval #PSA_ERROR_INSUFFICIENT_MEMORY
+ * \retval #PSA_ERROR_COMMUNICATION_FAILURE
+ * \retval #PSA_ERROR_HARDWARE_FAILURE
+ * \retval #PSA_ERROR_TAMPERING_DETECTED
+ * \retval #PSA_ERROR_BAD_STATE
+ *         The value of \p step is not valid given the state of \p generator.
+ * \retval #PSA_ERROR_BAD_STATE
+ *         The library has not been previously initialized by psa_crypto_init().
+ *         It is implementation-dependent whether a failure to initialize
+ *         results in this error code.
+ */
+psa_status_t psa_key_derivation_input_bytes(psa_crypto_generator_t *generator,
+                                            psa_key_derivation_step_t step,
+                                            const uint8_t *data,
+                                            size_t data_length);
+
+/** Provide an input for key derivation in the form of a key.
+ *
+ * See the descrition of psa_key_derivation_input_bytes() regarding
+ * what inputs are supported and in what order. An input step may only be
+ * a key if the descrition of psa_key_derivation_input_bytes() explicitly
+ * allows it.
+ *
+ * \param[in,out] generator       The generator object to use. It must
+ *                                have been set up with
+ *                                psa_key_derivation_setup() and must not
+ *                                have produced any output yet.
+ * \param step                    Which step the input data is for.
+ * \param handle                  Handle to the secret key.
+ *
+ * \retval #PSA_SUCCESS
+ *         Success.
+ * \retval #PSA_ERROR_INVALID_HANDLE
+ * \retval #PSA_ERROR_EMPTY_SLOT
+ * \retval #PSA_ERROR_NOT_PERMITTED
+ * \retval #PSA_ERROR_INVALID_ARGUMENT
+ *         \c step is not compatible with the generator's algorithm.
+ * \retval #PSA_ERROR_INSUFFICIENT_MEMORY
+ * \retval #PSA_ERROR_COMMUNICATION_FAILURE
+ * \retval #PSA_ERROR_HARDWARE_FAILURE
+ * \retval #PSA_ERROR_TAMPERING_DETECTED
+ * \retval #PSA_ERROR_BAD_STATE
+ *         The value of \p step is not valid given the state of \p generator.
+ * \retval #PSA_ERROR_BAD_STATE
+ *         The library has not been previously initialized by psa_crypto_init().
+ *         It is implementation-dependent whether a failure to initialize
+ *         results in this error code.
+ */
+psa_status_t psa_key_derivation_input_key(psa_crypto_generator_t *generator,
+                                          psa_key_derivation_step_t step,
+                                          psa_key_handle_t handle);
+
+/** Set up a key derivation operation.
+ *
  * A key derivation algorithm takes three inputs: a secret input \p key and
  * two non-secret inputs \p label and p salt.
  * The result of this function is a byte generator which can
diff --git a/include/psa/crypto_struct.h b/include/psa/crypto_struct.h
index ee3ecd7..bebc5c4 100644
--- a/include/psa/crypto_struct.h
+++ b/include/psa/crypto_struct.h
@@ -165,6 +165,8 @@
 #endif
     uint8_t offset_in_block;
     uint8_t block_number;
+    uint8_t state : 2;
+    uint8_t info_set : 1;
 } psa_hkdf_generator_t;
 #endif /* MBEDTLS_MD_C */
 
diff --git a/include/psa/crypto_types.h b/include/psa/crypto_types.h
index 9b44d6a..637e07c 100644
--- a/include/psa/crypto_types.h
+++ b/include/psa/crypto_types.h
@@ -98,4 +98,13 @@
 
 /**@}*/
 
+/** \defgroup derivation Key derivation
+ * @{
+ */
+
+/** \brief Encoding of the step of a key derivation. */
+typedef uint16_t psa_key_derivation_step_t;
+
+/**@}*/
+
 #endif /* PSA_CRYPTO_TYPES_H */
diff --git a/include/psa/crypto_values.h b/include/psa/crypto_values.h
index 4d25835..5c81acd 100644
--- a/include/psa/crypto_values.h
+++ b/include/psa/crypto_values.h
@@ -1417,4 +1417,16 @@
 
 /**@}*/
 
+/** \defgroup derivation Key derivation
+ * @{
+ */
+
+#define PSA_KDF_STEP_SECRET              ((psa_key_derivation_step_t)0x0101)
+#define PSA_KDF_STEP_LABEL               ((psa_key_derivation_step_t)0x0201)
+#define PSA_KDF_STEP_SALT                ((psa_key_derivation_step_t)0x0202)
+#define PSA_KDF_STEP_INFO                ((psa_key_derivation_step_t)0x0203)
+#define PSA_KDF_STEP_PEER_KEY            ((psa_key_derivation_step_t)0x0301)
+
+/**@}*/
+
 #endif /* PSA_CRYPTO_VALUES_H */
diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index fd76b27..916c52f 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -3313,6 +3313,11 @@
 /* Generators */
 /****************************************************************/
 
+#define HKDF_STATE_INIT 0 /* no input yet */
+#define HKDF_STATE_STARTED 1 /* got salt */
+#define HKDF_STATE_KEYED 2 /* got key */
+#define HKDF_STATE_OUTPUT 3 /* output started */
+
 psa_status_t psa_generator_abort( psa_crypto_generator_t *generator )
 {
     psa_status_t status = PSA_SUCCESS;
@@ -3366,7 +3371,6 @@
     return( status );
 }
 
-
 psa_status_t psa_get_generator_capacity(const psa_crypto_generator_t *generator,
                                         size_t *capacity)
 {
@@ -3374,6 +3378,17 @@
     return( PSA_SUCCESS );
 }
 
+psa_status_t psa_set_generator_capacity( psa_crypto_generator_t *generator,
+                                         size_t capacity )
+{
+    if( generator->alg == 0 )
+        return( PSA_ERROR_BAD_STATE );
+    if( capacity > generator->capacity )
+        return( PSA_ERROR_INVALID_ARGUMENT );
+    generator->capacity = capacity;
+    return( PSA_SUCCESS );
+}
+
 #if defined(MBEDTLS_MD_C)
 /* Read some bytes from an HKDF-based generator. This performs a chunk
  * of the expand phase of the HKDF algorithm. */
@@ -3385,6 +3400,10 @@
     uint8_t hash_length = PSA_HASH_SIZE( hash_alg );
     psa_status_t status;
 
+    if( hkdf->state < HKDF_STATE_KEYED || ! hkdf->info_set )
+        return( PSA_ERROR_BAD_STATE );
+    hkdf->state = HKDF_STATE_OUTPUT;
+
     while( output_length != 0 )
     {
         /* Copy what remains of the current block */
@@ -3755,6 +3774,8 @@
             return( PSA_ERROR_INSUFFICIENT_MEMORY );
         memcpy( hkdf->info, label, label_length );
     }
+    hkdf->state = HKDF_STATE_KEYED;
+    hkdf->info_set = 1;
     return( PSA_SUCCESS );
 }
 #endif /* MBEDTLS_MD_C */
@@ -3998,6 +4019,177 @@
     return( status );
 }
 
+psa_status_t psa_key_derivation_setup( psa_crypto_generator_t *generator,
+                                       psa_algorithm_t alg )
+{
+    if( generator->alg != 0 )
+        return( PSA_ERROR_BAD_STATE );
+    /* Make sure that alg is a supported key derivation algorithm.
+     * Key agreement algorithms and key selection algorithms are not
+     * supported by this function. */
+#if defined(MBEDTLS_MD_C)
+    if( PSA_ALG_IS_HKDF( alg ) ||
+        PSA_ALG_IS_TLS12_PRF( alg ) ||
+        PSA_ALG_IS_TLS12_PSK_TO_MS( alg ) )
+    {
+        psa_algorithm_t hash_alg = PSA_ALG_HKDF_GET_HASH( alg );
+        size_t hash_size = PSA_HASH_SIZE( hash_alg );
+        if( hash_size == 0 )
+            return( PSA_ERROR_NOT_SUPPORTED );
+        if( ( PSA_ALG_IS_TLS12_PRF( alg ) ||
+              PSA_ALG_IS_TLS12_PSK_TO_MS( alg ) ) &&
+            ! ( hash_alg == PSA_ALG_SHA_256 && hash_alg == PSA_ALG_SHA_384 ) )
+        {
+            return( PSA_ERROR_NOT_SUPPORTED );
+        }
+        generator->capacity = 255 * hash_size;
+    }
+#endif /* MBEDTLS_MD_C */
+    else if( PSA_ALG_IS_KEY_DERIVATION( alg ) )
+        return( PSA_ERROR_NOT_SUPPORTED );
+    else
+        return( PSA_ERROR_INVALID_ARGUMENT );
+    generator->alg = alg;
+    return( PSA_SUCCESS );
+}
+
+#if defined(MBEDTLS_MD_C)
+static psa_status_t psa_hkdf_input( psa_hkdf_generator_t *hkdf,
+                                    psa_algorithm_t hash_alg,
+                                    psa_key_derivation_step_t step,
+                                    const uint8_t *data,
+                                    size_t data_length )
+{
+    psa_status_t status;
+    switch( step )
+    {
+        case PSA_KDF_STEP_SALT:
+            if( hkdf->state == HKDF_STATE_INIT )
+            {
+                status = psa_hmac_setup_internal( &hkdf->hmac,
+                                                  data, data_length,
+                                                  hash_alg );
+                if( status != PSA_SUCCESS )
+                    return( status );
+                hkdf->state = HKDF_STATE_STARTED;
+                return( PSA_SUCCESS );
+            }
+            else
+                return( PSA_ERROR_BAD_STATE );
+            break;
+        case PSA_KDF_STEP_SECRET:
+            /* If no salt was provided, use an empty salt. */
+            if( hkdf->state == HKDF_STATE_INIT )
+            {
+                status = psa_hmac_setup_internal( &hkdf->hmac,
+                                                  NULL, 0,
+                                                  PSA_ALG_HMAC( hash_alg ) );
+                if( status != PSA_SUCCESS )
+                    return( status );
+                hkdf->state = HKDF_STATE_STARTED;
+            }
+            if( hkdf->state == HKDF_STATE_STARTED )
+            {
+                status = psa_hash_update( &hkdf->hmac.hash_ctx,
+                                          data, data_length );
+                if( status != PSA_SUCCESS )
+                    return( status );
+                status = psa_hmac_finish_internal( &hkdf->hmac,
+                                                   hkdf->prk,
+                                                   sizeof( hkdf->prk ) );
+                if( status != PSA_SUCCESS )
+                    return( status );
+                hkdf->offset_in_block = PSA_HASH_SIZE( hash_alg );
+                hkdf->block_number = 0;
+                hkdf->state = HKDF_STATE_KEYED;
+                return( PSA_SUCCESS );
+            }
+            else
+                return( PSA_ERROR_BAD_STATE );
+            break;
+        case PSA_KDF_STEP_INFO:
+            if( hkdf->state == HKDF_STATE_OUTPUT )
+                return( PSA_ERROR_BAD_STATE );
+            if( hkdf->info_set )
+                return( PSA_ERROR_BAD_STATE );
+            hkdf->info_length = data_length;
+            if( data_length != 0 )
+            {
+                hkdf->info = mbedtls_calloc( 1, data_length );
+                if( hkdf->info == NULL )
+                    return( PSA_ERROR_INSUFFICIENT_MEMORY );
+                memcpy( hkdf->info, data, data_length );
+            }
+            hkdf->info_set = 1;
+            return( PSA_SUCCESS );
+        default:
+            return( PSA_ERROR_INVALID_ARGUMENT );
+    }
+}
+#endif /* MBEDTLS_MD_C */
+
+psa_status_t psa_key_derivation_input_bytes( psa_crypto_generator_t *generator,
+                                             psa_key_derivation_step_t step,
+                                             const uint8_t *data,
+                                             size_t data_length )
+{
+    psa_status_t status;
+
+#if defined(MBEDTLS_MD_C)
+    if( PSA_ALG_IS_HKDF( generator->alg ) )
+    {
+        status = psa_hkdf_input( &generator->ctx.hkdf,
+                                 PSA_ALG_HKDF_GET_HASH( generator->alg ),
+                                 step, data, data_length );
+    }
+#endif /* MBEDTLS_MD_C */
+
+#if defined(MBEDTLS_MD_C)
+    /* TLS-1.2 PRF and TLS-1.2 PSK-to-MS are very similar, so share code. */
+    else if( PSA_ALG_IS_TLS12_PRF( generator->alg ) ||
+             PSA_ALG_IS_TLS12_PSK_TO_MS( generator->alg ) )
+    {
+        // TODO
+        status = PSA_ERROR_NOT_SUPPORTED;
+    }
+    else
+#endif /* MBEDTLS_MD_C */
+
+    {
+        /* This can't happen unless the generator object was not initialized */
+        return( PSA_ERROR_BAD_STATE );
+    }
+
+    if( status != PSA_SUCCESS )
+        psa_generator_abort( generator );
+    return( status );
+}
+
+psa_status_t psa_key_derivation_input_key( psa_crypto_generator_t *generator,
+                                           psa_key_derivation_step_t step,
+                                           psa_key_handle_t handle )
+{
+    psa_key_slot_t *slot;
+    psa_status_t status;
+    status = psa_get_key_from_slot( handle, &slot,
+                                    PSA_KEY_USAGE_DERIVE,
+                                    generator->alg );
+    if( status != PSA_SUCCESS )
+        return( status );
+    if( slot->type != PSA_KEY_TYPE_DERIVE )
+        return( PSA_ERROR_INVALID_ARGUMENT );
+    /* Don't allow a key to be used as an input that is usually public.
+     * This is debatable. It's ok from a cryptographic perspective to
+     * use secret material as an input that is usually public. However
+     * this is usually not intended, so be conservative at least for now. */
+    if( step != PSA_KDF_STEP_SECRET )
+        return( PSA_ERROR_INVALID_ARGUMENT );
+    return( psa_key_derivation_input_bytes( generator,
+                                            step,
+                                            slot->data.raw.data,
+                                            slot->data.raw.bytes ) );
+}
+
 
 
 /****************************************************************/
diff --git a/tests/suites/test_suite_psa_crypto.function b/tests/suites/test_suite_psa_crypto.function
index 6916bf4..9b8e01c 100644
--- a/tests/suites/test_suite_psa_crypto.function
+++ b/tests/suites/test_suite_psa_crypto.function
@@ -366,11 +366,30 @@
 
     if( usage & PSA_KEY_USAGE_DERIVE )
     {
-        PSA_ASSERT( psa_key_derivation( &generator,
-                                        handle, alg,
-                                        label, label_length,
-                                        seed, seed_length,
-                                        sizeof( output ) ) );
+        if( PSA_ALG_IS_HKDF( alg ) )
+        {
+            PSA_ASSERT( psa_key_derivation_setup( &generator, alg ) );
+            PSA_ASSERT( psa_key_derivation_input_bytes( &generator,
+                                                        PSA_KDF_STEP_SALT,
+                                                        label,
+                                                        label_length ) );
+            PSA_ASSERT( psa_key_derivation_input_key( &generator,
+                                                      PSA_KDF_STEP_SECRET,
+                                                      handle ) );
+            PSA_ASSERT( psa_key_derivation_input_bytes( &generator,
+                                                        PSA_KDF_STEP_INFO,
+                                                        seed,
+                                                        seed_length ) );
+        }
+        else
+        {
+            // legacy
+            PSA_ASSERT( psa_key_derivation( &generator,
+                                            handle, alg,
+                                            label, label_length,
+                                            seed, seed_length,
+                                            sizeof( output ) ) );
+        }
         PSA_ASSERT( psa_generator_read( &generator,
                                         output,
                                         sizeof( output ) ) );
@@ -3495,10 +3514,29 @@
                                 key_data->len ) );
 
     /* Extraction phase. */
-    PSA_ASSERT( psa_key_derivation( &generator, handle, alg,
-                                    salt->x, salt->len,
-                                    label->x, label->len,
-                                    requested_capacity ) );
+    if( PSA_ALG_IS_HKDF( alg ) )
+    {
+        PSA_ASSERT( psa_key_derivation_setup( &generator, alg ) );
+        PSA_ASSERT( psa_set_generator_capacity( &generator,
+                                                requested_capacity ) );
+        PSA_ASSERT( psa_key_derivation_input_bytes( &generator,
+                                                    PSA_KDF_STEP_SALT,
+                                                    salt->x, salt->len ) );
+        PSA_ASSERT( psa_key_derivation_input_key( &generator,
+                                                  PSA_KDF_STEP_SECRET,
+                                                  handle ) );
+        PSA_ASSERT( psa_key_derivation_input_bytes( &generator,
+                                                    PSA_KDF_STEP_INFO,
+                                                    label->x, label->len ) );
+    }
+    else
+    {
+        // legacy
+        PSA_ASSERT( psa_key_derivation( &generator, handle, alg,
+                                        salt->x, salt->len,
+                                        label->x, label->len,
+                                        requested_capacity ) );
+    }
     PSA_ASSERT( psa_get_generator_capacity( &generator,
                                             &current_capacity ) );
     TEST_EQUAL( current_capacity, requested_capacity );
@@ -3575,10 +3613,29 @@
                                 key_data->len ) );
 
     /* Extraction phase. */
-    PSA_ASSERT( psa_key_derivation( &generator, handle, alg,
-                                    salt->x, salt->len,
-                                    label->x, label->len,
-                                    requested_capacity ) );
+    if( PSA_ALG_IS_HKDF( alg ) )
+    {
+        PSA_ASSERT( psa_key_derivation_setup( &generator, alg ) );
+        PSA_ASSERT( psa_set_generator_capacity( &generator,
+                                                requested_capacity ) );
+        PSA_ASSERT( psa_key_derivation_input_bytes( &generator,
+                                                    PSA_KDF_STEP_SALT,
+                                                    salt->x, salt->len ) );
+        PSA_ASSERT( psa_key_derivation_input_key( &generator,
+                                                  PSA_KDF_STEP_SECRET,
+                                                  handle ) );
+        PSA_ASSERT( psa_key_derivation_input_bytes( &generator,
+                                                    PSA_KDF_STEP_INFO,
+                                                    label->x, label->len ) );
+    }
+    else
+    {
+        // legacy
+        PSA_ASSERT( psa_key_derivation( &generator, handle, alg,
+                                        salt->x, salt->len,
+                                        label->x, label->len,
+                                        requested_capacity ) );
+    }
     PSA_ASSERT( psa_get_generator_capacity( &generator,
                                             &current_capacity ) );
     TEST_EQUAL( current_capacity, expected_capacity );