set a threadpool number

pyca · Oct 29, 2024 · 9156762 · 9156762
1 parent 5a6a9fc
commit 9156762
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 3 deletions.
diff --git a/src/rust/src/backend/kdf.rs b/src/rust/src/backend/kdf.rs
@@ -201,7 +201,9 @@ impl Argon2id {
     ) -> CryptographyResult<Self> {
         cfg_if::cfg_if! {
             if #[cfg(not(CRYPTOGRAPHY_OPENSSL_320_OR_GREATER))] {
+                _ = py;
                 _ = salt;
+                _ = length;
                 _ = iterations;
                 _ = lanes;
                 _ = memory_cost;
@@ -274,7 +276,7 @@ impl Argon2id {
         }
     }
 
-    #[cfg(not(CRYPTOGRAPHY_IS_LIBRESSL))]
+    #[cfg(CRYPTOGRAPHY_OPENSSL_320_OR_GREATER)]
     fn derive<'p>(
         &mut self,
         py: pyo3::Python<'p>,
@@ -284,7 +286,6 @@ impl Argon2id {
             return Err(exceptions::already_finalized_error());
         }
         self.used = true;
-
         Ok(pyo3::types::PyBytes::new_bound_with(
             py,
             self.length,
@@ -306,7 +307,7 @@ impl Argon2id {
         )?)
     }
 
-    #[cfg(not(CRYPTOGRAPHY_IS_LIBRESSL))]
+    #[cfg(CRYPTOGRAPHY_OPENSSL_320_OR_GREATER)]
     fn verify(
         &mut self,
         py: pyo3::Python<'_>,

diff --git a/src/rust/src/lib.rs b/src/rust/src/lib.rs
@@ -11,6 +11,8 @@ use crate::error::CryptographyResult;
 use openssl::provider;
 #[cfg(CRYPTOGRAPHY_OPENSSL_300_OR_GREATER)]
 use std::env;
+#[cfg(CRYPTOGRAPHY_OPENSSL_320_OR_GREATER)]
+use std::ptr;
 
 mod asn1;
 mod backend;
@@ -74,6 +76,16 @@ fn _initialize_providers() -> CryptographyResult<LoadedProviders> {
     })
 }
 
+fn set_max_global_threads() {
+    // if we can't get available parallelism we do nothing
+    if let Ok(available) = std::thread::available_parallelism() {
+        // SAFETY: This sets a libctx provider limit, but we always use the same libctx by passing NULL.
+        unsafe {
+            openssl_sys::OSSL_set_max_threads(ptr::null_mut(), available.get() as u64);
+        }
+    };
+}
+
 fn _legacy_provider_error(success: bool) -> pyo3::PyResult<()> {
     if !success {
         return Err(pyo3::exceptions::PyRuntimeError::new_err(
@@ -210,6 +222,11 @@ mod _rust {
             cfg_if::cfg_if! {
                 if #[cfg(CRYPTOGRAPHY_OPENSSL_300_OR_GREATER)] {
                     let providers = super::super::_initialize_providers()?;
+                    cfg_if::cfg_if! {
+                        if #[cfg(CRYPTOGRAPHY_OPENSSL_320_OR_GREATER)] {
+                            super::super::set_max_global_threads();
+                        }
+                    }
                     if providers.legacy.is_some() {
                         openssl_mod.add("_legacy_provider_loaded", true)?;
                     } else {