diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index 3ecaebb..8c68077 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -5285,6 +5285,11 @@
             if( PSA_ALG_IS_HKDF_EXTRACT( kdf_alg ) )
                 return( PSA_ERROR_INVALID_ARGUMENT );
 #endif /* MBEDTLS_PSA_BUILTIN_ALG_HKDF_EXTRACT */
+#if defined(MBEDTLS_PSA_BUILTIN_ALG_HKDF_EXPAND)
+            if( PSA_ALG_IS_HKDF_EXPAND( kdf_alg ) &&
+                hkdf->state == HKDF_STATE_INIT )
+                return( PSA_ERROR_BAD_STATE );
+#endif /* MBEDTLS_PSA_BUILTIN_ALG_HKDF_EXTRACT */
             if( hkdf->state == HKDF_STATE_OUTPUT )
                 return( PSA_ERROR_BAD_STATE );
             if( hkdf->info_set )