From a5dd345be8bbda06fb3a963d950b616c63c3e85f Mon Sep 17 00:00:00 2001
From: Carson McManus <carson.mcmanus1@gmail.com>
Date: Mon, 10 Jul 2023 14:16:42 -0400
Subject: [PATCH] refactor to have the argon2 struct hold InitToken

---
 argon2/src/block.rs | 13 ++-----------
 argon2/src/lib.rs   | 24 +++++++++++++++++++++++-
 2 files changed, 25 insertions(+), 12 deletions(-)

diff --git a/argon2/src/block.rs b/argon2/src/block.rs
index bec3bb83..d9d0fd70 100644
--- 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
index 41b9f8c7..73b6bca2 100644
--- 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,9 @@ pub struct Argon2<'key> {
 
     /// Key array
     secret: Option<&'key [u8]>,
+
+    #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
+    cpu_feat_avx2: avx2_cpuid::InitToken,
 }
 
 impl Default for Argon2<'_> {
@@ -191,6 +197,8 @@ impl<'key> Argon2<'key> {
             version,
             params,
             secret: None,
+            #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
+            cpu_feat_avx2: avx2_cpuid::init(),
         }
     }
 
@@ -210,6 +218,8 @@ impl<'key> Argon2<'key> {
             version,
             params,
             secret: Some(secret),
+            #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
+            cpu_feat_avx2: avx2_cpuid::init(),
         })
     }
 
@@ -424,7 +434,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 +452,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 +599,8 @@ impl PasswordHasher for Argon2<'_> {
             algorithm,
             version,
             params,
+            #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
+            cpu_feat_avx2: self.cpu_feat_avx2,
         }
         .hash_password(password, salt)
     }