Move AEAD length checks to PSA core

Signed-off-by: Paul Elliott <paul.elliott@arm.com>
diff --git a/include/psa/crypto_builtin_primitives.h b/include/psa/crypto_builtin_primitives.h
index e3903bc..b28e0d7 100644
--- a/include/psa/crypto_builtin_primitives.h
+++ b/include/psa/crypto_builtin_primitives.h
@@ -130,7 +130,6 @@
     psa_algorithm_t alg;
     psa_key_type_t key_type;
 
-    unsigned int lengths_set : 1;
     unsigned int is_encrypt : 1;
     unsigned int ad_started : 1;
     unsigned int body_started : 1;
@@ -138,9 +137,6 @@
     uint8_t tag_length;
     uint8_t nonce_length;
 
-    size_t ad_remaining;
-    size_t body_remaining;
-
     /* Buffers for AD/data - only required until CCM gets proper multipart
        support. */
     uint8_t *ad_buffer;
@@ -172,7 +168,7 @@
 
 } mbedtls_psa_aead_operation_t;
 
-#define MBEDTLS_PSA_AEAD_OPERATION_INIT {0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, {0}, {0}}
+#define MBEDTLS_PSA_AEAD_OPERATION_INIT {0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, {0}, {0}}
 
 /*
  * BEYOND THIS POINT, TEST DRIVER DECLARATIONS ONLY.
diff --git a/include/psa/crypto_struct.h b/include/psa/crypto_struct.h
index 36503f9..0f74c54 100644
--- a/include/psa/crypto_struct.h
+++ b/include/psa/crypto_struct.h
@@ -165,6 +165,9 @@
     psa_algorithm_t alg;
     psa_key_type_t key_type;
 
+    size_t ad_remaining;
+    size_t body_remaining;
+
     unsigned int nonce_set : 1;
     unsigned int lengths_set : 1;
     unsigned int ad_started : 1;
@@ -173,7 +176,7 @@
     psa_driver_aead_context_t ctx;
 };
 
-#define PSA_AEAD_OPERATION_INIT {0, 0, 0, 0, 0, 0, 0, {0}}
+#define PSA_AEAD_OPERATION_INIT {0, 0, 0, 0, 0, 0, 0, 0, 0, {0}}
 static inline struct psa_aead_operation_s psa_aead_operation_init( void )
 {
     const struct psa_aead_operation_s v = PSA_AEAD_OPERATION_INIT;