allow multiple attempts for password entry

2 files changed, 94 insertions, 35 deletions

diff --git a/src/bin/rbw-agent/actions.rs b/src/bin/rbw-agent/actions.rs
index ac331ce..8823424 100644
--- a/src/bin/rbw-agent/actions.rs
+++ b/src/bin/rbw-agent/actions.rs
@@ -26,27 +26,58 @@ pub async fn login(
                 url_str
             ));
         };
-        let password = rbw::pinentry::getpin(
-            "Master Password",
-            &format!("Log in to {}", host),
-            tty,
-        )
-        .await
-        .context("failed to read password from pinentry")?;
-        let (access_token, refresh_token, iterations, protected_key, keys) =
-            rbw::actions::login(&email, &password)
-                .await
-                .context("failed to log in to bitwarden instance")?;
-
-        state.write().await.priv_key = Some(keys);
 
-        db.access_token = Some(access_token);
-        db.refresh_token = Some(refresh_token);
-        db.iterations = Some(iterations);
-        db.protected_key = Some(protected_key);
-        db.save_async(&email)
+        for i in 1_u8..=3 {
+            let err = if i > 1 {
+                Some(format!("Incorrect password (attempt {}/3)", i))
+            } else {
+                None
+            };
+            let password = rbw::pinentry::getpin(
+                "Master Password",
+                &format!("Log in to {}", host),
+                err.as_deref(),
+                tty,
+            )
             .await
-            .context("failed to save local database")?;
+            .context("failed to read password from pinentry")?;
+            let res = rbw::actions::login(&email, &password).await;
+            match res {
+                Ok((
+                    access_token,
+                    refresh_token,
+                    iterations,
+                    protected_key,
+                    keys,
+                )) => {
+                    state.write().await.priv_key = Some(keys);
+
+                    db.access_token = Some(access_token);
+                    db.refresh_token = Some(refresh_token);
+                    db.iterations = Some(iterations);
+                    db.protected_key = Some(protected_key);
+                    db.save_async(&email)
+                        .await
+                        .context("failed to save local database")?;
+
+                    break;
+                }
+                Err(rbw::error::Error::IncorrectPassword) => {
+                    if i == 3 {
+                        return Err(rbw::error::Error::IncorrectPassword)
+                            .context(
+                                "failed to log in to bitwarden instance",
+                            );
+                    } else {
+                        continue;
+                    }
+                }
+                Err(e) => {
+                    return Err(e)
+                        .context("failed to log in to bitwarden instance")
+                }
+            }
+        }
 
         sync(sock).await?;
     } else {
@@ -65,13 +96,6 @@ pub async fn unlock(
         let email = config_email()
             .await
             .context("failed to read email from config")?;
-        let password = rbw::pinentry::getpin(
-            "Master Password",
-            "Unlock the local database",
-            tty,
-        )
-        .await
-        .context("failed to read password from pinentry")?;
 
         let db = rbw::db::Db::load_async(&email)
             .await
@@ -91,16 +115,44 @@ pub async fn unlock(
                 "failed to find protected key in db"
             ));
         };
-        let keys = rbw::actions::unlock(
-            &email,
-            &password,
-            iterations,
-            &protected_key,
-        )
-        .await
-        .context("failed to unlock database")?;
 
-        state.write().await.priv_key = Some(keys);
+        for i in 1u8..=3 {
+            let err = if i > 1 {
+                Some(format!("Incorrect password (attempt {}/3)", i))
+            } else {
+                None
+            };
+            let password = rbw::pinentry::getpin(
+                "Master Password",
+                "Unlock the local database",
+                err.as_deref(),
+                tty,
+            )
+            .await
+            .context("failed to read password from pinentry")?;
+            let res = rbw::actions::unlock(
+                &email,
+                &password,
+                iterations,
+                &protected_key,
+            )
+            .await;
+            match res {
+                Ok(keys) => {
+                    state.write().await.priv_key = Some(keys);
+                    break;
+                }
+                Err(rbw::error::Error::IncorrectPassword) => {
+                    if i == 3 {
+                        return Err(rbw::error::Error::IncorrectPassword)
+                            .context("failed to unlock database");
+                    } else {
+                        continue;
+                    }
+                }
+                Err(e) => return Err(e).context("failed to unlock database"),
+            }
+        }
     }
 
     respond_ack(sock).await?;
diff --git a/src/pinentry.rs b/src/pinentry.rs
index c0a929b..707fb24 100644
--- a/src/pinentry.rs
+++ b/src/pinentry.rs
@@ -5,6 +5,7 @@ use tokio::io::AsyncWriteExt as _;
 pub async fn getpin(
     prompt: &str,
     desc: &str,
+    err: Option<&str>,
     tty: Option<&str>,
 ) -> Result<crate::locked::Password> {
     let mut opts = tokio::process::Command::new("pinentry");
@@ -33,6 +34,12 @@ pub async fn getpin(
         .write_all(format!("SETDESC {}\n", desc).as_bytes())
         .await
         .context(crate::error::WriteStdin)?;
+    if let Some(err) = err {
+        stdin
+            .write_all(format!("SETERROR {}\n", err).as_bytes())
+            .await
+            .context(crate::error::WriteStdin)?;
+    }
     stdin
         .write_all(b"GETPIN\n")
         .await