enforce AEAD add_data before update

6 files changed, 33 insertions, 14 deletions

diff --git a/cryptography/exceptions.py b/cryptography/exceptions.py
index 8b286679..d56db214 100644
--- a/cryptography/exceptions.py
+++ b/cryptography/exceptions.py
@@ -20,5 +20,9 @@ class AlreadyFinalized(Exception):
     pass
 
 
+class AlreadyUpdated(Exception):
+    pass
+
+
 class NotYetFinalized(Exception):
     pass
diff --git a/cryptography/hazmat/primitives/ciphers/base.py b/cryptography/hazmat/primitives/ciphers/base.py
index 89e56547..7c315898 100644
--- a/cryptography/hazmat/primitives/ciphers/base.py
+++ b/cryptography/hazmat/primitives/ciphers/base.py
@@ -14,7 +14,9 @@
 from __future__ import absolute_import, division, print_function
 
 from cryptography import utils
-from cryptography.exceptions import AlreadyFinalized, NotYetFinalized
+from cryptography.exceptions import (
+    AlreadyFinalized, NotYetFinalized, AlreadyUpdated,
+)
 from cryptography.hazmat.primitives import interfaces
 
 
@@ -80,10 +82,12 @@ class _AEADCipherContext(object):
     def __init__(self, ctx):
         self._ctx = ctx
         self._tag = None
+        self._updated = False
 
     def update(self, data):
         if self._ctx is None:
             raise AlreadyFinalized("Context was already finalized")
+        self._updated = True
         return self._ctx.update(data)
 
     def finalize(self):
@@ -97,6 +101,8 @@ class _AEADCipherContext(object):
     def add_data(self, data):
         if self._ctx is None:
             raise AlreadyFinalized("Context was already finalized")
+        if self._updated:
+            raise AlreadyUpdated("Update has been called on this context")
         self._ctx.add_data(data)
 
     @property
diff --git a/docs/exceptions.rst b/docs/exceptions.rst
index 7ec3cd27..087066b8 100644
--- a/docs/exceptions.rst
+++ b/docs/exceptions.rst
@@ -7,12 +7,19 @@ Exceptions
 
     This is raised when a context is used after being finalized.
 
-.. class:: NotFinalized
+
+.. class:: NotYetFinalized
 
     This is raised when the AEAD tag property is accessed on a context
     before it is finalized.
 
 
+.. class:: AlreadyUpdated
+
+    This is raised when additional data is added to a context after update
+    has already been called.
+
+
 .. class:: UnsupportedAlgorithm
 
     This is raised when a backend doesn't support the requested algorithm (or
diff --git a/tests/hazmat/primitives/test_block.py b/tests/hazmat/primitives/test_block.py
index 296821a4..2806efd5 100644
--- a/tests/hazmat/primitives/test_block.py
+++ b/tests/hazmat/primitives/test_block.py
@@ -26,7 +26,7 @@ from cryptography.hazmat.primitives.ciphers import (
     Cipher, algorithms, modes
 )
 
-from .utils import generate_aead_use_after_finalize_test
+from .utils import generate_aead_exception_test
 
 
 @utils.register_interface(interfaces.CipherAlgorithm)
@@ -127,7 +127,7 @@ class TestCipherContext(object):
 
 
 class TestAEADCipherContext(object):
-    test_use_after_finalize = generate_aead_use_after_finalize_test(
+    test_aead_exceptions = generate_aead_exception_test(
         algorithms.AES,
         modes.GCM,
         only_if=lambda backend: backend.cipher_supported(
diff --git a/tests/hazmat/primitives/test_utils.py b/tests/hazmat/primitives/test_utils.py
index f286e02d..ebb8b5c4 100644
--- a/tests/hazmat/primitives/test_utils.py
+++ b/tests/hazmat/primitives/test_utils.py
@@ -3,7 +3,7 @@ import pytest
 from .utils import (
     base_hash_test, encrypt_test, hash_test, long_string_hash_test,
     base_hmac_test, hmac_test, stream_encryption_test, aead_test,
-    aead_use_after_finalize_test,
+    aead_exception_test,
 )
 
 
@@ -32,7 +32,7 @@ class TestAEADTest(object):
 class TestAEADFinalizeTest(object):
     def test_skips_if_only_if_returns_false(self):
         with pytest.raises(pytest.skip.Exception) as exc_info:
-            aead_use_after_finalize_test(
+            aead_exception_test(
                 None, None, None,
                 only_if=lambda backend: False,
                 skip_message="message!"
diff --git a/tests/hazmat/primitives/utils.py b/tests/hazmat/primitives/utils.py
index 2a99cab9..8df02e78 100644
--- a/tests/hazmat/primitives/utils.py
+++ b/tests/hazmat/primitives/utils.py
@@ -7,7 +7,7 @@ from cryptography.hazmat.bindings import _ALL_BACKENDS
 from cryptography.hazmat.primitives import hashes, hmac
 from cryptography.hazmat.primitives.ciphers import Cipher
 from cryptography.exceptions import (
-    AlreadyFinalized, NotYetFinalized,
+    AlreadyFinalized, NotYetFinalized, AlreadyUpdated,
 )
 
 from ...utils import load_vectors_from_file
@@ -307,23 +307,23 @@ def base_hmac_test(backend, algorithm, only_if, skip_message):
     assert h._ctx != h_copy._ctx
 
 
-def generate_aead_use_after_finalize_test(cipher_factory, mode_factory,
-                                          only_if, skip_message):
-    def test_aead_use_after_finalize(self):
+def generate_aead_exception_test(cipher_factory, mode_factory,
+                                 only_if, skip_message):
+    def test_aead_exception(self):
         for backend in _ALL_BACKENDS:
             yield (
-                aead_use_after_finalize_test,
+                aead_exception_test,
                 backend,
                 cipher_factory,
                 mode_factory,
                 only_if,
                 skip_message
             )
-    return test_aead_use_after_finalize
+    return test_aead_exception
 
 
-def aead_use_after_finalize_test(backend, cipher_factory, mode_factory,
-                                 only_if, skip_message):
+def aead_exception_test(backend, cipher_factory, mode_factory,
+                        only_if, skip_message):
     if not only_if(backend):
         pytest.skip(skip_message)
     cipher = Cipher(
@@ -335,6 +335,8 @@ def aead_use_after_finalize_test(backend, cipher_factory, mode_factory,
     encryptor.update(b"a" * 16)
     with pytest.raises(NotYetFinalized):
         encryptor.tag
+    with pytest.raises(AlreadyUpdated):
+        encryptor.add_data(b"b" * 16)
     encryptor.finalize()
     with pytest.raises(AlreadyFinalized):
         encryptor.add_data(b"b" * 16)