Rustify PKCS7 unpadding (pyca#11556)

* refacto: Added rust PKCS7Unpadding

refacto: removed check_pkcs7_padding function

refacto: removed python _PKCS7Unpadding

* took comment into account

src/cryptography/hazmat/bindings/_rust/__init__.pyi

@@ -6,14 +6,18 @@ 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):
     def __init__(self, block_size: int) -> None: ...
     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

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:

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]

src/rust/src/padding.rs

@@ -20,7 +20,6 @@ fn constant_time_lt(a: u8, b: u8) -> u8 {
     duplicate_msb_to_all(a ^ ((a ^ b) | (a.wrapping_sub(b) ^ b)))
 }
 
-#[pyo3::pyfunction]
 pub(crate) fn check_pkcs7_padding(data: &[u8]) -> bool {
     let mut mismatch = 0;
     let pad_size = *data.last().unwrap();
@@ -111,6 +110,65 @@ impl PKCS7PaddingContext {
     }
 }
 
+#[pyo3::pyclass]
+pub(crate) struct PKCS7UnpaddingContext {
+    block_size: usize,
+    buffer: Option<Vec<u8>>,
+}
+
+#[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<pyo3::Bound<'a, pyo3::types::PyBytes>> {
+        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 = v.drain(..result_size);
+                Ok(pyo3::types::PyBytes::new_bound(py, result.as_slice()))
+            }
+            None => Err(exceptions::already_finalized_error()),
+        }
+    }
+
+    pub(crate) fn finalize<'p>(
+        &mut self,
+        py: pyo3::Python<'p>,
+    ) -> CryptographyResult<pyo3::Bound<'p, pyo3::types::PyBytes>> {
+        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];
+                Ok(pyo3::types::PyBytes::new_bound(py, result))
+            }
+            None => Err(exceptions::already_finalized_error()),
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::constant_time_lt;

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):