Remove the final call to into_gil_ref

pyca · Apr 15, 2024 · 0d80ec0 · 0d80ec0
1 parent e7a0023
commit 0d80ec0
Show file tree

Hide file tree

Showing 4 changed files with 39 additions and 17 deletions.
diff --git a/src/rust/src/backend/dsa.rs b/src/rust/src/backend/dsa.rs
@@ -77,7 +77,7 @@ impl DsaPrivateKey {
         let mut signer = openssl::pkey_ctx::PkeyCtx::new(&self.pkey)?;
         signer.sign_init()?;
         let mut sig = vec![];
-        signer.sign_to_vec(data, &mut sig)?;
+        signer.sign_to_vec(data.as_bytes(), &mut sig)?;
         Ok(pyo3::types::PyBytes::new_bound(py, &sig))
     }
 
@@ -162,7 +162,9 @@ impl DsaPublicKey {
 
         let mut verifier = openssl::pkey_ctx::PkeyCtx::new(&self.pkey)?;
         verifier.verify_init()?;
-        let valid = verifier.verify(data, signature.as_bytes()).unwrap_or(false);
+        let valid = verifier
+            .verify(data.as_bytes(), signature.as_bytes())
+            .unwrap_or(false);
         if !valid {
             return Err(CryptographyError::from(
                 exceptions::InvalidSignature::new_err(()),

diff --git a/src/rust/src/backend/ec.rs b/src/rust/src/backend/ec.rs
@@ -311,7 +311,7 @@ impl ECPrivateKey {
         // easily known a priori (if `r` or `s` has a leading 0, the signature
         // will be a byte or two shorter than the maximum possible length).
         let mut sig = vec![];
-        signer.sign_to_vec(data, &mut sig)?;
+        signer.sign_to_vec(data.as_bytes(), &mut sig)?;
         Ok(pyo3::types::PyBytes::new_bound(py, &sig))
     }
 
@@ -408,7 +408,9 @@ impl ECPublicKey {
 
         let mut verifier = openssl::pkey_ctx::PkeyCtx::new(&self.pkey)?;
         verifier.verify_init()?;
-        let valid = verifier.verify(data, signature.as_bytes()).unwrap_or(false);
+        let valid = verifier
+            .verify(data.as_bytes(), signature.as_bytes())
+            .unwrap_or(false);
         if !valid {
             return Err(CryptographyError::from(
                 exceptions::InvalidSignature::new_err(()),

diff --git a/src/rust/src/backend/rsa.rs b/src/rust/src/backend/rsa.rs
@@ -296,9 +296,9 @@ impl RsaPrivateKey {
         })?;
         setup_signature_ctx(py, &mut ctx, padding, &algorithm, self.pkey.size(), true)?;
 
-        let length = ctx.sign(data, None)?;
+        let length = ctx.sign(data.as_bytes(), None)?;
         Ok(pyo3::types::PyBytes::new_bound_with(py, length, |b| {
-            let length = ctx.sign(data, Some(b)).map_err(|_| {
+            let length = ctx.sign(data.as_bytes(), Some(b)).map_err(|_| {
                 pyo3::exceptions::PyValueError::new_err(
                     "Digest or salt length too long for key size. Use a larger key or shorter salt length if you are specifying a PSS salt",
                 )
@@ -434,7 +434,9 @@ impl RsaPublicKey {
         ctx.verify_init()?;
         setup_signature_ctx(py, &mut ctx, padding, &algorithm, self.pkey.size(), false)?;
 
-        let valid = ctx.verify(data, signature.as_bytes()).unwrap_or(false);
+        let valid = ctx
+            .verify(data.as_bytes(), signature.as_bytes())
+            .unwrap_or(false);
         if !valid {
             return Err(CryptographyError::from(
                 exceptions::InvalidSignature::new_err(()),

diff --git a/src/rust/src/backend/utils.rs b/src/rust/src/backend/utils.rs
@@ -5,7 +5,7 @@
 use crate::backend::hashes::Hash;
 use crate::error::{CryptographyError, CryptographyResult};
 use crate::{error, types};
-use pyo3::prelude::PyAnyMethods;
+use pyo3::prelude::{PyAnyMethods, PyBytesMethods};
 use pyo3::ToPyObject;
 
 pub(crate) fn py_int_to_bn(
@@ -354,31 +354,47 @@ pub(crate) fn pkey_public_bytes<'p>(
     ))
 }
 
+pub(crate) enum BytesOrPyBytes<'a> {
+    Bytes(&'a [u8]),
+    PyBytes(pyo3::Bound<'a, pyo3::types::PyBytes>),
+}
+
+impl BytesOrPyBytes<'_> {
+    pub(crate) fn as_bytes(&self) -> &[u8] {
+        match self {
+            BytesOrPyBytes::Bytes(v) => v,
+            BytesOrPyBytes::PyBytes(v) => v.as_bytes(),
+        }
+    }
+}
+
 pub(crate) fn calculate_digest_and_algorithm<'p>(
     py: pyo3::Python<'p>,
-    mut data: &'p [u8],
+    data: &'p [u8],
     algorithm: &pyo3::Bound<'p, pyo3::PyAny>,
-) -> CryptographyResult<(&'p [u8], pyo3::Bound<'p, pyo3::PyAny>)> {
-    let mut algorithm_result = algorithm.clone();
-    if algorithm.is_instance(&types::PREHASHED.get(py)?)? {
-        algorithm_result = algorithm.getattr("_algorithm")?;
+) -> CryptographyResult<(BytesOrPyBytes<'p>, pyo3::Bound<'p, pyo3::PyAny>)> {
+    let (algorithm, data) = if algorithm.is_instance(&types::PREHASHED.get(py)?)? {
+        (
+            algorithm.getattr("_algorithm")?,
+            BytesOrPyBytes::Bytes(data),
+        )
     } else {
         // Potential optimization: rather than allocate a PyBytes in
         // `h.finalize()`, have a way to get the `DigestBytes` directly.
         let mut h = Hash::new(py, algorithm, None)?;
         h.update_bytes(data)?;
-        data = h.finalize(py)?.into_gil_ref().as_bytes();
-    }
+        (algorithm.clone(), BytesOrPyBytes::PyBytes(h.finalize(py)?))
+    };
 
-    if data.len() != algorithm.getattr("digest_size")?.extract()? {
+    if data.as_bytes().len() != algorithm.getattr("digest_size")?.extract()? {
         return Err(CryptographyError::from(
             pyo3::exceptions::PyValueError::new_err(
                 "The provided data must be the same length as the hash algorithm's digest size.",
             ),
         ));
     }
 
-    Ok((data, algorithm_result))
+    Ok((data, algorithm))
 }
 
 pub(crate) enum PasswordCallbackStatus {