From e28828cc2bf565131122d1ef9f38541b144e6cba Mon Sep 17 00:00:00 2001
From: shimun <shimun@shimun.net>
Date: Mon, 22 Jun 2020 18:06:06 +0200
Subject: [PATCH] cleanup luks.rs

---
 src/cli.rs  | 24 ++++++++-----------
 src/luks.rs | 66 ++++++++++++-----------------------------------------
 2 files changed, 24 insertions(+), 66 deletions(-)

diff --git a/src/cli.rs b/src/cli.rs
index 5082aba..693ec43 100644
--- a/src/cli.rs
+++ b/src/cli.rs
@@ -1,5 +1,4 @@
 use crate::error::*;
-use crate::luks;
 use crate::*;
 
 use structopt::StructOpt;
@@ -445,20 +444,20 @@ pub fn run_cli() -> Fido2LuksResult<()> {
                     pin.as_deref(),
                 )
             };
+            let mut luks_dev = LuksDevice::load(&luks.device)?;
             // Non overlap
             match &args.command {
                 Command::AddKey { exclusive, .. } => {
                     let (existing_secret, _) = other_secret("Current password", false)?;
                     let (new_secret, cred) = secret(true)?;
-                    let added_slot = luks::add_key(
-                        &luks.device,
+                    let added_slot = luks_dev.add_key(
                         &new_secret,
                         &existing_secret[..],
                         luks_mod.kdf_time.or(Some(10)),
                         Some(&cred.id[..]).filter(|_| *token),
                     )?;
                     if *exclusive {
-                        let destroyed = luks::remove_keyslots(&luks.device, &[added_slot])?;
+                        let destroyed = luks_dev.remove_keyslots(&[added_slot])?;
                         println!(
                             "Added to key to device {}, slot: {}\nRemoved {} old keys",
                             luks.device.display(),
@@ -478,16 +477,14 @@ pub fn run_cli() -> Fido2LuksResult<()> {
                     let (existing_secret, _) = secret(false)?;
                     let (replacement_secret, cred) = other_secret("Replacement password", true)?;
                     let slot = if *add_password {
-                        luks::add_key(
-                            &luks.device,
+                        luks_dev.add_key(
                             &replacement_secret[..],
                             &existing_secret,
                             luks_mod.kdf_time,
                             cred.as_ref().filter(|_| *token).map(|cred| &cred.id[..]),
                         )
                     } else {
-                        luks::replace_key(
-                            &luks.device,
+                        luks_dev.replace_key(
                             &replacement_secret[..],
                             &existing_secret,
                             luks_mod.kdf_time,
@@ -545,14 +542,12 @@ pub fn run_cli() -> Fido2LuksResult<()> {
             };
 
             let mut retries = *retries;
+            let mut luks_dev = LuksDevice::load(&luks.device)?;
             loop {
                 let secret = match &args.command {
                     Command::Open { credentials, .. } => secret(Cow::Borrowed(&credentials.ids.0))
-                        .and_then(|(secret, _cred)| {
-                            luks::open_container(&luks.device, &name, &secret, luks.slot)
-                        }),
-                    Command::OpenToken { .. } => luks::open_container_token(
-                        &luks.device,
+                        .and_then(|(secret, _cred)| luks_dev.activate(&name, &secret, luks.slot)),
+                    Command::OpenToken { .. } => luks_dev.activate_token(
                         &name,
                         Box::new(|credentials: Vec<String>| {
                             let creds = credentials
@@ -562,6 +557,7 @@ pub fn run_cli() -> Fido2LuksResult<()> {
                             secret(Cow::Owned(creds))
                                 .map(|(secret, cred)| (secret, hex::encode(&cred.id)))
                         }),
+                        luks.slot,
                     ),
                     _ => unreachable!(),
                 };
@@ -577,7 +573,7 @@ pub fn run_cli() -> Fido2LuksResult<()> {
                         retries -= 1;
                         eprintln!("{}", e);
                     }
-                    res => break res,
+                    res => break res.map(|_| ()),
                 }
             }
         }
diff --git a/src/luks.rs b/src/luks.rs
index 019ffa2..e381a7e 100644
--- a/src/luks.rs
+++ b/src/luks.rs
@@ -262,12 +262,19 @@ impl LuksDevice {
             });
         }
         let (secret, credential) = secret(creds.keys().cloned().collect())?;
-        let slots = creds.get(&credential).unwrap();
-        let slots = slots
-            .iter()
-            .cloned()
-            .map(Option::Some)
-            .chain(std::iter::once(None).take(slots.is_empty() as usize));
+        let slots = if let Some(slots) = creds.get(&credential) {
+            slots
+        } else {
+            return Err(Fido2LuksError::LuksError {
+                cause: LuksError::NoToken,
+            });
+        };
+        //Try slots associated with the credential used
+        let slots = slots.iter().cloned().map(Option::Some).chain(
+            std::iter::once(slot_hint) // Try slot hint if there is one
+                .take(slot_hint.is_some() as usize)
+                .chain(std::iter::once(None).take(slots.is_empty() as usize)), // Try all slots as last resort
+        );
         for slot in slots {
             match self
                 .device
@@ -279,7 +286,7 @@ impl LuksDevice {
                 res => return res,
             }
         }
-        self.activate(name, &secret, slot_hint)
+        Err(Fido2LuksError::WrongSecret)
     }
 }
 
@@ -323,48 +330,3 @@ impl Default for Fido2LuksToken {
         }
     }
 }
-
-pub fn open_container<P: AsRef<Path>>(
-    path: P,
-    name: &str,
-    secret: &[u8],
-    slot_hint: Option<u32>,
-) -> Fido2LuksResult<()> {
-    LuksDevice::load(path)?
-        .activate(name, secret, slot_hint)
-        .map(|_| ())
-}
-
-pub fn open_container_token<P: AsRef<Path>>(
-    path: P,
-    name: &str,
-    secret: impl Fn(Vec<String>) -> Fido2LuksResult<([u8; 32], String)>,
-) -> Fido2LuksResult<()> {
-    LuksDevice::load(path)?
-        .activate_token(name, secret, None)
-        .map(|_| ())
-}
-
-pub fn add_key<P: AsRef<Path>>(
-    path: P,
-    secret: &[u8],
-    old_secret: &[u8],
-    iteration_time: Option<u64>,
-    credential_id: Option<&[u8]>,
-) -> Fido2LuksResult<u32> {
-    LuksDevice::load(path)?.add_key(secret, old_secret, iteration_time, credential_id)
-}
-
-pub fn remove_keyslots<P: AsRef<Path>>(path: P, exclude: &[u32]) -> Fido2LuksResult<u32> {
-    LuksDevice::load(path)?.remove_keyslots(exclude)
-}
-
-pub fn replace_key<P: AsRef<Path>>(
-    path: P,
-    secret: &[u8],
-    old_secret: &[u8],
-    iteration_time: Option<u64>,
-    credential_id: Option<&[u8]>,
-) -> Fido2LuksResult<u32> {
-    LuksDevice::load(path)?.replace_key(secret, old_secret, iteration_time, credential_id)
-}