From 4ed1ed18d2a5595164bbc2b08df7ef814235e289 Mon Sep 17 00:00:00 2001
From: Paul Elliott <paul.elliott@arm.com>
Date: Mon, 27 Sep 2021 18:09:28 +0100
Subject: [PATCH] Move nonce size checking to PSA Core

Signed-off-by: Paul Elliott <paul.elliott@arm.com>
---
 library/psa_crypto.c      | 42 ++++++++++++++++++++++++++++++++++-----
 library/psa_crypto_aead.c |  6 ------
 2 files changed, 37 insertions(+), 11 deletions(-)

diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index 65dc5c7fe..fd2069b79 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -3863,11 +3863,43 @@ psa_status_t psa_aead_set_nonce( psa_aead_operation_t *operation,
         goto exit;
     }
 
-    /* Not checking nonce size here as GCM spec allows almost arbitrarily
-     * large nonces. Please note that we do not generally recommend the usage
-     * of nonces of greater length than PSA_AEAD_NONCE_MAX_SIZE, as large
-     * nonces are hashed to a shorter size, which can then lead to collisions
-     * if you encrypt a very large number of messages.*/
+#if defined(MBEDTLS_PSA_BUILTIN_ALG_GCM)
+    if( operation->alg == PSA_ALG_GCM )
+    {
+        /* Not checking max nonce size here as GCM spec allows almost
+         * arbitrarily large nonces. Please note that we do not generally
+         * recommend the usage of nonces of greater length than
+         * PSA_AEAD_NONCE_MAX_SIZE, as large nonces are hashed to a shorter
+         * size, which can then lead to collisions if you encrypt a very
+         * large number of messages.*/
+        if( nonce_length == 0 )
+        {
+            status = PSA_ERROR_INVALID_ARGUMENT;
+            goto exit;
+        }
+    }
+#endif /* MBEDTLS_PSA_BUILTIN_ALG_GCM */
+#if defined(MBEDTLS_PSA_BUILTIN_ALG_CCM)
+    if( operation->alg == PSA_ALG_CCM )
+    {
+        if( nonce_length < 7 || nonce_length > 13 )
+        {
+            status = PSA_ERROR_INVALID_ARGUMENT;
+            goto exit;
+        }
+    }
+    else
+#endif /* MBEDTLS_PSA_BUILTIN_ALG_CCM */
+#if defined(MBEDTLS_PSA_BUILTIN_ALG_CHACHA20_POLY1305)
+        if( operation->alg == PSA_ALG_CHACHA20_POLY1305 )
+        {
+            if( nonce_length != 12 )
+            {
+                status = PSA_ERROR_INVALID_ARGUMENT;
+                goto exit;
+            }
+        }
+#endif /* MBEDTLS_PSA_BUILTIN_ALG_CHACHA20_POLY1305 */
 
     status = psa_driver_wrapper_aead_set_nonce( operation, nonce,
                                                 nonce_length );
diff --git a/library/psa_crypto_aead.c b/library/psa_crypto_aead.c
index d7317bd1a..4f6e70809 100644
--- a/library/psa_crypto_aead.c
+++ b/library/psa_crypto_aead.c
@@ -438,12 +438,6 @@ psa_status_t mbedtls_psa_aead_set_nonce(
 {
     psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
 
-    if( mbedtls_aead_check_nonce_length( operation, nonce_length )
-        != PSA_SUCCESS )
-    {
-        return( PSA_ERROR_INVALID_ARGUMENT );
-    }
-
 #if defined(MBEDTLS_PSA_BUILTIN_ALG_GCM)
     if( operation->alg == PSA_ALG_GCM )
     {