diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index 9e8f90b54..ef99403e3 100644
--- a/library/psa_crypto.c
+++ b/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;
diff --git a/tests/suites/test_suite_psa_crypto.data b/tests/suites/test_suite_psa_crypto.data
index 43b964730..8ffaa8779 100644
--- a/tests/suites/test_suite_psa_crypto.data
+++ b/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)
diff --git a/tests/suites/test_suite_psa_crypto.function b/tests/suites/test_suite_psa_crypto.function
index 1f1732e65..27bc4e1d5 100644
--- a/tests/suites/test_suite_psa_crypto.function
+++ b/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,