1use std::sync::Arc;
9use std::time::Duration;
10
11use rand::RngExt;
12use reqwest::{Method, StatusCode};
13use serde::Serialize;
14use serde::de::DeserializeOwned;
15use tokio::sync::Mutex;
16
17use crate::api::{ApiResponse, ResponseCode};
18use crate::config::{API_CONTENT_TYPE, ProtonClientConfiguration, RetryPolicy};
19use crate::error::{ProtonApiError, ProtonError, Result};
20use crate::ids::SessionId;
21use crate::telemetry::{NoopTelemetry, Telemetry, TelemetryExt};
22
23const SESSION_ID_HEADER: &str = "x-pm-uid";
24const APP_VERSION_HEADER: &str = "x-pm-appversion";
25const STORAGE_TOKEN_HEADER: &str = "pm-storage-token";
26
27#[derive(Debug, Clone)]
30pub struct Tokens {
31 pub access_token: String,
32 pub refresh_token: String,
33}
34
35#[derive(Clone)]
40pub struct ApiHttpClient {
41 inner: Arc<Inner>,
42 route_prefix: Arc<str>,
48}
49
50type TokensRefreshedCallback = Arc<dyn Fn(Tokens) + Send + Sync>;
52
53struct Inner {
54 http: reqwest::Client,
55 base_url: String,
56 config: ProtonClientConfiguration,
57 session_id: SessionId,
58 tokens: Mutex<Tokens>,
59 telemetry: std::sync::Mutex<Arc<dyn Telemetry>>,
65 on_tokens_refreshed: std::sync::Mutex<Option<TokensRefreshedCallback>>,
66}
67
68impl ApiHttpClient {
69 pub fn new(
71 config: ProtonClientConfiguration,
72 session_id: SessionId,
73 tokens: Tokens,
74 ) -> Result<Self> {
75 let http = reqwest::Client::builder()
76 .timeout(config.request_timeout)
77 .build()?;
78
79 let base_url = ensure_trailing_slash(&config.base_url);
80
81 Ok(Self {
82 inner: Arc::new(Inner {
83 http,
84 base_url,
85 config,
86 session_id,
87 tokens: Mutex::new(tokens),
88 telemetry: std::sync::Mutex::new(NoopTelemetry::shared()),
89 on_tokens_refreshed: std::sync::Mutex::new(None),
90 }),
91 route_prefix: Arc::from(""),
92 })
93 }
94
95 pub fn with_base_route(&self, route: impl Into<Arc<str>>) -> Self {
101 Self {
102 inner: Arc::clone(&self.inner),
103 route_prefix: route.into(),
104 }
105 }
106
107 pub async fn current_tokens(&self) -> Tokens {
109 self.inner.tokens.lock().await.clone()
110 }
111
112 pub fn set_telemetry(&self, telemetry: Arc<dyn Telemetry>) {
119 *self
120 .inner
121 .telemetry
122 .lock()
123 .expect("telemetry mutex poisoned") = telemetry;
124 }
125
126 pub fn set_on_tokens_refreshed(&self, callback: impl Fn(Tokens) + Send + Sync + 'static) {
129 *self
130 .inner
131 .on_tokens_refreshed
132 .lock()
133 .expect("on_tokens_refreshed mutex poisoned") = Some(Arc::new(callback));
134 }
135
136 fn telemetry(&self) -> Arc<dyn Telemetry> {
138 self.inner
139 .telemetry
140 .lock()
141 .expect("telemetry mutex poisoned")
142 .clone()
143 }
144
145 pub async fn get_storage_blob(&self, url: &str, token: &str) -> Result<Vec<u8>> {
154 let mut timer = self.telemetry().start("storage_download");
155 let response = send_retrying(&self.inner.config.retry_policy, || {
156 let mut request = self.inner.http.get(url).header(STORAGE_TOKEN_HEADER, token);
157 if !self.inner.config.user_agent.is_empty() {
158 request =
159 request.header(reqwest::header::USER_AGENT, &self.inner.config.user_agent);
160 }
161 request
162 })
163 .await?;
164 let status = response.status();
165 timer.attr("status", status.as_u16());
166 let bytes = response.bytes().await?;
167
168 if let Ok(envelope) = serde_json::from_slice::<ApiResponse>(&bytes) {
171 if !envelope.is_success() {
172 return Err(api_error(status, &bytes));
173 }
174 } else if !status.is_success() {
175 return Err(api_error(status, &bytes));
176 }
177
178 timer.success();
179 Ok(bytes.to_vec())
180 }
181
182 pub async fn post_storage_blob(&self, url: &str, token: &str, blob: Vec<u8>) -> Result<()> {
189 reqwest::multipart::Part::bytes(Vec::new())
192 .mime_str("application/octet-stream")
193 .map_err(ProtonError::from)?;
194
195 let mut timer = self.telemetry().start("storage_upload");
196 let response = send_retrying(&self.inner.config.retry_policy, || {
197 let part = reqwest::multipart::Part::bytes(blob.clone())
198 .file_name("blob")
199 .mime_str("application/octet-stream")
200 .expect("octet-stream is a valid MIME type");
201 let form = reqwest::multipart::Form::new().part("Block", part);
202
203 let mut request = self
204 .inner
205 .http
206 .post(url)
207 .header(STORAGE_TOKEN_HEADER, token)
208 .multipart(form);
209
210 if !self.inner.config.user_agent.is_empty() {
211 request =
212 request.header(reqwest::header::USER_AGENT, &self.inner.config.user_agent);
213 }
214 request
215 })
216 .await?;
217 let status = response.status();
218 timer.attr("status", status.as_u16());
219 let bytes = response.bytes().await?;
220
221 if let Ok(envelope) = serde_json::from_slice::<ApiResponse>(&bytes) {
222 if !envelope.is_success() {
223 return Err(api_error(status, &bytes));
224 }
225 } else if !status.is_success() {
226 return Err(api_error(status, &bytes));
227 }
228 timer.success();
229 Ok(())
230 }
231
232 pub fn session_id(&self) -> &SessionId {
233 &self.inner.session_id
234 }
235
236 pub async fn get<T: DeserializeOwned>(&self, path: &str) -> Result<T> {
238 self.send::<(), T>(Method::GET, path, None).await
239 }
240
241 pub async fn post<B: Serialize, T: DeserializeOwned>(&self, path: &str, body: &B) -> Result<T> {
243 self.send::<B, T>(Method::POST, path, Some(body)).await
244 }
245
246 pub async fn put<B: Serialize, T: DeserializeOwned>(&self, path: &str, body: &B) -> Result<T> {
248 self.send::<B, T>(Method::PUT, path, Some(body)).await
249 }
250
251 pub async fn delete<T: DeserializeOwned>(&self, path: &str) -> Result<T> {
253 self.send::<(), T>(Method::DELETE, path, None).await
254 }
255
256 async fn send<B: Serialize, T: DeserializeOwned>(
257 &self,
258 method: Method,
259 path: &str,
260 body: Option<&B>,
261 ) -> Result<T> {
262 let mut timer = self.telemetry().start("http_request");
263 timer.attr("method", method.as_str());
264
265 let access_token = self.inner.tokens.lock().await.access_token.clone();
266
267 let response = self
269 .send_with_token(method.clone(), path, body, &access_token)
270 .await?;
271
272 let response = if response.status() == StatusCode::UNAUTHORIZED {
273 self.handle_unauthorized(method, path, body, response, access_token)
274 .await?
275 } else {
276 response
277 };
278
279 timer.attr("status", response.status().as_u16());
280 let parsed = parse_response(response).await?;
281 timer.success();
282 Ok(parsed)
283 }
284
285 async fn handle_unauthorized<B: Serialize>(
286 &self,
287 method: Method,
288 path: &str,
289 body: Option<&B>,
290 response: reqwest::Response,
291 rejected_access_token: String,
292 ) -> Result<reqwest::Response> {
293 let bytes = response.bytes().await?;
295 if let Ok(envelope) = serde_json::from_slice::<ApiResponse>(&bytes)
296 && matches!(
297 envelope.code,
298 ResponseCode::AccountDeleted | ResponseCode::AccountDisabled
299 )
300 {
301 return Err(api_error(StatusCode::UNAUTHORIZED, &bytes));
302 }
303
304 let access_token = self.refresh_access_token(&rejected_access_token).await?;
305 self.send_with_token(method, path, body, &access_token)
306 .await
307 }
308
309 async fn send_with_token<B: Serialize>(
310 &self,
311 method: Method,
312 path: &str,
313 body: Option<&B>,
314 access_token: &str,
315 ) -> Result<reqwest::Response> {
316 let url = format!(
317 "{}{}{}",
318 self.inner.base_url,
319 self.route_prefix,
320 path.trim_start_matches('/')
321 );
322 send_retrying(&self.inner.config.retry_policy, || {
323 let mut request = self
324 .inner
325 .http
326 .request(method.clone(), &url)
327 .header(SESSION_ID_HEADER, self.inner.session_id.as_str())
328 .header(APP_VERSION_HEADER, &self.inner.config.app_version)
329 .header(reqwest::header::ACCEPT, API_CONTENT_TYPE)
330 .bearer_auth(access_token);
331
332 if !self.inner.config.user_agent.is_empty() {
333 request =
334 request.header(reqwest::header::USER_AGENT, &self.inner.config.user_agent);
335 }
336
337 if let Some(body) = body {
338 request = request.json(body);
339 }
340
341 request
342 })
343 .await
344 }
345
346 async fn refresh_access_token(&self, rejected_access_token: &str) -> Result<String> {
350 let mut guard = self.inner.tokens.lock().await;
351
352 if guard.access_token != rejected_access_token {
353 return Ok(guard.access_token.clone());
354 }
355
356 let refreshed = self.request_refresh(&guard.refresh_token).await?;
357 *guard = refreshed.clone();
358
359 if let Some(ref cb) = *self
361 .inner
362 .on_tokens_refreshed
363 .lock()
364 .expect("on_tokens_refreshed mutex poisoned")
365 {
366 cb(refreshed.clone());
367 }
368
369 Ok(refreshed.access_token)
370 }
371
372 async fn request_refresh(&self, refresh_token: &str) -> Result<Tokens> {
373 let url = format!("{}auth/v4/refresh", self.inner.base_url);
374 let body = SessionRefreshRequest {
375 response_type: "token",
376 grant_type: "refresh_token",
377 refresh_token,
378 redirect_uri: &self.inner.config.refresh_redirect_uri,
379 };
380
381 let response = send_retrying(&self.inner.config.retry_policy, || {
384 self.inner
385 .http
386 .post(&url)
387 .header(SESSION_ID_HEADER, self.inner.session_id.as_str())
388 .header(APP_VERSION_HEADER, &self.inner.config.app_version)
389 .header(reqwest::header::ACCEPT, API_CONTENT_TYPE)
390 .json(&body)
391 })
392 .await?;
393
394 let refreshed: SessionRefreshResponse = parse_response(response).await?;
395 Ok(Tokens {
396 access_token: refreshed.access_token,
397 refresh_token: refreshed.refresh_token,
398 })
399 }
400}
401
402#[derive(Serialize)]
403struct SessionRefreshRequest<'a> {
404 #[serde(rename = "ResponseType")]
405 response_type: &'a str,
406 #[serde(rename = "GrantType")]
407 grant_type: &'a str,
408 #[serde(rename = "RefreshToken")]
409 refresh_token: &'a str,
410 #[serde(rename = "RedirectURI")]
411 redirect_uri: &'a str,
412}
413
414#[derive(serde::Deserialize)]
415struct SessionRefreshResponse {
416 #[serde(rename = "AccessToken")]
417 access_token: String,
418 #[serde(rename = "RefreshToken")]
419 refresh_token: String,
420}
421
422pub async fn post_unauthenticated<B: Serialize, T: DeserializeOwned>(
428 config: &ProtonClientConfiguration,
429 path: &str,
430 body: &B,
431) -> Result<T> {
432 let http = reqwest::Client::builder()
433 .timeout(config.request_timeout)
434 .build()?;
435
436 let base_url = ensure_trailing_slash(&config.base_url);
437 let url = format!("{}{}", base_url, path.trim_start_matches('/'));
438
439 let response = send_retrying(&config.retry_policy, || {
440 let mut request = http
441 .post(&url)
442 .header(APP_VERSION_HEADER, &config.app_version)
443 .header(reqwest::header::ACCEPT, API_CONTENT_TYPE)
444 .json(body);
445 if !config.user_agent.is_empty() {
446 request = request.header(reqwest::header::USER_AGENT, &config.user_agent);
447 }
448 request
449 })
450 .await?;
451
452 parse_response(response).await
453}
454
455async fn parse_response<T: DeserializeOwned>(response: reqwest::Response) -> Result<T> {
458 let status = response.status();
459 let bytes = response.bytes().await?;
460
461 if let Ok(envelope) = serde_json::from_slice::<ApiResponse>(&bytes) {
466 if !envelope.is_success() && envelope.code != ResponseCode::MultipleResponses {
467 return Err(api_error(status, &bytes));
468 }
469 } else if !status.is_success() {
470 return Err(api_error(status, &bytes));
471 }
472
473 Ok(serde_json::from_slice::<T>(&bytes)?)
474}
475
476fn api_error(status: StatusCode, bytes: &[u8]) -> ProtonError {
477 let envelope = serde_json::from_slice::<ApiResponse>(bytes).ok();
478 let code = envelope
479 .as_ref()
480 .map(|e| e.code)
481 .unwrap_or(ResponseCode::Unknown);
482 let message = envelope.and_then(|e| e.error_message).unwrap_or_else(|| {
483 status
484 .canonical_reason()
485 .unwrap_or("unknown error")
486 .to_owned()
487 });
488
489 ProtonError::Api(ProtonApiError {
490 code,
491 http_status: status.as_u16(),
492 message,
493 })
494}
495
496async fn send_retrying<F>(policy: &RetryPolicy, build: F) -> Result<reqwest::Response>
508where
509 F: Fn() -> reqwest::RequestBuilder,
510{
511 let mut attempt: u32 = 0;
512 loop {
513 match build().send().await {
514 Ok(response) => {
515 if attempt < policy.max_retries && is_retryable_status(response.status()) {
516 let delay = retry_after(&response).unwrap_or_else(|| backoff(policy, attempt));
517 tokio::time::sleep(delay).await;
518 attempt += 1;
519 continue;
520 }
521 return Ok(response);
522 }
523 Err(err) => {
524 if attempt < policy.max_retries && is_retryable_error(&err) {
525 tokio::time::sleep(backoff(policy, attempt)).await;
526 attempt += 1;
527 continue;
528 }
529 return Err(err.into());
530 }
531 }
532 }
533}
534
535fn is_retryable_status(status: StatusCode) -> bool {
538 matches!(status.as_u16(), 408 | 429 | 502 | 503 | 504)
539}
540
541fn is_retryable_error(err: &reqwest::Error) -> bool {
545 err.is_timeout() || err.is_connect()
546}
547
548fn retry_after(response: &reqwest::Response) -> Option<Duration> {
551 let value = response
552 .headers()
553 .get(reqwest::header::RETRY_AFTER)?
554 .to_str()
555 .ok()?;
556 parse_retry_after_secs(value)
557}
558
559fn parse_retry_after_secs(value: &str) -> Option<Duration> {
562 value.trim().parse().ok().map(Duration::from_secs)
563}
564
565fn backoff(policy: &RetryPolicy, attempt: u32) -> Duration {
568 let ceiling = policy
569 .base_delay
570 .saturating_mul(1u32.checked_shl(attempt).unwrap_or(u32::MAX))
571 .min(policy.max_delay);
572 let ceiling_ms = ceiling.as_millis() as u64;
573 if ceiling_ms == 0 {
574 return Duration::ZERO;
575 }
576 Duration::from_millis(rand::rng().random_range(0..=ceiling_ms))
577}
578
579fn ensure_trailing_slash(url: &str) -> String {
580 if url.ends_with('/') {
581 url.to_owned()
582 } else {
583 format!("{url}/")
584 }
585}
586
587#[cfg(test)]
588mod tests {
589 use super::*;
590
591 #[test]
592 fn retryable_statuses() {
593 for code in [408u16, 429, 502, 503, 504] {
594 assert!(is_retryable_status(StatusCode::from_u16(code).unwrap()));
595 }
596 for code in [200u16, 400, 401, 403, 404, 500] {
597 assert!(!is_retryable_status(StatusCode::from_u16(code).unwrap()));
598 }
599 }
600
601 #[test]
602 fn retry_after_parses_seconds_only() {
603 assert_eq!(parse_retry_after_secs("5"), Some(Duration::from_secs(5)));
604 assert_eq!(
605 parse_retry_after_secs(" 12 "),
606 Some(Duration::from_secs(12))
607 );
608 assert_eq!(parse_retry_after_secs("0"), Some(Duration::ZERO));
609 assert_eq!(
611 parse_retry_after_secs("Wed, 21 Oct 2015 07:28:00 GMT"),
612 None
613 );
614 assert_eq!(parse_retry_after_secs(""), None);
615 }
616
617 #[test]
618 fn backoff_grows_then_caps_within_jitter_bounds() {
619 let policy = RetryPolicy {
620 max_retries: 5,
621 base_delay: Duration::from_millis(100),
622 max_delay: Duration::from_millis(1000),
623 };
624 for attempt in 0..8u32 {
627 let ceiling = Duration::from_millis(100u64.saturating_mul(1 << attempt.min(20)))
628 .min(policy.max_delay);
629 for _ in 0..64 {
630 assert!(backoff(&policy, attempt) <= ceiling);
631 }
632 }
633 }
634
635 #[test]
636 fn backoff_handles_large_attempt_without_overflow() {
637 let policy = RetryPolicy::default();
638 assert!(backoff(&policy, 64) <= policy.max_delay);
640 }
641
642 #[test]
643 fn disabled_policy_has_no_retries() {
644 assert_eq!(RetryPolicy::disabled().max_retries, 0);
645 }
646
647 struct Capture(std::sync::Mutex<Vec<crate::telemetry::TelemetryEvent>>);
649
650 impl Telemetry for Capture {
651 fn record(&self, event: &crate::telemetry::TelemetryEvent) {
652 self.0.lock().unwrap().push(event.clone());
653 }
654 }
655
656 #[tokio::test]
657 async fn http_request_records_telemetry_event() {
658 use crate::telemetry::Outcome;
659 use tokio::io::{AsyncReadExt, AsyncWriteExt};
660
661 let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
664 let addr = listener.local_addr().unwrap();
665 let server = tokio::spawn(async move {
666 let (mut sock, _) = listener.accept().await.unwrap();
667 let mut buf = [0u8; 2048];
668 let _ = sock.read(&mut buf).await.unwrap();
669 let body = br#"{"Code":1000}"#;
670 let head = format!(
671 "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n",
672 body.len()
673 );
674 sock.write_all(head.as_bytes()).await.unwrap();
675 sock.write_all(body).await.unwrap();
676 sock.flush().await.unwrap();
677 });
678
679 let config = ProtonClientConfiguration::new("test@1.0")
680 .with_base_url(format!("http://{addr}/"))
681 .with_retry_policy(RetryPolicy::disabled());
682 let client = ApiHttpClient::new(
683 config,
684 SessionId::from("test-session"),
685 Tokens {
686 access_token: "access".into(),
687 refresh_token: "refresh".into(),
688 },
689 )
690 .unwrap();
691
692 let capture = Arc::new(Capture(std::sync::Mutex::new(Vec::new())));
693 client.set_telemetry(capture.clone());
694
695 let _: ApiResponse = client.get("some/path").await.unwrap();
696 server.await.unwrap();
697
698 let events = capture.0.lock().unwrap();
699 assert_eq!(events.len(), 1, "exactly one http_request event");
700 let event = &events[0];
701 assert_eq!(event.operation, "http_request");
702 assert_eq!(event.outcome, Outcome::Success);
703 assert!(
704 event
705 .attributes
706 .iter()
707 .any(|(k, v)| *k == "method" && v == "GET")
708 );
709 assert!(
710 event
711 .attributes
712 .iter()
713 .any(|(k, v)| *k == "status" && v == "200")
714 );
715 }
716}