From d2f44d511c15cee0a92fe5003218eebf6c38a1a2 Mon Sep 17 00:00:00 2001 From: Nils Pukropp Date: Thu, 25 Jun 2026 15:54:03 +0200 Subject: [PATCH] feat(sdk/http): implement automatic retries with exponential backoff and jitter --- Cargo.toml | 4 +- crates/proton-sdk/src/api.rs | 4 + crates/proton-sdk/src/config.rs | 55 ++++++ crates/proton-sdk/src/http.rs | 316 ++++++++++++++++++++++++++------ 4 files changed, 317 insertions(+), 62 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index e2f8570..3df0c7b 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -10,7 +10,7 @@ repository = "https://github.com/ProtonDriveApps/sdk" rust-version = "1.80" [workspace.dependencies] -tokio = { version = "1", features = ["rt-multi-thread", "macros", "sync", "fs", "io-util"] } +tokio = { version = "1", features = ["rt-multi-thread", "macros", "sync", "fs", "io-util", "time"] } reqwest = { version = "0.12", default-features = false, features = ["json", "rustls-tls", "stream", "multipart"] } serde = { version = "1", features = ["derive"] } serde_json = "1" @@ -24,6 +24,8 @@ hex = "0.4" hmac = "0.12" sha1 = "0.10" sha2 = "0.10" +aes-gcm = "0.10" +hkdf = "0.12" bcrypt = "0.17" pgp = "0.16" tracing = "0.1" diff --git a/crates/proton-sdk/src/api.rs b/crates/proton-sdk/src/api.rs index 5857178..37a1928 100644 --- a/crates/proton-sdk/src/api.rs +++ b/crates/proton-sdk/src/api.rs @@ -33,6 +33,8 @@ pub enum ResponseCode { Unauthorized = 401, Forbidden = 403, RequestTimeout = 408, + TooManyRequests = 429, + ServiceUnavailable = 503, Success = 1000, MultipleResponses = 1001, @@ -70,6 +72,8 @@ impl ResponseCode { 401 => Self::Unauthorized, 403 => Self::Forbidden, 408 => Self::RequestTimeout, + 429 => Self::TooManyRequests, + 503 => Self::ServiceUnavailable, 1000 => Self::Success, 1001 => Self::MultipleResponses, 2000 => Self::InvalidRequirements, diff --git a/crates/proton-sdk/src/config.rs b/crates/proton-sdk/src/config.rs index ed7ee63..0459447 100644 --- a/crates/proton-sdk/src/config.rs +++ b/crates/proton-sdk/src/config.rs @@ -11,6 +11,54 @@ pub const DEFAULT_REFRESH_REDIRECT_URI: &str = "https://proton.me"; /// Default per-request timeout, matching `ProtonApiDefaults.DefaultTimeoutSeconds`. pub const DEFAULT_TIMEOUT: Duration = Duration::from_secs(30); +/// Default number of automatic retries on a retryable failure (rate-limit, +/// transient transport error). The original attempt is not counted, so the +/// total request budget is `1 + max_retries`. +pub const DEFAULT_MAX_RETRIES: u32 = 3; + +/// Default base delay for exponential backoff between retries. +pub const DEFAULT_RETRY_BASE_DELAY: Duration = Duration::from_millis(500); + +/// Default upper bound for a single backoff sleep. A server-supplied +/// `Retry-After` is honoured even if it exceeds this. +pub const DEFAULT_RETRY_MAX_DELAY: Duration = Duration::from_secs(30); + +/// Controls automatic retries on retryable responses (HTTP 408/429/502/503/504) +/// and transient transport errors (timeout, connect). +/// +/// Mirrors the Polly-style retry pipeline in the C# SDK: a server-supplied +/// `Retry-After` header wins; otherwise the delay is exponential backoff +/// (`base_delay * 2^attempt`, capped at `max_delay`) with full jitter. +#[derive(Debug, Clone)] +pub struct RetryPolicy { + /// Maximum number of retries after the initial attempt. `0` disables retry. + pub max_retries: u32, + /// Base delay for exponential backoff. + pub base_delay: Duration, + /// Cap on a single computed backoff sleep. + pub max_delay: Duration, +} + +impl Default for RetryPolicy { + fn default() -> Self { + Self { + max_retries: DEFAULT_MAX_RETRIES, + base_delay: DEFAULT_RETRY_BASE_DELAY, + max_delay: DEFAULT_RETRY_MAX_DELAY, + } + } +} + +impl RetryPolicy { + /// A policy that performs no automatic retries. + pub fn disabled() -> Self { + Self { + max_retries: 0, + ..Self::default() + } + } +} + /// Content type used for the `Accept` header on API requests. pub const API_CONTENT_TYPE: &str = "application/vnd.protonmail.api+json"; @@ -26,6 +74,7 @@ pub struct ProtonClientConfiguration { pub user_agent: String, pub refresh_redirect_uri: String, pub request_timeout: Duration, + pub retry_policy: RetryPolicy, } impl ProtonClientConfiguration { @@ -37,6 +86,7 @@ impl ProtonClientConfiguration { user_agent: String::new(), refresh_redirect_uri: DEFAULT_REFRESH_REDIRECT_URI.to_owned(), request_timeout: DEFAULT_TIMEOUT, + retry_policy: RetryPolicy::default(), } } @@ -54,4 +104,9 @@ impl ProtonClientConfiguration { self.request_timeout = timeout; self } + + pub fn with_retry_policy(mut self, retry_policy: RetryPolicy) -> Self { + self.retry_policy = retry_policy; + self + } } diff --git a/crates/proton-sdk/src/http.rs b/crates/proton-sdk/src/http.rs index 3d47f56..da257d1 100644 --- a/crates/proton-sdk/src/http.rs +++ b/crates/proton-sdk/src/http.rs @@ -6,16 +6,19 @@ //! no `DelegatingHandler` pipeline. use std::sync::Arc; +use std::time::Duration; +use rand::Rng; use reqwest::{Method, StatusCode}; use serde::de::DeserializeOwned; use serde::Serialize; use tokio::sync::Mutex; use crate::api::{ApiResponse, ResponseCode}; -use crate::config::{ProtonClientConfiguration, API_CONTENT_TYPE}; +use crate::config::{ProtonClientConfiguration, RetryPolicy, API_CONTENT_TYPE}; use crate::error::{ProtonApiError, ProtonError, Result}; use crate::ids::SessionId; +use crate::telemetry::{NoopTelemetry, Telemetry, TelemetryExt}; const SESSION_ID_HEADER: &str = "x-pm-uid"; const APP_VERSION_HEADER: &str = "x-pm-appversion"; @@ -44,6 +47,12 @@ struct Inner { config: ProtonClientConfiguration, session_id: SessionId, tokens: Mutex, + /// Telemetry sink for per-request events. Interior-mutable because the + /// client is already shared (cloned into the Drive client) by the time a + /// caller attaches a sink via [`ApiHttpClient::set_telemetry`]. Defaults to + /// a no-op. `std::sync::Mutex` (not tokio's) — held only for the cheap + /// clone/replace, never across an await. + telemetry: std::sync::Mutex>, } impl ApiHttpClient { @@ -66,6 +75,7 @@ impl ApiHttpClient { config, session_id, tokens: Mutex::new(tokens), + telemetry: std::sync::Mutex::new(NoopTelemetry::shared()), }), }) } @@ -75,6 +85,25 @@ impl ApiHttpClient { self.inner.tokens.lock().await.clone() } + /// Attach a telemetry sink to receive a per-request + /// [`TelemetryEvent`](crate::telemetry::TelemetryEvent) (operation + /// `http_request` for API calls, `storage_download` / `storage_upload` for + /// block storage; attributes carry the HTTP method and status). Replaces any + /// previous sink. Takes effect for every clone of this client, since they + /// share state. + pub fn set_telemetry(&self, telemetry: Arc) { + *self.inner.telemetry.lock().expect("telemetry mutex poisoned") = telemetry; + } + + /// Snapshot the current telemetry sink. + fn telemetry(&self) -> Arc { + self.inner + .telemetry + .lock() + .expect("telemetry mutex poisoned") + .clone() + } + /// `GET {url}` against block storage, returning the raw (still-encrypted) /// blob bytes. /// @@ -84,18 +113,18 @@ impl ApiHttpClient { /// .GetBlobStreamAsync`. A successful response is raw binary; an error /// response is the usual JSON envelope. pub async fn get_storage_blob(&self, url: &str, token: &str) -> Result> { - let mut request = self - .inner - .http - .get(url) - .header(STORAGE_TOKEN_HEADER, token); - - if !self.inner.config.user_agent.is_empty() { - request = request.header(reqwest::header::USER_AGENT, &self.inner.config.user_agent); - } - - let response = request.send().await?; + let mut timer = self.telemetry().start("storage_download"); + let response = send_retrying(&self.inner.config.retry_policy, || { + let mut request = self.inner.http.get(url).header(STORAGE_TOKEN_HEADER, token); + if !self.inner.config.user_agent.is_empty() { + request = + request.header(reqwest::header::USER_AGENT, &self.inner.config.user_agent); + } + request + }) + .await?; let status = response.status(); + timer.attr("status", status.as_u16()); let bytes = response.bytes().await?; // Success bodies are raw block bytes (not JSON); only error responses @@ -108,6 +137,7 @@ impl ApiHttpClient { return Err(api_error(status, &bytes)); } + timer.success(); Ok(bytes.to_vec()) } @@ -118,25 +148,36 @@ impl ApiHttpClient { /// authorized by the per-block `pm-storage-token` header rather than the /// session bearer. The response is the usual JSON envelope. pub async fn post_storage_blob(&self, url: &str, token: &str, blob: Vec) -> Result<()> { - let part = reqwest::multipart::Part::bytes(blob) - .file_name("blob") + // Validate the part once up front; the multipart body itself is rebuilt + // per attempt inside the retry closure (a stream body can't be cloned). + reqwest::multipart::Part::bytes(Vec::new()) .mime_str("application/octet-stream") .map_err(ProtonError::from)?; - let form = reqwest::multipart::Form::new().part("Block", part); - let mut request = self - .inner - .http - .post(url) - .header(STORAGE_TOKEN_HEADER, token) - .multipart(form); + let mut timer = self.telemetry().start("storage_upload"); + let response = send_retrying(&self.inner.config.retry_policy, || { + let part = reqwest::multipart::Part::bytes(blob.clone()) + .file_name("blob") + .mime_str("application/octet-stream") + .expect("octet-stream is a valid MIME type"); + let form = reqwest::multipart::Form::new().part("Block", part); - if !self.inner.config.user_agent.is_empty() { - request = request.header(reqwest::header::USER_AGENT, &self.inner.config.user_agent); - } + let mut request = self + .inner + .http + .post(url) + .header(STORAGE_TOKEN_HEADER, token) + .multipart(form); - let response = request.send().await?; + if !self.inner.config.user_agent.is_empty() { + request = + request.header(reqwest::header::USER_AGENT, &self.inner.config.user_agent); + } + request + }) + .await?; let status = response.status(); + timer.attr("status", status.as_u16()); let bytes = response.bytes().await?; if let Ok(envelope) = serde_json::from_slice::(&bytes) { @@ -146,6 +187,7 @@ impl ApiHttpClient { } else if !status.is_success() { return Err(api_error(status, &bytes)); } + timer.success(); Ok(()) } @@ -179,8 +221,12 @@ impl ApiHttpClient { path: &str, body: Option<&B>, ) -> Result { + let mut timer = self.telemetry().start("http_request"); + timer.attr("method", method.as_str()); + let access_token = self.inner.tokens.lock().await.access_token.clone(); + // An early `?` here records the op as a failure (OpTimer defaults to it). let response = self .send_with_token(method.clone(), path, body, &access_token) .await?; @@ -192,7 +238,10 @@ impl ApiHttpClient { response }; - parse_response(response).await + timer.attr("status", response.status().as_u16()); + let parsed = parse_response(response).await?; + timer.success(); + Ok(parsed) } async fn handle_unauthorized( @@ -226,24 +275,28 @@ impl ApiHttpClient { access_token: &str, ) -> Result { let url = format!("{}{}", self.inner.base_url, path.trim_start_matches('/')); - let mut request = self - .inner - .http - .request(method, url) - .header(SESSION_ID_HEADER, self.inner.session_id.as_str()) - .header(APP_VERSION_HEADER, &self.inner.config.app_version) - .header(reqwest::header::ACCEPT, API_CONTENT_TYPE) - .bearer_auth(access_token); + send_retrying(&self.inner.config.retry_policy, || { + let mut request = self + .inner + .http + .request(method.clone(), &url) + .header(SESSION_ID_HEADER, self.inner.session_id.as_str()) + .header(APP_VERSION_HEADER, &self.inner.config.app_version) + .header(reqwest::header::ACCEPT, API_CONTENT_TYPE) + .bearer_auth(access_token); - if !self.inner.config.user_agent.is_empty() { - request = request.header(reqwest::header::USER_AGENT, &self.inner.config.user_agent); - } + if !self.inner.config.user_agent.is_empty() { + request = + request.header(reqwest::header::USER_AGENT, &self.inner.config.user_agent); + } - if let Some(body) = body { - request = request.json(body); - } + if let Some(body) = body { + request = request.json(body); + } - Ok(request.send().await?) + request + }) + .await } /// Refresh the session tokens, deduplicating concurrent refreshes: if the @@ -272,16 +325,16 @@ impl ApiHttpClient { // The refresh call carries the session id but, deliberately, no bearer // token (the access token is the thing being replaced). - let response = self - .inner - .http - .post(url) - .header(SESSION_ID_HEADER, self.inner.session_id.as_str()) - .header(APP_VERSION_HEADER, &self.inner.config.app_version) - .header(reqwest::header::ACCEPT, API_CONTENT_TYPE) - .json(&body) - .send() - .await?; + let response = send_retrying(&self.inner.config.retry_policy, || { + self.inner + .http + .post(&url) + .header(SESSION_ID_HEADER, self.inner.session_id.as_str()) + .header(APP_VERSION_HEADER, &self.inner.config.app_version) + .header(reqwest::header::ACCEPT, API_CONTENT_TYPE) + .json(&body) + }) + .await?; let refreshed: SessionRefreshResponse = parse_response(response).await?; Ok(Tokens { @@ -328,17 +381,20 @@ pub async fn post_unauthenticated( let base_url = ensure_trailing_slash(&config.base_url); let url = format!("{}{}", base_url, path.trim_start_matches('/')); - let mut request = http - .post(url) - .header(APP_VERSION_HEADER, &config.app_version) - .header(reqwest::header::ACCEPT, API_CONTENT_TYPE) - .json(body); + let response = send_retrying(&config.retry_policy, || { + let mut request = http + .post(&url) + .header(APP_VERSION_HEADER, &config.app_version) + .header(reqwest::header::ACCEPT, API_CONTENT_TYPE) + .json(body); + if !config.user_agent.is_empty() { + request = request.header(reqwest::header::USER_AGENT, &config.user_agent); + } + request + }) + .await?; - if !config.user_agent.is_empty() { - request = request.header(reqwest::header::USER_AGENT, &config.user_agent); - } - - parse_response(request.send().await?).await + parse_response(response).await } /// Read a response body, enforce the Proton success envelope, and deserialize @@ -379,6 +435,89 @@ fn api_error(status: StatusCode, bytes: &[u8]) -> ProtonError { }) } +/// Send a request, transparently retrying retryable failures per `policy`. +/// +/// `build` is called once per attempt to produce a fresh `RequestBuilder` +/// (reqwest builders are consumed by `send`, and streaming bodies like +/// `multipart` can't be cloned), so every retry resends the full request. +/// +/// Retryable = HTTP 408/429/502/503/504 or a transient transport error +/// (timeout / connect). A `Retry-After` header (delta-seconds) is honoured; +/// otherwise the delay is exponential backoff with full jitter. Non-retryable +/// responses and errors — including ordinary 4xx and the 401 that drives token +/// refresh — pass straight through to the caller untouched. +async fn send_retrying(policy: &RetryPolicy, build: F) -> Result +where + F: Fn() -> reqwest::RequestBuilder, +{ + let mut attempt: u32 = 0; + loop { + match build().send().await { + Ok(response) => { + if attempt < policy.max_retries && is_retryable_status(response.status()) { + let delay = retry_after(&response).unwrap_or_else(|| backoff(policy, attempt)); + tokio::time::sleep(delay).await; + attempt += 1; + continue; + } + return Ok(response); + } + Err(err) => { + if attempt < policy.max_retries && is_retryable_error(&err) { + tokio::time::sleep(backoff(policy, attempt)).await; + attempt += 1; + continue; + } + return Err(err.into()); + } + } + } +} + +/// Status codes Proton (or an intermediary) returns for transient conditions: +/// request timeout, rate limit, and the gateway/unavailable family. +fn is_retryable_status(status: StatusCode) -> bool { + matches!(status.as_u16(), 408 | 429 | 502 | 503 | 504) +} + +/// A transport error worth retrying: a timeout or a failure to connect. A +/// mid-body error (`is_body`) is not retried — the request may have been +/// applied server-side. +fn is_retryable_error(err: &reqwest::Error) -> bool { + err.is_timeout() || err.is_connect() +} + +/// Parse a `Retry-After` header expressed as delta-seconds. The HTTP-date form +/// is not emitted by the Proton API, so it is ignored (falls back to backoff). +fn retry_after(response: &reqwest::Response) -> Option { + let value = response + .headers() + .get(reqwest::header::RETRY_AFTER)? + .to_str() + .ok()?; + parse_retry_after_secs(value) +} + +/// Parse a `Retry-After` delta-seconds value into a delay. Non-numeric values +/// (the HTTP-date form, which Proton does not emit) yield `None`. +fn parse_retry_after_secs(value: &str) -> Option { + value.trim().parse().ok().map(Duration::from_secs) +} + +/// Exponential backoff with full jitter: a uniformly random delay in +/// `[0, base_delay * 2^attempt]`, capped at `max_delay`. +fn backoff(policy: &RetryPolicy, attempt: u32) -> Duration { + let ceiling = policy + .base_delay + .saturating_mul(1u32.checked_shl(attempt).unwrap_or(u32::MAX)) + .min(policy.max_delay); + let ceiling_ms = ceiling.as_millis() as u64; + if ceiling_ms == 0 { + return Duration::ZERO; + } + Duration::from_millis(rand::thread_rng().gen_range(0..=ceiling_ms)) +} + fn ensure_trailing_slash(url: &str) -> String { if url.ends_with('/') { url.to_owned() @@ -386,3 +525,58 @@ fn ensure_trailing_slash(url: &str) -> String { format!("{url}/") } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn retryable_statuses() { + for code in [408u16, 429, 502, 503, 504] { + assert!(is_retryable_status(StatusCode::from_u16(code).unwrap())); + } + for code in [200u16, 400, 401, 403, 404, 500] { + assert!(!is_retryable_status(StatusCode::from_u16(code).unwrap())); + } + } + + #[test] + fn retry_after_parses_seconds_only() { + assert_eq!(parse_retry_after_secs("5"), Some(Duration::from_secs(5))); + assert_eq!(parse_retry_after_secs(" 12 "), Some(Duration::from_secs(12))); + assert_eq!(parse_retry_after_secs("0"), Some(Duration::ZERO)); + // HTTP-date form is unsupported -> falls back to backoff. + assert_eq!(parse_retry_after_secs("Wed, 21 Oct 2015 07:28:00 GMT"), None); + assert_eq!(parse_retry_after_secs(""), None); + } + + #[test] + fn backoff_grows_then_caps_within_jitter_bounds() { + let policy = RetryPolicy { + max_retries: 5, + base_delay: Duration::from_millis(100), + max_delay: Duration::from_millis(1000), + }; + // Full jitter: every sample stays within [0, ceiling] where the + // ceiling is base*2^attempt capped at max_delay. + for attempt in 0..8u32 { + let ceiling = Duration::from_millis(100u64.saturating_mul(1 << attempt.min(20))) + .min(policy.max_delay); + for _ in 0..64 { + assert!(backoff(&policy, attempt) <= ceiling); + } + } + } + + #[test] + fn backoff_handles_large_attempt_without_overflow() { + let policy = RetryPolicy::default(); + // attempt >= 32 would overflow a naive shift; must saturate to max_delay. + assert!(backoff(&policy, 64) <= policy.max_delay); + } + + #[test] + fn disabled_policy_has_no_retries() { + assert_eq!(RetryPolicy::disabled().max_retries, 0); + } +}