DynamicAuthenticationProvider supports refreshing authentication credentials (#821)

* `DynamicAuthenticationProvider` supports refreshing authentication credentials

* Fix a compiling issue in the kotlin wrapper

* Update API doc

Co-authored-by: Oguz Kocer <[email protected]>

---------

Co-authored-by: Oguz Kocer <[email protected]>

Author: crazytonyli

native/kotlin/api/kotlin/src/integrationTest/kotlin/AuthProviderTest.kt

@@ -1,5 +1,8 @@
 package rs.wordpress.api.kotlin
 
+import kotlinx.coroutines.CoroutineDispatcher
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.withContext
 import kotlinx.coroutines.test.runTest
 import org.junit.jupiter.api.Test
 import uniffi.wp_api.ModifiableAuthenticationProvider
@@ -31,7 +34,10 @@ class AuthProviderTest {
 
     @Test
     fun testDynamicAuthProvider() = runTest {
-        class DynamicAuthProvider(var isAuthorized: Boolean = false): WpDynamicAuthenticationProvider {
+        class DynamicAuthProvider(
+            var isAuthorized: Boolean = false,
+            private val dispatcher: CoroutineDispatcher = Dispatchers.IO
+        ): WpDynamicAuthenticationProvider {
             override fun auth(): WpAuthentication =
                 if (isAuthorized) {
                     wpAuthenticationFromUsernameAndPassword(
@@ -41,6 +47,7 @@ class AuthProviderTest {
                 } else {
                     WpAuthentication.None
                 }
+            override suspend fun refresh(): Boolean = withContext(dispatcher) { false }
         }
 
         val dynamicAuthProvider = DynamicAuthProvider()

wp_api/src/auth.rs

@@ -1,4 +1,5 @@
 use http::HeaderValue;
+use std::fmt::Debug;
 use std::sync::{Arc, RwLock};
 
 #[derive(Debug, Clone, uniffi::Enum)]
@@ -70,8 +71,17 @@ impl ModifiableAuthenticationProvider {
 }
 
 #[uniffi::export(with_foreign)]
-pub trait WpDynamicAuthenticationProvider: Send + Sync {
+#[async_trait::async_trait]
+pub trait WpDynamicAuthenticationProvider: Send + Sync + Debug {
     fn auth(&self) -> WpAuthentication;
+
+    /// Refresh the authentication token. The implementation should only return true
+    /// if the authentication was successfully refreshed.
+    ///
+    /// **Concurrency:** This method may be called concurrently by multiple request
+    /// executors. Implementations must handle concurrent calls safely and avoid
+    /// unnecessary duplicate refresh operations.
+    async fn refresh(&self) -> bool;
 }
 
 #[derive(uniffi::Object)]
@@ -135,4 +145,14 @@ impl WpAuthenticationProvider {
             WpAuthenticationProvider::Modifiable { inner } => inner.header_value(),
         }
     }
+
+    pub(crate) async fn refresh(&self) -> bool {
+        match self {
+            WpAuthenticationProvider::StaticAuthenticationProvider { .. } => false,
+            WpAuthenticationProvider::DynamicAuthenticationProvider { inner } => {
+                inner.refresh().await
+            }
+            WpAuthenticationProvider::Modifiable { .. } => false,
+        }
+    }
 }

wp_api_integration_tests/tests/test_auth_provider_immut.rs

@@ -1,4 +1,4 @@
-use std::sync::atomic::{AtomicBool, Ordering};
+use std::sync::RwLock;
 use wp_api::auth::{ModifiableAuthenticationProvider, WpDynamicAuthenticationProvider};
 use wp_api_integration_tests::prelude::*;
 
@@ -20,27 +20,79 @@ async fn test_static_auth_provider() {
     assert_eq!(FIRST_USER_ID, user.id);
 }
 
-#[tokio::test]
-#[parallel]
-async fn test_dynamic_auth_provider() {
-    #[derive(Default)]
-    struct DynamicAuthProvider {
-        is_authorized: AtomicBool,
+#[derive(Debug)]
+struct DynamicAuthProvider {
+    auth: RwLock<WpAuthentication>,
+    refreshed_auth: RwLock<Option<WpAuthentication>>,
+}
+
+impl DynamicAuthProvider {
+    fn new() -> Self {
+        Self {
+            auth: RwLock::new(WpAuthentication::None),
+            refreshed_auth: RwLock::new(None),
+        }
+    }
+
+    fn authenticate_with(&self, auth: WpAuthentication) {
+        *(self.auth.write().unwrap()) = auth;
+    }
+
+    fn allow_refresh_auth(&self, auth: WpAuthentication) {
+        *(self.refreshed_auth.write().unwrap()) = Some(auth);
+    }
+}
+
+#[async_trait::async_trait]
+impl WpDynamicAuthenticationProvider for DynamicAuthProvider {
+    fn auth(&self) -> WpAuthentication {
+        self.auth.read().unwrap().clone()
     }
 
-    impl WpDynamicAuthenticationProvider for DynamicAuthProvider {
-        fn auth(&self) -> WpAuthentication {
-            if self.is_authorized.load(Ordering::Relaxed) {
-                WpAuthentication::from_username_and_password(
-                    TestCredentials::instance().admin_username.to_string(),
-                    TestCredentials::instance().admin_password.to_string(),
-                )
-            } else {
-                WpAuthentication::None
-            }
+    async fn refresh(&self) -> bool {
+        if let Some(ref auth) = *(self.refreshed_auth.read().unwrap()) {
+            *(self.auth.write().unwrap()) = auth.clone();
+            return true;
+        } else {
+            return false;
         }
     }
-    let dynamic_auth_provider = Arc::new(DynamicAuthProvider::default());
+}
+
+#[tokio::test]
+#[parallel]
+async fn test_dynamic_auth_provider() {
+    let dynamic_auth_provider = Arc::new(DynamicAuthProvider::new());
+    let client: WpApiClient = api_client_with_auth_provider(Arc::new(
+        WpAuthenticationProvider::dynamic(dynamic_auth_provider.clone()),
+    ));
+
+    // Assert that initial unauthorized request fails
+    client
+        .users()
+        .retrieve_me_with_edit_context()
+        .await
+        .assert_wp_error(WpErrorCode::Unauthorized);
+
+    // Assert that request succeeds after providing a valid authentication
+    dynamic_auth_provider.authenticate_with(WpAuthentication::from_username_and_password(
+        TestCredentials::instance().admin_username.to_string(),
+        TestCredentials::instance().admin_password.to_string(),
+    ));
+    let user = client
+        .users()
+        .retrieve_me_with_edit_context()
+        .await
+        .assert_response()
+        .data;
+    // FIRST_USER_ID is the current user's id
+    assert_eq!(FIRST_USER_ID, user.id);
+}
+
+#[tokio::test]
+#[parallel]
+async fn test_refresh_dynamic_auth_provider() {
+    let dynamic_auth_provider = Arc::new(DynamicAuthProvider::new());
     let client: WpApiClient = api_client_with_auth_provider(Arc::new(
         WpAuthenticationProvider::dynamic(dynamic_auth_provider.clone()),
     ));
@@ -52,10 +104,13 @@ async fn test_dynamic_auth_provider() {
         .await
         .assert_wp_error(WpErrorCode::Unauthorized);
 
-    // Assert that request succeeds after setting `is_authorized = true`
-    dynamic_auth_provider
-        .is_authorized
-        .store(true, Ordering::Relaxed);
+    // Set the refreshed authentication
+    dynamic_auth_provider.allow_refresh_auth(WpAuthentication::from_username_and_password(
+        TestCredentials::instance().admin_username.to_string(),
+        TestCredentials::instance().admin_password.to_string(),
+    ));
+
+    // Assert that request succeeds after refresh
     let user = client
         .users()
         .retrieve_me_with_edit_context()

wp_derive_request_builder/src/generate.rs

@@ -74,13 +74,33 @@ fn generate_async_request_executor(
                 pub async #fn_signature -> Result<#response_type_ident, #error_type> {
                     use #crate_ident::api_error::MaybeWpError;
                     use #crate_ident::middleware::PerformsRequests;
-                    #request_from_request_builder
-                    let response = self.perform(std::sync::Arc::new(request)).await?;
-                    let response_status_code = response.status_code;
-                    let parsed_response = response.parse();
-                    if parsed_response.is_unauthorized_error().unwrap_or_default() || (response_status_code == 401 && self.fetch_authentication_state().await.map(|auth_state| auth_state.is_unauthorized()).unwrap_or_default()) {
+                    let perform_request = async || {
+                        #request_from_request_builder
+                        let response = self.perform(std::sync::Arc::new(request)).await?;
+                        let response_status_code = response.status_code;
+                        let parsed_response = response.parse();
+                        let unauthorized = parsed_response.is_unauthorized_error().unwrap_or_default() || (response_status_code == 401 && self.fetch_authentication_state().await.map(|auth_state| auth_state.is_unauthorized()).unwrap_or_default());
+
+                        Ok((parsed_response, unauthorized))
+                    };
+
+                    let mut parsed_response;
+                    let mut unauthorized;
+
+                    // The Rust compiler has trouble figuring out the return type. The statement below helps it.
+                    let result: Result<_, #error_type> = perform_request().await;
+                    (parsed_response, unauthorized) = result?;
+
+                    // If the auth provider successfully refreshed the authentication, we can retry the request.
+                    if unauthorized && self.delegate.auth_provider.refresh().await {
+                        // The response of the retried request will be returned instead of the original one.
+                        (parsed_response, unauthorized) = perform_request().await?;
+                    }
+
+                    if unauthorized {
                         self.delegate.app_notifier.requested_with_invalid_authentication().await;
                     }
+
                     parsed_response
                }
             }