diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index 31520b8..7b1d16b 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -4851,17 +4851,11 @@
                                              const uint8_t *data,
                                              size_t data_length )
 {
-    switch( step )
-    {
-        case PSA_KEY_DERIVATION_INPUT_LABEL:
-        case PSA_KEY_DERIVATION_INPUT_SALT:
-        case PSA_KEY_DERIVATION_INPUT_INFO:
-        case PSA_KEY_DERIVATION_INPUT_SEED:
-            return( psa_key_derivation_input_internal( operation, step,
-                                                       data, data_length ) );
-        default:
-            return( PSA_ERROR_INVALID_ARGUMENT );
-    }
+    if( step == PSA_KEY_DERIVATION_INPUT_SECRET )
+        return( PSA_ERROR_INVALID_ARGUMENT );
+
+    return( psa_key_derivation_input_internal( operation, step,
+                                               data, data_length ) );
 }
 
 psa_status_t psa_key_derivation_input_key( psa_key_derivation_operation_t *operation,