diff --git a/src/bin/rbw-agent/actions.rs b/src/bin/rbw-agent/actions.rs index 9ddd2ad9..94a9f033 100644 --- a/src/bin/rbw-agent/actions.rs +++ b/src/bin/rbw-agent/actions.rs @@ -688,6 +688,38 @@ pub async fn decrypt( Ok(()) } +pub async fn decrypt_batch( + sock: &mut crate::sock::Sock, + state: std::sync::Arc>, + environment: &rbw::protocol::Environment, + entries: &[rbw::protocol::DecryptRequest], +) -> anyhow::Result<()> { + let mut results = Vec::with_capacity(entries.len()); + for entry in entries { + let result = decrypt_cipher( + state.clone(), + environment, + &entry.cipherstring, + entry.entry_key.as_deref(), + entry.org_id.as_deref(), + ) + .await; + results.push(match result { + Ok(plaintext) => { + rbw::protocol::DecryptResult::Success { plaintext } + } + Err(e) => rbw::protocol::DecryptResult::Failure { + error: format!("{e:#}"), + }, + }); + } + + sock.send(&rbw::protocol::Response::DecryptBatch { results }) + .await?; + + Ok(()) +} + pub async fn encrypt( sock: &mut crate::sock::Sock, state: std::sync::Arc>, diff --git a/src/bin/rbw-agent/agent.rs b/src/bin/rbw-agent/agent.rs index 1691ed51..4821e4a5 100644 --- a/src/bin/rbw-agent/agent.rs +++ b/src/bin/rbw-agent/agent.rs @@ -162,6 +162,16 @@ async fn handle_request( .await?; true } + rbw::protocol::Action::DecryptBatch { entries } => { + crate::actions::decrypt_batch( + sock, + state.clone(), + &environment, + entries, + ) + .await?; + true + } rbw::protocol::Action::Encrypt { plaintext, org_id } => { crate::actions::encrypt( sock, diff --git a/src/bin/rbw/actions.rs b/src/bin/rbw/actions.rs index a0a34e8c..8364aba1 100644 --- a/src/bin/rbw/actions.rs +++ b/src/bin/rbw/actions.rs @@ -105,6 +105,25 @@ pub fn decrypt( } } +pub fn decrypt_batch( + requests: Vec, +) -> anyhow::Result> { + let mut sock = connect()?; + sock.send(&rbw::protocol::Request::new( + get_environment(), + rbw::protocol::Action::DecryptBatch { entries: requests }, + ))?; + + let res = sock.recv()?; + match res { + rbw::protocol::Response::DecryptBatch { results } => Ok(results), + rbw::protocol::Response::Error { error } => { + Err(anyhow::anyhow!("failed to decrypt: {error}")) + } + _ => Err(anyhow::anyhow!("unexpected message: {res:?}")), + } +} + pub fn encrypt( plaintext: &str, org_id: Option<&str>, diff --git a/src/bin/rbw/commands.rs b/src/bin/rbw/commands.rs index bddf0efe..c3cafd57 100644 --- a/src/bin/rbw/commands.rs +++ b/src/bin/rbw/commands.rs @@ -1371,10 +1371,27 @@ pub fn list(fields: &[String], raw: bool) -> anyhow::Result<()> { unlock()?; let db = load_db()?; - let mut entries: Vec = db + + // Gather every cipherstring that needs decrypting across all entries, then + // decrypt them in a single batch request to the agent. This avoids a + // separate socket round-trip per field per entry, which dominates the + // runtime of `list` on large vaults. + let mut requests = BatchRequests::new(); + let plans: Vec = db .entries .iter() - .map(|entry| decrypt_list_cipher(entry, &fields)) + .map(|entry| ListCipherPlan::build(entry, &fields, &mut requests)) + .collect(); + + let results = if requests.is_empty() { + Vec::new() + } else { + crate::actions::decrypt_batch(requests.into_vec())? + }; + + let mut entries: Vec = plans + .into_iter() + .map(|plan| plan.resolve(&results)) .collect::>()?; entries.sort_unstable_by(|a, b| a.name.cmp(&b.name)); @@ -1500,10 +1517,24 @@ pub fn search( let db = load_db()?; - let mut entries: Vec = db + // As in `list`, decrypt every entry's searchable fields in a single batch + // request rather than one socket round-trip per field per entry. + let mut requests = BatchRequests::new(); + let plans: Vec = db .entries .iter() - .map(decrypt_search_cipher) + .map(|entry| SearchCipherPlan::build(entry, &mut requests)) + .collect(); + + let results = if requests.is_empty() { + Vec::new() + } else { + crate::actions::decrypt_batch(requests.into_vec())? + }; + + let mut entries: Vec = plans + .into_iter() + .map(|plan| plan.resolve(&results)) .filter(|entry| { entry .as_ref() @@ -2026,11 +2057,23 @@ fn find_entry( needle = Needle::Name(s); } + let mut requests = BatchRequests::new(); + let plans: Vec = db + .entries + .iter() + .map(|entry| SearchCipherPlan::build(entry, &mut requests)) + .collect(); + let results = if requests.is_empty() { + Vec::new() + } else { + crate::actions::decrypt_batch(requests.into_vec())? + }; let ciphers: Vec<(rbw::db::Entry, DecryptedSearchCipher)> = db .entries .iter() - .map(|entry| { - decrypt_search_cipher(entry) + .zip(plans) + .map(|(entry, plan)| { + plan.resolve(&results) .map(|decrypted| (entry.clone(), decrypted)) }) .collect::>()?; @@ -2122,179 +2165,345 @@ fn decrypt_field( } } -fn decrypt_list_cipher( - entry: &rbw::db::Entry, - fields: &[ListField], -) -> anyhow::Result { - let id = entry.id.clone(); - let name = if fields.contains(&ListField::Name) { - Some(crate::actions::decrypt( +// Accumulates the cipherstrings to be decrypted in a single `decrypt_batch` +// call. `push` returns the index at which the corresponding plaintext will +// appear in the results vector, which the cipher plans record and later +// resolve. +struct BatchRequests(Vec); + +impl BatchRequests { + fn new() -> Self { + Self(Vec::new()) + } + + fn push( + &mut self, + cipherstring: &str, + entry_key: Option<&str>, + org_id: Option<&str>, + ) -> usize { + let index = self.0.len(); + self.0.push(rbw::protocol::DecryptRequest { + cipherstring: cipherstring.to_string(), + entry_key: entry_key.map(std::string::ToString::to_string), + org_id: org_id.map(std::string::ToString::to_string), + }); + index + } + + fn is_empty(&self) -> bool { + self.0.is_empty() + } + + fn into_vec(self) -> Vec { + self.0 + } +} + +fn entry_type_name(data: &rbw::db::EntryData) -> &'static str { + match data { + rbw::db::EntryData::Login { .. } => "Login", + rbw::db::EntryData::Identity { .. } => "Identity", + rbw::db::EntryData::SshKey { .. } => "SSH Key", + rbw::db::EntryData::SecureNote => "Note", + rbw::db::EntryData::Card { .. } => "Card", + } +} + +// A plan describing which batch-decrypt results make up a single list entry. +// The `usize` fields are indices into the flat results vector returned by +// `decrypt_batch`; `entry_type` needs no decryption so it is resolved up front. +struct ListCipherPlan { + id: String, + name: Option, + user: Option, + folder: Option, + uris: Option>, + entry_type: Option, +} + +impl ListCipherPlan { + fn build( + entry: &rbw::db::Entry, + fields: &[ListField], + requests: &mut BatchRequests, + ) -> Self { + let name = fields.contains(&ListField::Name).then(|| { + requests.push( + &entry.name, + entry.key.as_deref(), + entry.org_id.as_deref(), + ) + }); + + let user = if fields.contains(&ListField::User) { + match &entry.data { + rbw::db::EntryData::Login { + username: Some(username), + .. + } => Some(requests.push( + username, + entry.key.as_deref(), + entry.org_id.as_deref(), + )), + _ => None, + } + } else { + None + }; + + let folder = if fields.contains(&ListField::Folder) { + // folder name should always be decrypted with the local key + // because folders are local to a specific user's vault, not the + // organization + entry + .folder + .as_ref() + .map(|folder| requests.push(folder, None, None)) + } else { + None + }; + + let uris = if fields.contains(&ListField::Uri) { + match &entry.data { + rbw::db::EntryData::Login { uris, .. } => Some( + uris.iter() + .map(|s| { + requests.push( + &s.uri, + entry.key.as_deref(), + entry.org_id.as_deref(), + ) + }) + .collect(), + ), + _ => None, + } + } else { + None + }; + + let entry_type = fields + .contains(&ListField::EntryType) + .then(|| entry_type_name(&entry.data).to_string()); + + Self { + id: entry.id.clone(), + name, + user, + folder, + uris, + entry_type, + } + } + + fn resolve( + self, + results: &[rbw::protocol::DecryptResult], + ) -> anyhow::Result { + // entry name and folder are required, so a decryption failure is fatal + let name = self + .name + .map(|index| strict_result(&results[index])) + .transpose()?; + let folder = self + .folder + .map(|index| strict_result(&results[index])) + .transpose()?; + // optional login fields are skipped (with a warning) on failure, to + // match the previous best-effort behavior of `decrypt_field` + let user = self.user.and_then(|index| { + lenient_result(&results[index], Field::Username) + }); + let uris = self.uris.map(|indices| { + indices + .iter() + .filter_map(|&index| { + lenient_result(&results[index], Field::Uris) + }) + .collect() + }); + + Ok(DecryptedListCipher { + id: self.id, + name, + user, + folder, + uris, + entry_type: self.entry_type, + }) + } +} + +fn strict_result( + result: &rbw::protocol::DecryptResult, +) -> anyhow::Result { + match result { + rbw::protocol::DecryptResult::Success { plaintext } => { + Ok(plaintext.clone()) + } + rbw::protocol::DecryptResult::Failure { error } => { + Err(anyhow::anyhow!("{error}")) + } + } +} + +fn lenient_result( + result: &rbw::protocol::DecryptResult, + name: Field, +) -> Option { + match result { + rbw::protocol::DecryptResult::Success { plaintext } => { + Some(plaintext.clone()) + } + rbw::protocol::DecryptResult::Failure { error } => { + log::warn!("failed to decrypt {name}: {error}"); + None + } + } +} + +// A plan describing which batch-decrypt results make up a single search entry. +// Like `ListCipherPlan`, the `usize` fields index into the flat results vector +// returned by `decrypt_batch`. Search decrypts more per entry than list (notes +// and the custom field values), because those are searchable too. +struct SearchCipherPlan { + id: String, + entry_type: String, + name: usize, + user: Option, + folder: Option, + notes: Option, + uris: Vec<(usize, Option)>, + fields: Vec, +} + +impl SearchCipherPlan { + fn build(entry: &rbw::db::Entry, requests: &mut BatchRequests) -> Self { + let name = requests.push( &entry.name, entry.key.as_deref(), entry.org_id.as_deref(), - )?) - } else { - None - }; - let user = if fields.contains(&ListField::User) { - match &entry.data { - rbw::db::EntryData::Login { username, .. } => decrypt_field( - Field::Username, - username.as_deref(), + ); + + let user = match &entry.data { + rbw::db::EntryData::Login { + username: Some(username), + .. + } => Some(requests.push( + username, entry.key.as_deref(), entry.org_id.as_deref(), - ), + )), _ => None, - } - } else { - None - }; - let folder = if fields.contains(&ListField::Folder) { + }; + // folder name should always be decrypted with the local key because // folders are local to a specific user's vault, not the organization - entry + let folder = entry .folder .as_ref() - .map(|folder| crate::actions::decrypt(folder, None, None)) - .transpose()? - } else { - None - }; - let uris = if fields.contains(&ListField::Uri) { - match &entry.data { - rbw::db::EntryData::Login { uris, .. } => Some( - uris.iter() - .filter_map(|s| { - decrypt_field( - Field::Uris, - Some(&s.uri), - entry.key.as_deref(), - entry.org_id.as_deref(), - ) - }) - .collect(), - ), - _ => None, - } - } else { - None - }; - let entry_type = fields - .contains(&ListField::EntryType) - .then_some(match &entry.data { - rbw::db::EntryData::Login { .. } => "Login", - rbw::db::EntryData::Identity { .. } => "Identity", - rbw::db::EntryData::SshKey { .. } => "SSH Key", - rbw::db::EntryData::SecureNote => "Note", - rbw::db::EntryData::Card { .. } => "Card", - }) - .map(str::to_string); - - Ok(DecryptedListCipher { - id, - name, - user, - folder, - uris, - entry_type, - }) -} + .map(|folder| requests.push(folder, None, None)); -fn decrypt_search_cipher( - entry: &rbw::db::Entry, -) -> anyhow::Result { - let id = entry.id.clone(); - let name = crate::actions::decrypt( - &entry.name, - entry.key.as_deref(), - entry.org_id.as_deref(), - )?; - let user = match &entry.data { - rbw::db::EntryData::Login { username, .. } => decrypt_field( - Field::Username, - username.as_deref(), - entry.key.as_deref(), - entry.org_id.as_deref(), - ), - _ => None, - }; - // folder name should always be decrypted with the local key because - // folders are local to a specific user's vault, not the organization - let folder = entry - .folder - .as_ref() - .map(|folder| crate::actions::decrypt(folder, None, None)) - .transpose()?; - let notes = entry - .notes - .as_ref() - .map(|notes| { - crate::actions::decrypt( + let notes = entry.notes.as_ref().map(|notes| { + requests.push( notes, entry.key.as_deref(), entry.org_id.as_deref(), ) - }) - .transpose(); - let uris = if let rbw::db::EntryData::Login { uris, .. } = &entry.data { - uris.iter() - .filter_map(|s| { - decrypt_field( - Field::Uris, - Some(&s.uri), + }); + + let uris = match &entry.data { + rbw::db::EntryData::Login { uris, .. } => uris + .iter() + .map(|s| { + ( + requests.push( + &s.uri, + entry.key.as_deref(), + entry.org_id.as_deref(), + ), + s.match_type, + ) + }) + .collect(), + _ => vec![], + }; + + let fields = entry + .fields + .iter() + .filter_map(|field| { + if field.ty == Some(rbw::api::FieldType::Hidden) { + None + } else { + field.value.as_ref() + } + }) + .map(|value| { + requests.push( + value, entry.key.as_deref(), entry.org_id.as_deref(), ) - .map(|uri| (uri, s.match_type)) }) - .collect() - } else { - vec![] - }; - let fields = entry - .fields - .iter() - .filter_map(|field| { - if field.ty == Some(rbw::api::FieldType::Hidden) { - None - } else { - field.value.as_ref() - } - }) - .map(|value| { - crate::actions::decrypt( - value, - entry.key.as_deref(), - entry.org_id.as_deref(), - ) - }) - .collect::>()?; - let notes = match notes { - Ok(notes) => notes, - Err(e) => { - log::warn!("failed to decrypt notes: {e}"); - None + .collect(); + + Self { + id: entry.id.clone(), + entry_type: entry_type_name(&entry.data).to_string(), + name, + user, + folder, + notes, + uris, + fields, } - }; - let entry_type = (match &entry.data { - rbw::db::EntryData::Login { .. } => "Login", - rbw::db::EntryData::Identity { .. } => "Identity", - rbw::db::EntryData::SshKey { .. } => "SSH Key", - rbw::db::EntryData::SecureNote => "Note", - rbw::db::EntryData::Card { .. } => "Card", - }) - .to_string(); + } - Ok(DecryptedSearchCipher { - id, - entry_type, - folder, - name, - user, - uris, - fields, - notes, - }) + fn resolve( + self, + results: &[rbw::protocol::DecryptResult], + ) -> anyhow::Result { + // name, folder, and the (non-hidden) custom fields were previously + // decrypted with `?`, so their failures stay fatal; user, uris, and + // notes were best-effort and are skipped (with a warning) on failure + let name = strict_result(&results[self.name])?; + let folder = self + .folder + .map(|index| strict_result(&results[index])) + .transpose()?; + let fields = self + .fields + .iter() + .map(|&index| strict_result(&results[index])) + .collect::>()?; + let user = self.user.and_then(|index| { + lenient_result(&results[index], Field::Username) + }); + let notes = self + .notes + .and_then(|index| lenient_result(&results[index], Field::Notes)); + let uris = self + .uris + .into_iter() + .filter_map(|(index, match_type)| { + lenient_result(&results[index], Field::Uris) + .map(|uri| (uri, match_type)) + }) + .collect(); + + Ok(DecryptedSearchCipher { + id: self.id, + entry_type: self.entry_type, + folder, + name, + user, + uris, + fields, + notes, + }) + } } fn decrypt_cipher(entry: &rbw::db::Entry) -> anyhow::Result { diff --git a/src/protocol.rs b/src/protocol.rs index ec0c06eb..2f9ec4d9 100644 --- a/src/protocol.rs +++ b/src/protocol.rs @@ -179,6 +179,9 @@ pub enum Action { entry_key: Option, org_id: Option, }, + DecryptBatch { + entries: Vec, + }, Encrypt { plaintext: String, org_id: Option, @@ -196,6 +199,28 @@ pub enum Response { Ack, Error { error: String }, Decrypt { plaintext: String }, + DecryptBatch { results: Vec }, Encrypt { cipherstring: String }, Version { version: u32 }, } + +// A single cipherstring to decrypt as part of an `Action::DecryptBatch`. Each +// entry carries its own keys so that fields encrypted with different keys +// (e.g. organization items vs. local folders) can be batched together. +#[derive(serde::Serialize, serde::Deserialize, Debug, Clone)] +pub struct DecryptRequest { + pub cipherstring: String, + pub entry_key: Option, + pub org_id: Option, +} + +// The result of decrypting a single `DecryptRequest`. Failures are reported +// per entry rather than failing the whole batch, so the caller can decide +// whether a given field is fatal (e.g. an entry name) or skippable (e.g. an +// optional login field). +#[derive(serde::Serialize, serde::Deserialize, Debug, Clone)] +#[serde(tag = "type")] +pub enum DecryptResult { + Success { plaintext: String }, + Failure { error: String }, +}