enforce AEAD add_data before update
diff --git a/cryptography/exceptions.py b/cryptography/exceptions.py
index 8b28667..d56db21 100644
--- a/cryptography/exceptions.py
+++ b/cryptography/exceptions.py
@@ -20,5 +20,9 @@
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 89e5654..7c31589 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 @@
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 @@
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 7ec3cd2..087066b 100644
--- a/docs/exceptions.rst
+++ b/docs/exceptions.rst
@@ -7,12 +7,19 @@
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 296821a..2806efd 100644
--- a/tests/hazmat/primitives/test_block.py
+++ b/tests/hazmat/primitives/test_block.py
@@ -26,7 +26,7 @@
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 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 f286e02..ebb8b5c 100644
--- a/tests/hazmat/primitives/test_utils.py
+++ b/tests/hazmat/primitives/test_utils.py
@@ -3,7 +3,7 @@
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 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 2a99cab..8df02e7 100644
--- a/tests/hazmat/primitives/utils.py
+++ b/tests/hazmat/primitives/utils.py
@@ -7,7 +7,7 @@
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 @@
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 @@
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)