HKDF: be more robust if we reach the maximum ouptut length

In psa_generator_hkdf_read, return BAD_STATE if we're trying to
construct more output than the algorithm allows. This can't happen
through the API due to the capacity limit, but it could potentially
happen in an internal call.

Also add a test case that verifies that we can set up HKDF with its
maximum capacity and read up to the maximum capacity.

library/psa_crypto.c

@@ -3047,8 +3047,15 @@ static psa_status_t psa_generator_hkdf_read( psa_hkdf_generator_t *hkdf,
         output += n;
         output_length -= n;
         hkdf->offset_in_block += n;
-        if( output_length == 0 || hkdf->block_number == 0xff )
+        if( output_length == 0 )
             break;
+        /* We can't be wanting more output after block 0xff, otherwise
+         * the capacity check in psa_generator_read() would have
+         * prevented this call. It could happen only if the generator
+         * object was corrupted or if this function is called directly
+         * inside the library. */
+        if( hkdf->block_number == 0xff )
+            return( PSA_ERROR_BAD_STATE );
 
         /* We need a new block */
         ++hkdf->block_number;

tests/suites/test_suite_psa_crypto.data

@@ -853,6 +853,14 @@ PSA key derivation: over capacity 42: output 43+1
 depends_on:MBEDTLS_MD_C:MBEDTLS_SHA256_C
 derive_output:PSA_ALG_HKDF(PSA_ALG_SHA_256):"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":"000102030405060708090a0b0c":"f0f1f2f3f4f5f6f7f8f9":42:"3cb25f25faacd57a90434f64d0362f2a2d2d0a90cf1a5a4c5db02d56ecc4c5bf34007208d5b887185865ff":"ff"
 
+PSA key derivation: HKDF SHA-256, read maximum capacity minus 1
+depends_on:MBEDTLS_MD_C:MBEDTLS_SHA256_C
+derive_full:PSA_ALG_HKDF(PSA_ALG_SHA_256):"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":"000102030405060708090a0b0c":"f0f1f2f3f4f5f6f7f8f9":255 * 32 - 1
+
+PSA key derivation: HKDF SHA-256, read maximum capacity
+depends_on:MBEDTLS_MD_C:MBEDTLS_SHA256_C
+derive_full:PSA_ALG_HKDF(PSA_ALG_SHA_256):"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":"000102030405060708090a0b0c":"f0f1f2f3f4f5f6f7f8f9":255 * 32
+
 PSA key derivation: HKDF SHA-256, exercise HMAC-SHA-256
 depends_on:MBEDTLS_MD_C:MBEDTLS_SHA256_C
 derive_key_exercise:PSA_ALG_HKDF(PSA_ALG_SHA_256):"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":"000102030405060708090a0b0c":"f0f1f2f3f4f5f6f7f8f9":PSA_KEY_TYPE_HMAC:256:PSA_KEY_USAGE_SIGN:PSA_ALG_HMAC(PSA_ALG_SHA_256)

tests/suites/test_suite_psa_crypto.function

@@ -2583,6 +2583,72 @@ exit:
 }
 /* END_CASE */
 
+/* BEGIN_CASE */
+void derive_full( int alg_arg,
+                  data_t *key_data,
+                  data_t *salt,
+                  data_t *label,
+                  int requested_capacity_arg )
+{
+    psa_key_slot_t slot = 1;
+    psa_algorithm_t alg = alg_arg;
+    size_t requested_capacity = requested_capacity_arg;
+    psa_crypto_generator_t generator = PSA_CRYPTO_GENERATOR_INIT;
+    unsigned char output_buffer[16];
+    size_t expected_capacity = requested_capacity;
+    size_t current_capacity;
+    psa_key_policy_t policy;
+
+    TEST_ASSERT( psa_crypto_init( ) == PSA_SUCCESS );
+
+    psa_key_policy_init( &policy );
+    psa_key_policy_set_usage( &policy, PSA_KEY_USAGE_DERIVE, alg );
+    TEST_ASSERT( psa_set_key_policy( slot, &policy ) == PSA_SUCCESS );
+
+    TEST_ASSERT( psa_import_key( slot, PSA_KEY_TYPE_DERIVE,
+                                 key_data->x,
+                                 key_data->len ) == PSA_SUCCESS );
+
+    /* Extraction phase. */
+    TEST_ASSERT( psa_key_derivation( &generator, slot, alg,
+                                     salt->x, salt->len,
+                                     label->x, label->len,
+                                     requested_capacity ) == PSA_SUCCESS );
+    TEST_ASSERT( psa_get_generator_capacity( &generator,
+                                             &current_capacity ) ==
+                 PSA_SUCCESS );
+    TEST_ASSERT( current_capacity == expected_capacity );
+
+    /* Expansion phase. */
+    while( current_capacity > 0 )
+    {
+        size_t read_size = sizeof( output_buffer );
+        if( read_size > current_capacity )
+            read_size = current_capacity;
+        TEST_ASSERT( psa_generator_read( &generator,
+                                         output_buffer,
+                                         read_size ) == PSA_SUCCESS );
+        expected_capacity -= read_size;
+        TEST_ASSERT( psa_get_generator_capacity( &generator,
+                                                 &current_capacity ) ==
+                     PSA_SUCCESS );
+        TEST_ASSERT( current_capacity == expected_capacity );
+    }
+
+    /* Check that the generator refuses to go over capacity. */
+    TEST_ASSERT( psa_generator_read( &generator,
+                                     output_buffer,
+                                     1 ) == PSA_ERROR_INSUFFICIENT_CAPACITY );
+
+    TEST_ASSERT( psa_generator_abort( &generator ) == PSA_SUCCESS );
+
+exit:
+    psa_generator_abort( &generator );
+    psa_destroy_key( slot );
+    mbedtls_psa_crypto_free( );
+}
+/* END_CASE */
+
 /* BEGIN_CASE */
 void derive_key_exercise( int alg_arg,
                           data_t *key_data,