2 files changed, 42 insertions, 19 deletions

diff --git a/nitrocli/src/commands.rs b/nitrocli/src/commands.rs
index 6285c90..ccbff69 100644
--- a/nitrocli/src/commands.rs
+++ b/nitrocli/src/commands.rs
@@ -30,6 +30,7 @@ use nitrokey::Device;
 use nitrokey::GenerateOtp;
 
 use crate::args;
+use crate::error;
 use crate::error::Error;
 use crate::pinentry;
 use crate::Result;
@@ -188,7 +189,7 @@ fn get_volume_status(status: &nitrokey::VolumeStatus) -> &'static str {
 /// the first try, this function will call `op` with `data`.  At the
 /// second or third try, it will call `op` with the data returned by the
 /// previous call to `op`.
-fn try_with_pin_and_data_with_pinentry<D, F, R>(
+fn try_with_pin_and_data_with_pinentry<D, F, R, E>(
   ctx: &mut args::ExecCtx<'_>,
   pin_entry: &pinentry::PinEntry,
   msg: &'static str,
@@ -196,7 +197,8 @@ fn try_with_pin_and_data_with_pinentry<D, F, R>(
   mut op: F,
 ) -> Result<R>
 where
-  F: FnMut(D, &str) -> result::Result<R, (D, nitrokey::CommandError)>,
+  F: FnMut(D, &str) -> result::Result<R, (D, E)>,
+  E: error::TryInto<nitrokey::CommandError>,
 {
   let mut data = data;
   let mut retry = 3;
@@ -205,26 +207,29 @@ where
     let pin = pinentry::inquire(ctx, pin_entry, pinentry::Mode::Query, error_msg)?;
     match op(data, &pin) {
       Ok(result) => return Ok(result),
-      Err((new_data, err)) => match err {
-        nitrokey::CommandError::WrongPassword => {
-          pinentry::clear(pin_entry)?;
-          retry -= 1;
-
-          if retry > 0 {
-            error_msg = Some("Wrong password, please reenter");
-            data = new_data;
-            continue;
+      Err((new_data, err)) => match err.try_into() {
+        Ok(err) => match err {
+          nitrokey::CommandError::WrongPassword => {
+            pinentry::clear(pin_entry)?;
+            retry -= 1;
+
+            if retry > 0 {
+              error_msg = Some("Wrong password, please reenter");
+              data = new_data;
+              continue;
+            }
+            return Err(get_error(msg, err));
           }
-          return Err(get_error(msg, err));
-        }
-        err => return Err(get_error(msg, err)),
+          err => return Err(get_error(msg, err)),
+        },
+        Err(err) => return Err(err),
       },
     };
   }
 }
 
 /// Try to execute the given function with a PIN.
-fn try_with_pin_and_data<D, F, R>(
+fn try_with_pin_and_data<D, F, R, E>(
   ctx: &mut args::ExecCtx<'_>,
   pin_entry: &pinentry::PinEntry,
   msg: &'static str,
@@ -232,7 +237,8 @@ fn try_with_pin_and_data<D, F, R>(
   mut op: F,
 ) -> Result<R>
 where
-  F: FnMut(D, &str) -> result::Result<R, (D, nitrokey::CommandError)>,
+  F: FnMut(D, &str) -> result::Result<R, (D, E)>,
+  E: Into<Error> + error::TryInto<nitrokey::CommandError>,
 {
   let pin = match pin_entry.pin_type() {
     pinentry::PinType::Admin => &ctx.admin_pin,
@@ -246,7 +252,7 @@ where
         msg
       ))
     })?;
-    op(data, &pin).map_err(|(_, err)| get_error(msg, err))
+    op(data, &pin).map_err(|(_, err)| err.into())
   } else {
     try_with_pin_and_data_with_pinentry(ctx, pin_entry, msg, data, op)
   }
@@ -256,14 +262,15 @@ where
 ///
 /// This function behaves exactly as `try_with_pin_and_data`, but
 /// it refrains from passing any data to it.
-fn try_with_pin<F>(
+fn try_with_pin<F, E>(
   ctx: &mut args::ExecCtx<'_>,
   pin_entry: &pinentry::PinEntry,
   msg: &'static str,
   mut op: F,
 ) -> Result<()>
 where
-  F: FnMut(&str) -> result::Result<(), nitrokey::CommandError>,
+  F: FnMut(&str) -> result::Result<(), E>,
+  E: Into<Error> + error::TryInto<nitrokey::CommandError>,
 {
   try_with_pin_and_data(ctx, pin_entry, msg, (), |data, pin| {
     op(pin).map_err(|err| (data, err))
diff --git a/nitrocli/src/error.rs b/nitrocli/src/error.rs
index 1346526..d1eb2eb 100644
--- a/nitrocli/src/error.rs
+++ b/nitrocli/src/error.rs
@@ -22,6 +22,22 @@ use std::io;
 use std::str;
 use std::string;
 
+/// A trait used to simplify error handling in conjunction with the
+/// try_with_* functions we use for repeatedly asking the user for a
+/// secret.
+pub trait TryInto<T> {
+  fn try_into(self) -> Result<T, Error>;
+}
+
+impl<T, U> TryInto<U> for T
+where
+  T: Into<U>,
+{
+  fn try_into(self) -> Result<U, Error> {
+    Ok(self.into())
+  }
+}
+
 #[derive(Debug)]
 pub enum Error {
   ArgparseError(i32),