refactor to have the argon2 struct hold InitToken

RustCrypto · Jul 10, 2023 · 22ead03 · 22ead03
1 parent 3faaff1
commit 22ead03
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 12 deletions.
diff --git a/argon2/src/block.rs b/argon2/src/block.rs
@@ -44,9 +44,6 @@ macro_rules! permute {
     };
 }
 
-#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
-cpufeatures::new!(avx2_cpuid, "avx2");
-
 /// Structure for the (1 KiB) memory block implemented as 128 64-bit words.
 #[derive(Copy, Clone, Debug)]
 #[repr(align(64))]
@@ -70,17 +67,11 @@ impl Block {
     }
 
     pub(crate) fn compress(rhs: &Self, lhs: &Self) -> Self {
-        #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
-        {
-            if avx2_cpuid::get() {
-                return unsafe { Self::compress_avx2(rhs, lhs) };
-            }
-        }
         Self::compress_soft(rhs, lhs)
     }
 
     #[inline(always)]
-    fn compress_soft(rhs: &Self, lhs: &Self) -> Self {
+    pub(crate) fn compress_soft(rhs: &Self, lhs: &Self) -> Self {
         let r = *rhs ^ lhs;
 
         // Apply permutations rowwise
@@ -118,7 +109,7 @@ impl Block {
 
     #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
     #[target_feature(enable = "avx2")]
-    unsafe fn compress_avx2(rhs: &Self, lhs: &Self) -> Self {
+    pub(crate) unsafe fn compress_avx2(rhs: &Self, lhs: &Self) -> Self {
         Self::compress_soft(rhs, lhs)
     }
 }

diff --git a/argon2/src/lib.rs b/argon2/src/lib.rs
@@ -144,6 +144,9 @@ pub(crate) const SYNC_POINTS: usize = 4;
 /// To generate reference block positions
 const ADDRESSES_IN_BLOCK: usize = 128;
 
+#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
+cpufeatures::new!(avx2_cpuid, "avx2");
+
 /// Argon2 context.
 ///
 /// This is the primary type of this crate's API, and contains the following:
@@ -165,6 +168,8 @@ pub struct Argon2<'key> {
 
     /// Key array
     secret: Option<&'key [u8]>,
+
+    cpu_feat_avx2: avx2_cpuid::InitToken,
 }
 
 impl Default for Argon2<'_> {
@@ -191,6 +196,7 @@ impl<'key> Argon2<'key> {
             version,
             params,
             secret: None,
+            cpu_feat_avx2: avx2_cpuid::init(),
         }
     }
 
@@ -210,6 +216,7 @@ impl<'key> Argon2<'key> {
             version,
             params,
             secret: Some(secret),
+            cpu_feat_avx2: avx2_cpuid::init(),
         })
     }
 
@@ -424,7 +431,7 @@ impl<'key> Argon2<'key> {
 
                         // Calculate new block
                         let result =
-                            Block::compress(&memory_blocks[prev_index], &memory_blocks[ref_index]);
+                            self.compress(&memory_blocks[prev_index], &memory_blocks[ref_index]);
 
                         if self.version == Version::V0x10 || pass == 0 {
                             memory_blocks[cur_index] = result;
@@ -442,6 +449,16 @@ impl<'key> Argon2<'key> {
         Ok(())
     }
 
+    fn compress(&self, rhs: &Block, lhs: &Block) -> Block {
+        #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
+        {
+            if self.cpu_feat_avx2.get() {
+                return unsafe { Block::compress_avx2(rhs, lhs) };
+            }
+        }
+        Block::compress(rhs, lhs)
+    }
+
     /// Get default configured [`Params`].
     pub fn params(&self) -> &Params {
         &self.params
@@ -579,6 +596,7 @@ impl PasswordHasher for Argon2<'_> {
             algorithm,
             version,
             params,
+            cpu_feat_avx2: self.cpu_feat_avx2,
         }
         .hash_password(password, salt)
     }