fix unlocking org keys on initial login

2 files changed, 53 insertions, 15 deletions

diff --git a/src/bin/rbw-agent/actions.rs b/src/bin/rbw-agent/actions.rs
index 4e76d4b..e1124e0 100644
--- a/src/bin/rbw-agent/actions.rs
+++ b/src/bin/rbw-agent/actions.rs
@@ -45,16 +45,50 @@ pub async fn login(
                     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.protected_key = Some(protected_key.clone());
                     save_db(&db).await?;
 
+                    sync(sock, false).await?;
+                    db = load_db().await?;
+
+                    let protected_private_key =
+                        if let Some(protected_private_key) =
+                            db.protected_private_key
+                        {
+                            protected_private_key
+                        } else {
+                            return Err(anyhow::anyhow!(
+                                "failed to find protected private key in db"
+                            ));
+                        };
+
+                    let res = rbw::actions::unlock(
+                        &email,
+                        &password,
+                        iterations,
+                        &protected_key,
+                        &protected_private_key,
+                        &db.protected_org_keys,
+                    )
+                    .await;
+
+                    match res {
+                        Ok((keys, org_keys)) => {
+                            let mut state = state.write().await;
+                            state.priv_key = Some(keys);
+                            state.org_keys = Some(org_keys);
+                        }
+                        Err(e) => {
+                            return Err(e)
+                                .context("failed to unlock database")
+                        }
+                    }
+
                     break;
                 }
                 Err(rbw::error::Error::IncorrectPassword) => {
@@ -73,12 +107,10 @@ pub async fn login(
                 }
             }
         }
-
-        sync(sock).await?;
-    } else {
-        respond_ack(sock).await?;
     }
 
+    respond_ack(sock).await?;
+
     Ok(())
 }
 
@@ -142,7 +174,7 @@ pub async fn unlock(
                 Ok((keys, org_keys)) => {
                     let mut state = state.write().await;
                     state.priv_key = Some(keys);
-                    state.org_keys = org_keys;
+                    state.org_keys = Some(org_keys);
                     break;
                 }
                 Err(rbw::error::Error::IncorrectPassword) => {
@@ -174,7 +206,10 @@ pub async fn lock(
     Ok(())
 }
 
-pub async fn sync(sock: &mut crate::sock::Sock) -> anyhow::Result<()> {
+pub async fn sync(
+    sock: &mut crate::sock::Sock,
+    ack: bool,
+) -> anyhow::Result<()> {
     let mut db = load_db().await?;
 
     let access_token = if let Some(access_token) = &db.access_token {
@@ -202,7 +237,9 @@ pub async fn sync(sock: &mut crate::sock::Sock) -> anyhow::Result<()> {
     db.entries = entries;
     save_db(&db).await?;
 
-    respond_ack(sock).await?;
+    if ack {
+        respond_ack(sock).await?;
+    }
 
     Ok(())
 }
diff --git a/src/bin/rbw-agent/agent.rs b/src/bin/rbw-agent/agent.rs
index 0094f7d..655c3dc 100644
--- a/src/bin/rbw-agent/agent.rs
+++ b/src/bin/rbw-agent/agent.rs
@@ -9,20 +9,21 @@ pub enum TimeoutEvent {
 
 pub struct State {
     pub priv_key: Option<rbw::locked::Keys>,
-    pub org_keys: std::collections::HashMap<String, rbw::locked::Keys>,
+    pub org_keys:
+        Option<std::collections::HashMap<String, rbw::locked::Keys>>,
     pub timeout_chan: tokio::sync::mpsc::UnboundedSender<TimeoutEvent>,
 }
 
 impl State {
     pub fn key(&self, org_id: Option<&str>) -> Option<&rbw::locked::Keys> {
         match org_id {
-            Some(id) => self.org_keys.get(id),
+            Some(id) => self.org_keys.as_ref().and_then(|h| h.get(id)),
             None => self.priv_key.as_ref(),
         }
     }
 
     pub fn needs_unlock(&self) -> bool {
-        self.priv_key.is_none()
+        self.priv_key.is_none() || self.org_keys.is_none()
     }
 
     pub fn set_timeout(&mut self) {
@@ -153,7 +154,7 @@ async fn handle_request(
             false
         }
         rbw::protocol::Action::Sync => {
-            crate::actions::sync(sock).await?;
+            crate::actions::sync(sock, true).await?;
             false
         }
         rbw::protocol::Action::Decrypt {