Optimize + test cases

2 files changed, 24 insertions, 8 deletions

diff --git a/cryptography/hazmat/primitives/padding.py b/cryptography/hazmat/primitives/padding.py
index 27a1face..5c7c7e91 100644
--- a/cryptography/hazmat/primitives/padding.py
+++ b/cryptography/hazmat/primitives/padding.py
@@ -48,11 +48,12 @@ class _PKCS7PaddingContext(object):
             raise ValueError("Context was already finalized")
 
         self._buffer += data
-        result = b""
 
-        while len(self._buffer) >= self.block_size // 8:
-            result += self._buffer[:self.block_size // 8]
-            self._buffer = self._buffer[self.block_size // 8:]
+        finished_blocks = len(self._buffer) // (self.block_size // 8)
+
+        result = self._buffer[:finished_blocks * (self.block_size // 8)]
+        self._buffer = self._buffer[finished_blocks * (self.block_size // 8):]
+
         return result
 
     def finalize(self):
@@ -79,10 +80,15 @@ class _PKCS7UnpaddingContext(object):
             raise ValueError("Context was already finalized")
 
         self._buffer += data
-        result = b""
-        while len(self._buffer) >= 2 * (self.block_size // 8):
-            result += self._buffer[:self.block_size // 8]
-            self._buffer = self._buffer[self.block_size // 8:]
+
+        finished_blocks = max(
+            len(self._buffer) // (self.block_size // 8) - 1,
+            0
+        )
+
+        result = self._buffer[:finished_blocks * (self.block_size // 8)]
+        self._buffer = self._buffer[finished_blocks * (self.block_size // 8):]
+
         return result
 
     def finalize(self):
diff --git a/tests/hazmat/primitives/test_padding.py b/tests/hazmat/primitives/test_padding.py
index c47b6c5c..f689dda0 100644
--- a/tests/hazmat/primitives/test_padding.py
+++ b/tests/hazmat/primitives/test_padding.py
@@ -47,6 +47,16 @@ class TestPKCS7(object):
             b"111111111111111122222222222222",
             b"111111111111111122222222222222\x02\x02",
         ),
+        (
+            128,
+            b"1" * 16,
+            b"1" * 16 + b"\x10" * 16,
+        ),
+        (
+            128,
+            b"1" * 17,
+            b"1" * 17 + b"\x0F" * 15,
+        )
     ])
     def test_pad(self, size, unpadded, padded):
         padder = padding.PKCS7(size)