From 8bae83abaeb6566f74eb23be52af5234594d7eb9 Mon Sep 17 00:00:00 2001 From: Jannis Morgenstern <49486580+nyarthan@users.noreply.github.com> Date: Tue, 26 Mar 2024 14:34:43 +0100 Subject: [PATCH] feat: retry bitwarden requests (#37) Retry requests to bitwarden up to 4 times using an exponential backoff strategy & jitter to avoid rate-limits. --- Cargo.lock | 30 +++++++++++++- Cargo.toml | 2 +- bwenv.yaml | 2 +- lib/Cargo.toml | 2 + lib/src/bitwarden.rs | 94 +++++++++++++++++++++++++++++++------------- lib/src/config.rs | 4 +- lib/src/data.rs | 13 +----- src/main.rs | 1 + 8 files changed, 105 insertions(+), 43 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index e3dd9ba..8f80e28 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -130,6 +130,15 @@ version = "1.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9b34d609dfbaf33d6889b2b7106d3ca345eacad44200913df5ba02bfd31d2ba9" +[[package]] +name = "async-mutex" +version = "1.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "479db852db25d9dbf6204e6cb6253698f175c15726470f78af0d918e99d6156e" +dependencies = [ + "event-listener", +] + [[package]] name = "atty" version = "0.2.14" @@ -298,7 +307,7 @@ checksum = "a3e2c3daef883ecc1b5d58c15adae93470a91d425f3532ba1695849656af3fc1" [[package]] name = "bwenv" -version = "1.2.3" +version = "1.2.4" dependencies = [ "assert_cmd", "atty", @@ -321,6 +330,7 @@ name = "bwenv-lib" version = "0.1.0" dependencies = [ "anyhow", + "async-mutex", "bitwarden", "colored", "derived-deref", @@ -335,6 +345,7 @@ dependencies = [ "tabular", "tempfile", "tokio", + "tokio-retry", "toml", "tracing", "uuid", @@ -627,6 +638,12 @@ dependencies = [ "windows-sys 0.52.0", ] +[[package]] +name = "event-listener" +version = "2.5.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0206175f82b8d6bf6652ff7d71a1e27fd2e4efde587fd368662814d6ec1d9ce0" + [[package]] name = "fastrand" version = "2.0.1" @@ -2312,6 +2329,17 @@ dependencies = [ "tokio", ] +[[package]] +name = "tokio-retry" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7f57eb36ecbe0fc510036adff84824dd3c24bb781e21bfa67b69d556aa85214f" +dependencies = [ + "pin-project", + "rand", + "tokio", +] + [[package]] name = "tokio-util" version = "0.7.8" diff --git a/Cargo.toml b/Cargo.toml index b06c953..f50a679 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "bwenv" -version = "1.2.3" +version = "1.2.4" edition = "2021" [dependencies] diff --git a/bwenv.yaml b/bwenv.yaml index c2a198c..c9e5271 100644 --- a/bwenv.yaml +++ b/bwenv.yaml @@ -1,7 +1,7 @@ version: 1.2 cache: - path: ../node_modules/.cache/bwenv + path: ./node_modules/.cache/bwenv global: overrides: diff --git a/lib/Cargo.toml b/lib/Cargo.toml index d22e51b..206ffd1 100644 --- a/lib/Cargo.toml +++ b/lib/Cargo.toml @@ -20,6 +20,8 @@ derived-deref = "2.1.0" dirs = "5.0.1" reqwest = { version = "0.12.1", features = ["json"] } once_cell = "1.19.0" +tokio-retry = "0.3.0" +async-mutex = "1.4.0" [dev-dependencies] tempfile = "3.10.1" diff --git a/lib/src/bitwarden.rs b/lib/src/bitwarden.rs index ac3907f..0086e46 100644 --- a/lib/src/bitwarden.rs +++ b/lib/src/bitwarden.rs @@ -1,12 +1,19 @@ -use bitwarden::secrets_manager::secrets::{SecretIdentifiersByProjectRequest, SecretsGetRequest}; +use async_mutex::Mutex; + +use bitwarden::secrets_manager::secrets::{ + SecretIdentifiersByProjectRequest, SecretIdentifiersResponse, SecretsGetRequest, +}; use bitwarden::{ auth::login::AccessTokenLoginRequest, client::client_settings::{ClientSettings, DeviceType}, Client, }; +use tokio_retry::strategy::{jitter, ExponentialBackoff}; +use tokio_retry::Retry; use uuid::Uuid; use crate::config_yaml::Secrets; +use tracing::{error, info}; pub struct BitwardenClient { _identity_url: String, @@ -14,7 +21,7 @@ pub struct BitwardenClient { _user_agent: String, _device_type: DeviceType, _access_token: String, - client: Client, + client: Mutex, } impl BitwardenClient { @@ -36,10 +43,13 @@ impl BitwardenClient { access_token: access_token.to_owned(), }) .await - .unwrap(); + .unwrap_or_else(|_| { + error!(message = "Failed to login using access token"); + std::process::exit(1); + }); Self { - client, + client: Mutex::new(client), _access_token: access_token, _identity_url: identity_url, _api_url: api_url, @@ -51,34 +61,64 @@ impl BitwardenClient { pub async fn get_secrets_by_project_id<'a, T: AsRef>( &mut self, project_id: T, - ) -> Secrets<'a> { - let secrets_by_project_request = SecretIdentifiersByProjectRequest { - project_id: Uuid::parse_str(project_id.as_ref()).unwrap(), - }; + ) -> Result, Box> { + let secret_identifiers = async { + let retry_strategy = ExponentialBackoff::from_millis(10).map(jitter).take(4); + let request = || async { + let secrets_by_project_request = SecretIdentifiersByProjectRequest { + project_id: Uuid::parse_str(project_id.as_ref())?, + }; - let secret_identifiers = self - .client - .secrets() - .list_by_project(&secrets_by_project_request) - .await - .unwrap(); - - let secrets_get_request = SecretsGetRequest { - ids: secret_identifiers - .data - .into_iter() - .map(|ident| ident.id) - .collect(), + info!(message = "Fetching secret IDs"); + + Ok(self + .client + .lock() + .await + .secrets() + .list_by_project(&secrets_by_project_request) + .await?) + }; + + let result: Result> = + Retry::spawn(retry_strategy, request).await; + + result.unwrap() }; - self.client - .secrets() - .get_by_ids(secrets_get_request) + let ids: Vec = secret_identifiers .await - .unwrap() .data .into_iter() - .map(|secret| (secret.key, secret.value)) - .collect() + .map(|ident| ident.id) + .collect(); + + let secrets = async { + let retry_strategy = ExponentialBackoff::from_millis(10).map(jitter).take(4); + let request = || async { + let secrets_get_request = SecretsGetRequest { ids: ids.clone() }; + + info!(message = "Fetching secrets"); + + Ok(self + .client + .lock() + .await + .secrets() + .get_by_ids(secrets_get_request) + .await? + .data + .into_iter() + .map(|secret| (secret.key, secret.value)) + .collect()) + }; + + let result: Result, Box> = + Retry::spawn(retry_strategy, request).await; + + result + }; + + secrets.await } } diff --git a/lib/src/config.rs b/lib/src/config.rs index b126682..ad406ec 100644 --- a/lib/src/config.rs +++ b/lib/src/config.rs @@ -53,7 +53,7 @@ mod tests { #[test] fn finds_yaml_config_in_current_dir() { let temp_dir = tempdir().unwrap(); - create_config_file(&temp_dir.path().to_path_buf(), "bwenv.yaml"); + create_config_file(temp_dir.path(), "bwenv.yaml"); let result = find_local_config(Some(temp_dir.path())); assert!(result.is_ok()); @@ -66,7 +66,7 @@ mod tests { let temp_dir = tempdir().unwrap(); let child_dir = temp_dir.path().join("child"); std::fs::create_dir(&child_dir).unwrap(); - create_config_file(&temp_dir.path().to_path_buf(), "bwenv.toml"); + create_config_file(temp_dir.path(), "bwenv.toml"); let result = find_local_config(Some(&child_dir)); assert!(result.is_ok()); diff --git a/lib/src/data.rs b/lib/src/data.rs index 6953c29..22e629a 100644 --- a/lib/src/data.rs +++ b/lib/src/data.rs @@ -17,21 +17,12 @@ impl Default for Data { } } -#[derive(Serialize, Deserialize, Debug)] +#[derive(Serialize, Deserialize, Debug, Default)] pub struct DataContent { pub last_update_check: u64, pub last_checked_version: Option, } -impl Default for DataContent { - fn default() -> Self { - Self { - last_update_check: 0, - last_checked_version: None, - } - } -} - impl Data { pub fn new() -> Self { let path = dirs::data_dir() @@ -67,7 +58,7 @@ impl Data { } pub fn get_content(&self) -> DataContent { - self.read().unwrap_or(DataContent::default()) + self.read().unwrap_or_default() } pub fn set_content( diff --git a/src/main.rs b/src/main.rs index 9145cb1..964801b 100644 --- a/src/main.rs +++ b/src/main.rs @@ -164,6 +164,7 @@ async fn run_with<'a>( bitwarden_client .get_secrets_by_project_id(&project_id) .await + .unwrap() }) .await .unwrap();