add support for update_into on CipherContext (#3190)

* add support for update_into on CipherContext

This allows you to provide your own buffer (like recv_into) to improve
performance when repeatedly calling encrypt/decrypt on large payloads.

* another skip_if

* more skip_if complexity

* maybe do this right

* correct number of args

* coverage for the coverage gods

* add a cffi minimum test tox target and travis builder

This tests against macOS so we capture some commoncrypto branches

* extra arg

* need to actually install py35

* fix

* coverage for GCM decrypt in CC

* no longer relevant

* 1.8 now

* pep8

* dramatically simplify

* update docs

* remove unneeded test

* changelog entry

* test improvements

* coverage fix

* add some comments to example

* move the comments to their own line

* fix and move comment

9 files changed, 316 insertions, 6 deletions

diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index 787e3072..1616947c 100644
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -11,7 +11,9 @@ Changelog
 * Changed ASN.1 dependency from ``pyasn1`` to ``asn1crypto`` resulting in a
   general performance increase when encoding/decoding ASN.1 structures. Also,
   the ``pyasn1_modules`` test dependency is no longer required.
-
+* Added support for
+  :meth:`~cryptography.hazmat.primitives.ciphers.CipherContext.update_into` on
+  :class:`~cryptography.hazmat.primitives.ciphers.CipherContext`.
 * Added
   :meth:`~cryptography.hazmat.primitives.asymmetric.dh.DHPrivateKeyWithSerialization.private_bytes`
   to
diff --git a/docs/hazmat/primitives/symmetric-encryption.rst b/docs/hazmat/primitives/symmetric-encryption.rst
index 24b2c045..1fd5a546 100644
--- a/docs/hazmat/primitives/symmetric-encryption.rst
+++ b/docs/hazmat/primitives/symmetric-encryption.rst
@@ -456,6 +456,49 @@ Interfaces
         return bytes immediately, however in other modes it will return chunks
         whose size is determined by the cipher's block size.
 
+    .. method:: update_into(data, buf)
+
+        .. versionadded:: 1.8
+
+        .. warning::
+
+            This method allows you to avoid a memory copy by passing a writable
+            buffer and reading the resulting data. You are responsible for
+            correctly sizing the buffer and properly handling the data. This
+            method should only be used when extremely high performance is a
+            requirement and you will be making many small calls to
+            ``update_into``.
+
+        :param bytes data: The data you wish to pass into the context.
+        :param buf: A writable Python buffer that the data will be written
+            into. This buffer should be ``len(data) + n - 1`` bytes where ``n``
+            is the block size (in bytes) of the cipher and mode combination.
+        :return int: Number of bytes written.
+        :raises NotImplementedError: This is raised if the version of ``cffi``
+            used is too old (this can happen on older PyPy releases).
+        :raises ValueError: This is raised if the supplied buffer is too small.
+
+        .. doctest::
+
+            >>> import os
+            >>> from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
+            >>> from cryptography.hazmat.backends import default_backend
+            >>> backend = default_backend()
+            >>> key = os.urandom(32)
+            >>> iv = os.urandom(16)
+            >>> cipher = Cipher(algorithms.AES(key), modes.CBC(iv), backend=backend)
+            >>> encryptor = cipher.encryptor()
+            >>> # the buffer needs to be at least len(data) + n - 1 where n is cipher/mode block size in bytes
+            >>> buf = bytearray(31)
+            >>> len_encrypted = encryptor.update_into(b"a secret message", buf)
+            >>> # get the ciphertext from the buffer reading only the bytes written to it (len_encrypted)
+            >>> ct = bytes(buf[:len_encrypted]) + encryptor.finalize()
+            >>> decryptor = cipher.decryptor()
+            >>> len_decrypted = decryptor.update_into(ct, buf)
+            >>> # get the plaintext from the buffer reading only the bytes written (len_decrypted)
+            >>> bytes(buf[:len_decrypted]) + decryptor.finalize()
+            'a secret message'
+
     .. method:: finalize()
 
         :return bytes: Returns the remainder of the data.
diff --git a/setup.py b/setup.py
index 1b1ff608..1ec3a354 100644
--- a/setup.py
+++ b/setup.py
@@ -36,6 +36,7 @@ VECTORS_DEPENDENCY = "cryptography_vectors=={0}".format(about['__version__'])
 requirements = [
     "idna>=2.0",
     "asn1crypto>=0.21.0",
+    "packaging",
     "six>=1.4.1",
     "setuptools>=11.3",
 ]
diff --git a/src/cryptography/hazmat/backends/commoncrypto/ciphers.py b/src/cryptography/hazmat/backends/commoncrypto/ciphers.py
index 1ce8aec5..b59381cb 100644
--- a/src/cryptography/hazmat/backends/commoncrypto/ciphers.py
+++ b/src/cryptography/hazmat/backends/commoncrypto/ciphers.py
@@ -86,6 +86,24 @@ class _CipherContext(object):
         self._backend._check_cipher_response(res)
         return self._backend._ffi.buffer(buf)[:outlen[0]]
 
+    def update_into(self, data, buf):
+        if len(buf) < (len(data) + self._byte_block_size - 1):
+            raise ValueError(
+                "buffer must be at least {0} bytes for this "
+                "payload".format(len(data) + self._byte_block_size - 1)
+            )
+        # Count bytes processed to handle block alignment.
+        self._bytes_processed += len(data)
+        outlen = self._backend._ffi.new("size_t *")
+        buf = self._backend._ffi.cast(
+            "unsigned char *", self._backend._ffi.from_buffer(buf)
+        )
+        res = self._backend._lib.CCCryptorUpdate(
+            self._ctx[0], data, len(data), buf,
+            len(data) + self._byte_block_size - 1, outlen)
+        self._backend._check_cipher_response(res)
+        return outlen[0]
+
     def finalize(self):
         # Raise error if block alignment is wrong.
         if self._bytes_processed % self._byte_block_size:
@@ -161,6 +179,24 @@ class _GCMCipherContext(object):
         self._backend._check_cipher_response(res)
         return self._backend._ffi.buffer(buf)[:]
 
+    def update_into(self, data, buf):
+        if len(buf) < len(data):
+            raise ValueError(
+                "buffer must be at least {0} bytes".format(len(data))
+            )
+
+        buf = self._backend._ffi.cast(
+            "unsigned char *", self._backend._ffi.from_buffer(buf)
+        )
+        args = (self._ctx[0], data, len(data), buf)
+        if self._operation == self._backend._lib.kCCEncrypt:
+            res = self._backend._lib.CCCryptorGCMEncrypt(*args)
+        else:
+            res = self._backend._lib.CCCryptorGCMDecrypt(*args)
+
+        self._backend._check_cipher_response(res)
+        return len(data)
+
     def finalize(self):
         # CommonCrypto has a yet another bug where you must make at least one
         # call to update. If you pass just AAD and call finalize without a call
diff --git a/src/cryptography/hazmat/backends/openssl/ciphers.py b/src/cryptography/hazmat/backends/openssl/ciphers.py
index 898b3497..0e0918af 100644
--- a/src/cryptography/hazmat/backends/openssl/ciphers.py
+++ b/src/cryptography/hazmat/backends/openssl/ciphers.py
@@ -109,6 +109,22 @@ class _CipherContext(object):
         self._backend.openssl_assert(res != 0)
         return self._backend._ffi.buffer(buf)[:outlen[0]]
 
+    def update_into(self, data, buf):
+        if len(buf) < (len(data) + self._block_size_bytes - 1):
+            raise ValueError(
+                "buffer must be at least {0} bytes for this "
+                "payload".format(len(data) + self._block_size_bytes - 1)
+            )
+
+        buf = self._backend._ffi.cast(
+            "unsigned char *", self._backend._ffi.from_buffer(buf)
+        )
+        outlen = self._backend._ffi.new("int *")
+        res = self._backend._lib.EVP_CipherUpdate(self._ctx, buf, outlen,
+                                                  data, len(data))
+        self._backend.openssl_assert(res != 0)
+        return outlen[0]
+
     def finalize(self):
         # OpenSSL 1.0.1 on Ubuntu 12.04 (and possibly other distributions)
         # appears to have a bug where you must make at least one call to update
diff --git a/src/cryptography/hazmat/primitives/ciphers/base.py b/src/cryptography/hazmat/primitives/ciphers/base.py
index 496975ae..502d9804 100644
--- a/src/cryptography/hazmat/primitives/ciphers/base.py
+++ b/src/cryptography/hazmat/primitives/ciphers/base.py
@@ -6,6 +6,8 @@ from __future__ import absolute_import, division, print_function
 
 import abc
 
+import cffi
+
 import six
 
 from cryptography import utils
@@ -51,6 +53,13 @@ class CipherContext(object):
         """
 
     @abc.abstractmethod
+    def update_into(self, data, buf):
+        """
+        Processes the provided bytes and writes the resulting data into the
+        provided buffer. Returns the number of bytes written.
+        """
+
+    @abc.abstractmethod
     def finalize(self):
         """
         Returns the results of processing the final block as bytes.
@@ -136,6 +145,20 @@ class _CipherContext(object):
             raise AlreadyFinalized("Context was already finalized.")
         return self._ctx.update(data)
 
+    # cffi 1.7 supports from_buffer on bytearray, which is required. We can
+    # remove this check in the future when we raise our minimum PyPy version.
+    if utils._version_check(cffi.__version__, "1.7"):
+        def update_into(self, data, buf):
+            if self._ctx is None:
+                raise AlreadyFinalized("Context was already finalized.")
+            return self._ctx.update_into(data, buf)
+    else:
+        def update_into(self, data, buf):
+            raise NotImplementedError(
+                "update_into requires cffi 1.7+. To use this method please "
+                "update cffi."
+            )
+
     def finalize(self):
         if self._ctx is None:
             raise AlreadyFinalized("Context was already finalized.")
@@ -154,11 +177,11 @@ class _AEADCipherContext(object):
         self._tag = None
         self._updated = False
 
-    def update(self, data):
+    def _check_limit(self, data_size):
         if self._ctx is None:
             raise AlreadyFinalized("Context was already finalized.")
         self._updated = True
-        self._bytes_processed += len(data)
+        self._bytes_processed += data_size
         if self._bytes_processed > self._ctx._mode._MAX_ENCRYPTED_BYTES:
             raise ValueError(
                 "{0} has a maximum encrypted byte limit of {1}".format(
@@ -166,8 +189,23 @@ class _AEADCipherContext(object):
                 )
             )
 
+    def update(self, data):
+        self._check_limit(len(data))
         return self._ctx.update(data)
 
+    # cffi 1.7 supports from_buffer on bytearray, which is required. We can
+    # remove this check in the future when we raise our minimum PyPy version.
+    if utils._version_check(cffi.__version__, "1.7"):
+        def update_into(self, data, buf):
+            self._check_limit(len(data))
+            return self._ctx.update_into(data, buf)
+    else:
+        def update_into(self, data, buf):
+            raise NotImplementedError(
+                "update_into requires cffi 1.7+. To use this method please "
+                "update cffi."
+            )
+
     def finalize(self):
         if self._ctx is None:
             raise AlreadyFinalized("Context was already finalized.")
diff --git a/src/cryptography/utils.py b/src/cryptography/utils.py
index f16b7efa..8183bdaf 100644
--- a/src/cryptography/utils.py
+++ b/src/cryptography/utils.py
@@ -10,6 +10,8 @@ import inspect
 import sys
 import warnings
 
+from packaging.version import parse
+
 
 # the functions deprecated in 1.0 and 1.4 are on an arbitrarily extended
 # deprecation cycle and should not be removed until we agree on when that cycle
@@ -98,6 +100,11 @@ else:
         return len(bin(x)) - (2 + (x <= 0))
 
 
+def _version_check(version, required_version):
+    # This is used to check if we support update_into on CipherContext.
+    return parse(version) >= parse(required_version)
+
+
 class _DeprecatedValue(object):
     def __init__(self, value, message, warning_class):
         self.value = value
diff --git a/tests/hazmat/primitives/test_block.py b/tests/hazmat/primitives/test_block.py
index eb0a2c3b..11a70195 100644
--- a/tests/hazmat/primitives/test_block.py
+++ b/tests/hazmat/primitives/test_block.py
@@ -6,6 +6,8 @@ from __future__ import absolute_import, division, print_function
 
 import binascii
 
+import cffi
+
 import pytest
 
 from cryptography.exceptions import (
@@ -15,6 +17,7 @@ from cryptography.hazmat.backends.interfaces import CipherBackend
 from cryptography.hazmat.primitives.ciphers import (
     Cipher, algorithms, base, modes
 )
+from cryptography.utils import _version_check
 
 from .utils import (
     generate_aead_exception_test, generate_aead_tag_exception_test
@@ -70,6 +73,23 @@ class TestCipherContext(object):
         with pytest.raises(AlreadyFinalized):
             decryptor.finalize()
 
+    @pytest.mark.skipif(
+        not _version_check(cffi.__version__, '1.7'),
+        reason="cffi version too old"
+    )
+    def test_use_update_into_after_finalize(self, backend):
+        cipher = Cipher(
+            algorithms.AES(binascii.unhexlify(b"0" * 32)),
+            modes.CBC(binascii.unhexlify(b"0" * 32)),
+            backend
+        )
+        encryptor = cipher.encryptor()
+        encryptor.update(b"a" * 16)
+        encryptor.finalize()
+        with pytest.raises(AlreadyFinalized):
+            buf = bytearray(31)
+            encryptor.update_into(b"b" * 16, buf)
+
     def test_unaligned_block_encryption(self, backend):
         cipher = Cipher(
             algorithms.AES(binascii.unhexlify(b"0" * 32)),
diff --git a/tests/hazmat/primitives/test_ciphers.py b/tests/hazmat/primitives/test_ciphers.py
index d9a07ff6..83952a87 100644
--- a/tests/hazmat/primitives/test_ciphers.py
+++ b/tests/hazmat/primitives/test_ciphers.py
@@ -5,17 +5,24 @@
 from __future__ import absolute_import, division, print_function
 
 import binascii
+import os
+
+import cffi
 
 import pytest
 
 from cryptography.exceptions import _Reasons
+from cryptography.hazmat.backends.interfaces import CipherBackend
 from cryptography.hazmat.primitives import ciphers
+from cryptography.hazmat.primitives.ciphers import modes
 from cryptography.hazmat.primitives.ciphers.algorithms import (
     AES, ARC4, Blowfish, CAST5, Camellia, IDEA, SEED, TripleDES
 )
-from cryptography.hazmat.primitives.ciphers.modes import ECB
+from cryptography.utils import _version_check
 
-from ...utils import raises_unsupported_algorithm
+from ...utils import (
+    load_nist_vectors, load_vectors_from_file, raises_unsupported_algorithm
+)
 
 
 class TestAES(object):
@@ -132,4 +139,144 @@ def test_invalid_backend():
     pretend_backend = object()
 
     with raises_unsupported_algorithm(_Reasons.BACKEND_MISSING_INTERFACE):
-        ciphers.Cipher(AES(b"AAAAAAAAAAAAAAAA"), ECB, pretend_backend)
+        ciphers.Cipher(AES(b"AAAAAAAAAAAAAAAA"), modes.ECB, pretend_backend)
+
+
+@pytest.mark.skipif(
+    not _version_check(cffi.__version__, '1.7'),
+    reason="cffi version too old"
+)
+@pytest.mark.supported(
+    only_if=lambda backend: backend.cipher_supported(
+        AES(b"\x00" * 16), modes.ECB()
+    ),
+    skip_message="Does not support AES ECB",
+)
+@pytest.mark.requires_backend_interface(interface=CipherBackend)
+class TestCipherUpdateInto(object):
+    @pytest.mark.parametrize(
+        "params",
+        load_vectors_from_file(
+            os.path.join("ciphers", "AES", "ECB", "ECBGFSbox128.rsp"),
+            load_nist_vectors
+        )
+    )
+    def test_update_into(self, params, backend):
+        key = binascii.unhexlify(params["key"])
+        pt = binascii.unhexlify(params["plaintext"])
+        ct = binascii.unhexlify(params["ciphertext"])
+        c = ciphers.Cipher(AES(key), modes.ECB(), backend)
+        encryptor = c.encryptor()
+        buf = bytearray(len(pt) + 15)
+        res = encryptor.update_into(pt, buf)
+        assert res == len(pt)
+        assert bytes(buf)[:res] == ct
+
+    @pytest.mark.supported(
+        only_if=lambda backend: backend.cipher_supported(
+            AES(b"\x00" * 16), modes.GCM(b"0" * 12)
+        ),
+        skip_message="Does not support AES GCM",
+    )
+    def test_update_into_gcm(self, backend):
+        key = binascii.unhexlify(b"e98b72a9881a84ca6b76e0f43e68647a")
+        iv = binascii.unhexlify(b"8b23299fde174053f3d652ba")
+        ct = binascii.unhexlify(b"5a3c1cf1985dbb8bed818036fdd5ab42")
+        pt = binascii.unhexlify(b"28286a321293253c3e0aa2704a278032")
+        c = ciphers.Cipher(AES(key), modes.GCM(iv), backend)
+        encryptor = c.encryptor()
+        buf = bytearray(len(pt) + 15)
+        res = encryptor.update_into(pt, buf)
+        assert res == len(pt)
+        assert bytes(buf)[:res] == ct
+        encryptor.finalize()
+        c = ciphers.Cipher(AES(key), modes.GCM(iv, encryptor.tag), backend)
+        decryptor = c.decryptor()
+        res = decryptor.update_into(ct, buf)
+        decryptor.finalize()
+        assert res == len(pt)
+        assert bytes(buf)[:res] == pt
+
+    @pytest.mark.parametrize(
+        "params",
+        load_vectors_from_file(
+            os.path.join("ciphers", "AES", "ECB", "ECBGFSbox128.rsp"),
+            load_nist_vectors
+        )
+    )
+    def test_update_into_multiple_calls(self, params, backend):
+        key = binascii.unhexlify(params["key"])
+        pt = binascii.unhexlify(params["plaintext"])
+        ct = binascii.unhexlify(params["ciphertext"])
+        c = ciphers.Cipher(AES(key), modes.ECB(), backend)
+        encryptor = c.encryptor()
+        buf = bytearray(len(pt) + 15)
+        res = encryptor.update_into(pt[:3], buf)
+        assert res == 0
+        res = encryptor.update_into(pt[3:], buf)
+        assert res == len(pt)
+        assert bytes(buf)[:res] == ct
+
+    def test_update_into_buffer_too_small(self, backend):
+        key = b"\x00" * 16
+        c = ciphers.Cipher(AES(key), modes.ECB(), backend)
+        encryptor = c.encryptor()
+        buf = bytearray(16)
+        with pytest.raises(ValueError):
+            encryptor.update_into(b"testing", buf)
+
+    @pytest.mark.supported(
+        only_if=lambda backend: backend.cipher_supported(
+            AES(b"\x00" * 16), modes.GCM(b"\x00" * 12)
+        ),
+        skip_message="Does not support AES GCM",
+    )
+    def test_update_into_buffer_too_small_gcm(self, backend):
+        key = b"\x00" * 16
+        c = ciphers.Cipher(AES(key), modes.GCM(b"\x00" * 12), backend)
+        encryptor = c.encryptor()
+        buf = bytearray(5)
+        with pytest.raises(ValueError):
+            encryptor.update_into(b"testing", buf)
+
+
+@pytest.mark.skipif(
+    _version_check(cffi.__version__, '1.7'),
+    reason="cffi version too new"
+)
+@pytest.mark.requires_backend_interface(interface=CipherBackend)
+class TestCipherUpdateIntoUnsupported(object):
+    def _too_old(self, mode, backend):
+        key = b"\x00" * 16
+        c = ciphers.Cipher(AES(key), mode, backend)
+        encryptor = c.encryptor()
+        buf = bytearray(32)
+        with pytest.raises(NotImplementedError):
+            encryptor.update_into(b"\x00" * 16, buf)
+
+    @pytest.mark.supported(
+        only_if=lambda backend: backend.cipher_supported(
+            AES(b"\x00" * 16), modes.ECB()
+        ),
+        skip_message="Does not support AES ECB",
+    )
+    def test_cffi_too_old_ecb(self, backend):
+        self._too_old(modes.ECB(), backend)
+
+    @pytest.mark.supported(
+        only_if=lambda backend: backend.cipher_supported(
+            AES(b"\x00" * 16), modes.CTR(b"0" * 16)
+        ),
+        skip_message="Does not support AES CTR",
+    )
+    def test_cffi_too_old_ctr(self, backend):
+        self._too_old(modes.CTR(b"0" * 16), backend)
+
+    @pytest.mark.supported(
+        only_if=lambda backend: backend.cipher_supported(
+            AES(b"\x00" * 16), modes.GCM(b"0" * 16)
+        ),
+        skip_message="Does not support AES GCM",
+    )
+    def test_cffi_too_old_gcm(self, backend):
+        self._too_old(modes.GCM(b"0" * 16), backend)