Move buffer size checks up to psa_crypto layer

Signed-off-by: Paul Elliott <paul.elliott@arm.com>
diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index a9026e4..a5027f3 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -3749,8 +3749,11 @@
     return( status );
 }
 
-static psa_status_t psa_aead_final_checks( psa_aead_operation_t *operation )
+static psa_status_t psa_aead_final_checks( psa_aead_operation_t *operation,
+                                           size_t output_size )
 {
+    size_t finish_output_size;
+
     if( operation->id == 0 || !operation->nonce_set )
         return( PSA_ERROR_BAD_STATE );
 
@@ -3758,6 +3761,13 @@
                                    operation->body_remaining != 0 ) )
         return( PSA_ERROR_INVALID_ARGUMENT );
 
+    finish_output_size = operation->is_encrypt ?
+         PSA_AEAD_FINISH_OUTPUT_SIZE( operation->key_type, operation->alg ) :
+         PSA_AEAD_VERIFY_OUTPUT_SIZE( operation->key_type, operation->alg );
+
+    if( output_size < finish_output_size )
+        return( PSA_ERROR_BUFFER_TOO_SMALL );
+
     return( PSA_SUCCESS );
 }
 
@@ -3775,7 +3785,7 @@
     *ciphertext_length = 0;
     *tag_length = tag_size;
 
-    status = psa_aead_final_checks( operation );
+    status = psa_aead_final_checks( operation, ciphertext_size );
 
     if( status != PSA_SUCCESS )
         goto exit;
@@ -3816,7 +3826,7 @@
 
     *plaintext_length = 0;
 
-    status = psa_aead_final_checks( operation );
+    status = psa_aead_final_checks( operation, plaintext_size );
 
     if( status != PSA_SUCCESS )
         goto exit;