diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index c53020a..fcc22e1 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -3429,6 +3429,12 @@
         goto exit;
     }
 
+    /* Not checking nonce size here as GCM spec allows almost abitrarily 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. */
+
     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 bbfc927..10849b2 100644
--- a/library/psa_crypto_aead.c
+++ b/library/psa_crypto_aead.c
@@ -388,11 +388,16 @@
     #if defined(MBEDTLS_PSA_BUILTIN_ALG_GCM)
     if( operation->alg == PSA_ALG_GCM )
     {
+        operation->nonce = mbedtls_calloc( 1, nonce_length );
+
+        if( operation->nonce == NULL )
+            return( PSA_ERROR_INSUFFICIENT_MEMORY );
+
         /* GCM sets nonce once additional data has been supplied */
         memcpy( operation->nonce, nonce, nonce_length );
 
         /* We know that nonce size cannot exceed the uint8_t size */
-        operation->nonce_length = ( uint8_t ) nonce_length;
+        operation->nonce_length = nonce_length;
         status = PSA_SUCCESS;
     }
     else
@@ -400,12 +405,17 @@
 #if defined(MBEDTLS_PSA_BUILTIN_ALG_CCM)
     if( operation->alg == PSA_ALG_CCM )
     {
+        operation->nonce = mbedtls_calloc( 1, nonce_length );
+
+        if( operation->nonce == NULL )
+            return( PSA_ERROR_INSUFFICIENT_MEMORY );
+
         /* Multipart CCM not supported as yet, so CCM is basically operating
            in oneshot mode. Store the nonce as we need this later */
         memcpy( operation->nonce, nonce, nonce_length );
 
         /* We know that nonce size cannot exceed the uint8_t size */
-        operation->nonce_length = ( uint8_t ) nonce_length;
+        operation->nonce_length = nonce_length;
         status = PSA_SUCCESS;
     }
     else
@@ -919,6 +929,10 @@
     mbedtls_free( operation->tag_buffer );
     operation->tag_buffer = NULL;
 
+    mbedtls_free( operation->nonce );
+    operation->nonce = NULL;
+    operation->nonce_length = 0;
+
     return( PSA_SUCCESS );
 }
 
diff --git a/library/psa_crypto_aead.h b/library/psa_crypto_aead.h
index fcac5ca..ef4842e 100644
--- a/library/psa_crypto_aead.h
+++ b/library/psa_crypto_aead.h
@@ -263,6 +263,8 @@
  * \retval #PSA_ERROR_NOT_SUPPORTED
  *         Algorithm previously set is not supported in this configuration of
  *         the library.
+ * \retval #PSA_ERROR_INSUFFICIENT_MEMORY
+ *         (GCM and CCM only) Unable to allocate buffer for nonce.
  */
 psa_status_t mbedtls_psa_aead_set_nonce(
     mbedtls_psa_aead_operation_t *operation,