diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index 9e8f90b..ef99403 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -3047,8 +3047,15 @@
         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;