Improve nonce length checks

Add the missing nonce length checks (this function is being used by
oneshot functions as well as multipart, and thus all cipher suites are
being used) and cover the case where a NULL buffer gets passed in.
Extended the set nonce test to cover this.

Signed-off-by: Paul Elliott <paul.elliott@arm.com>
diff --git a/library/psa_crypto_aead.c b/library/psa_crypto_aead.c
index d877638..92c5ccf 100644
--- a/library/psa_crypto_aead.c
+++ b/library/psa_crypto_aead.c
@@ -141,6 +141,21 @@
     mbedtls_psa_aead_operation_t *operation,
     size_t nonce_length )
 {
+#if defined(MBEDTLS_PSA_BUILTIN_ALG_GCM)
+    if( operation->alg == PSA_ALG_GCM )
+    {
+        if( nonce_length == 0 )
+            return( PSA_ERROR_NOT_SUPPORTED );
+    }
+#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 )
+            return( PSA_ERROR_NOT_SUPPORTED );
+    }
+    else
+#endif /* MBEDTLS_PSA_BUILTIN_ALG_CCM */
 #if defined(MBEDTLS_PSA_BUILTIN_ALG_CHACHA20_POLY1305)
     if( operation->alg == PSA_ALG_CHACHA20_POLY1305 )
     {
@@ -428,7 +443,7 @@
     psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
 
     if( mbedtls_aead_check_nonce_length( operation, nonce_length )
-        != PSA_SUCCESS )
+        != PSA_SUCCESS || nonce == NULL )
     {
         return( PSA_ERROR_INVALID_ARGUMENT );
     }
diff --git a/tests/suites/test_suite_psa_crypto.data b/tests/suites/test_suite_psa_crypto.data
index d223537..f2355d6 100644
--- a/tests/suites/test_suite_psa_crypto.data
+++ b/tests/suites/test_suite_psa_crypto.data
@@ -2598,33 +2598,41 @@
 depends_on:PSA_WANT_ALG_CHACHA20_POLY1305:PSA_WANT_KEY_TYPE_CHACHA20
 aead_multipart_generate_nonce:PSA_KEY_TYPE_CHACHA20:"808182838485868788898a8b8c8d8e8f909192939495969798999a9b9c9d9e9f":PSA_ALG_CHACHA20_POLY1305:16:12:"":"":PSA_SUCCESS
 
+PSA Multipart Set Nonce, AES - GCM, IV = 0 (NULL)
+depends_on:PSA_WANT_ALG_GCM:PSA_WANT_KEY_TYPE_AES
+aead_multipart_set_nonce:PSA_KEY_TYPE_AES:"aa740abfadcda779220d3b406c5d7ec09a77fe9d94104539":PSA_ALG_AEAD_WITH_SHORTENED_TAG(PSA_ALG_GCM,16):0:1:"":"":PSA_ERROR_INVALID_ARGUMENT
+
 PSA Multipart Set Nonce, AES - GCM, IV = 0
 depends_on:PSA_WANT_ALG_GCM:PSA_WANT_KEY_TYPE_AES
-aead_multipart_set_nonce:PSA_KEY_TYPE_AES:"aa740abfadcda779220d3b406c5d7ec09a77fe9d94104539":PSA_ALG_AEAD_WITH_SHORTENED_TAG(PSA_ALG_GCM,16):0:"":"":PSA_ERROR_INVALID_ARGUMENT
+aead_multipart_set_nonce:PSA_KEY_TYPE_AES:"aa740abfadcda779220d3b406c5d7ec09a77fe9d94104539":PSA_ALG_AEAD_WITH_SHORTENED_TAG(PSA_ALG_GCM,16):0:0:"":"":PSA_ERROR_INVALID_ARGUMENT
 
 PSA Multipart Set Nonce, AES - GCM, IV = 16
 depends_on:PSA_WANT_ALG_GCM:PSA_WANT_KEY_TYPE_AES
-aead_multipart_set_nonce:PSA_KEY_TYPE_AES:"aa740abfadcda779220d3b406c5d7ec09a77fe9d94104539":PSA_ALG_AEAD_WITH_SHORTENED_TAG(PSA_ALG_GCM,16):16:"290322092d57479e20f6281e331d95a9":"e7fb0631eebf9bdba87045b33650c4ce":PSA_SUCCESS
+aead_multipart_set_nonce:PSA_KEY_TYPE_AES:"aa740abfadcda779220d3b406c5d7ec09a77fe9d94104539":PSA_ALG_AEAD_WITH_SHORTENED_TAG(PSA_ALG_GCM,16):16:0:"290322092d57479e20f6281e331d95a9":"e7fb0631eebf9bdba87045b33650c4ce":PSA_SUCCESS
 
 PSA Multipart Set Nonce, AES - GCM, IV = 20
 depends_on:PSA_WANT_ALG_GCM:PSA_WANT_KEY_TYPE_AES
-aead_multipart_set_nonce:PSA_KEY_TYPE_AES:"aa740abfadcda779220d3b406c5d7ec09a77fe9d94104539":PSA_ALG_AEAD_WITH_SHORTENED_TAG(PSA_ALG_GCM,16):20:"290322092d57479e20f6281e331d95a9":"e7fb0631eebf9bdba87045b33650c4ce":PSA_SUCCESS
+aead_multipart_set_nonce:PSA_KEY_TYPE_AES:"aa740abfadcda779220d3b406c5d7ec09a77fe9d94104539":PSA_ALG_AEAD_WITH_SHORTENED_TAG(PSA_ALG_GCM,16):20:0:"290322092d57479e20f6281e331d95a9":"e7fb0631eebf9bdba87045b33650c4ce":PSA_SUCCESS
 
 PSA Multipart Set Nonce: ChaCha20 - Poly1305, IV = 12
 depends_on:PSA_WANT_ALG_CHACHA20_POLY1305:PSA_WANT_KEY_TYPE_CHACHA20
-aead_multipart_set_nonce:PSA_KEY_TYPE_CHACHA20:"808182838485868788898a8b8c8d8e8f909192939495969798999a9b9c9d9e9f":PSA_ALG_CHACHA20_POLY1305:12:"50515253c0c1c2c3c4c5c6c7":"4c616469657320616e642047656e746c656d656e206f662074686520636c617373206f66202739393a204966204920636f756c64206f6666657220796f75206f6e6c79206f6e652074697020666f7220746865206675747572652c2073756e73637265656e20776f756c642062652069742e":PSA_SUCCESS
+aead_multipart_set_nonce:PSA_KEY_TYPE_CHACHA20:"808182838485868788898a8b8c8d8e8f909192939495969798999a9b9c9d9e9f":PSA_ALG_CHACHA20_POLY1305:12:0:"50515253c0c1c2c3c4c5c6c7":"4c616469657320616e642047656e746c656d656e206f662074686520636c617373206f66202739393a204966204920636f756c64206f6666657220796f75206f6e6c79206f6e652074697020666f7220746865206675747572652c2073756e73637265656e20776f756c642062652069742e":PSA_SUCCESS
 
 PSA Multipart Set Nonce: ChaCha20 - Poly1305, IV = 8
 depends_on:PSA_WANT_ALG_CHACHA20_POLY1305:PSA_WANT_KEY_TYPE_CHACHA20
-aead_multipart_set_nonce:PSA_KEY_TYPE_CHACHA20:"808182838485868788898a8b8c8d8e8f909192939495969798999a9b9c9d9e9f":PSA_ALG_CHACHA20_POLY1305:8:"":"":PSA_ERROR_INVALID_ARGUMENT
+aead_multipart_set_nonce:PSA_KEY_TYPE_CHACHA20:"808182838485868788898a8b8c8d8e8f909192939495969798999a9b9c9d9e9f":PSA_ALG_CHACHA20_POLY1305:8:0:"":"":PSA_ERROR_INVALID_ARGUMENT
+
+PSA Multipart Set Nonce: ChaCha20 - Poly1305, IV = 0 (NULL)
+depends_on:PSA_WANT_ALG_CHACHA20_POLY1305:PSA_WANT_KEY_TYPE_CHACHA20
+aead_multipart_set_nonce:PSA_KEY_TYPE_CHACHA20:"808182838485868788898a8b8c8d8e8f909192939495969798999a9b9c9d9e9f":PSA_ALG_CHACHA20_POLY1305:0:1:"":"":PSA_ERROR_INVALID_ARGUMENT
 
 PSA Multipart Set Nonce: ChaCha20 - Poly1305, IV = 0
 depends_on:PSA_WANT_ALG_CHACHA20_POLY1305:PSA_WANT_KEY_TYPE_CHACHA20
-aead_multipart_set_nonce:PSA_KEY_TYPE_CHACHA20:"808182838485868788898a8b8c8d8e8f909192939495969798999a9b9c9d9e9f":PSA_ALG_CHACHA20_POLY1305:0:"":"":PSA_ERROR_INVALID_ARGUMENT
+aead_multipart_set_nonce:PSA_KEY_TYPE_CHACHA20:"808182838485868788898a8b8c8d8e8f909192939495969798999a9b9c9d9e9f":PSA_ALG_CHACHA20_POLY1305:0:0:"":"":PSA_ERROR_INVALID_ARGUMENT
 
 PSA Multipart Set Nonce: ChaCha20 - Poly1305, IV = 16
 depends_on:PSA_WANT_ALG_CHACHA20_POLY1305:PSA_WANT_KEY_TYPE_CHACHA20
-aead_multipart_set_nonce:PSA_KEY_TYPE_CHACHA20:"808182838485868788898a8b8c8d8e8f909192939495969798999a9b9c9d9e9f":PSA_ALG_CHACHA20_POLY1305:16:"":"":PSA_ERROR_INVALID_ARGUMENT
+aead_multipart_set_nonce:PSA_KEY_TYPE_CHACHA20:"808182838485868788898a8b8c8d8e8f909192939495969798999a9b9c9d9e9f":PSA_ALG_CHACHA20_POLY1305:16:0:"":"":PSA_ERROR_INVALID_ARGUMENT
 
 PSA AEAD output buffer test: AES - GCM, IN = 16, BUF = 10
 depends_on:PSA_WANT_ALG_GCM:PSA_WANT_KEY_TYPE_AES
diff --git a/tests/suites/test_suite_psa_crypto.function b/tests/suites/test_suite_psa_crypto.function
index b8023ee..58e4387 100644
--- a/tests/suites/test_suite_psa_crypto.function
+++ b/tests/suites/test_suite_psa_crypto.function
@@ -3771,6 +3771,7 @@
 void aead_multipart_set_nonce( int key_type_arg, data_t *key_data,
                                int alg_arg,
                                int nonce_len,
+                               int allow_null_nonce_buffer,
                                data_t *additional_data,
                                data_t *input_data,
                                int expected_status_arg )
@@ -3829,11 +3830,22 @@
 
     PSA_ASSERT( status );
 
-    ASSERT_ALLOC( nonce_buffer, nonce_len );
-
-    for( index = 0; index < nonce_len - 1; ++index)
+    if( nonce_len == 0 )
     {
-        nonce_buffer[index] = 'a' + index;
+        if( !allow_null_nonce_buffer )
+        {
+            /* Arbitrary size buffer, to test zero length valid buffer. */
+            ASSERT_ALLOC( nonce_buffer, 4 );
+        }
+    }
+    else
+    {
+        ASSERT_ALLOC( nonce_buffer, nonce_len );
+
+        for( index = 0; index < nonce_len - 1; ++index)
+        {
+            nonce_buffer[index] = 'a' + index;
+        }
     }
 
     status = psa_aead_set_nonce( &operation, nonce_buffer, nonce_len );