Convert unpadding code to Rust (#5668)

5 files changed, 90 insertions, 105 deletions

diff --git a/src/_cffi_src/build_padding.py b/src/_cffi_src/build_padding.py
deleted file mode 100644
index 61f36ef69..000000000
--- a/src/_cffi_src/build_padding.py
+++ /dev/null
@@ -1,26 +0,0 @@
-# This file is dual licensed under the terms of the Apache License, Version
-# 2.0, and the BSD License. See the LICENSE file in the root of this repository
-# for complete details.
-
-
-import os
-
-from _cffi_src.utils import build_ffi, compiler_type, extra_link_args
-
-
-with open(
-    os.path.join(os.path.dirname(__file__), "hazmat_src/padding.h")
-) as f:
-    types = f.read()
-
-with open(
-    os.path.join(os.path.dirname(__file__), "hazmat_src/padding.c")
-) as f:
-    functions = f.read()
-
-ffi = build_ffi(
-    module_name="_padding",
-    cdef_source=types,
-    verify_source=functions,
-    extra_link_args=extra_link_args(compiler_type()),
-)
diff --git a/src/_cffi_src/hazmat_src/padding.c b/src/_cffi_src/hazmat_src/padding.c
deleted file mode 100644
index a6e05dee1..000000000
--- a/src/_cffi_src/hazmat_src/padding.c
+++ /dev/null
@@ -1,65 +0,0 @@
-// This file is dual licensed under the terms of the Apache License, Version
-// 2.0, and the BSD License. See the LICENSE file in the root of this
-// repository for complete details.
-
-/* Returns the value of the input with the most-significant-bit copied to all
-   of the bits. */
-static uint16_t Cryptography_DUPLICATE_MSB_TO_ALL(uint16_t a) {
-    return (1 - (a >> (sizeof(uint16_t) * 8 - 1))) - 1;
-}
-
-/* This returns 0xFFFF if a < b else 0x0000, but does so in a constant time
-   fashion */
-static uint16_t Cryptography_constant_time_lt(uint16_t a, uint16_t b) {
-    a -= b;
-    return Cryptography_DUPLICATE_MSB_TO_ALL(a);
-}
-
-uint8_t Cryptography_check_pkcs7_padding(const uint8_t *data,
-                                         uint16_t block_len) {
-    uint16_t i;
-    uint16_t pad_size = data[block_len - 1];
-    uint16_t mismatch = 0;
-    for (i = 0; i < block_len; i++) {
-        unsigned int mask = Cryptography_constant_time_lt(i, pad_size);
-        uint16_t b = data[block_len - 1 - i];
-        mismatch |= (mask & (pad_size ^ b));
-    }
-
-    /* Check to make sure the pad_size was within the valid range. */
-    mismatch |= ~Cryptography_constant_time_lt(0, pad_size);
-    mismatch |= Cryptography_constant_time_lt(block_len, pad_size);
-
-    /* Make sure any bits set are copied to the lowest bit */
-    mismatch |= mismatch >> 8;
-    mismatch |= mismatch >> 4;
-    mismatch |= mismatch >> 2;
-    mismatch |= mismatch >> 1;
-    /* Now check the low bit to see if it's set */
-    return (mismatch & 1) == 0;
-}
-
-uint8_t Cryptography_check_ansix923_padding(const uint8_t *data,
-                                            uint16_t block_len) {
-    uint16_t i;
-    uint16_t pad_size = data[block_len - 1];
-    uint16_t mismatch = 0;
-    /* Skip the first one with the pad size */
-    for (i = 1; i < block_len; i++) {
-        unsigned int mask = Cryptography_constant_time_lt(i, pad_size);
-        uint16_t b = data[block_len - 1 - i];
-        mismatch |= (mask & b);
-    }
-
-    /* Check to make sure the pad_size was within the valid range. */
-    mismatch |= ~Cryptography_constant_time_lt(0, pad_size);
-    mismatch |= Cryptography_constant_time_lt(block_len, pad_size);
-
-    /* Make sure any bits set are copied to the lowest bit */
-    mismatch |= mismatch >> 8;
-    mismatch |= mismatch >> 4;
-    mismatch |= mismatch >> 2;
-    mismatch |= mismatch >> 1;
-    /* Now check the low bit to see if it's set */
-    return (mismatch & 1) == 0;
-}
diff --git a/src/_cffi_src/hazmat_src/padding.h b/src/_cffi_src/hazmat_src/padding.h
deleted file mode 100644
index fb023c171..000000000
--- a/src/_cffi_src/hazmat_src/padding.h
+++ /dev/null
@@ -1,6 +0,0 @@
-// This file is dual licensed under the terms of the Apache License, Version
-// 2.0, and the BSD License. See the LICENSE file in the root of this
-// repository for complete details.
-
-uint8_t Cryptography_check_pkcs7_padding(const uint8_t *, uint8_t);
-uint8_t Cryptography_check_ansix923_padding(const uint8_t *, uint8_t);
diff --git a/src/cryptography/hazmat/primitives/padding.py b/src/cryptography/hazmat/primitives/padding.py
index e6f46eb4f..ccfde7404 100644
--- a/src/cryptography/hazmat/primitives/padding.py
+++ b/src/cryptography/hazmat/primitives/padding.py
@@ -8,7 +8,10 @@ import typing
 
 from cryptography import utils
 from cryptography.exceptions import AlreadyFinalized
-from cryptography.hazmat.bindings._padding import lib
+from cryptography.hazmat.bindings._rust import (
+    check_ansix923_padding,
+    check_pkcs7_padding,
+)
 
 
 class PaddingContext(metaclass=abc.ABCMeta):
@@ -84,7 +87,7 @@ def _byte_unpadding_update(
 def _byte_unpadding_check(
     buffer_: typing.Optional[bytes],
     block_size: int,
-    checkfn: typing.Callable[[bytes, int], int],
+    checkfn: typing.Callable[[bytes], int],
 ) -> bytes:
     if buffer_ is None:
         raise AlreadyFinalized("Context was already finalized.")
@@ -92,7 +95,7 @@ def _byte_unpadding_check(
     if len(buffer_) != block_size // 8:
         raise ValueError("Invalid padding bytes.")
 
-    valid = checkfn(buffer_, block_size // 8)
+    valid = checkfn(buffer_)
 
     if not valid:
         raise ValueError("Invalid padding bytes.")
@@ -154,7 +157,7 @@ class _PKCS7UnpaddingContext(PaddingContext):
 
     def finalize(self) -> bytes:
         result = _byte_unpadding_check(
-            self._buffer, self.block_size, lib.Cryptography_check_pkcs7_padding
+            self._buffer, self.block_size, check_pkcs7_padding
         )
         self._buffer = None
         return result
@@ -215,7 +218,7 @@ class _ANSIX923UnpaddingContext(PaddingContext):
         result = _byte_unpadding_check(
             self._buffer,
             self.block_size,
-            lib.Cryptography_check_ansix923_padding,
+            check_ansix923_padding,
         )
         self._buffer = None
         return result
diff --git a/src/rust/src/lib.rs b/src/rust/src/lib.rs
index 1580ca4fc..3257b35e1 100644
--- a/src/rust/src/lib.rs
+++ b/src/rust/src/lib.rs
@@ -2,9 +2,88 @@
 // 2.0, and the BSD License. See the LICENSE file in the root of this repository
 // for complete details.
 
+use std::convert::TryInto;
+
+/// 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 {
+    0u8.wrapping_sub(a >> 7)
+}
+
+/// This returns 0xFF if a < b else 0x00, but does so in a constant time
+/// fashion.
+fn constant_time_lt(a: u8, b: u8) -> u8 {
+    // Derived from:
+    // https://github.com/openssl/openssl/blob/OpenSSL_1_1_1i/include/internal/constant_time.h#L120
+    duplicate_msb_to_all(a ^ ((a ^ b) | (a.wrapping_sub(b) ^ b)))
+}
+
+#[pyo3::prelude::pyfunction]
+fn check_pkcs7_padding(data: &[u8]) -> bool {
+    let mut mismatch = 0;
+    let pad_size = *data.last().unwrap();
+    let len: u8 = data.len().try_into().expect("data too long");
+    for (i, b) in (0..len).zip(data.iter().rev()) {
+        let mask = constant_time_lt(i, pad_size);
+        mismatch |= mask & (pad_size ^ b);
+    }
+
+    // Check to make sure the pad_size was within the valid range.
+    mismatch |= !constant_time_lt(0, pad_size);
+    mismatch |= constant_time_lt(len, pad_size);
+
+    // Make sure any bits set are copied to the lowest bit
+    mismatch |= mismatch >> 4;
+    mismatch |= mismatch >> 2;
+    mismatch |= mismatch >> 1;
+
+    // Now check the low bit to see if it's set
+    (mismatch & 1) == 0
+}
+
+#[pyo3::prelude::pyfunction]
+fn check_ansix923_padding(data: &[u8]) -> bool {
+    let mut mismatch = 0;
+    let pad_size = *data.last().unwrap();
+    let len: u8 = data.len().try_into().expect("data too long");
+    // Skip the first one with the pad size
+    for (i, b) in (1..len).zip(data[..data.len() - 1].iter().rev()) {
+        let mask = constant_time_lt(i, pad_size);
+        mismatch |= mask & b;
+    }
+
+    // Check to make sure the pad_size was within the valid range.
+    mismatch |= !constant_time_lt(0, pad_size);
+    mismatch |= constant_time_lt(len, pad_size);
+
+    // Make sure any bits set are copied to the lowest bit
+    mismatch |= mismatch >> 4;
+    mismatch |= mismatch >> 2;
+    mismatch |= mismatch >> 1;
+
+    // Now check the low bit to see if it's set
+    (mismatch & 1) == 0
+}
+
 #[pyo3::prelude::pymodule]
-// False positive: https://github.com/rust-lang/rust-clippy/issues/6721
-#[allow(clippy::unnecessary_wraps)]
-fn _rust(_py: pyo3::Python<'_>, _m: &pyo3::types::PyModule) -> pyo3::PyResult<()> {
+fn _rust(_py: pyo3::Python<'_>, m: &pyo3::types::PyModule) -> pyo3::PyResult<()> {
+    m.add_function(pyo3::wrap_pyfunction!(check_pkcs7_padding, m)?)?;
+    m.add_function(pyo3::wrap_pyfunction!(check_ansix923_padding, m)?)?;
+
     Ok(())
 }
+
+#[cfg(test)]
+mod tests {
+    use super::constant_time_lt;
+
+    #[test]
+    fn test_constant_time_lt() {
+        for a in 0..=255 {
+            for b in 0..=255 {
+                let expected = if a < b { 0xff } else { 0 };
+                assert_eq!(constant_time_lt(a, b), expected);
+            }
+        }
+    }
+}