diff --git a/src/cryptography/hazmat/bindings/_rust/__init__.pyi b/src/cryptography/hazmat/bindings/_rust/__init__.pyi index c0ea0a5405ca..30b67d85597e 100644 --- a/src/cryptography/hazmat/bindings/_rust/__init__.pyi +++ b/src/cryptography/hazmat/bindings/_rust/__init__.pyi @@ -6,7 +6,6 @@ import typing from cryptography.hazmat.primitives import padding -def check_pkcs7_padding(data: bytes) -> bool: ... def check_ansix923_padding(data: bytes) -> bool: ... class PKCS7PaddingContext(padding.PaddingContext): @@ -14,6 +13,11 @@ class PKCS7PaddingContext(padding.PaddingContext): def update(self, data: bytes) -> bytes: ... def finalize(self) -> bytes: ... +class PKCS7UnpaddingContext(padding.PaddingContext): + def __init__(self, block_size: int) -> None: ... + def update(self, data: bytes) -> bytes: ... + def finalize(self) -> bytes: ... + class ObjectIdentifier: def __init__(self, val: str) -> None: ... @property diff --git a/src/cryptography/hazmat/primitives/padding.py b/src/cryptography/hazmat/primitives/padding.py index d1ca775f33d0..b2a3f1cfffaa 100644 --- a/src/cryptography/hazmat/primitives/padding.py +++ b/src/cryptography/hazmat/primitives/padding.py @@ -11,8 +11,8 @@ from cryptography.exceptions import AlreadyFinalized from cryptography.hazmat.bindings._rust import ( PKCS7PaddingContext, + PKCS7UnpaddingContext, check_ansix923_padding, - check_pkcs7_padding, ) @@ -115,32 +115,11 @@ def padder(self) -> PaddingContext: return PKCS7PaddingContext(self.block_size) def unpadder(self) -> PaddingContext: - return _PKCS7UnpaddingContext(self.block_size) - - -class _PKCS7UnpaddingContext(PaddingContext): - _buffer: bytes | None - - def __init__(self, block_size: int): - self.block_size = block_size - # TODO: more copies than necessary, we should use zero-buffer (#193) - self._buffer = b"" - - def update(self, data: bytes) -> bytes: - self._buffer, result = _byte_unpadding_update( - self._buffer, data, self.block_size - ) - return result - - def finalize(self) -> bytes: - result = _byte_unpadding_check( - self._buffer, self.block_size, check_pkcs7_padding - ) - self._buffer = None - return result + return PKCS7UnpaddingContext(self.block_size) PaddingContext.register(PKCS7PaddingContext) +PaddingContext.register(PKCS7UnpaddingContext) class ANSIX923: diff --git a/src/rust/src/lib.rs b/src/rust/src/lib.rs index cd7b99f1570a..e15fffa6d32e 100644 --- a/src/rust/src/lib.rs +++ b/src/rust/src/lib.rs @@ -102,7 +102,7 @@ mod _rust { #[pymodule_export] use crate::oid::ObjectIdentifier; #[pymodule_export] - use crate::padding::{check_ansix923_padding, check_pkcs7_padding, PKCS7PaddingContext}; + use crate::padding::{check_ansix923_padding, PKCS7PaddingContext, PKCS7UnpaddingContext}; #[pymodule_export] use crate::pkcs12::pkcs12; #[pymodule_export] diff --git a/src/rust/src/padding.rs b/src/rust/src/padding.rs index 3a55039d3385..12737825f735 100644 --- a/src/rust/src/padding.rs +++ b/src/rust/src/padding.rs @@ -111,6 +111,65 @@ impl PKCS7PaddingContext { } } +#[pyo3::pyclass] +pub(crate) struct PKCS7UnpaddingContext { + block_size: usize, + buffer: Option>, +} + +#[pyo3::pymethods] +impl PKCS7UnpaddingContext { + #[new] + pub(crate) fn new(block_size: usize) -> PKCS7UnpaddingContext { + PKCS7UnpaddingContext { + block_size: block_size / 8, + buffer: Some(Vec::new()), + } + } + + pub(crate) fn update<'a>( + &mut self, + py: pyo3::Python<'a>, + buf: CffiBuf<'a>, + ) -> CryptographyResult> { + match self.buffer.as_mut() { + Some(v) => { + v.extend_from_slice(buf.as_bytes()); + let finished_blocks = (v.len() / self.block_size).saturating_sub(1); + let result_size = finished_blocks * self.block_size; + let result: Vec = v.drain(..result_size).collect(); + Ok(pyo3::types::PyBytes::new_bound(py, &result)) + } + None => Err(exceptions::already_finalized_error()), + } + } + + pub(crate) fn finalize<'p>( + &mut self, + py: pyo3::Python<'p>, + ) -> CryptographyResult> { + match self.buffer.take() { + Some(v) => { + if v.len() != self.block_size { + return Err( + pyo3::exceptions::PyValueError::new_err("Invalid padding bytes.").into(), + ); + } + if !check_pkcs7_padding(&v) { + return Err( + pyo3::exceptions::PyValueError::new_err("Invalid padding bytes.").into(), + ); + } + + let pad_size = *v.last().unwrap(); + let result = v[..v.len() - pad_size as usize].to_vec(); + Ok(pyo3::types::PyBytes::new_bound(py, &result)) + } + None => Err(exceptions::already_finalized_error()), + } + } +} + #[cfg(test)] mod tests { use super::constant_time_lt; diff --git a/tests/hazmat/primitives/test_padding.py b/tests/hazmat/primitives/test_padding.py index 0ab1125f5bfb..df1ee4ec1131 100644 --- a/tests/hazmat/primitives/test_padding.py +++ b/tests/hazmat/primitives/test_padding.py @@ -80,6 +80,8 @@ def test_pad(self, size, unpadded, padded): 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_unpad(self, size, unpadded, padded):