put the AAD and encrypted byte limit checks in the parent context

2 files changed, 21 insertions, 0 deletions

diff --git a/src/cryptography/hazmat/primitives/ciphers/base.py b/src/cryptography/hazmat/primitives/ciphers/base.py
index 8f3028fc..dae93655 100644
--- a/src/cryptography/hazmat/primitives/ciphers/base.py
+++ b/src/cryptography/hazmat/primitives/ciphers/base.py
@@ -149,6 +149,8 @@ class _CipherContext(object):
 class _AEADCipherContext(object):
     def __init__(self, ctx):
         self._ctx = ctx
+        self._bytes_processed = 0
+        self._aad_bytes_processed = 0
         self._tag = None
         self._updated = False
 
@@ -156,6 +158,14 @@ class _AEADCipherContext(object):
         if self._ctx is None:
             raise AlreadyFinalized("Context was already finalized.")
         self._updated = True
+        self._bytes_processed += len(data)
+        if self._bytes_processed > self._ctx._mode._MAX_ENCRYPTED_BYTES:
+            raise ValueError(
+                "{0} has a maximum encrypted byte limit of {1}".format(
+                    self._ctx._mode.name, self._ctx._mode._MAX_ENCRYPTED_BYTES
+                )
+            )
+
         return self._ctx.update(data)
 
     def finalize(self):
@@ -171,6 +181,15 @@ class _AEADCipherContext(object):
             raise AlreadyFinalized("Context was already finalized.")
         if self._updated:
             raise AlreadyUpdated("Update has been called on this context.")
+
+        self._aad_bytes_processed += len(data)
+        if self._aad_bytes_processed > self._ctx._mode._MAX_AAD_BYTES:
+            raise ValueError(
+                "{0} has a maximum AAD byte limit of {0}".format(
+                    self._ctx._mode.name, self._ctx._mode._MAX_AAD_BYTES
+                )
+            )
+
         self._ctx.authenticate_additional_data(data)
 
 
diff --git a/src/cryptography/hazmat/primitives/ciphers/modes.py b/src/cryptography/hazmat/primitives/ciphers/modes.py
index e31c9060..4284042d 100644
--- a/src/cryptography/hazmat/primitives/ciphers/modes.py
+++ b/src/cryptography/hazmat/primitives/ciphers/modes.py
@@ -139,6 +139,8 @@ class CTR(object):
 @utils.register_interface(ModeWithAuthenticationTag)
 class GCM(object):
     name = "GCM"
+    _MAX_ENCRYPTED_BYTES = (2 ** 39 - 256) // 8
+    _MAX_AAD_BYTES = (2 ** 64) // 8
 
     def __init__(self, initialization_vector, tag=None, min_tag_length=16):
         # len(initialization_vector) must in [1, 2 ** 64), but it's impossible