Convert `PKCS7PaddingContext` to Rust (#11089)

7 files changed, 73 insertions, 39 deletions

diff --git a/docs/hazmat/primitives/padding.rst b/docs/hazmat/primitives/padding.rst
index ecd70e6..a1be2ab 100644
--- a/docs/hazmat/primitives/padding.rst
+++ b/docs/hazmat/primitives/padding.rst
@@ -24,16 +24,13 @@ multiple of the block size.
         >>> from cryptography.hazmat.primitives import padding
         >>> padder = padding.PKCS7(128).padder()
         >>> padded_data = padder.update(b"11111111111111112222222222")
-        >>> padded_data
-        b'1111111111111111'
         >>> padded_data += padder.finalize()
         >>> padded_data
         b'11111111111111112222222222\x06\x06\x06\x06\x06\x06'
         >>> unpadder = padding.PKCS7(128).unpadder()
         >>> data = unpadder.update(padded_data)
+        >>> data += unpadder.finalize()
         >>> data
-        b'1111111111111111'
-        >>> data + unpadder.finalize()
         b'11111111111111112222222222'
 
     :param block_size: The size of the block in :term:`bits` that the data is
@@ -67,16 +64,13 @@ multiple of the block size.
 
         >>> padder = padding.ANSIX923(128).padder()
         >>> padded_data = padder.update(b"11111111111111112222222222")
-        >>> padded_data
-        b'1111111111111111'
         >>> padded_data += padder.finalize()
         >>> padded_data
         b'11111111111111112222222222\x00\x00\x00\x00\x00\x06'
         >>> unpadder = padding.ANSIX923(128).unpadder()
         >>> data = unpadder.update(padded_data)
+        >>> data += unpadder.finalize()
         >>> data
-        b'1111111111111111'
-        >>> data + unpadder.finalize()
         b'11111111111111112222222222'
 
     :param block_size: The size of the block in :term:`bits` that the data is
diff --git a/src/cryptography/hazmat/bindings/_rust/__init__.pyi b/src/cryptography/hazmat/bindings/_rust/__init__.pyi
index 18a6fb8..c0ea0a5 100644
--- a/src/cryptography/hazmat/bindings/_rust/__init__.pyi
+++ b/src/cryptography/hazmat/bindings/_rust/__init__.pyi
@@ -4,9 +4,16 @@
 
 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 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 baceaf3..d1ca775 100644
--- a/src/cryptography/hazmat/primitives/padding.py
+++ b/src/cryptography/hazmat/primitives/padding.py
@@ -10,6 +10,7 @@ import typing
 from cryptography import utils
 from cryptography.exceptions import AlreadyFinalized
 from cryptography.hazmat.bindings._rust import (
+    PKCS7PaddingContext,
     check_ansix923_padding,
     check_pkcs7_padding,
 )
@@ -111,37 +112,12 @@ class PKCS7:
         self.block_size = block_size
 
     def padder(self) -> PaddingContext:
-        return _PKCS7PaddingContext(self.block_size)
+        return PKCS7PaddingContext(self.block_size)
 
     def unpadder(self) -> PaddingContext:
         return _PKCS7UnpaddingContext(self.block_size)
 
 
-class _PKCS7PaddingContext(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_padding_update(
-            self._buffer, data, self.block_size
-        )
-        return result
-
-    def _padding(self, size: int) -> bytes:
-        return bytes([size]) * size
-
-    def finalize(self) -> bytes:
-        result = _byte_padding_pad(
-            self._buffer, self.block_size, self._padding
-        )
-        self._buffer = None
-        return result
-
-
 class _PKCS7UnpaddingContext(PaddingContext):
     _buffer: bytes | None
 
@@ -164,6 +140,9 @@ class _PKCS7UnpaddingContext(PaddingContext):
         return result
 
 
+PaddingContext.register(PKCS7PaddingContext)
+
+
 class ANSIX923:
     def __init__(self, block_size: int):
         _byte_padding_check(block_size)
diff --git a/src/rust/src/buf.rs b/src/rust/src/buf.rs
index e077932..ff9ca0c 100644
--- a/src/rust/src/buf.rs
+++ b/src/rust/src/buf.rs
@@ -8,7 +8,7 @@ use pyo3::types::IntoPyDict;
 use std::slice;
 
 pub(crate) struct CffiBuf<'p> {
-    _pyobj: pyo3::Bound<'p, pyo3::PyAny>,
+    pyobj: pyo3::Bound<'p, pyo3::PyAny>,
     _bufobj: pyo3::Bound<'p, pyo3::PyAny>,
     buf: &'p [u8],
 }
@@ -34,10 +34,14 @@ fn _extract_buffer_length<'p>(
     Ok((bufobj, ptrval))
 }
 
-impl CffiBuf<'_> {
+impl<'a> CffiBuf<'a> {
     pub(crate) fn as_bytes(&self) -> &[u8] {
         self.buf
     }
+
+    pub(crate) fn into_pyobj(self) -> pyo3::Bound<'a, pyo3::PyAny> {
+        self.pyobj
+    }
 }
 
 impl<'a> pyo3::conversion::FromPyObject<'a> for CffiBuf<'a> {
@@ -59,7 +63,7 @@ impl<'a> pyo3::conversion::FromPyObject<'a> for CffiBuf<'a> {
         };
 
         Ok(CffiBuf {
-            _pyobj: pyobj.clone(),
+            pyobj: pyobj.clone(),
             _bufobj: bufobj,
             buf,
         })
diff --git a/src/rust/src/lib.rs b/src/rust/src/lib.rs
index ac076e6..da929fe 100644
--- a/src/rust/src/lib.rs
+++ b/src/rust/src/lib.rs
@@ -100,6 +100,7 @@ fn _rust(py: pyo3::Python<'_>, m: &pyo3::Bound<'_, pyo3::types::PyModule>) -> py
         padding::check_ansix923_padding,
         m
     )?)?;
+    m.add_class::<padding::PKCS7PaddingContext>()?;
     m.add_class::<oid::ObjectIdentifier>()?;
 
     m.add_submodule(&asn1::create_submodule(py)?)?;
diff --git a/src/rust/src/padding.rs b/src/rust/src/padding.rs
index 523fe85..c4396c2 100644
--- a/src/rust/src/padding.rs
+++ b/src/rust/src/padding.rs
@@ -2,6 +2,10 @@
 // 2.0, and the BSD License. See the LICENSE file in the root of this repository
 // for complete details.
 
+use crate::buf::CffiBuf;
+use crate::error::{CryptographyError, CryptographyResult};
+use crate::exceptions;
+
 /// Returns the value of the input with the most-significant-bit copied to all
 /// of the bits.
 fn duplicate_msb_to_all(a: u8) -> u8 {
@@ -63,6 +67,51 @@ pub(crate) fn check_ansix923_padding(data: &[u8]) -> bool {
     (mismatch & 1) == 0
 }
 
+#[pyo3::prelude::pyclass]
+pub(crate) struct PKCS7PaddingContext {
+    block_size: usize,
+    length_seen: Option<usize>,
+}
+
+#[pyo3::prelude::pymethods]
+impl PKCS7PaddingContext {
+    #[new]
+    fn new(block_size: usize) -> PKCS7PaddingContext {
+        PKCS7PaddingContext {
+            block_size: block_size / 8,
+            length_seen: Some(0),
+        }
+    }
+
+    fn update<'a>(&mut self, buf: CffiBuf<'a>) -> CryptographyResult<pyo3::Bound<'a, pyo3::PyAny>> {
+        match self.length_seen.as_mut() {
+            Some(v) => {
+                *v += buf.as_bytes().len();
+                Ok(buf.into_pyobj())
+            }
+            None => Err(CryptographyError::from(
+                exceptions::AlreadyFinalized::new_err("Context was already finalized."),
+            )),
+        }
+    }
+
+    fn finalize<'p>(
+        &mut self,
+        py: pyo3::Python<'p>,
+    ) -> CryptographyResult<pyo3::Bound<'p, pyo3::types::PyBytes>> {
+        match self.length_seen.take() {
+            Some(v) => {
+                let pad_size = self.block_size - (v % self.block_size);
+                let pad = vec![pad_size as u8; pad_size];
+                Ok(pyo3::types::PyBytes::new_bound(py, &pad))
+            }
+            None => Err(CryptographyError::from(
+                exceptions::AlreadyFinalized::new_err("Context was already finalized."),
+            )),
+        }
+    }
+}
+
 #[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 2e20363..0ab1125 100644
--- a/tests/hazmat/primitives/test_padding.py
+++ b/tests/hazmat/primitives/test_padding.py
@@ -47,9 +47,9 @@ class TestPKCS7:
 
         str(mybytes())
         padder = padding.PKCS7(128).padder()
-        padder.update(mybytes(b"abc"))
+        data = padder.update(mybytes(b"abc")) + padder.finalize()
         unpadder = padding.PKCS7(128).unpadder()
-        unpadder.update(mybytes(padder.finalize()))
+        unpadder.update(mybytes(data))
         assert unpadder.finalize() == b"abc"
 
     @pytest.mark.parametrize(