From fd22e1d472d53c406d2e8c6b5b2dc00e20a9625a Mon Sep 17 00:00:00 2001
From: Zack Fu Zi Xiang <speed2exe@live.com.sg>
Date: Mon, 2 Sep 2024 21:13:33 +0800
Subject: [PATCH] chore: preserve provider token

---
 libs/client-api/src/http.rs         | 14 ++++++++++++--
 libs/client-api/src/native/retry.rs | 26 +++++++++++++++++---------
 2 files changed, 29 insertions(+), 11 deletions(-)

diff --git a/libs/client-api/src/http.rs b/libs/client-api/src/http.rs
index acd532ff..74981c38 100644
--- a/libs/client-api/src/http.rs
+++ b/libs/client-api/src/http.rs
@@ -273,13 +273,18 @@ impl Client {
       .split('&');
 
     let mut refresh_token: Option<&str> = None;
+    let mut provider_token: Option<String> = None;
+    let mut provider_refresh_token: Option<String> = None;
     for param in key_value_pairs {
       match param.split_once('=') {
         Some(pair) => {
           let (k, v) = pair;
           if k == "refresh_token" {
             refresh_token = Some(v);
-            break;
+          } else if k == "provider_token" {
+            provider_token = Some(v.to_string());
+          } else if k == "provider_refresh_token" {
+            provider_refresh_token = Some(v.to_string());
           }
         },
         None => warn!("param is not in key=value format: {}", param),
@@ -287,13 +292,18 @@ impl Client {
     }
     let refresh_token = refresh_token.ok_or(url_missing_param("refresh_token"))?;
 
-    let new_token = self
+    let mut new_token = self
       .gotrue_client
       .token(&Grant::RefreshToken(RefreshTokenGrant {
         refresh_token: refresh_token.to_owned(),
       }))
       .await?;
 
+    // refresh endpoint does not return provider token
+    // so we need to set it manually to preserve this information
+    new_token.provider_access_token = provider_token;
+    new_token.provider_refresh_token = provider_refresh_token;
+
     let (_user, new) = self.verify_token(&new_token.access_token).await?;
     self.token.write().set(new_token);
     Ok(new)
diff --git a/libs/client-api/src/native/retry.rs b/libs/client-api/src/native/retry.rs
index 3454de4f..02078e7c 100644
--- a/libs/client-api/src/native/retry.rs
+++ b/libs/client-api/src/native/retry.rs
@@ -47,18 +47,26 @@ impl Action for RefreshTokenAction {
       if let (Some(token), Some(gotrue_client)) =
         (weak_token.upgrade(), weak_gotrue_client.upgrade())
       {
-        let refresh_token = token
-          .read()
-          .as_ref()
-          .ok_or(GoTrueError::NotLoggedIn(
+        let (refresh_token, provider_access_token, provider_refresh_token) = {
+          let mut token_write = token.write();
+          let gotrue_resp_token = token_write.as_mut().ok_or(GoTrueError::NotLoggedIn(
             "fail to refresh user token".to_owned(),
-          ))?
-          .refresh_token
-          .as_str()
-          .to_owned();
-        let access_token_resp = gotrue_client
+          ))?;
+          let refresh_token = gotrue_resp_token.refresh_token.as_str().to_owned();
+          let provider_access_token = gotrue_resp_token.provider_access_token.take();
+          let provider_refresh_token = gotrue_resp_token.provider_refresh_token.take();
+          (refresh_token, provider_access_token, provider_refresh_token)
+        };
+
+        let mut access_token_resp = gotrue_client
           .token(&Grant::RefreshToken(RefreshTokenGrant { refresh_token }))
           .await?;
+
+        // refresh does not preserve provider token and refresh token
+        // so we need to set it manually to preserve this information
+        access_token_resp.provider_access_token = provider_access_token;
+        access_token_resp.provider_refresh_token = provider_refresh_token;
+
         token.write().set(access_token_resp);
       }
       Ok(())