Add support for non-default-tag-size AEAD (CCM and GCM)
diff --git a/include/psa/crypto.h b/include/psa/crypto.h
index 0269be9..c22e853 100644
--- a/include/psa/crypto.h
+++ b/include/psa/crypto.h
@@ -884,8 +884,34 @@
  */
 #define PSA_ALG_CBC_PKCS7                       ((psa_algorithm_t)0x04600101)
 
-#define PSA_ALG_CCM                             ((psa_algorithm_t)0x06000001)
-#define PSA_ALG_GCM                             ((psa_algorithm_t)0x06000002)
+#define PSA_ALG_CCM                             ((psa_algorithm_t)0x06001001)
+#define PSA_ALG_GCM                             ((psa_algorithm_t)0x06001002)
+
+#define PSA_ALG_AEAD_TAG_LENGTH_MASK            ((psa_algorithm_t)0x00003f00)
+#define PSA_AEAD_TAG_LENGTH_OFFSET 8
+
+/** Macro to build a shortened AEAD algorithm.
+ *
+ * A shortened AEAD algorithm is similar to the corresponding AEAD
+ * algorithm, but has an authentication tag that consists of fewer bytes.
+ * Depending on the algorithm, the tag length may affect the calculation
+ * of the ciphertext.
+ *
+ * \param alg           A AEAD algorithm identifier (value of type
+ *                      #psa_algorithm_t such that #PSA_ALG_IS_AEAD(\p alg)
+ *                      is true).
+ * \param mac_length    Desired length of the authentication tag in bytes.
+ *
+ * \return              The corresponding AEAD algorithm with the specified
+ *                      length.
+ * \return              Unspecified if \p alg is not a supported
+ *                      AEAD algorithm or if \p tag_length is not valid
+ *                      for the specified AEAD algorithm.
+ */
+#define PSA_ALG_AEAD_WITH_TAG_LENGTH(alg, tag_length)                   \
+    (((alg) & ~PSA_ALG_AEAD_TAG_LENGTH_MASK) |                          \
+     ((tag_length) << PSA_AEAD_TAG_LENGTH_OFFSET &                      \
+      PSA_ALG_AEAD_TAG_LENGTH_MASK))
 
 #define PSA_ALG_RSA_PKCS1V15_SIGN_BASE          ((psa_algorithm_t)0x10020000)
 /** RSA PKCS#1 v1.5 signature with hashing.
@@ -2432,9 +2458,9 @@
  *                            correct size for an AEAD algorithm that it
  *                            recognizes, but does not support.
  */
-#define PSA_AEAD_TAG_SIZE(alg)             \
-    ((alg) == PSA_ALG_GCM ? 16 :           \
-     (alg) == PSA_ALG_CCM ? 16 :           \
+#define PSA_AEAD_TAG_LENGTH(alg)                                        \
+    (PSA_ALG_IS_AEAD(alg) ?                                             \
+     (((alg) & PSA_ALG_AEAD_TAG_LENGTH_MASK) >> PSA_AEAD_TAG_LENGTH_OFFSET) : \
      0)
 
 /** Process an authenticated encryption operation.
diff --git a/include/psa/crypto_sizes.h b/include/psa/crypto_sizes.h
index c058afc..169566e 100644
--- a/include/psa/crypto_sizes.h
+++ b/include/psa/crypto_sizes.h
@@ -188,9 +188,9 @@
  *                            correct size for an AEAD algorithm that it
  *                            recognizes, but does not support.
  */
-#define PSA_AEAD_ENCRYPT_OUTPUT_SIZE(alg, plaintext_length)     \
-    (PSA_AEAD_TAG_SIZE(alg) != 0 ?                              \
-     (plaintext_length) + PSA_AEAD_TAG_SIZE(alg) :              \
+#define PSA_AEAD_ENCRYPT_OUTPUT_SIZE(alg, plaintext_length)       \
+    (PSA_AEAD_TAG_LENGTH(alg) != 0 ?                              \
+     (plaintext_length) + PSA_AEAD_TAG_LENGTH(alg) :              \
      0)
 
 /** The maximum size of the output of psa_aead_decrypt(), in bytes.
@@ -212,9 +212,9 @@
  *                            correct size for an AEAD algorithm that it
  *                            recognizes, but does not support.
  */
-#define PSA_AEAD_DECRYPT_OUTPUT_SIZE(alg, ciphertext_length)    \
-    (PSA_AEAD_TAG_SIZE(alg) != 0 ?                              \
-     (plaintext_length) - PSA_AEAD_TAG_SIZE(alg) :              \
+#define PSA_AEAD_DECRYPT_OUTPUT_SIZE(alg, ciphertext_length)      \
+    (PSA_AEAD_TAG_LENGTH(alg) != 0 ?                              \
+     (plaintext_length) - PSA_AEAD_TAG_LENGTH(alg) :              \
      0)
 
 /** Safe signature buffer size for psa_asymmetric_sign().
diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index 82af920..4486242 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -1224,6 +1224,9 @@
     mbedtls_cipher_mode_t mode;
     mbedtls_cipher_id_t cipher_id_tmp;
 
+    if( PSA_ALG_IS_AEAD( alg ) )
+        alg &= ~PSA_ALG_AEAD_TAG_LENGTH_MASK;
+
     if( PSA_ALG_IS_CIPHER( alg ) || PSA_ALG_IS_AEAD( alg ) )
     {
         switch( alg )
@@ -1246,10 +1249,10 @@
             case PSA_ALG_CBC_PKCS7:
                 mode = MBEDTLS_MODE_CBC;
                 break;
-            case PSA_ALG_CCM:
+            case PSA_ALG_CCM & ~PSA_ALG_AEAD_TAG_LENGTH_MASK:
                 mode = MBEDTLS_MODE_CCM;
                 break;
-            case PSA_ALG_GCM:
+            case PSA_ALG_GCM & ~PSA_ALG_AEAD_TAG_LENGTH_MASK:
                 mode = MBEDTLS_MODE_GCM;
                 break;
             default:
@@ -2834,6 +2837,8 @@
         mbedtls_gcm_context gcm;
 #endif /* MBEDTLS_GCM_C */
     } ctx;
+    psa_algorithm_t core_alg;
+    uint8_t full_tag_length;
     uint8_t tag_length;
 } aead_operation_t;
 
@@ -2876,11 +2881,12 @@
     if( operation->cipher_info == NULL )
         return( PSA_ERROR_NOT_SUPPORTED );
 
-    switch( alg )
+    switch( PSA_ALG_AEAD_WITH_TAG_LENGTH( alg, 0 ) )
     {
 #if defined(MBEDTLS_CCM_C)
-        case PSA_ALG_CCM:
-            operation->tag_length = 16;
+        case PSA_ALG_AEAD_WITH_TAG_LENGTH( PSA_ALG_CCM, 0 ):
+            operation->core_alg = PSA_ALG_CCM;
+            operation->full_tag_length = 16;
             if( PSA_BLOCK_CIPHER_BLOCK_SIZE( operation->slot->type ) != 16 )
                 return( PSA_ERROR_INVALID_ARGUMENT );
             mbedtls_ccm_init( &operation->ctx.ccm );
@@ -2894,8 +2900,9 @@
 #endif /* MBEDTLS_CCM_C */
 
 #if defined(MBEDTLS_GCM_C)
-        case PSA_ALG_GCM:
-            operation->tag_length = 16;
+        case PSA_ALG_AEAD_WITH_TAG_LENGTH( PSA_ALG_GCM, 0 ):
+            operation->core_alg = PSA_ALG_GCM;
+            operation->full_tag_length = 16;
             if( PSA_BLOCK_CIPHER_BLOCK_SIZE( operation->slot->type ) != 16 )
                 return( PSA_ERROR_INVALID_ARGUMENT );
             mbedtls_gcm_init( &operation->ctx.gcm );
@@ -2910,6 +2917,16 @@
             return( PSA_ERROR_NOT_SUPPORTED );
     }
 
+    if( PSA_AEAD_TAG_LENGTH( alg ) > operation->full_tag_length )
+    {
+        status = PSA_ERROR_INVALID_ARGUMENT;
+        goto cleanup;
+    }
+    operation->tag_length = PSA_AEAD_TAG_LENGTH( alg );
+    /* CCM allows the following tag lengths: 4, 6, 8, 10, 12, 14, 16.
+     * GCM allows the following tag lengths: 4, 8, 12, 13, 14, 15, 16.
+     * In both cases, mbedtls_xxx will validate the tag length below. */
+
     return( PSA_SUCCESS );
 
 cleanup:
@@ -2948,7 +2965,7 @@
     }
     tag = ciphertext + plaintext_length;
 
-    if( alg == PSA_ALG_GCM )
+    if( operation.core_alg == PSA_ALG_GCM )
     {
         status = mbedtls_to_psa_error(
             mbedtls_gcm_crypt_and_tag( &operation.ctx.gcm,
@@ -2959,7 +2976,7 @@
                                        plaintext, ciphertext,
                                        operation.tag_length, tag ) );
     }
-    else if( alg == PSA_ALG_CCM )
+    else if( operation.core_alg == PSA_ALG_CCM )
     {
         status = mbedtls_to_psa_error(
             mbedtls_ccm_encrypt_and_tag( &operation.ctx.ccm,
@@ -3028,7 +3045,7 @@
     if( status != PSA_SUCCESS )
         return( status );
 
-    if( alg == PSA_ALG_GCM )
+    if( operation.core_alg == PSA_ALG_GCM )
     {
         status = psa_aead_unpadded_locate_tag( operation.tag_length,
                                                ciphertext, ciphertext_length,
@@ -3045,7 +3062,7 @@
                                       tag, operation.tag_length,
                                       ciphertext, plaintext ) );
     }
-    else if( alg == PSA_ALG_CCM )
+    else if( operation.core_alg == PSA_ALG_CCM )
     {
         status = psa_aead_unpadded_locate_tag( operation.tag_length,
                                                ciphertext, ciphertext_length,
diff --git a/tests/suites/test_suite_psa_crypto_metadata.function b/tests/suites/test_suite_psa_crypto_metadata.function
index ca9d557..a264389 100644
--- a/tests/suites/test_suite_psa_crypto_metadata.function
+++ b/tests/suites/test_suite_psa_crypto_metadata.function
@@ -227,7 +227,7 @@
     algorithm_classification( alg, classification_flags );
 
     /* Tag length */
-    TEST_ASSERT( tag_length == PSA_AEAD_TAG_SIZE( alg ) );
+    TEST_ASSERT( tag_length == PSA_AEAD_TAG_LENGTH( alg ) );
 }
 /* END_CASE */