Fix psa_key_derivation_input_integer() not detecting bad state

Signed-off-by: Waleed Elmelegy <waleed.elmelegy@arm.com>
diff --git a/ChangeLog.d/fix-key-derive-bad-state-error.txt b/ChangeLog.d/fix-key-derive-bad-state-error.txt
new file mode 100644
index 0000000..0bccf77
--- /dev/null
+++ b/ChangeLog.d/fix-key-derive-bad-state-error.txt
@@ -0,0 +1,3 @@
+Bugfix
+   * Fix issue where psa_key_derivation_input_integer() is not detecting
+     bad state after an operation has been aborted.
diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index ec5934e..69d037b 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -4751,6 +4751,12 @@
     psa_status_t status;
     psa_algorithm_t kdf_alg = psa_key_derivation_get_kdf_alg(operation);
 
+    if (kdf_alg == 0) {
+       /* This is a blank or aborted operation. */
+        status = PSA_ERROR_BAD_STATE;
+        goto exit;
+    }
+
     status = psa_key_derivation_check_input_type(step, key_type);
     if (status != PSA_SUCCESS) {
         goto exit;
diff --git a/tests/suites/test_suite_psa_crypto.function b/tests/suites/test_suite_psa_crypto.function
index 21b768b..838717e 100644
--- a/tests/suites/test_suite_psa_crypto.function
+++ b/tests/suites/test_suite_psa_crypto.function
@@ -4566,6 +4566,16 @@
     }
     TEST_EQUAL(actual_output_status, expected_output_status);
 
+    /* Test calling input functions after operation has been aborted
+       result in PSA_ERROR_BAD_STATE error.
+    */
+    psa_key_derivation_abort(&operation);
+
+    TEST_EQUAL(psa_key_derivation_input_bytes(
+                    &operation, steps[0],
+                    inputs[0]->x, inputs[0]->len),
+                PSA_ERROR_BAD_STATE);
+
 exit:
     psa_key_derivation_abort(&operation);
     for (i = 0; i < ARRAY_LENGTH(keys); i++) {