Allow extending the iterations number when restoring from a checkpoint

lbeder · Oct 21, 2024 · 2708644 · 2708644
1 parent 47fa6ff
commit 2708644
Show file tree

Hide file tree

Showing 2 changed files with 38 additions and 6 deletions.
diff --git a/src/main.rs b/src/main.rs
@@ -380,7 +380,19 @@ fn main() {
                     path: path.clone(),
                 });
 
-                slowkey_opts = checkpoint_data.data.slowkey.clone();
+                if iterations <= checkpoint_data.data.iteration {
+                    panic!(
+                        "Invalid iterations number {} for checkpoint {}",
+                        iterations, checkpoint_data.data.iteration
+                    );
+                }
+
+                slowkey_opts = SlowKeyOptions {
+                    iterations,
+                    length: checkpoint_data.data.slowkey.length,
+                    scrypt: checkpoint_data.data.slowkey.scrypt,
+                    argon2id: checkpoint_data.data.slowkey.argon2id,
+                };
 
                 offset = checkpoint_data.data.iteration + 1;
                 offset_data.clone_from(&checkpoint_data.data.data);
@@ -614,9 +626,8 @@ fn main() {
             let slowkey_opts = checkpoint_data.data.slowkey.clone();
 
             println!(
-                "{}: iterations: {}, length: {}, {}: (n: {}, r: {}, p: {}), {}: (version: {}, m_cost: {}, t_cost: {})",
+                "{}: length: {}, {}: (n: {}, r: {}, p: {}), {}: (version: {}, m_cost: {}, t_cost: {})",
                 "SlowKey Parameters".yellow(),
-                &slowkey_opts.iterations.to_string().cyan(),
                 &slowkey_opts.length.to_string().cyan(),
                 "Scrypt".green(),
                 &slowkey_opts.scrypt.n.to_string().cyan(),

diff --git a/src/utils/checkpoints/checkpoint.rs b/src/utils/checkpoints/checkpoint.rs
@@ -1,6 +1,10 @@
 use crate::{
     slowkey::SlowKeyOptions,
-    utils::chacha20poly1305::{ChaCha20Poly1305, Nonce},
+    utils::{
+        argon2id::Argon2idOptions,
+        chacha20poly1305::{ChaCha20Poly1305, Nonce},
+        scrypt::ScryptOptions,
+    },
 };
 
 use glob::{glob_with, MatchOptions};
@@ -34,11 +38,28 @@ pub struct OpenCheckpointOptions {
     pub key: Vec<u8>,
 }
 
+#[derive(PartialEq, Debug, Clone, Serialize, Deserialize)]
+pub struct CheckpointSlowKeyOptions {
+    pub length: usize,
+    pub scrypt: ScryptOptions,
+    pub argon2id: Argon2idOptions,
+}
+
+impl From<SlowKeyOptions> for CheckpointSlowKeyOptions {
+    fn from(options: SlowKeyOptions) -> Self {
+        CheckpointSlowKeyOptions {
+            length: options.length,
+            scrypt: options.scrypt,
+            argon2id: options.argon2id,
+        }
+    }
+}
+
 #[derive(Serialize, Deserialize)]
 pub struct SlowKeyData {
     pub iteration: usize,
     pub data: Vec<u8>,
-    pub slowkey: SlowKeyOptions,
+    pub slowkey: CheckpointSlowKeyOptions,
 }
 
 #[derive(Serialize, Deserialize)]
@@ -104,7 +125,7 @@ impl Checkpoint {
                 data: SlowKeyData {
                     iteration: 0,
                     data: Vec::new(),
-                    slowkey: opts.slowkey.clone(),
+                    slowkey: opts.slowkey.clone().into(),
                 },
             },
             checkpoint_extension_padding: (opts.iterations as f64).log10().round() as usize + 1,