Merge pull request #4243 from bensze01/psa_vararg

PSA: Update AEAD output buffer macros to PSA API version 1.0
diff --git a/ChangeLog.d/psa-aead-output-size-macros-1.0.txt b/ChangeLog.d/psa-aead-output-size-macros-1.0.txt
new file mode 100644
index 0000000..94a66a5
--- /dev/null
+++ b/ChangeLog.d/psa-aead-output-size-macros-1.0.txt
@@ -0,0 +1,12 @@
+API changes
+   * Update AEAD output size macros to bring them in line with the PSA Crypto
+     API version 1.0 spec. This version of the spec parameterizes them on the
+     key type used, as well as the key bit-size in the case of
+     PSA_AEAD_TAG_LENGTH.
+     The old versions of these macros were renamed and deprecated as follows:
+     - PSA_AEAD_TAG_LENGTH          -> PSA_AEAD_TAG_LENGTH_1_ARG
+     - PSA_AEAD_ENCRYPT_OUTPUT_SIZE -> PSA_AEAD_ENCRYPT_OUTPUT_SIZE_2_ARG
+     - PSA_AEAD_DECRYPT_OUTPUT_SIZE -> PSA_AEAD_DECRYPT_OUTPUT_SIZE_2_ARG
+     - PSA_AEAD_UPDATE_OUTPUT_SIZE  -> PSA_AEAD_UPDATE_OUTPUT_SIZE_2_ARG
+     - PSA_AEAD_FINISH_OUTPUT_SIZE  -> PSA_AEAD_FINISH_OUTPUT_SIZE_1_ARG
+     - PSA_AEAD_VERIFY_OUTPUT_SIZE  -> PSA_AEAD_VERIFY_OUTPUT_SIZE_1_ARG
diff --git a/include/psa/crypto.h b/include/psa/crypto.h
index 81e1f28..94b8f99 100644
--- a/include/psa/crypto.h
+++ b/include/psa/crypto.h
@@ -2113,9 +2113,16 @@
  *                                authentication tag is appended to the
  *                                encrypted data.
  * \param ciphertext_size         Size of the \p ciphertext buffer in bytes.
- *                                This must be at least
- *                                #PSA_AEAD_ENCRYPT_OUTPUT_SIZE(\p alg,
- *                                \p plaintext_length).
+ *                                This must be appropriate for the selected
+ *                                algorithm and key:
+ *                                - A sufficient output size is
+ *                                  #PSA_AEAD_ENCRYPT_OUTPUT_SIZE(\c key_type,
+ *                                  \p alg, \p plaintext_length) where
+ *                                  \c key_type is the type of \p key.
+ *                                - #PSA_AEAD_ENCRYPT_OUTPUT_MAX_SIZE(\p
+ *                                  plaintext_length) evaluates to the maximum
+ *                                  ciphertext size of any supported AEAD
+ *                                  encryption.
  * \param[out] ciphertext_length  On success, the size of the output
  *                                in the \p ciphertext buffer.
  *
@@ -2129,7 +2136,11 @@
  *         \p alg is not supported or is not an AEAD algorithm.
  * \retval #PSA_ERROR_INSUFFICIENT_MEMORY
  * \retval #PSA_ERROR_BUFFER_TOO_SMALL
- *         \p ciphertext_size is too small
+ *         \p ciphertext_size is too small.
+ *         #PSA_AEAD_ENCRYPT_OUTPUT_SIZE(\c key_type, \p alg,
+ *         \p plaintext_length) or
+ *         #PSA_AEAD_ENCRYPT_OUTPUT_MAX_SIZE(\p plaintext_length) can be used to
+ *         determine the required buffer size.
  * \retval #PSA_ERROR_COMMUNICATION_FAILURE
  * \retval #PSA_ERROR_HARDWARE_FAILURE
  * \retval #PSA_ERROR_CORRUPTION_DETECTED
@@ -2173,9 +2184,16 @@
  * \param ciphertext_length       Size of \p ciphertext in bytes.
  * \param[out] plaintext          Output buffer for the decrypted data.
  * \param plaintext_size          Size of the \p plaintext buffer in bytes.
- *                                This must be at least
- *                                #PSA_AEAD_DECRYPT_OUTPUT_SIZE(\p alg,
- *                                \p ciphertext_length).
+ *                                This must be appropriate for the selected
+ *                                algorithm and key:
+ *                                - A sufficient output size is
+ *                                  #PSA_AEAD_DECRYPT_OUTPUT_SIZE(\c key_type,
+ *                                  \p alg, \p ciphertext_length) where
+ *                                  \c key_type is the type of \p key.
+ *                                - #PSA_AEAD_DECRYPT_OUTPUT_MAX_SIZE(\p
+ *                                  ciphertext_length) evaluates to the maximum
+ *                                  plaintext size of any supported AEAD
+ *                                  decryption.
  * \param[out] plaintext_length   On success, the size of the output
  *                                in the \p plaintext buffer.
  *
@@ -2191,7 +2209,11 @@
  *         \p alg is not supported or is not an AEAD algorithm.
  * \retval #PSA_ERROR_INSUFFICIENT_MEMORY
  * \retval #PSA_ERROR_BUFFER_TOO_SMALL
- *         \p plaintext_size or \p nonce_length is too small
+ *         \p plaintext_size is too small.
+ *         #PSA_AEAD_DECRYPT_OUTPUT_SIZE(\c key_type, \p alg,
+ *         \p ciphertext_length) or
+ *         #PSA_AEAD_DECRYPT_OUTPUT_MAX_SIZE(\p ciphertext_length) can be used
+ *         to determine the required buffer size.
  * \retval #PSA_ERROR_COMMUNICATION_FAILURE
  * \retval #PSA_ERROR_HARDWARE_FAILURE
  * \retval #PSA_ERROR_CORRUPTION_DETECTED
@@ -2612,10 +2634,18 @@
  * \param input_length          Size of the \p input buffer in bytes.
  * \param[out] output           Buffer where the output is to be written.
  * \param output_size           Size of the \p output buffer in bytes.
- *                              This must be at least
- *                              #PSA_AEAD_UPDATE_OUTPUT_SIZE(\c alg,
- *                              \p input_length) where \c alg is the
- *                              algorithm that is being calculated.
+ *                              This must be appropriate for the selected
+ *                                algorithm and key:
+ *                                - A sufficient output size is
+ *                                  #PSA_AEAD_UPDATE_OUTPUT_SIZE(\c key_type,
+ *                                  \c alg, \p input_length) where
+ *                                  \c key_type is the type of key and \c alg is
+ *                                  the algorithm that were used to set up the
+ *                                  operation.
+ *                                - #PSA_AEAD_UPDATE_OUTPUT_MAX_SIZE(\p
+ *                                  input_length) evaluates to the maximum
+ *                                  output size of any supported AEAD
+ *                                  algorithm.
  * \param[out] output_length    On success, the number of bytes
  *                              that make up the returned output.
  *
@@ -2626,9 +2656,9 @@
  *         set, and have lengths set if required by the algorithm).
  * \retval #PSA_ERROR_BUFFER_TOO_SMALL
  *         The size of the \p output buffer is too small.
- *         You can determine a sufficient buffer size by calling
- *         #PSA_AEAD_UPDATE_OUTPUT_SIZE(\c alg, \p input_length)
- *         where \c alg is the algorithm that is being calculated.
+ *         #PSA_AEAD_UPDATE_OUTPUT_SIZE(\c key_type, \c alg, \p input_length) or
+ *         #PSA_AEAD_UPDATE_OUTPUT_MAX_SIZE(\p input_length) can be used to
+ *         determine the required buffer size.
  * \retval #PSA_ERROR_INVALID_ARGUMENT
  *         The total length of input to psa_aead_update_ad() so far is
  *         less than the additional data length that was previously
@@ -2665,9 +2695,7 @@
  * This function has two output buffers:
  * - \p ciphertext contains trailing ciphertext that was buffered from
  *   preceding calls to psa_aead_update().
- * - \p tag contains the authentication tag. Its length is always
- *   #PSA_AEAD_TAG_LENGTH(\c alg) where \c alg is the AEAD algorithm
- *   that the operation performs.
+ * - \p tag contains the authentication tag.
  *
  * When this function returns successfuly, the operation becomes inactive.
  * If this function returns an error status, the operation enters an error
@@ -2677,18 +2705,32 @@
  * \param[out] ciphertext       Buffer where the last part of the ciphertext
  *                              is to be written.
  * \param ciphertext_size       Size of the \p ciphertext buffer in bytes.
- *                              This must be at least
- *                              #PSA_AEAD_FINISH_OUTPUT_SIZE(\c alg) where
- *                              \c alg is the algorithm that is being
- *                              calculated.
+ *                              This must be appropriate for the selected
+ *                              algorithm and key:
+ *                              - A sufficient output size is
+ *                                #PSA_AEAD_FINISH_OUTPUT_SIZE(\c key_type,
+ *                                \c alg) where \c key_type is the type of key
+ *                                and \c alg is the algorithm that were used to
+ *                                set up the operation.
+ *                              - #PSA_AEAD_FINISH_OUTPUT_MAX_SIZE evaluates to
+ *                                the maximum output size of any supported AEAD
+ *                                algorithm.
  * \param[out] ciphertext_length On success, the number of bytes of
  *                              returned ciphertext.
  * \param[out] tag              Buffer where the authentication tag is
  *                              to be written.
  * \param tag_size              Size of the \p tag buffer in bytes.
- *                              This must be at least
- *                              #PSA_AEAD_TAG_LENGTH(\c alg) where \c alg is
- *                              the algorithm that is being calculated.
+ *                              This must be appropriate for the selected
+ *                              algorithm and key:
+ *                              - The exact tag size is #PSA_AEAD_TAG_LENGTH(\c
+ *                                key_type, \c key_bits, \c alg) where
+ *                                \c key_type and \c key_bits are the type and
+ *                                bit-size of the key, and \c alg is the
+ *                                algorithm that were used in the call to
+ *                                psa_aead_encrypt_setup().
+ *                              - #PSA_AEAD_TAG_MAX_SIZE evaluates to the
+ *                                maximum tag size of any supported AEAD
+ *                                algorithm.
  * \param[out] tag_length       On success, the number of bytes
  *                              that make up the returned tag.
  *
@@ -2699,11 +2741,11 @@
  *         operation with a nonce set).
  * \retval #PSA_ERROR_BUFFER_TOO_SMALL
  *         The size of the \p ciphertext or \p tag buffer is too small.
- *         You can determine a sufficient buffer size for \p ciphertext by
- *         calling #PSA_AEAD_FINISH_OUTPUT_SIZE(\c alg)
- *         where \c alg is the algorithm that is being calculated.
- *         You can determine a sufficient buffer size for \p tag by
- *         calling #PSA_AEAD_TAG_LENGTH(\c alg).
+ *         #PSA_AEAD_FINISH_OUTPUT_SIZE(\c key_type, \c alg) or
+ *         #PSA_AEAD_FINISH_OUTPUT_MAX_SIZE can be used to determine the
+ *         required \p ciphertext buffer size. #PSA_AEAD_TAG_LENGTH(\c key_type,
+ *         \c key_bits, \c alg) or #PSA_AEAD_TAG_MAX_SIZE can be used to
+ *         determine the required \p tag buffer size.
  * \retval #PSA_ERROR_INVALID_ARGUMENT
  *         The total length of input to psa_aead_update_ad() so far is
  *         less than the additional data length that was previously
@@ -2762,10 +2804,15 @@
  *                              that could not be processed until the end
  *                              of the input.
  * \param plaintext_size        Size of the \p plaintext buffer in bytes.
- *                              This must be at least
- *                              #PSA_AEAD_VERIFY_OUTPUT_SIZE(\c alg) where
- *                              \c alg is the algorithm that is being
- *                              calculated.
+ *                              This must be appropriate for the selected algorithm and key:
+ *                              - A sufficient output size is
+ *                                #PSA_AEAD_VERIFY_OUTPUT_SIZE(\c key_type,
+ *                                \c alg) where \c key_type is the type of key
+ *                                and \c alg is the algorithm that were used to
+ *                                set up the operation.
+ *                              - #PSA_AEAD_VERIFY_OUTPUT_MAX_SIZE evaluates to
+ *                                the maximum output size of any supported AEAD
+ *                                algorithm.
  * \param[out] plaintext_length On success, the number of bytes of
  *                              returned plaintext.
  * \param[in] tag               Buffer containing the authentication tag.
@@ -2781,9 +2828,9 @@
  *         operation with a nonce set).
  * \retval #PSA_ERROR_BUFFER_TOO_SMALL
  *         The size of the \p plaintext buffer is too small.
- *         You can determine a sufficient buffer size for \p plaintext by
- *         calling #PSA_AEAD_VERIFY_OUTPUT_SIZE(\c alg)
- *         where \c alg is the algorithm that is being calculated.
+ *         #PSA_AEAD_VERIFY_OUTPUT_SIZE(\c key_type, \c alg) or
+ *         #PSA_AEAD_VERIFY_OUTPUT_MAX_SIZE can be used to determine the
+ *         required buffer size.
  * \retval #PSA_ERROR_INVALID_ARGUMENT
  *         The total length of input to psa_aead_update_ad() so far is
  *         less than the additional data length that was previously
diff --git a/include/psa/crypto_compat.h b/include/psa/crypto_compat.h
index ae09a70..5dabbd2 100644
--- a/include/psa/crypto_compat.h
+++ b/include/psa/crypto_compat.h
@@ -269,6 +269,153 @@
 #define PSA_ALG_AEAD_WITH_TAG_LENGTH( aead_alg, tag_length ) \
     MBEDTLS_DEPRECATED_CONSTANT( psa_algorithm_t, PSA_ALG_AEAD_WITH_SHORTENED_TAG( aead_alg, tag_length ) )
 
+/*
+ * Deprecated PSA AEAD output size macros (PSA Crypto API  <= 1.0 beta3)
+ */
+
+/** The tag size for an AEAD algorithm, in bytes.
+ *
+ * \param alg                 An AEAD algorithm
+ *                            (\c PSA_ALG_XXX value such that
+ *                            #PSA_ALG_IS_AEAD(\p alg) is true).
+ *
+ * \return                    The tag size for the specified algorithm.
+ *                            If the AEAD algorithm does not have an identified
+ *                            tag that can be distinguished from the rest of
+ *                            the ciphertext, return 0.
+ *                            If the AEAD algorithm is not recognized, return 0.
+ */
+#define PSA_AEAD_TAG_LENGTH_1_ARG( alg )     \
+    MBEDTLS_DEPRECATED_CONSTANT( size_t,     \
+        PSA_ALG_IS_AEAD( alg ) ?             \
+        PSA_ALG_AEAD_GET_TAG_LENGTH( alg ) : \
+        0 )
+
+/** The maximum size of the output of psa_aead_encrypt(), in bytes.
+ *
+ * If the size of the ciphertext buffer is at least this large, it is
+ * guaranteed that psa_aead_encrypt() will not fail due to an
+ * insufficient buffer size. Depending on the algorithm, the actual size of
+ * the ciphertext may be smaller.
+ *
+ * \warning This macro may evaluate its arguments multiple times or
+ *          zero times, so you should not pass arguments that contain
+ *          side effects.
+ *
+ * \param alg                 An AEAD algorithm
+ *                            (\c PSA_ALG_XXX value such that
+ *                            #PSA_ALG_IS_AEAD(\p alg) is true).
+ * \param plaintext_length    Size of the plaintext in bytes.
+ *
+ * \return                    The AEAD ciphertext size for the specified
+ *                            algorithm.
+ *                            If the AEAD algorithm is not recognized, return 0.
+ */
+#define PSA_AEAD_ENCRYPT_OUTPUT_SIZE_2_ARG( alg, plaintext_length ) \
+    MBEDTLS_DEPRECATED_CONSTANT( size_t,                            \
+        PSA_ALG_IS_AEAD( alg ) ?                                    \
+        (plaintext_length) + PSA_ALG_AEAD_GET_TAG_LENGTH( alg ) :   \
+        0 )
+
+/** The maximum size of the output of psa_aead_decrypt(), in bytes.
+ *
+ * If the size of the plaintext buffer is at least this large, it is
+ * guaranteed that psa_aead_decrypt() will not fail due to an
+ * insufficient buffer size. Depending on the algorithm, the actual size of
+ * the plaintext may be smaller.
+ *
+ * \warning This macro may evaluate its arguments multiple times or
+ *          zero times, so you should not pass arguments that contain
+ *          side effects.
+ *
+ * \param alg                 An AEAD algorithm
+ *                            (\c PSA_ALG_XXX value such that
+ *                            #PSA_ALG_IS_AEAD(\p alg) is true).
+ * \param ciphertext_length   Size of the plaintext in bytes.
+ *
+ * \return                    The AEAD ciphertext size for the specified
+ *                            algorithm.
+ *                            If the AEAD algorithm is not recognized, return 0.
+ */
+#define PSA_AEAD_DECRYPT_OUTPUT_SIZE_2_ARG( alg, ciphertext_length )   \
+    MBEDTLS_DEPRECATED_CONSTANT( size_t,                               \
+        PSA_ALG_IS_AEAD( alg ) &&                                      \
+            (ciphertext_length) > PSA_ALG_AEAD_GET_TAG_LENGTH( alg ) ? \
+            (ciphertext_length) - PSA_ALG_AEAD_GET_TAG_LENGTH( alg ) : \
+        0 )
+
+/** A sufficient output buffer size for psa_aead_update().
+ *
+ * If the size of the output buffer is at least this large, it is
+ * guaranteed that psa_aead_update() will not fail due to an
+ * insufficient buffer size. The actual size of the output may be smaller
+ * in any given call.
+ *
+ * \warning This macro may evaluate its arguments multiple times or
+ *          zero times, so you should not pass arguments that contain
+ *          side effects.
+ *
+ * \param alg                 An AEAD algorithm
+ *                            (\c PSA_ALG_XXX value such that
+ *                            #PSA_ALG_IS_AEAD(\p alg) is true).
+ * \param input_length        Size of the input in bytes.
+ *
+ * \return                    A sufficient output buffer size for the specified
+ *                            algorithm.
+ *                            If the AEAD algorithm is not recognized, return 0.
+ */
+/* For all the AEAD modes defined in this specification, it is possible
+ * to emit output without delay. However, hardware may not always be
+ * capable of this. So for modes based on a block cipher, allow the
+ * implementation to delay the output until it has a full block. */
+#define PSA_AEAD_UPDATE_OUTPUT_SIZE_2_ARG( alg, input_length )                        \
+    MBEDTLS_DEPRECATED_CONSTANT( size_t,                                              \
+        PSA_ALG_IS_AEAD_ON_BLOCK_CIPHER( alg ) ?                                      \
+        PSA_ROUND_UP_TO_MULTIPLE( PSA_BLOCK_CIPHER_BLOCK_MAX_SIZE, (input_length) ) : \
+        (input_length) )
+
+/** A sufficient ciphertext buffer size for psa_aead_finish().
+ *
+ * If the size of the ciphertext buffer is at least this large, it is
+ * guaranteed that psa_aead_finish() will not fail due to an
+ * insufficient ciphertext buffer size. The actual size of the output may
+ * be smaller in any given call.
+ *
+ * \param alg                 An AEAD algorithm
+ *                            (\c PSA_ALG_XXX value such that
+ *                            #PSA_ALG_IS_AEAD(\p alg) is true).
+ *
+ * \return                    A sufficient ciphertext buffer size for the
+ *                            specified algorithm.
+ *                            If the AEAD algorithm is not recognized, return 0.
+ */
+#define PSA_AEAD_FINISH_OUTPUT_SIZE_1_ARG( alg )                        \
+    MBEDTLS_DEPRECATED_CONSTANT( size_t,                                \
+        PSA_ALG_IS_AEAD_ON_BLOCK_CIPHER( alg ) ?                        \
+        PSA_BLOCK_CIPHER_BLOCK_MAX_SIZE :                               \
+        0 )
+
+/** A sufficient plaintext buffer size for psa_aead_verify().
+ *
+ * If the size of the plaintext buffer is at least this large, it is
+ * guaranteed that psa_aead_verify() will not fail due to an
+ * insufficient plaintext buffer size. The actual size of the output may
+ * be smaller in any given call.
+ *
+ * \param alg                 An AEAD algorithm
+ *                            (\c PSA_ALG_XXX value such that
+ *                            #PSA_ALG_IS_AEAD(\p alg) is true).
+ *
+ * \return                    A sufficient plaintext buffer size for the
+ *                            specified algorithm.
+ *                            If the AEAD algorithm is not recognized, return 0.
+ */
+#define PSA_AEAD_VERIFY_OUTPUT_SIZE_1_ARG( alg )                        \
+    MBEDTLS_DEPRECATED_CONSTANT( size_t,                                \
+        PSA_ALG_IS_AEAD_ON_BLOCK_CIPHER( alg ) ?                        \
+        PSA_BLOCK_CIPHER_BLOCK_MAX_SIZE :                               \
+        0 )
+
 #endif /* MBEDTLS_DEPRECATED_REMOVED */
 
 /** Open a handle to an existing persistent key.
diff --git a/include/psa/crypto_sizes.h b/include/psa/crypto_sizes.h
index c9de062..79f9673 100644
--- a/include/psa/crypto_sizes.h
+++ b/include/psa/crypto_sizes.h
@@ -117,26 +117,35 @@
  */
 #define PSA_MAC_MAX_SIZE PSA_HASH_MAX_SIZE
 
-/** The tag size for an AEAD algorithm, in bytes.
+/** The length of a tag for an AEAD algorithm, in bytes.
  *
+ * This macro can be used to allocate a buffer of sufficient size to store the
+ * tag output from psa_aead_finish().
+ *
+ * See also #PSA_AEAD_TAG_MAX_SIZE.
+ *
+ * \param key_type            The type of the AEAD key.
+ * \param key_bits            The size of the AEAD key in bits.
  * \param alg                 An AEAD algorithm
  *                            (\c PSA_ALG_XXX value such that
  *                            #PSA_ALG_IS_AEAD(\p alg) is true).
  *
- * \return                    The tag size for the specified algorithm.
+ * \return                    The tag length for the specified algorithm and key.
  *                            If the AEAD algorithm does not have an identified
  *                            tag that can be distinguished from the rest of
  *                            the ciphertext, return 0.
- *                            If the AEAD algorithm is not recognized, return 0.
+ *                            If the key type or AEAD algorithm is not
+ *                            recognized, or the parameters are incompatible,
+ *                            return 0.
  */
-#define PSA_AEAD_TAG_LENGTH(alg)                                        \
-    (PSA_ALG_IS_AEAD(alg) ?                                             \
-     (((alg) & PSA_ALG_AEAD_TAG_LENGTH_MASK) >> PSA_AEAD_TAG_LENGTH_OFFSET) : \
-     0)
+#define PSA_AEAD_TAG_LENGTH(key_type, key_bits, alg)                        \
+    (PSA_AEAD_NONCE_LENGTH(key_type, alg) != 0 ?                            \
+     PSA_ALG_AEAD_GET_TAG_LENGTH(alg) :                                     \
+     ((void) (key_bits), 0))
 
 /** The maximum tag size for all supported AEAD algorithms, in bytes.
  *
- * See also #PSA_AEAD_TAG_LENGTH(\p alg).
+ * See also #PSA_AEAD_TAG_LENGTH(\p key_type, \p key_bits, \p alg).
  */
 #define PSA_AEAD_TAG_MAX_SIZE       16
 
@@ -241,10 +250,14 @@
  * insufficient buffer size. Depending on the algorithm, the actual size of
  * the ciphertext may be smaller.
  *
+ * See also #PSA_AEAD_ENCRYPT_OUTPUT_MAX_SIZE(\p plaintext_length).
+ *
  * \warning This macro may evaluate its arguments multiple times or
  *          zero times, so you should not pass arguments that contain
  *          side effects.
  *
+ * \param key_type            A symmetric key type that is
+ *                            compatible with algorithm \p alg.
  * \param alg                 An AEAD algorithm
  *                            (\c PSA_ALG_XXX value such that
  *                            #PSA_ALG_IS_AEAD(\p alg) is true).
@@ -252,11 +265,13 @@
  *
  * \return                    The AEAD ciphertext size for the specified
  *                            algorithm.
- *                            If the AEAD algorithm is not recognized, return 0.
+ *                            If the key type or AEAD algorithm is not
+ *                            recognized, or the parameters are incompatible,
+ *                            return 0.
  */
-#define PSA_AEAD_ENCRYPT_OUTPUT_SIZE(alg, plaintext_length)       \
-    (PSA_AEAD_TAG_LENGTH(alg) != 0 ?                              \
-     (plaintext_length) + PSA_AEAD_TAG_LENGTH(alg) :              \
+#define PSA_AEAD_ENCRYPT_OUTPUT_SIZE(key_type, alg, plaintext_length) \
+    (PSA_AEAD_NONCE_LENGTH(key_type, alg) != 0 ?                      \
+     (plaintext_length) + PSA_ALG_AEAD_GET_TAG_LENGTH(alg) :          \
      0)
 
 /** A sufficient output buffer size for psa_aead_encrypt(), for any of the
@@ -268,7 +283,8 @@
  * \note This macro returns a compile-time constant if its arguments are
  *       compile-time constants.
  *
- * See also #PSA_AEAD_ENCRYPT_OUTPUT_SIZE(\p alg, \p plaintext_length).
+ * See also #PSA_AEAD_ENCRYPT_OUTPUT_SIZE(\p key_type, \p alg,
+ * \p plaintext_length).
  *
  * \param plaintext_length    Size of the plaintext in bytes.
  *
@@ -287,10 +303,14 @@
  * insufficient buffer size. Depending on the algorithm, the actual size of
  * the plaintext may be smaller.
  *
+ * See also #PSA_AEAD_DECRYPT_OUTPUT_MAX_SIZE(\p ciphertext_length).
+ *
  * \warning This macro may evaluate its arguments multiple times or
  *          zero times, so you should not pass arguments that contain
  *          side effects.
  *
+ * \param key_type            A symmetric key type that is
+ *                            compatible with algorithm \p alg.
  * \param alg                 An AEAD algorithm
  *                            (\c PSA_ALG_XXX value such that
  *                            #PSA_ALG_IS_AEAD(\p alg) is true).
@@ -298,11 +318,14 @@
  *
  * \return                    The AEAD ciphertext size for the specified
  *                            algorithm.
- *                            If the AEAD algorithm is not recognized, return 0.
+ *                            If the key type or AEAD algorithm is not
+ *                            recognized, or the parameters are incompatible,
+ *                            return 0.
  */
-#define PSA_AEAD_DECRYPT_OUTPUT_SIZE(alg, ciphertext_length)      \
-    (PSA_AEAD_TAG_LENGTH(alg) != 0 ?                              \
-     (ciphertext_length) - PSA_AEAD_TAG_LENGTH(alg) :             \
+#define PSA_AEAD_DECRYPT_OUTPUT_SIZE(key_type, alg, ciphertext_length) \
+    (PSA_AEAD_NONCE_LENGTH(key_type, alg) != 0 &&                      \
+         (ciphertext_length) > PSA_ALG_AEAD_GET_TAG_LENGTH(alg) ?      \
+         (ciphertext_length) - PSA_ALG_AEAD_GET_TAG_LENGTH(alg) :      \
      0)
 
 /** A sufficient output buffer size for psa_aead_decrypt(), for any of the
@@ -314,7 +337,8 @@
  * \note This macro returns a compile-time constant if its arguments are
  *       compile-time constants.
  *
- * See also #PSA_AEAD_DECRYPT_OUTPUT_SIZE(\p alg, \p ciphertext_length).
+ * See also #PSA_AEAD_DECRYPT_OUTPUT_SIZE(\p key_type, \p alg,
+ * \p ciphertext_length).
  *
  * \param ciphertext_length   Size of the ciphertext in bytes.
  *
@@ -352,11 +376,11 @@
  */
 #define PSA_AEAD_NONCE_LENGTH(key_type, alg) \
     (PSA_BLOCK_CIPHER_BLOCK_LENGTH(key_type) == 16 ? \
-          PSA_ALG_AEAD_WITH_DEFAULT_LENGTH_TAG(alg) == PSA_ALG_CCM ? 13 : \
-          PSA_ALG_AEAD_WITH_DEFAULT_LENGTH_TAG(alg) == PSA_ALG_GCM ? 12 : \
+          MBEDTLS_PSA_ALG_AEAD_EQUAL(alg, PSA_ALG_CCM) ? 13 : \
+          MBEDTLS_PSA_ALG_AEAD_EQUAL(alg, PSA_ALG_GCM) ? 12 : \
           0 : \
      (key_type) == PSA_KEY_TYPE_CHACHA20 && \
-          PSA_ALG_AEAD_WITH_DEFAULT_LENGTH_TAG(alg) == PSA_ALG_CHACHA20_POLY1305 ? 12 : \
+          MBEDTLS_PSA_ALG_AEAD_EQUAL(alg, PSA_ALG_CHACHA20_POLY1305) ? 12 : \
      0)
 
 /** The maximum default nonce size among all supported pairs of key types and
@@ -379,10 +403,14 @@
  * insufficient buffer size. The actual size of the output may be smaller
  * in any given call.
  *
+ * See also #PSA_AEAD_UPDATE_OUTPUT_MAX_SIZE(\p input_length).
+ *
  * \warning This macro may evaluate its arguments multiple times or
  *          zero times, so you should not pass arguments that contain
  *          side effects.
  *
+ * \param key_type            A symmetric key type that is
+ *                            compatible with algorithm \p alg.
  * \param alg                 An AEAD algorithm
  *                            (\c PSA_ALG_XXX value such that
  *                            #PSA_ALG_IS_AEAD(\p alg) is true).
@@ -390,16 +418,20 @@
  *
  * \return                    A sufficient output buffer size for the specified
  *                            algorithm.
- *                            If the AEAD algorithm is not recognized, return 0.
+ *                            If the key type or AEAD algorithm is not
+ *                            recognized, or the parameters are incompatible,
+ *                            return 0.
  */
 /* For all the AEAD modes defined in this specification, it is possible
  * to emit output without delay. However, hardware may not always be
  * capable of this. So for modes based on a block cipher, allow the
  * implementation to delay the output until it has a full block. */
-#define PSA_AEAD_UPDATE_OUTPUT_SIZE(alg, input_length)                              \
-    (PSA_ALG_IS_AEAD_ON_BLOCK_CIPHER(alg) ?                                         \
-     PSA_ROUND_UP_TO_MULTIPLE(PSA_BLOCK_CIPHER_BLOCK_MAX_SIZE, (input_length)) :    \
-     (input_length))
+#define PSA_AEAD_UPDATE_OUTPUT_SIZE(key_type, alg, input_length)                             \
+    (PSA_AEAD_NONCE_LENGTH(key_type, alg) != 0 ?                                             \
+         PSA_ALG_IS_AEAD_ON_BLOCK_CIPHER(alg) ?                                              \
+         PSA_ROUND_UP_TO_MULTIPLE(PSA_BLOCK_CIPHER_BLOCK_LENGTH(key_type), (input_length)) : \
+         (input_length) : \
+     0)
 
 /** A sufficient output buffer size for psa_aead_update(), for any of the
  *  supported key types and AEAD algorithms.
@@ -407,7 +439,7 @@
  * If the size of the output buffer is at least this large, it is guaranteed
  * that psa_aead_update() will not fail due to an insufficient buffer size.
  *
- * See also #PSA_AEAD_UPDATE_OUTPUT_SIZE(\p alg, \p input_length).
+ * See also #PSA_AEAD_UPDATE_OUTPUT_SIZE(\p key_type, \p alg, \p input_length).
  *
  * \param input_length      Size of the input in bytes.
  */
@@ -421,23 +453,30 @@
  * insufficient ciphertext buffer size. The actual size of the output may
  * be smaller in any given call.
  *
+ * See also #PSA_AEAD_FINISH_OUTPUT_MAX_SIZE.
+ *
+ * \param key_type            A symmetric key type that is
+                              compatible with algorithm \p alg.
  * \param alg                 An AEAD algorithm
  *                            (\c PSA_ALG_XXX value such that
  *                            #PSA_ALG_IS_AEAD(\p alg) is true).
  *
  * \return                    A sufficient ciphertext buffer size for the
  *                            specified algorithm.
- *                            If the AEAD algorithm is not recognized, return 0.
+ *                            If the key type or AEAD algorithm is not
+ *                            recognized, or the parameters are incompatible,
+ *                            return 0.
  */
-#define PSA_AEAD_FINISH_OUTPUT_SIZE(alg)                                \
-    (PSA_ALG_IS_AEAD_ON_BLOCK_CIPHER(alg) ?                             \
-     PSA_BLOCK_CIPHER_BLOCK_MAX_SIZE :                                  \
+#define PSA_AEAD_FINISH_OUTPUT_SIZE(key_type, alg) \
+    (PSA_AEAD_NONCE_LENGTH(key_type, alg) != 0 &&  \
+         PSA_ALG_IS_AEAD_ON_BLOCK_CIPHER(alg) ?    \
+         PSA_BLOCK_CIPHER_BLOCK_LENGTH(key_type) : \
      0)
 
 /** A sufficient ciphertext buffer size for psa_aead_finish(), for any of the
  *  supported key types and AEAD algorithms.
  *
- * See also #PSA_AEAD_FINISH_OUTPUT_SIZE(\p alg).
+ * See also #PSA_AEAD_FINISH_OUTPUT_SIZE(\p key_type, \p alg).
  */
 #define PSA_AEAD_FINISH_OUTPUT_MAX_SIZE     (PSA_BLOCK_CIPHER_BLOCK_MAX_SIZE)
 
@@ -448,23 +487,30 @@
  * insufficient plaintext buffer size. The actual size of the output may
  * be smaller in any given call.
  *
+ * See also #PSA_AEAD_VERIFY_OUTPUT_MAX_SIZE.
+ *
+ * \param key_type            A symmetric key type that is
+ *                            compatible with algorithm \p alg.
  * \param alg                 An AEAD algorithm
  *                            (\c PSA_ALG_XXX value such that
  *                            #PSA_ALG_IS_AEAD(\p alg) is true).
  *
  * \return                    A sufficient plaintext buffer size for the
  *                            specified algorithm.
- *                            If the AEAD algorithm is not recognized, return 0.
+ *                            If the key type or AEAD algorithm is not
+ *                            recognized, or the parameters are incompatible,
+ *                            return 0.
  */
-#define PSA_AEAD_VERIFY_OUTPUT_SIZE(alg)                                \
-    (PSA_ALG_IS_AEAD_ON_BLOCK_CIPHER(alg) ?                             \
-     PSA_BLOCK_CIPHER_BLOCK_MAX_SIZE :                                  \
+#define PSA_AEAD_VERIFY_OUTPUT_SIZE(key_type, alg) \
+    (PSA_AEAD_NONCE_LENGTH(key_type, alg) != 0 &&  \
+         PSA_ALG_IS_AEAD_ON_BLOCK_CIPHER(alg) ?    \
+         PSA_BLOCK_CIPHER_BLOCK_LENGTH(key_type) : \
      0)
 
 /** A sufficient plaintext buffer size for psa_aead_verify(), for any of the
  *  supported key types and AEAD algorithms.
  *
- * See also #PSA_AEAD_VERIFY_OUTPUT_SIZE(\p alg).
+ * See also #PSA_AEAD_VERIFY_OUTPUT_SIZE(\p key_type, \p alg).
  */
 #define PSA_AEAD_VERIFY_OUTPUT_MAX_SIZE     (PSA_BLOCK_CIPHER_BLOCK_MAX_SIZE)
 
diff --git a/include/psa/crypto_values.h b/include/psa/crypto_values.h
index 5e865c9..9bfd5ab 100644
--- a/include/psa/crypto_values.h
+++ b/include/psa/crypto_values.h
@@ -2137,4 +2137,27 @@
 
 /**@}*/
 
+/** \defgroup helper_macros Helper macros
+ * @{
+ */
+
+/* Helper macros */
+
+/** Check if two AEAD algorithm identifiers refer to the same AEAD algorithm
+ *  regardless of the tag length they encode.
+ *
+ * \param aead_alg_1 An AEAD algorithm identifier.
+ * \param aead_alg_2 An AEAD algorithm identifier.
+ *
+ * \return           1 if both identifiers refer to the same AEAD algorithm,
+ *                   0 otherwise.
+ *                   Unspecified if neither \p aead_alg_1 nor \p aead_alg_2 are
+ *                   a supported AEAD algorithm.
+ */
+#define MBEDTLS_PSA_ALG_AEAD_EQUAL(aead_alg_1, aead_alg_2) \
+    (!(((aead_alg_1) ^ (aead_alg_2)) & \
+       ~(PSA_ALG_AEAD_TAG_LENGTH_MASK | PSA_ALG_AEAD_AT_LEAST_THIS_LENGTH_FLAG)))
+
+/**@}*/
+
 #endif /* PSA_CRYPTO_VALUES_H */
diff --git a/library/psa_crypto_aead.c b/library/psa_crypto_aead.c
index 2632830..356679c 100644
--- a/library/psa_crypto_aead.c
+++ b/library/psa_crypto_aead.c
@@ -154,10 +154,14 @@
             return( PSA_ERROR_NOT_SUPPORTED );
     }
 
-    if( PSA_AEAD_TAG_LENGTH( alg ) > full_tag_length )
+    if( PSA_AEAD_TAG_LENGTH( attributes->core.type,
+                             key_bits, alg )
+        > full_tag_length )
         return( PSA_ERROR_INVALID_ARGUMENT );
 
-    operation->tag_length = PSA_AEAD_TAG_LENGTH( alg );
+    operation->tag_length = PSA_AEAD_TAG_LENGTH( attributes->core.type,
+                                                 key_bits,
+                                                 alg );
 
     return( PSA_SUCCESS );
 }
diff --git a/programs/psa/key_ladder_demo.c b/programs/psa/key_ladder_demo.c
index 47d5de6..5d64349 100644
--- a/programs/psa/key_ladder_demo.c
+++ b/programs/psa/key_ladder_demo.c
@@ -365,6 +365,8 @@
     psa_status_t status;
     FILE *input_file = NULL;
     FILE *output_file = NULL;
+    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
+    psa_key_type_t key_type;
     long input_position;
     size_t input_size;
     size_t buffer_size = 0;
@@ -385,7 +387,10 @@
     }
 #endif
     input_size = input_position;
-    buffer_size = PSA_AEAD_ENCRYPT_OUTPUT_SIZE( WRAPPING_ALG, input_size );
+    PSA_CHECK( psa_get_key_attributes( wrapping_key, &attributes ) );
+    key_type = psa_get_key_type( &attributes );
+    buffer_size =
+        PSA_AEAD_ENCRYPT_OUTPUT_SIZE( key_type, WRAPPING_ALG, input_size );
     /* Check for integer overflow. */
     if( buffer_size < input_size )
     {
@@ -442,6 +447,8 @@
     psa_status_t status;
     FILE *input_file = NULL;
     FILE *output_file = NULL;
+    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
+    psa_key_type_t key_type;
     unsigned char *buffer = NULL;
     size_t ciphertext_size = 0;
     size_t plaintext_size;
@@ -465,8 +472,10 @@
         status = DEMO_ERROR;
         goto exit;
     }
+    PSA_CHECK( psa_get_key_attributes( wrapping_key, &attributes) );
+    key_type = psa_get_key_type( &attributes);
     ciphertext_size =
-        PSA_AEAD_ENCRYPT_OUTPUT_SIZE( WRAPPING_ALG, header.payload_size );
+        PSA_AEAD_ENCRYPT_OUTPUT_SIZE( key_type, WRAPPING_ALG, header.payload_size );
     /* Check for integer overflow. */
     if( ciphertext_size < header.payload_size )
     {
diff --git a/programs/psa/psa_constant_names_generated.c b/programs/psa/psa_constant_names_generated.c
index 2175af9..dcbe87f 100644
--- a/programs/psa/psa_constant_names_generated.c
+++ b/programs/psa/psa_constant_names_generated.c
@@ -169,11 +169,11 @@
         } else if (alg & PSA_ALG_AEAD_AT_LEAST_THIS_LENGTH_FLAG) {
             append(&buffer, buffer_size, &required_size,
                    "PSA_ALG_AEAD_WITH_AT_LEAST_THIS_LENGTH_TAG(", 43);
-            length_modifier = PSA_AEAD_TAG_LENGTH(alg);
+            length_modifier = PSA_ALG_AEAD_GET_TAG_LENGTH(alg);
         } else if (core_alg != alg) {
             append(&buffer, buffer_size, &required_size,
                    "PSA_ALG_AEAD_WITH_SHORTENED_TAG(", 32);
-            length_modifier = PSA_AEAD_TAG_LENGTH(alg);
+            length_modifier = PSA_ALG_AEAD_GET_TAG_LENGTH(alg);
         }
     } else if (PSA_ALG_IS_KEY_AGREEMENT(alg) &&
                !PSA_ALG_IS_RAW_KEY_AGREEMENT(alg)) {
diff --git a/scripts/generate_psa_constants.py b/scripts/generate_psa_constants.py
index ff07ecd..71afd02 100755
--- a/scripts/generate_psa_constants.py
+++ b/scripts/generate_psa_constants.py
@@ -117,11 +117,11 @@
         } else if (alg & PSA_ALG_AEAD_AT_LEAST_THIS_LENGTH_FLAG) {
             append(&buffer, buffer_size, &required_size,
                    "PSA_ALG_AEAD_WITH_AT_LEAST_THIS_LENGTH_TAG(", 43);
-            length_modifier = PSA_AEAD_TAG_LENGTH(alg);
+            length_modifier = PSA_ALG_AEAD_GET_TAG_LENGTH(alg);
         } else if (core_alg != alg) {
             append(&buffer, buffer_size, &required_size,
                    "PSA_ALG_AEAD_WITH_SHORTENED_TAG(", 32);
-            length_modifier = PSA_AEAD_TAG_LENGTH(alg);
+            length_modifier = PSA_ALG_AEAD_GET_TAG_LENGTH(alg);
         }
     } else if (PSA_ALG_IS_KEY_AGREEMENT(alg) &&
                !PSA_ALG_IS_RAW_KEY_AGREEMENT(alg)) {
diff --git a/tests/suites/test_suite_psa_crypto.function b/tests/suites/test_suite_psa_crypto.function
index 4e56800..310b2a7 100644
--- a/tests/suites/test_suite_psa_crypto.function
+++ b/tests/suites/test_suite_psa_crypto.function
@@ -2985,24 +2985,16 @@
     mbedtls_svc_key_id_t key = MBEDTLS_SVC_KEY_ID_INIT;
     psa_key_type_t key_type = key_type_arg;
     psa_algorithm_t alg = alg_arg;
+    size_t key_bits;
     unsigned char *output_data = NULL;
     size_t output_size = 0;
     size_t output_length = 0;
     unsigned char *output_data2 = NULL;
     size_t output_length2 = 0;
-    size_t tag_length = PSA_AEAD_TAG_LENGTH( alg );
     psa_status_t status = PSA_ERROR_GENERIC_ERROR;
     psa_status_t expected_result = expected_result_arg;
     psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
 
-    output_size = input_data->len + tag_length;
-    /* For all currently defined algorithms, PSA_AEAD_ENCRYPT_OUTPUT_SIZE
-     * should be exact. */
-    if( expected_result != PSA_ERROR_INVALID_ARGUMENT )
-        TEST_EQUAL( output_size,
-                    PSA_AEAD_ENCRYPT_OUTPUT_SIZE( alg, input_data->len ) );
-    ASSERT_ALLOC( output_data, output_size );
-
     PSA_ASSERT( psa_crypto_init( ) );
 
     psa_set_key_usage_flags( &attributes, PSA_KEY_USAGE_ENCRYPT | PSA_KEY_USAGE_DECRYPT );
@@ -3011,6 +3003,22 @@
 
     PSA_ASSERT( psa_import_key( &attributes, key_data->x, key_data->len,
                                 &key ) );
+    PSA_ASSERT( psa_get_key_attributes( key, &attributes ) );
+    key_bits = psa_get_key_bits( &attributes );
+
+    output_size = input_data->len + PSA_AEAD_TAG_LENGTH( key_type, key_bits,
+                                                         alg );
+    /* For all currently defined algorithms, PSA_AEAD_ENCRYPT_OUTPUT_SIZE
+     * should be exact. */
+    if( expected_result != PSA_ERROR_INVALID_ARGUMENT &&
+        expected_result != PSA_ERROR_NOT_SUPPORTED )
+    {
+        TEST_EQUAL( output_size,
+                    PSA_AEAD_ENCRYPT_OUTPUT_SIZE( key_type, alg, input_data->len ) );
+        TEST_ASSERT( output_size <=
+                     PSA_AEAD_ENCRYPT_OUTPUT_MAX_SIZE( input_data->len ) );
+    }
+    ASSERT_ALLOC( output_data, output_size );
 
     status = psa_aead_encrypt( key, alg,
                                nonce->x, nonce->len,
@@ -3038,7 +3046,7 @@
         /* For all currently defined algorithms, PSA_AEAD_DECRYPT_OUTPUT_SIZE
          * should be exact. */
         TEST_EQUAL( input_data->len,
-                    PSA_AEAD_DECRYPT_OUTPUT_SIZE( alg, output_length ) );
+                    PSA_AEAD_DECRYPT_OUTPUT_SIZE( key_type, alg, output_length ) );
 
         TEST_ASSERT( input_data->len <=
                      PSA_AEAD_DECRYPT_OUTPUT_MAX_SIZE( output_length ) );
@@ -3075,22 +3083,13 @@
     mbedtls_svc_key_id_t key = MBEDTLS_SVC_KEY_ID_INIT;
     psa_key_type_t key_type = key_type_arg;
     psa_algorithm_t alg = alg_arg;
+    size_t key_bits;
     unsigned char *output_data = NULL;
     size_t output_size = 0;
     size_t output_length = 0;
-    size_t tag_length = PSA_AEAD_TAG_LENGTH( alg );
     psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
     psa_status_t status = PSA_ERROR_GENERIC_ERROR;
 
-    output_size = input_data->len + tag_length;
-    /* For all currently defined algorithms, PSA_AEAD_ENCRYPT_OUTPUT_SIZE
-     * should be exact. */
-    TEST_EQUAL( output_size,
-                PSA_AEAD_ENCRYPT_OUTPUT_SIZE( alg, input_data->len ) );
-    TEST_ASSERT( output_size <=
-                 PSA_AEAD_ENCRYPT_OUTPUT_MAX_SIZE( input_data->len ) );
-    ASSERT_ALLOC( output_data, output_size );
-
     PSA_ASSERT( psa_crypto_init( ) );
 
     psa_set_key_usage_flags( &attributes, PSA_KEY_USAGE_ENCRYPT  );
@@ -3099,6 +3098,18 @@
 
     PSA_ASSERT( psa_import_key( &attributes, key_data->x, key_data->len,
                                 &key ) );
+    PSA_ASSERT( psa_get_key_attributes( key, &attributes ) );
+    key_bits = psa_get_key_bits( &attributes );
+
+    output_size = input_data->len + PSA_AEAD_TAG_LENGTH( key_type, key_bits,
+                                                         alg );
+    /* For all currently defined algorithms, PSA_AEAD_ENCRYPT_OUTPUT_SIZE
+     * should be exact. */
+    TEST_EQUAL( output_size,
+                PSA_AEAD_ENCRYPT_OUTPUT_SIZE( key_type, alg, input_data->len ) );
+    TEST_ASSERT( output_size <=
+                 PSA_AEAD_ENCRYPT_OUTPUT_MAX_SIZE( input_data->len ) );
+    ASSERT_ALLOC( output_data, output_size );
 
     status = psa_aead_encrypt( key, alg,
                                nonce->x, nonce->len,
@@ -3139,26 +3150,14 @@
     mbedtls_svc_key_id_t key = MBEDTLS_SVC_KEY_ID_INIT;
     psa_key_type_t key_type = key_type_arg;
     psa_algorithm_t alg = alg_arg;
+    size_t key_bits;
     unsigned char *output_data = NULL;
     size_t output_size = 0;
     size_t output_length = 0;
-    size_t tag_length = PSA_AEAD_TAG_LENGTH( alg );
     psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
     psa_status_t expected_result = expected_result_arg;
     psa_status_t status = PSA_ERROR_GENERIC_ERROR;
 
-    output_size = input_data->len - tag_length;
-    if( expected_result != PSA_ERROR_INVALID_ARGUMENT )
-    {
-        /* For all currently defined algorithms, PSA_AEAD_DECRYPT_OUTPUT_SIZE
-         * should be exact. */
-        TEST_EQUAL( output_size,
-                    PSA_AEAD_DECRYPT_OUTPUT_SIZE( alg, input_data->len ) );
-        TEST_ASSERT( output_size <=
-                     PSA_AEAD_DECRYPT_OUTPUT_MAX_SIZE( input_data->len ) );
-    }
-    ASSERT_ALLOC( output_data, output_size );
-
     PSA_ASSERT( psa_crypto_init( ) );
 
     psa_set_key_usage_flags( &attributes, PSA_KEY_USAGE_DECRYPT  );
@@ -3167,6 +3166,22 @@
 
     PSA_ASSERT( psa_import_key( &attributes, key_data->x, key_data->len,
                                 &key ) );
+    PSA_ASSERT( psa_get_key_attributes( key, &attributes ) );
+    key_bits = psa_get_key_bits( &attributes );
+
+    output_size = input_data->len - PSA_AEAD_TAG_LENGTH( key_type, key_bits,
+                                                         alg );
+    if( expected_result != PSA_ERROR_INVALID_ARGUMENT &&
+        expected_result != PSA_ERROR_NOT_SUPPORTED )
+    {
+        /* For all currently defined algorithms, PSA_AEAD_DECRYPT_OUTPUT_SIZE
+         * should be exact. */
+        TEST_EQUAL( output_size,
+                    PSA_AEAD_DECRYPT_OUTPUT_SIZE( key_type, alg, input_data->len ) );
+        TEST_ASSERT( output_size <=
+                     PSA_AEAD_DECRYPT_OUTPUT_MAX_SIZE( input_data->len ) );
+    }
+    ASSERT_ALLOC( output_data, output_size );
 
     status = psa_aead_decrypt( key, alg,
                                nonce->x, nonce->len,
diff --git a/tests/suites/test_suite_psa_crypto_driver_wrappers.function b/tests/suites/test_suite_psa_crypto_driver_wrappers.function
index ad5b6c5..8ccce06 100644
--- a/tests/suites/test_suite_psa_crypto_driver_wrappers.function
+++ b/tests/suites/test_suite_psa_crypto_driver_wrappers.function
@@ -822,24 +822,15 @@
     mbedtls_svc_key_id_t key = MBEDTLS_SVC_KEY_ID_INIT;
     psa_key_type_t key_type = key_type_arg;
     psa_algorithm_t alg = alg_arg;
+    size_t key_bits;
     psa_status_t forced_status = forced_status_arg;
     unsigned char *output_data = NULL;
     size_t output_size = 0;
     size_t output_length = 0;
-    size_t tag_length = PSA_AEAD_TAG_LENGTH( alg );
     psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
     psa_status_t status = PSA_ERROR_GENERIC_ERROR;
     test_driver_aead_hooks = test_driver_aead_hooks_init();
 
-    output_size = input_data->len + tag_length;
-    /* For all currently defined algorithms, PSA_AEAD_ENCRYPT_OUTPUT_SIZE
-     * should be exact. */
-    TEST_EQUAL( output_size,
-                PSA_AEAD_ENCRYPT_OUTPUT_SIZE( alg, input_data->len ) );
-    TEST_ASSERT( output_size <=
-                 PSA_AEAD_ENCRYPT_OUTPUT_MAX_SIZE( input_data->len ) );
-    ASSERT_ALLOC( output_data, output_size );
-
     PSA_ASSERT( psa_crypto_init( ) );
 
     psa_set_key_usage_flags( &attributes, PSA_KEY_USAGE_ENCRYPT  );
@@ -848,6 +839,18 @@
 
     PSA_ASSERT( psa_import_key( &attributes, key_data->x, key_data->len,
                                 &key ) );
+    PSA_ASSERT( psa_get_key_attributes( key, &attributes ) );
+    key_bits = psa_get_key_bits( &attributes );
+
+    output_size = input_data->len + PSA_AEAD_TAG_LENGTH( key_type, key_bits,
+                                                         alg );
+    /* For all currently defined algorithms, PSA_AEAD_ENCRYPT_OUTPUT_SIZE
+     * should be exact. */
+    TEST_EQUAL( output_size,
+                PSA_AEAD_ENCRYPT_OUTPUT_SIZE( key_type, alg, input_data->len ) );
+    TEST_ASSERT( output_size <=
+                 PSA_AEAD_ENCRYPT_OUTPUT_MAX_SIZE( input_data->len ) );
+    ASSERT_ALLOC( output_data, output_size );
 
     test_driver_aead_hooks.forced_status = forced_status;
     status = psa_aead_encrypt( key, alg,
@@ -888,18 +891,15 @@
     mbedtls_svc_key_id_t key = MBEDTLS_SVC_KEY_ID_INIT;
     psa_key_type_t key_type = key_type_arg;
     psa_algorithm_t alg = alg_arg;
+    size_t key_bits;
     psa_status_t forced_status = forced_status_arg;
     unsigned char *output_data = NULL;
     size_t output_size = 0;
     size_t output_length = 0;
-    size_t tag_length = PSA_AEAD_TAG_LENGTH( alg );
     psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
     psa_status_t status = PSA_ERROR_GENERIC_ERROR;
     test_driver_aead_hooks = test_driver_aead_hooks_init();
 
-    output_size = input_data->len - tag_length;
-    ASSERT_ALLOC( output_data, output_size );
-
     PSA_ASSERT( psa_crypto_init( ) );
 
     psa_set_key_usage_flags( &attributes, PSA_KEY_USAGE_DECRYPT  );
@@ -908,6 +908,12 @@
 
     PSA_ASSERT( psa_import_key( &attributes, key_data->x, key_data->len,
                                 &key ) );
+    PSA_ASSERT( psa_get_key_attributes( key, &attributes ) );
+    key_bits = psa_get_key_bits( &attributes );
+
+    output_size = input_data->len - PSA_AEAD_TAG_LENGTH( key_type, key_bits,
+                                                         alg );
+    ASSERT_ALLOC( output_data, output_size );
 
     test_driver_aead_hooks.forced_status = forced_status;
     status = psa_aead_decrypt( key, alg,
diff --git a/tests/suites/test_suite_psa_crypto_metadata.data b/tests/suites/test_suite_psa_crypto_metadata.data
index bd98a76..4e2f4d5 100644
--- a/tests/suites/test_suite_psa_crypto_metadata.data
+++ b/tests/suites/test_suite_psa_crypto_metadata.data
@@ -134,17 +134,57 @@
 depends_on:PSA_WANT_ALG_XTS:MBEDTLS_CIPHER_C
 cipher_algorithm:PSA_ALG_XTS:0
 
-AEAD: CCM
-depends_on:PSA_WANT_ALG_CCM
-aead_algorithm:PSA_ALG_CCM:ALG_IS_AEAD_ON_BLOCK_CIPHER:16
+AEAD: CCM-AES-128
+depends_on:PSA_WANT_KEY_TYPE_AES:PSA_WANT_ALG_CCM
+aead_algorithm:PSA_ALG_CCM:ALG_IS_AEAD_ON_BLOCK_CIPHER:16:PSA_KEY_TYPE_AES:128
 
-AEAD: GCM
-depends_on:PSA_WANT_ALG_GCM
-aead_algorithm:PSA_ALG_GCM:ALG_IS_AEAD_ON_BLOCK_CIPHER:16
+AEAD: CCM-AES-192
+depends_on:PSA_WANT_KEY_TYPE_AES:PSA_WANT_ALG_CCM
+aead_algorithm:PSA_ALG_CCM:ALG_IS_AEAD_ON_BLOCK_CIPHER:16:PSA_KEY_TYPE_AES:192
+
+AEAD: CCM-AES-256
+depends_on:PSA_WANT_KEY_TYPE_AES:PSA_WANT_ALG_CCM
+aead_algorithm:PSA_ALG_CCM:ALG_IS_AEAD_ON_BLOCK_CIPHER:16:PSA_KEY_TYPE_AES:256
+
+AEAD: CCM-CAMELLIA-128
+depends_on:PSA_WANT_KEY_TYPE_CAMELLIA:PSA_WANT_ALG_CCM
+aead_algorithm:PSA_ALG_CCM:ALG_IS_AEAD_ON_BLOCK_CIPHER:16:PSA_KEY_TYPE_CAMELLIA:128
+
+AEAD: CCM-CAMELLIA-192
+depends_on:PSA_WANT_KEY_TYPE_CAMELLIA:PSA_WANT_ALG_CCM
+aead_algorithm:PSA_ALG_CCM:ALG_IS_AEAD_ON_BLOCK_CIPHER:16:PSA_KEY_TYPE_CAMELLIA:192
+
+AEAD: CCM-CAMELLIA-256
+depends_on:PSA_WANT_KEY_TYPE_CAMELLIA:PSA_WANT_ALG_CCM
+aead_algorithm:PSA_ALG_CCM:ALG_IS_AEAD_ON_BLOCK_CIPHER:16:PSA_KEY_TYPE_CAMELLIA:256
+
+AEAD: GCM-AES-128
+depends_on:PSA_WANT_KEY_TYPE_AES:PSA_WANT_ALG_GCM
+aead_algorithm:PSA_ALG_GCM:ALG_IS_AEAD_ON_BLOCK_CIPHER:16:PSA_KEY_TYPE_AES:128
+
+AEAD: GCM-AES-192
+depends_on:PSA_WANT_KEY_TYPE_AES:PSA_WANT_ALG_GCM
+aead_algorithm:PSA_ALG_GCM:ALG_IS_AEAD_ON_BLOCK_CIPHER:16:PSA_KEY_TYPE_AES:192
+
+AEAD: GCM-AES-256
+depends_on:PSA_WANT_KEY_TYPE_AES:PSA_WANT_ALG_GCM
+aead_algorithm:PSA_ALG_GCM:ALG_IS_AEAD_ON_BLOCK_CIPHER:16:PSA_KEY_TYPE_AES:256
+
+AEAD: GCM-CAMELLIA-128
+depends_on:PSA_WANT_KEY_TYPE_CAMELLIA:PSA_WANT_ALG_GCM
+aead_algorithm:PSA_ALG_GCM:ALG_IS_AEAD_ON_BLOCK_CIPHER:16:PSA_KEY_TYPE_CAMELLIA:128
+
+AEAD: GCM-CAMELLIA-192
+depends_on:PSA_WANT_KEY_TYPE_CAMELLIA:PSA_WANT_ALG_GCM
+aead_algorithm:PSA_ALG_GCM:ALG_IS_AEAD_ON_BLOCK_CIPHER:16:PSA_KEY_TYPE_CAMELLIA:192
+
+AEAD: GCM-CAMELLIA-256
+depends_on:PSA_WANT_KEY_TYPE_CAMELLIA:PSA_WANT_ALG_GCM
+aead_algorithm:PSA_ALG_GCM:ALG_IS_AEAD_ON_BLOCK_CIPHER:16:PSA_KEY_TYPE_CAMELLIA:256
 
 AEAD: ChaCha20_Poly1305
 depends_on:PSA_WANT_ALG_CHACHA20_POLY1305
-aead_algorithm:PSA_ALG_CHACHA20_POLY1305:0:16
+aead_algorithm:PSA_ALG_CHACHA20_POLY1305:0:16:PSA_KEY_TYPE_CHACHA20:256
 
 Asymmetric signature: RSA PKCS#1 v1.5 raw
 depends_on:PSA_WANT_ALG_RSA_PKCS1V15_SIGN
diff --git a/tests/suites/test_suite_psa_crypto_metadata.function b/tests/suites/test_suite_psa_crypto_metadata.function
index 8acbe44..8134f44 100644
--- a/tests/suites/test_suite_psa_crypto_metadata.function
+++ b/tests/suites/test_suite_psa_crypto_metadata.function
@@ -169,6 +169,7 @@
 }
 
 void aead_algorithm_core( psa_algorithm_t alg, int classification_flags,
+                          psa_key_type_t key_type, size_t key_bits,
                           size_t tag_length )
 {
     /* Algorithm classification */
@@ -183,7 +184,7 @@
     algorithm_classification( alg, classification_flags );
 
     /* Tag length */
-    TEST_EQUAL( tag_length, PSA_AEAD_TAG_LENGTH( alg ) );
+    TEST_EQUAL( tag_length, PSA_AEAD_TAG_LENGTH( key_type, key_bits, alg ) );
 
 exit: ;
 }
@@ -367,19 +368,24 @@
 
 /* BEGIN_CASE */
 void aead_algorithm( int alg_arg, int classification_flags,
-                     int tag_length_arg )
+                     int tag_length_arg,
+                     int key_type_arg, int key_bits_arg )
 {
     psa_algorithm_t alg = alg_arg;
     size_t tag_length = tag_length_arg;
     size_t n;
+    psa_key_type_t key_type = key_type_arg;
+    size_t key_bits = key_bits_arg;
 
-    aead_algorithm_core( alg, classification_flags, tag_length );
+    aead_algorithm_core( alg, classification_flags,
+                         key_type, key_bits, tag_length );
 
     /* Truncated versions */
     for( n = 1; n <= tag_length; n++ )
     {
         psa_algorithm_t truncated_alg = PSA_ALG_AEAD_WITH_SHORTENED_TAG( alg, n );
-        aead_algorithm_core( truncated_alg, classification_flags, n );
+        aead_algorithm_core( truncated_alg, classification_flags,
+                             key_type, key_bits, n );
         TEST_EQUAL( PSA_ALG_AEAD_WITH_DEFAULT_LENGTH_TAG( truncated_alg ),
                     alg );
         /* Check that calling PSA_ALG_AEAD_WITH_SHORTENED_TAG twice gives
@@ -411,7 +417,8 @@
     for( n = 1; n <= tag_length; n++ )
     {
         psa_algorithm_t policy_alg = PSA_ALG_AEAD_WITH_AT_LEAST_THIS_LENGTH_TAG( alg, n );
-        aead_algorithm_core( policy_alg, classification_flags | ALG_IS_WILDCARD, n );
+        aead_algorithm_core( policy_alg, classification_flags | ALG_IS_WILDCARD,
+                             key_type, key_bits, n );
         TEST_EQUAL( PSA_ALG_AEAD_WITH_DEFAULT_LENGTH_TAG( policy_alg ),
                     alg );
         /* Check that calling PSA_ALG_AEAD_WITH_AT_LEAST_THIS_LENGTH_TAG twice