diff --git a/library/psa_crypto_aead.c b/library/psa_crypto_aead.c
index f2096ce..9ac2646 100644
--- a/library/psa_crypto_aead.c
+++ b/library/psa_crypto_aead.c
@@ -136,6 +136,22 @@
     return( PSA_SUCCESS );
 }
 
+/* Perform common nonce length checks */
+static psa_status_t mbedtls_aead_check_nonce_length(
+    mbedtls_psa_aead_operation_t *operation,
+    size_t nonce_length )
+{
+#if defined(MBEDTLS_PSA_BUILTIN_ALG_CHACHA20_POLY1305)
+    if( operation->alg == PSA_ALG_CHACHA20_POLY1305 )
+    {
+        if( nonce_length != 12 )
+            return( PSA_ERROR_NOT_SUPPORTED );
+    }
+#endif /* MBEDTLS_PSA_BUILTIN_ALG_CHACHA20_POLY1305 */
+
+    return PSA_SUCCESS;
+}
+
 psa_status_t mbedtls_psa_aead_encrypt(
     const psa_key_attributes_t *attributes,
     const uint8_t *key_buffer, size_t key_buffer_size,
@@ -164,6 +180,13 @@
     }
     tag = ciphertext + plaintext_length;
 
+    if( mbedtls_aead_check_nonce_length( &operation, nonce_length )
+        != PSA_SUCCESS )
+    {
+        status = PSA_ERROR_NOT_SUPPORTED;
+        goto exit;
+    }
+
 #if defined(MBEDTLS_PSA_BUILTIN_ALG_CCM)
     if( operation.alg == PSA_ALG_CCM )
     {
@@ -195,7 +218,7 @@
 #if defined(MBEDTLS_PSA_BUILTIN_ALG_CHACHA20_POLY1305)
     if( operation.alg == PSA_ALG_CHACHA20_POLY1305 )
     {
-        if( nonce_length != 12 || operation.tag_length != 16 )
+        if( operation.tag_length != 16 )
         {
             status = PSA_ERROR_NOT_SUPPORTED;
             goto exit;
@@ -247,21 +270,6 @@
     return( PSA_SUCCESS );
 }
 
-static psa_status_t mbedtls_aead_check_nonce_length(
-    mbedtls_psa_aead_operation_t *operation,
-    size_t nonce_length )
-{
-#if defined(MBEDTLS_PSA_BUILTIN_ALG_CHACHA20_POLY1305)
-    if( operation->alg == PSA_ALG_CHACHA20_POLY1305 )
-    {
-        if( nonce_length != 12 )
-            return( PSA_ERROR_NOT_SUPPORTED );
-    }
-#endif /* MBEDTLS_PSA_BUILTIN_ALG_CHACHA20_POLY1305 */
-
-    return PSA_SUCCESS;
-}
-
 psa_status_t mbedtls_psa_aead_decrypt(
     const psa_key_attributes_t *attributes,
     const uint8_t *key_buffer, size_t key_buffer_size,