Merge pull request #8198 from silabs-Kusumit/kdf_incorrect_initial_capacity

KDF incorrect initial capacity
diff --git a/ChangeLog.d/fix_kdf_incorrect_initial_capacity.txt b/ChangeLog.d/fix_kdf_incorrect_initial_capacity.txt
new file mode 100644
index 0000000..10e2795
--- /dev/null
+++ b/ChangeLog.d/fix_kdf_incorrect_initial_capacity.txt
@@ -0,0 +1,3 @@
+Bugfix
+   * Correct initial capacities for key derivation algorithms:TLS12_PRF,
+     TLS12_PSK_TO_MS, PBKDF2-HMAC, PBKDF2-CMAC
diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index e6d3851..e203e16 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -6096,6 +6096,91 @@
     return status;
 }
 
+static psa_status_t psa_key_derivation_set_maximum_capacity(
+    psa_key_derivation_operation_t *operation,
+    psa_algorithm_t kdf_alg)
+{
+#if defined(PSA_WANT_ALG_TLS12_ECJPAKE_TO_PMS)
+    if (kdf_alg == PSA_ALG_TLS12_ECJPAKE_TO_PMS) {
+        operation->capacity = PSA_HASH_LENGTH(PSA_ALG_SHA_256);
+        return PSA_SUCCESS;
+    }
+#endif
+#if defined(PSA_WANT_ALG_PBKDF2_AES_CMAC_PRF_128)
+    if (kdf_alg == PSA_ALG_PBKDF2_AES_CMAC_PRF_128) {
+#if (SIZE_MAX > UINT32_MAX)
+        operation->capacity = UINT32_MAX * (size_t) PSA_MAC_LENGTH(
+            PSA_KEY_TYPE_AES,
+            128U,
+            PSA_ALG_CMAC);
+#else
+        operation->capacity = SIZE_MAX;
+#endif
+        return PSA_SUCCESS;
+    }
+#endif /* PSA_WANT_ALG_PBKDF2_AES_CMAC_PRF_128 */
+
+    /* After this point, if kdf_alg is not valid then value of hash_alg may be
+     * invalid or meaningless but it does not affect this function */
+    psa_algorithm_t hash_alg = PSA_ALG_GET_HASH(kdf_alg);
+    size_t hash_size = PSA_HASH_LENGTH(hash_alg);
+    if (hash_size == 0) {
+        return PSA_ERROR_NOT_SUPPORTED;
+    }
+
+    /* Make sure that hash_alg is a supported hash algorithm. Otherwise
+     * we might fail later, which is somewhat unfriendly and potentially
+     * risk-prone. */
+    psa_status_t status = psa_hash_try_support(hash_alg);
+    if (status != PSA_SUCCESS) {
+        return status;
+    }
+
+#if defined(PSA_WANT_ALG_HKDF)
+    if (PSA_ALG_IS_HKDF(kdf_alg)) {
+        operation->capacity = 255 * hash_size;
+    } else
+#endif
+#if defined(PSA_WANT_ALG_HKDF_EXTRACT)
+    if (PSA_ALG_IS_HKDF_EXTRACT(kdf_alg)) {
+        operation->capacity = hash_size;
+    } else
+#endif
+#if defined(PSA_WANT_ALG_HKDF_EXPAND)
+    if (PSA_ALG_IS_HKDF_EXPAND(kdf_alg)) {
+        operation->capacity = 255 * hash_size;
+    } else
+#endif
+#if defined(PSA_WANT_ALG_TLS12_PRF)
+    if (PSA_ALG_IS_TLS12_PRF(kdf_alg) &&
+        (hash_alg == PSA_ALG_SHA_256 || hash_alg == PSA_ALG_SHA_384)) {
+        operation->capacity = SIZE_MAX;
+    } else
+#endif
+#if defined(PSA_WANT_ALG_TLS12_PSK_TO_MS)
+    if (PSA_ALG_IS_TLS12_PSK_TO_MS(kdf_alg) &&
+        (hash_alg == PSA_ALG_SHA_256 || hash_alg == PSA_ALG_SHA_384)) {
+        /* Master Secret is always 48 bytes
+         * https://datatracker.ietf.org/doc/html/rfc5246.html#section-8.1 */
+        operation->capacity = 48U;
+    } else
+#endif
+#if defined(PSA_WANT_ALG_PBKDF2_HMAC)
+    if (PSA_ALG_IS_PBKDF2_HMAC(kdf_alg)) {
+#if (SIZE_MAX > UINT32_MAX)
+        operation->capacity = UINT32_MAX * hash_size;
+#else
+        operation->capacity = SIZE_MAX;
+#endif
+    } else
+#endif /* PSA_WANT_ALG_PBKDF2_HMAC */
+    {
+        (void) hash_size;
+        status = PSA_ERROR_NOT_SUPPORTED;
+    }
+    return status;
+}
+
 static psa_status_t psa_key_derivation_setup_kdf(
     psa_key_derivation_operation_t *operation,
     psa_algorithm_t kdf_alg)
@@ -6109,43 +6194,9 @@
         return PSA_ERROR_NOT_SUPPORTED;
     }
 
-    /* All currently supported key derivation algorithms (apart from
-     * ecjpake to pms and pbkdf2_aes_cmac_128) are based on a hash algorithm. */
-    psa_algorithm_t hash_alg = PSA_ALG_HKDF_GET_HASH(kdf_alg);
-    size_t hash_size = PSA_HASH_LENGTH(hash_alg);
-    if (kdf_alg == PSA_ALG_TLS12_ECJPAKE_TO_PMS) {
-        hash_size = PSA_HASH_LENGTH(PSA_ALG_SHA_256);
-    } else if (kdf_alg == PSA_ALG_PBKDF2_AES_CMAC_PRF_128) {
-        hash_size = PSA_MAC_LENGTH(PSA_KEY_TYPE_AES, 128U, PSA_ALG_CMAC);
-    } else {
-        if (hash_size == 0) {
-            return PSA_ERROR_NOT_SUPPORTED;
-        }
-
-        /* Make sure that hash_alg is a supported hash algorithm. Otherwise
-         * we might fail later, which is somewhat unfriendly and potentially
-         * risk-prone. */
-        psa_status_t status = psa_hash_try_support(hash_alg);
-        if (status != PSA_SUCCESS) {
-            return status;
-        }
-    }
-
-    if ((PSA_ALG_IS_TLS12_PRF(kdf_alg) ||
-         PSA_ALG_IS_TLS12_PSK_TO_MS(kdf_alg)) &&
-        !(hash_alg == PSA_ALG_SHA_256 || hash_alg == PSA_ALG_SHA_384)) {
-        return PSA_ERROR_NOT_SUPPORTED;
-    }
-#if defined(MBEDTLS_PSA_BUILTIN_ALG_HKDF_EXTRACT) || \
-    defined(MBEDTLS_PSA_BUILTIN_ALG_TLS12_ECJPAKE_TO_PMS)
-    if (PSA_ALG_IS_HKDF_EXTRACT(kdf_alg) ||
-        (kdf_alg == PSA_ALG_TLS12_ECJPAKE_TO_PMS)) {
-        operation->capacity = hash_size;
-    } else
-#endif /* MBEDTLS_PSA_BUILTIN_ALG_HKDF_EXTRACT ||
-          MBEDTLS_PSA_BUILTIN_ALG_TLS12_ECJPAKE_TO_PMS */
-    operation->capacity = 255 * hash_size;
-    return PSA_SUCCESS;
+    psa_status_t status = psa_key_derivation_set_maximum_capacity(operation,
+                                                                  kdf_alg);
+    return status;
 }
 
 static psa_status_t psa_key_agreement_try_support(psa_algorithm_t alg)
diff --git a/tests/src/psa_exercise_key.c b/tests/src/psa_exercise_key.c
index f8b36e1..560b711 100644
--- a/tests/src/psa_exercise_key.c
+++ b/tests/src/psa_exercise_key.c
@@ -414,6 +414,21 @@
                                                   PSA_KEY_DERIVATION_INPUT_INFO,
                                                   input2,
                                                   input2_length));
+    } else if (PSA_ALG_IS_HKDF_EXTRACT(alg)) {
+        PSA_ASSERT(psa_key_derivation_input_bytes(operation,
+                                                  PSA_KEY_DERIVATION_INPUT_SALT,
+                                                  input1, input1_length));
+        PSA_ASSERT(psa_key_derivation_input_key(operation,
+                                                PSA_KEY_DERIVATION_INPUT_SECRET,
+                                                key));
+    } else if (PSA_ALG_IS_HKDF_EXPAND(alg)) {
+        PSA_ASSERT(psa_key_derivation_input_key(operation,
+                                                PSA_KEY_DERIVATION_INPUT_SECRET,
+                                                key));
+        PSA_ASSERT(psa_key_derivation_input_bytes(operation,
+                                                  PSA_KEY_DERIVATION_INPUT_INFO,
+                                                  input2,
+                                                  input2_length));
     } else if (PSA_ALG_IS_TLS12_PRF(alg) ||
                PSA_ALG_IS_TLS12_PSK_TO_MS(alg)) {
         PSA_ASSERT(psa_key_derivation_input_bytes(operation,
@@ -436,6 +451,10 @@
         PSA_ASSERT(psa_key_derivation_input_key(operation,
                                                 PSA_KEY_DERIVATION_INPUT_PASSWORD,
                                                 key));
+    } else if (alg == PSA_ALG_TLS12_ECJPAKE_TO_PMS) {
+        PSA_ASSERT(psa_key_derivation_input_bytes(operation,
+                                                  PSA_KEY_DERIVATION_INPUT_SECRET,
+                                                  input1, input1_length));
     } else {
         TEST_FAIL("Key derivation algorithm not supported");
     }
diff --git a/tests/suites/test_suite_psa_crypto.data b/tests/suites/test_suite_psa_crypto.data
index 3dd3866..0882e37 100644
--- a/tests/suites/test_suite_psa_crypto.data
+++ b/tests/suites/test_suite_psa_crypto.data
@@ -5456,19 +5456,19 @@
 derive_input:PSA_ALG_KEY_AGREEMENT(PSA_ALG_ECDH, PSA_ALG_HKDF(PSA_ALG_SHA_256)):PSA_KEY_DERIVATION_INPUT_SALT:PSA_KEY_TYPE_NONE:"":PSA_SUCCESS:PSA_KEY_DERIVATION_INPUT_SECRET:PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_SECP_R1):"c88f01f510d9ac3f70a292daa2316de544e9aab8afe84049c62a9c57862d1433":PSA_SUCCESS:0:UNUSED:"":UNUSED:PSA_KEY_TYPE_NONE:PSA_ERROR_BAD_STATE
 
 PSA key derivation: TLS12_ECJPAKE_TO_PMS, good input, output too short
-depends_on:PSA_WANT_ALG_TLS12_ECJPAKE_TO_PMS
+depends_on:PSA_WANT_ALG_TLS12_ECJPAKE_TO_PMS:PSA_WANT_ALG_SHA_256
 derive_input:PSA_ALG_TLS12_ECJPAKE_TO_PMS:PSA_KEY_DERIVATION_INPUT_SECRET:PSA_KEY_TYPE_NONE:"04aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa":PSA_SUCCESS:0:UNUSED:"":UNUSED:0:UNUSED:"":UNUSED:PSA_KEY_TYPE_NONE:PSA_ERROR_INVALID_ARGUMENT
 
 PSA key derivation: TLS12_ECJPAKE_TO_PMS, input[0]=0x02
-depends_on:PSA_WANT_ALG_TLS12_ECJPAKE_TO_PMS
+depends_on:PSA_WANT_ALG_TLS12_ECJPAKE_TO_PMS:PSA_WANT_ALG_SHA_256
 derive_input:PSA_ALG_TLS12_ECJPAKE_TO_PMS:PSA_KEY_DERIVATION_INPUT_SECRET:PSA_KEY_TYPE_NONE:"02aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa":PSA_ERROR_INVALID_ARGUMENT:0:UNUSED:"":UNUSED:0:UNUSED:"":UNUSED:PSA_KEY_TYPE_NONE:PSA_ERROR_BAD_STATE
 
 PSA key derivation: TLS12_ECJPAKE_TO_PMS, input too short
-depends_on:PSA_WANT_ALG_TLS12_ECJPAKE_TO_PMS
+depends_on:PSA_WANT_ALG_TLS12_ECJPAKE_TO_PMS:PSA_WANT_ALG_SHA_256
 derive_input:PSA_ALG_TLS12_ECJPAKE_TO_PMS:PSA_KEY_DERIVATION_INPUT_SECRET:PSA_KEY_TYPE_NONE:"04aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa":PSA_ERROR_INVALID_ARGUMENT:0:UNUSED:"":UNUSED:0:UNUSED:"":UNUSED:PSA_KEY_TYPE_NONE:PSA_ERROR_BAD_STATE
 
 PSA key derivation: TLS12_ECJPAKE_TO_PMS, input too long
-depends_on:PSA_WANT_ALG_TLS12_ECJPAKE_TO_PMS
+depends_on:PSA_WANT_ALG_TLS12_ECJPAKE_TO_PMS:PSA_WANT_ALG_SHA_256
 derive_input:PSA_ALG_TLS12_ECJPAKE_TO_PMS:PSA_KEY_DERIVATION_INPUT_SECRET:PSA_KEY_TYPE_NONE:"04aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa":PSA_ERROR_INVALID_ARGUMENT:0:UNUSED:"":UNUSED:0:UNUSED:"":UNUSED:PSA_KEY_TYPE_NONE:PSA_ERROR_BAD_STATE
 
 PSA key derivation: PBKDF2-HMAC-SHA256, good case, direct output
@@ -6324,6 +6324,53 @@
 depends_on:PSA_WANT_ALG_HKDF_EXPAND:PSA_WANT_ALG_SHA_1
 derive_set_capacity:PSA_ALG_HKDF_EXPAND(PSA_ALG_SHA_1):255 * PSA_HASH_LENGTH(PSA_ALG_SHA_1) + 1:PSA_ERROR_INVALID_ARGUMENT
 
+# TLS 1.2 PRF does not have a maximum capacity therefore
+# derive_set_capacity negative test case is not added
+
+PSA key derivation: TLS 1.2 PSK-to-MS SHA-256, request too much capacity
+depends_on:PSA_WANT_ALG_SHA_256:PSA_WANT_ALG_TLS12_PSK_TO_MS
+derive_set_capacity:PSA_ALG_TLS12_PSK_TO_MS(PSA_ALG_SHA_256):48U + 1U:PSA_ERROR_INVALID_ARGUMENT
+
+PSA key derivation: TLS 1.2 PSK-to-MS SHA-384, request too much capacity
+depends_on:PSA_WANT_ALG_SHA_384:PSA_WANT_ALG_TLS12_PSK_TO_MS
+derive_set_capacity:PSA_ALG_TLS12_PSK_TO_MS(PSA_ALG_SHA_384):48U + 1U:PSA_ERROR_INVALID_ARGUMENT
+
+PSA key derivation: TLS 1.2 ECJPAKE-to-PMS, request too much capacity
+depends_on:PSA_WANT_ALG_TLS12_ECJPAKE_TO_PMS:PSA_WANT_ALG_SHA_256
+derive_set_capacity:PSA_ALG_TLS12_ECJPAKE_TO_PMS:PSA_HASH_LENGTH(PSA_ALG_SHA_256) + 1:PSA_ERROR_INVALID_ARGUMENT
+
+PSA key derivation: PBKDF2-HMAC-SHA256, request too much capacity
+depends_on:PSA_WANT_ALG_PBKDF2_HMAC:PSA_WANT_ALG_SHA_256:SIZE_MAX>=0xffffffffffffffff
+derive_set_capacity:PSA_ALG_PBKDF2_HMAC(PSA_ALG_SHA_256):4294967295ULL * PSA_HASH_LENGTH(PSA_ALG_SHA_256) + 1:PSA_ERROR_INVALID_ARGUMENT
+
+PSA key derivation: PBKDF2-HMAC-SHA512, request too much capacity
+depends_on:PSA_WANT_ALG_PBKDF2_HMAC:PSA_WANT_ALG_SHA_512:SIZE_MAX>=0xffffffffffffffff
+derive_set_capacity:PSA_ALG_PBKDF2_HMAC(PSA_ALG_SHA_512):4294967295ULL * PSA_HASH_LENGTH(PSA_ALG_SHA_512) + 1:PSA_ERROR_INVALID_ARGUMENT
+
+PSA key derivation: PBKDF2-AES-CMAC-PRF-128, request too much capacity
+depends_on:PSA_WANT_ALG_PBKDF2_AES_CMAC_PRF_128:PSA_WANT_ALG_CMAC:PSA_WANT_KEY_TYPE_AES:SIZE_MAX>=0xffffffffffffffff
+derive_set_capacity:PSA_ALG_PBKDF2_AES_CMAC_PRF_128:4294967295ULL * 16 + 1:PSA_ERROR_INVALID_ARGUMENT
+
+PSA key derivation: TLS 1.2 PRF SHA-256, request maximum capacity
+depends_on:PSA_WANT_ALG_TLS12_PRF:PSA_WANT_ALG_SHA_256
+derive_set_capacity:PSA_ALG_TLS12_PRF(PSA_ALG_SHA_256):SIZE_MAX:PSA_SUCCESS
+
+PSA key derivation: TLS 1.2 PRF SHA-384, request maximum capacity
+depends_on:PSA_WANT_ALG_TLS12_PRF:PSA_WANT_ALG_SHA_384
+derive_set_capacity:PSA_ALG_TLS12_PRF(PSA_ALG_SHA_384):SIZE_MAX:PSA_SUCCESS
+
+PSA key derivation: PBKDF2-HMAC-SHA256, request maximum capacity
+depends_on:PSA_WANT_ALG_PBKDF2_HMAC:PSA_WANT_ALG_SHA_256:SIZE_MAX>=0xffffffffffffffff
+derive_set_capacity:PSA_ALG_PBKDF2_HMAC(PSA_ALG_SHA_256):4294967295ULL * PSA_HASH_LENGTH(PSA_ALG_SHA_256):PSA_SUCCESS
+
+PSA key derivation: PBKDF2-HMAC-SHA512, request maximum capacity
+depends_on:PSA_WANT_ALG_PBKDF2_HMAC:PSA_WANT_ALG_SHA_512:SIZE_MAX>=0xffffffffffffffff
+derive_set_capacity:PSA_ALG_PBKDF2_HMAC(PSA_ALG_SHA_512):4294967295ULL * PSA_HASH_LENGTH(PSA_ALG_SHA_512):PSA_SUCCESS
+
+PSA key derivation: PBKDF2-AES-CMAC-PRF-128, request maximum capacity
+depends_on:PSA_WANT_ALG_PBKDF2_AES_CMAC_PRF_128:PSA_WANT_ALG_CMAC:PSA_WANT_KEY_TYPE_AES:SIZE_MAX>=0xffffffffffffffff
+derive_set_capacity:PSA_ALG_PBKDF2_AES_CMAC_PRF_128:4294967295ULL * 16:PSA_SUCCESS
+
 PSA key derivation: over capacity 42: output 42+1
 depends_on:PSA_WANT_ALG_HKDF:PSA_WANT_ALG_SHA_256
 derive_output:PSA_ALG_HKDF(PSA_ALG_SHA_256):PSA_KEY_DERIVATION_INPUT_SALT:"000102030405060708090a0b0c":PSA_SUCCESS:PSA_KEY_DERIVATION_INPUT_SECRET:"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":PSA_SUCCESS:PSA_KEY_DERIVATION_INPUT_INFO:"f0f1f2f3f4f5f6f7f8f9":PSA_SUCCESS:0:"":PSA_SUCCESS:"":42:"3cb25f25faacd57a90434f64d0362f2a2d2d0a90cf1a5a4c5db02d56ecc4c5bf34007208d5b887185865":"ff":0:1:0
@@ -6520,17 +6567,69 @@
 depends_on:PSA_WANT_ALG_HKDF:PSA_WANT_ALG_SHA_256
 derive_full:PSA_ALG_HKDF(PSA_ALG_SHA_256):"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":"000102030405060708090a0b0c":"f0f1f2f3f4f5f6f7f8f9":255 * PSA_HASH_LENGTH(PSA_ALG_SHA_256) - 1
 
+PSA key derivation: HKDF SHA-512, read maximum capacity minus 1
+depends_on:PSA_WANT_ALG_HKDF:PSA_WANT_ALG_SHA_512
+derive_full:PSA_ALG_HKDF(PSA_ALG_SHA_512):"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":"000102030405060708090a0b0c":"f0f1f2f3f4f5f6f7f8f9":255 * PSA_HASH_LENGTH(PSA_ALG_SHA_512) - 1
+
 PSA key derivation: HKDF SHA-256, read maximum capacity
 depends_on:PSA_WANT_ALG_HKDF:PSA_WANT_ALG_SHA_256
 derive_full:PSA_ALG_HKDF(PSA_ALG_SHA_256):"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":"000102030405060708090a0b0c":"f0f1f2f3f4f5f6f7f8f9":255 * PSA_HASH_LENGTH(PSA_ALG_SHA_256)
 
-PSA key derivation: TLS 1.2 PRF SHA-256, read maximum capacity minus 1
-depends_on:PSA_WANT_ALG_SHA_256:PSA_WANT_ALG_TLS12_PRF
-derive_full:PSA_ALG_TLS12_PRF(PSA_ALG_SHA_256):"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":"000102030405060708090a0b0c":"f0f1f2f3f4f5f6f7f8f9":255 * PSA_HASH_LENGTH(PSA_ALG_SHA_256) - 1
+PSA key derivation: HKDF SHA-512, read maximum capacity
+depends_on:PSA_WANT_ALG_HKDF:PSA_WANT_ALG_SHA_512
+derive_full:PSA_ALG_HKDF(PSA_ALG_SHA_512):"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":"000102030405060708090a0b0c":"f0f1f2f3f4f5f6f7f8f9":255 * PSA_HASH_LENGTH(PSA_ALG_SHA_512)
 
-PSA key derivation: TLS 1.2 PRF SHA-256, read maximum capacity
-depends_on:PSA_WANT_ALG_SHA_256:PSA_WANT_ALG_TLS12_PRF
-derive_full:PSA_ALG_TLS12_PRF(PSA_ALG_SHA_256):"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":"000102030405060708090a0b0c":"f0f1f2f3f4f5f6f7f8f9":255 * PSA_HASH_LENGTH(PSA_ALG_SHA_256)
+PSA key derivation: HKDF-Extract SHA-256, read maximum capacity minus 1
+depends_on:PSA_WANT_ALG_HKDF_EXTRACT:PSA_WANT_ALG_SHA_256
+derive_full:PSA_ALG_HKDF_EXTRACT(PSA_ALG_SHA_256):"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":"000102030405060708090a0b0c":"":PSA_HASH_LENGTH(PSA_ALG_SHA_256) - 1
+
+PSA key derivation: HKDF-Extract SHA-512, read maximum capacity minus 1
+depends_on:PSA_WANT_ALG_HKDF_EXTRACT:PSA_WANT_ALG_SHA_512
+derive_full:PSA_ALG_HKDF_EXTRACT(PSA_ALG_SHA_512):"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":"000102030405060708090a0b0c":"":PSA_HASH_LENGTH(PSA_ALG_SHA_512) - 1
+
+PSA key derivation: HKDF-Extract SHA-256, read maximum capacity
+depends_on:PSA_WANT_ALG_HKDF_EXTRACT:PSA_WANT_ALG_SHA_256
+derive_full:PSA_ALG_HKDF_EXTRACT(PSA_ALG_SHA_256):"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":"000102030405060708090a0b0c":"":PSA_HASH_LENGTH(PSA_ALG_SHA_256)
+
+PSA key derivation: HKDF-Extract SHA-512, read maximum capacity
+depends_on:PSA_WANT_ALG_HKDF_EXTRACT:PSA_WANT_ALG_SHA_512
+derive_full:PSA_ALG_HKDF_EXTRACT(PSA_ALG_SHA_512):"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":"000102030405060708090a0b0c":"":PSA_HASH_LENGTH(PSA_ALG_SHA_512)
+
+PSA key derivation: HKDF-Expand SHA-256, read maximum capacity minus 1
+depends_on:PSA_WANT_ALG_HKDF_EXPAND:PSA_WANT_ALG_SHA_256
+derive_full:PSA_ALG_HKDF_EXPAND(PSA_ALG_SHA_256):"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":"":"f0f1f2f3f4f5f6f7f8f9":255 * PSA_HASH_LENGTH(PSA_ALG_SHA_256) - 1
+
+PSA key derivation: HKDF-Expand SHA-512, read maximum capacity minus 1
+depends_on:PSA_WANT_ALG_HKDF_EXPAND:PSA_WANT_ALG_SHA_512
+derive_full:PSA_ALG_HKDF_EXPAND(PSA_ALG_SHA_512):"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":"":"f0f1f2f3f4f5f6f7f8f9":255 * PSA_HASH_LENGTH(PSA_ALG_SHA_512) - 1
+
+PSA key derivation: HKDF-Expand SHA-256, read maximum capacity
+depends_on:PSA_WANT_ALG_HKDF_EXPAND:PSA_WANT_ALG_SHA_256
+derive_full:PSA_ALG_HKDF_EXPAND(PSA_ALG_SHA_256):"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":"":"f0f1f2f3f4f5f6f7f8f9":255 * PSA_HASH_LENGTH(PSA_ALG_SHA_256)
+
+PSA key derivation: HKDF-Expand SHA-512, read maximum capacity
+depends_on:PSA_WANT_ALG_HKDF_EXPAND:PSA_WANT_ALG_SHA_512
+derive_full:PSA_ALG_HKDF_EXPAND(PSA_ALG_SHA_512):"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":"":"f0f1f2f3f4f5f6f7f8f9":255 * PSA_HASH_LENGTH(PSA_ALG_SHA_512)
+
+PSA key derivation: TLS 1.2 PSK-to-MS SHA-256, read maximum capacity minus 1
+depends_on:PSA_WANT_ALG_SHA_256:PSA_WANT_ALG_TLS12_PSK_TO_MS
+derive_full:PSA_ALG_TLS12_PSK_TO_MS(PSA_ALG_SHA_256):"01020304":"5bc0b19b4a8b24b07afe7ec65c471e94a7d518fcef06c3574315255c52afe21b5bc0b19b872b9b26508458f03603744d575f463a11ae7f1b090c012606fd3e9f":"6d617374657220736563726574":47
+
+PSA key derivation: TLS 1.2 PSK-to-MS SHA-384, read maximum capacity minus 1
+depends_on:PSA_WANT_ALG_SHA_384:PSA_WANT_ALG_TLS12_PSK_TO_MS
+derive_full:PSA_ALG_TLS12_PSK_TO_MS(PSA_ALG_SHA_384):"01020304":"5bc0b19b4a8b24b07afe7ec65c471e94a7d518fcef06c3574315255c52afe21b5bc0b19b872b9b26508458f03603744d575f463a11ae7f1b090c012606fd3e9f":"6d617374657220736563726574":47
+
+PSA key derivation: TLS 1.2 PSK-to-MS SHA-256, read maximum capacity
+depends_on:PSA_WANT_ALG_SHA_256:PSA_WANT_ALG_TLS12_PSK_TO_MS
+derive_full:PSA_ALG_TLS12_PSK_TO_MS(PSA_ALG_SHA_256):"01020304":"5bc0b19b4a8b24b07afe7ec65c471e94a7d518fcef06c3574315255c52afe21b5bc0b19b872b9b26508458f03603744d575f463a11ae7f1b090c012606fd3e9f":"6d617374657220736563726574":48
+
+PSA key derivation: TLS 1.2 PSK-to-MS SHA-384, read maximum capacity
+depends_on:PSA_WANT_ALG_SHA_384:PSA_WANT_ALG_TLS12_PSK_TO_MS
+derive_full:PSA_ALG_TLS12_PSK_TO_MS(PSA_ALG_SHA_384):"01020304":"5bc0b19b4a8b24b07afe7ec65c471e94a7d518fcef06c3574315255c52afe21b5bc0b19b872b9b26508458f03603744d575f463a11ae7f1b090c012606fd3e9f":"6d617374657220736563726574":48
+
+PSA key derivation: TLS 1.2 ECJPAKE-to-PMS, read maximum capacity
+depends_on:PSA_WANT_ALG_TLS12_ECJPAKE_TO_PMS:PSA_WANT_ALG_SHA_256
+derive_full:PSA_ALG_TLS12_ECJPAKE_TO_PMS:"deadbeef":"0409fc1accc230a205e4a208e64a8f204291f581a12756392da4b8c0cf5ef02b950000000000000000000000000000000000000000000000000000000000000000":"":PSA_HASH_LENGTH(PSA_ALG_SHA_256)
 
 PSA key derivation: HKDF SHA-256, exercise AES128-CTR
 depends_on:PSA_WANT_ALG_CTR:PSA_WANT_ALG_HKDF:PSA_WANT_ALG_SHA_256:PSA_WANT_KEY_TYPE_AES
diff --git a/tests/suites/test_suite_psa_crypto.function b/tests/suites/test_suite_psa_crypto.function
index 889a3b0..95e7a2d 100644
--- a/tests/suites/test_suite_psa_crypto.function
+++ b/tests/suites/test_suite_psa_crypto.function
@@ -8459,7 +8459,7 @@
 /* END_CASE */
 
 /* BEGIN_CASE */
-void derive_set_capacity(int alg_arg, int capacity_arg,
+void derive_set_capacity(int alg_arg, int64_t capacity_arg,
                          int expected_status_arg)
 {
     psa_algorithm_t alg = alg_arg;
@@ -8940,7 +8940,7 @@
     psa_algorithm_t alg = alg_arg;
     size_t requested_capacity = requested_capacity_arg;
     psa_key_derivation_operation_t operation = PSA_KEY_DERIVATION_OPERATION_INIT;
-    unsigned char output_buffer[16];
+    unsigned char output_buffer[32];
     size_t expected_capacity = requested_capacity;
     size_t current_capacity;
     psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;