refactor: Rewrite session policies as request repeater

The behavior is expected for all proton clients, there's no real benefit
in leaving this an optional configuration parameter.

Author: LeanderBB

examples/user_id.rs

@@ -1,5 +1,5 @@
-use proton_api_rs::SessionType;
-use proton_api_rs::{http, ping_async, DefaultSession};
+use proton_api_rs::{http, ping_async};
+use proton_api_rs::{Session, SessionType};
 pub use tokio;
 use tokio::io::{AsyncBufReadExt, AsyncWriteExt};
 
@@ -16,7 +16,7 @@ async fn main() {
 
     ping_async(&client).await.unwrap();
 
-    let session = match DefaultSession::login_async(&client, &user_email, &user_password)
+    let session = match Session::login_async(&client, &user_email, &user_password, None)
         .await
         .unwrap()
     {

examples/user_id_sync.rs

@@ -1,5 +1,5 @@
 use proton_api_rs::clientv2::{ping, SessionType};
-use proton_api_rs::{http, DefaultSession};
+use proton_api_rs::{http, Session};
 use std::io::{BufRead, Write};
 
 fn main() {
@@ -16,7 +16,7 @@ fn main() {
 
     ping(&client).unwrap();
 
-    let session = match DefaultSession::login(&client, &user_email, &user_password).unwrap() {
+    let session = match Session::login(&client, &user_email, &user_password, None).unwrap() {
         SessionType::Authenticated(s) => s,
         SessionType::AwaitingTotp(mut t) => {
             let mut line_reader = std::io::BufReader::new(std::io::stdin());

src/clientv2/mod.rs

@@ -1,9 +1,9 @@
 mod client;
+mod request_repeater;
 mod session;
-mod session_policies;
 mod totp;
 
 pub use client::*;
+pub use request_repeater::*;
 pub use session::*;
-pub use session_policies::*;
 pub use totp::*;

src/clientv2/request_repeater.rs

@@ -0,0 +1,157 @@
+//! Automatic request repeater based on the expectations Proton has for their clients.
+
+use crate::domain::UserUid;
+use crate::http::{
+    ClientAsync, ClientSync, DefaultRequestFactory, Method, Request, RequestData, RequestFactory,
+};
+use crate::requests::{AuthRefreshRequest, UserAuth};
+use crate::{http, SessionRefreshData};
+use secrecy::ExposeSecret;
+
+pub type OnAuthRefreshedCallback = Box<dyn Fn(&UserUid, &str)>;
+
+pub struct RequestRepeater {
+    user_auth: parking_lot::RwLock<UserAuth>,
+    on_auth_refreshed: Option<OnAuthRefreshedCallback>,
+}
+
+impl std::fmt::Debug for RequestRepeater {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(
+            f,
+            "RequestRepeater{{user_auth:{:?} on_auth_refreshed:{}}}",
+            self.user_auth,
+            if self.on_auth_refreshed.is_some() {
+                "Some"
+            } else {
+                "None"
+            }
+        )
+    }
+}
+
+impl RequestRepeater {
+    pub fn new(user_auth: UserAuth, on_auth_refreshed: Option<OnAuthRefreshedCallback>) -> Self {
+        Self {
+            user_auth: parking_lot::RwLock::new(user_auth),
+            on_auth_refreshed,
+        }
+    }
+
+    fn refresh_auth<C: ClientSync>(&self, client: &C) -> http::Result<()> {
+        let borrow = self.user_auth.read();
+        match AuthRefreshRequest::new(
+            borrow.uid.expose_secret(),
+            borrow.refresh_token.expose_secret(),
+        )
+        .execute_sync(client, &DefaultRequestFactory {})
+        {
+            Ok(s) => {
+                let mut borrow = self.user_auth.write();
+                *borrow = UserAuth::from_auth_refresh_response(&s);
+                if let Some(cb) = &self.on_auth_refreshed {
+                    (cb)(
+                        borrow.uid.expose_secret(),
+                        borrow.access_token.expose_secret(),
+                    );
+                }
+                Ok(())
+            }
+            Err(e) => Err(e),
+        }
+    }
+
+    async fn refresh_auth_async<C: ClientAsync>(&self, client: &C) -> http::Result<()> {
+        // Have to clone here due to async boundaries.
+        let user_auth = self.user_auth.read().clone();
+        match AuthRefreshRequest::new(
+            user_auth.uid.expose_secret(),
+            user_auth.refresh_token.expose_secret(),
+        )
+        .execute_async(client, &DefaultRequestFactory {})
+        .await
+        {
+            Ok(s) => {
+                let mut borrow = self.user_auth.write();
+                *borrow = UserAuth::from_auth_refresh_response(&s);
+                if let Some(cb) = &self.on_auth_refreshed {
+                    (cb)(
+                        borrow.uid.expose_secret(),
+                        borrow.access_token.expose_secret(),
+                    );
+                }
+                Ok(())
+            }
+            Err(e) => Err(e),
+        }
+    }
+
+    pub fn execute<C: ClientSync, R: Request>(
+        &self,
+        client: &C,
+        request: R,
+    ) -> http::Result<R::Output> {
+        match request.execute_sync(client, self) {
+            Ok(r) => Ok(r),
+            Err(original_error) => {
+                if let http::Error::API(api_err) = &original_error {
+                    if api_err.http_code == 401 {
+                        log::debug!("Account session expired, attempting refresh");
+                        // Session expired/not authorized, try auth refresh.
+                        if let Err(e) = self.refresh_auth(client) {
+                            log::error!("Failed to refresh account {e}");
+                            return Err(original_error);
+                        }
+
+                        // Execute request again
+                        return request.execute_sync(client, self);
+                    }
+                }
+                Err(original_error)
+            }
+        }
+    }
+
+    pub async fn execute_async<'a, C: ClientAsync, R: Request + 'a>(
+        &'a self,
+        client: &'a C,
+        request: R,
+    ) -> http::Result<R::Output> {
+        match request.execute_async(client, self).await {
+            Ok(r) => Ok(r),
+            Err(original_error) => {
+                if let http::Error::API(api_err) = &original_error {
+                    log::debug!("Account session expired, attempting refresh");
+                    if api_err.http_code == 401 {
+                        // Session expired/not authorized, try auth refresh.
+                        if let Err(e) = self.refresh_auth_async(client).await {
+                            log::error!("Failed to refresh account {e}");
+                            return Err(original_error);
+                        }
+
+                        // Execute request again
+                        return request.execute_async(client, self).await;
+                    }
+                }
+                Err(original_error)
+            }
+        }
+    }
+
+    pub fn get_refresh_data(&self) -> SessionRefreshData {
+        let borrow = self.user_auth.read();
+        SessionRefreshData {
+            user_uid: borrow.uid.clone(),
+            token: borrow.refresh_token.clone(),
+        }
+    }
+}
+
+impl RequestFactory for RequestRepeater {
+    fn new_request(&self, method: Method, url: &str) -> RequestData {
+        let accessor = self.user_auth.read();
+        RequestData::new(method, url)
+            .header(http::X_PM_UID_HEADER, &accessor.uid.expose_secret().0)
+            .bearer_token(accessor.access_token.expose_secret())
+    }
+}