diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index e97cbaf..c53020a 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -3467,7 +3467,11 @@
 exit:
 
     if( status == PSA_SUCCESS )
+    {
+        operation->ad_remaining = ad_length;
+        operation->body_remaining = plaintext_length;
         operation->lengths_set = 1;
+    }
     else
         psa_aead_abort( operation );
 
@@ -3492,6 +3496,17 @@
         goto exit;
     }
 
+    if( operation->lengths_set )
+    {
+        if ( operation->ad_remaining < input_length )
+        {
+            status = PSA_ERROR_INVALID_ARGUMENT;
+            goto exit;
+        }
+
+        operation->ad_remaining -= input_length;
+    }
+
     status = psa_driver_wrapper_aead_update_ad( operation, input,
                                                 input_length );
 
@@ -3530,6 +3545,26 @@
         goto exit;
     }
 
+    if( operation->lengths_set )
+    {
+        /* Additional data length was supplied, but not all the additional
+           data was supplied.*/
+        if( operation->ad_remaining != 0 )
+        {
+            status = PSA_ERROR_INVALID_ARGUMENT;
+            goto exit;
+        }
+
+        /* Too much data provided. */
+        if( operation->body_remaining < input_length )
+        {
+            status = PSA_ERROR_INVALID_ARGUMENT;
+            goto exit;
+        }
+
+        operation->body_remaining -= input_length;
+    }
+
     status = psa_driver_wrapper_aead_update( operation, input, input_length,
                                              output, output_size,
                                              output_length );
@@ -3571,6 +3606,13 @@
         goto exit;
     }
 
+    if( operation->lengths_set && (operation->ad_remaining != 0 ||
+                                   operation->body_remaining != 0 ) )
+    {
+        status = PSA_ERROR_BAD_STATE;
+        goto exit;
+    }
+
     status = psa_driver_wrapper_aead_finish( operation, ciphertext,
                                              ciphertext_size,
                                              ciphertext_length,
@@ -3609,6 +3651,13 @@
         goto exit;
     }
 
+    if( operation->lengths_set && (operation->ad_remaining != 0 ||
+                                   operation->body_remaining != 0 ) )
+    {
+        status = PSA_ERROR_BAD_STATE;
+        goto exit;
+    }
+
     status = psa_driver_wrapper_aead_verify( operation, plaintext,
                                              plaintext_size,
                                              plaintext_length,