diff options
| -rw-r--r-- | cryptography/hazmat/primitives/padding.py | 22 | ||||
| -rw-r--r-- | tests/hazmat/primitives/test_padding.py | 10 |
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) |