Merge pull request #260 from dreid/hmac-already-finalized

Make HMAC methods raise AlreadyFinalized.

3 files changed, 36 insertions, 3 deletions

diff --git a/cryptography/hazmat/primitives/hmac.py b/cryptography/hazmat/primitives/hmac.py
index 1a67b332..1bbe39c7 100644
--- a/cryptography/hazmat/primitives/hmac.py
+++ b/cryptography/hazmat/primitives/hmac.py
@@ -15,6 +15,7 @@ from __future__ import absolute_import, division, print_function
 
 import six
 
+from cryptography.exceptions import AlreadyFinalized
 from cryptography.hazmat.primitives import interfaces
 
 
@@ -37,11 +38,15 @@ class HMAC(object):
             self._ctx = ctx
 
     def update(self, msg):
+        if self._ctx is None:
+            raise AlreadyFinalized("Context was already finalized")
         if isinstance(msg, six.text_type):
             raise TypeError("Unicode-objects must be encoded before hashing")
         self._ctx.update(msg)
 
     def copy(self):
+        if self._ctx is None:
+            raise AlreadyFinalized("Context was already finalized")
         return HMAC(
             self._key,
             self.algorithm,
@@ -50,4 +55,8 @@ class HMAC(object):
         )
 
     def finalize(self):
-        return self._ctx.finalize()
+        if self._ctx is None:
+            raise AlreadyFinalized("Context was already finalized")
+        digest = self._ctx.finalize()
+        self._ctx = None
+        return digest
diff --git a/docs/hazmat/primitives/hmac.rst b/docs/hazmat/primitives/hmac.rst
index bd1a4934..cff2dbf1 100644
--- a/docs/hazmat/primitives/hmac.rst
+++ b/docs/hazmat/primitives/hmac.rst
@@ -36,15 +36,25 @@ message.
     .. method:: update(msg)
 
         :param bytes msg: The bytes to hash and authenticate.
+        :raises cryptography.exceptions.AlreadyFinalized: See :meth:`finalize`
 
     .. method:: copy()
 
-        :return: a new instance of this object with a copied internal state.
+        Copy this :class:`HMAC` instance, usually so that we may call
+        :meth:`finalize` and get an intermediate digest value while we continue
+        to call :meth:`update` on the original.
+
+        :return: A new instance of :class:`HMAC` which can be updated
+            and finalized independently of the original instance.
+        :raises cryptography.exceptions.AlreadyFinalized: See :meth:`finalize`
 
     .. method:: finalize()
 
         Finalize the current context and return the message digest as bytes.
 
-        Once ``finalize`` is called this object can no longer be used.
+        Once ``finalize`` is called this object can no longer be used and
+        :meth:`update`, :meth:`copy`, and :meth:`finalize` will raise
+        :class:`~cryptography.exceptions.AlreadyFinalized`.
 
         :return bytes: The message digest as bytes.
+        :raises cryptography.exceptions.AlreadyFinalized:
diff --git a/tests/hazmat/primitives/test_hmac.py b/tests/hazmat/primitives/test_hmac.py
index 4186967a..d17049e3 100644
--- a/tests/hazmat/primitives/test_hmac.py
+++ b/tests/hazmat/primitives/test_hmac.py
@@ -19,6 +19,7 @@ import pytest
 
 import six
 
+from cryptography.exceptions import AlreadyFinalized
 from cryptography.hazmat.primitives import hashes, hmac
 
 from .utils import generate_base_hmac_test
@@ -49,3 +50,16 @@ class TestHMAC(object):
     def test_hmac_algorithm_instance(self):
         with pytest.raises(TypeError):
             hmac.HMAC(b"key", hashes.SHA1)
+
+    def test_raises_after_finalize(self):
+        h = hmac.HMAC(b"key", hashes.SHA1())
+        h.finalize()
+
+        with pytest.raises(AlreadyFinalized):
+            h.update(b"foo")
+
+        with pytest.raises(AlreadyFinalized):
+            h.copy()
+
+        with pytest.raises(AlreadyFinalized):
+            h.finalize()