From 7db45a4b23329183f0979aee66dcdcc3ecb5bf99 Mon Sep 17 00:00:00 2001
From: Benjamin Bouvier <public@benj.me>
Date: Tue, 5 Sep 2023 16:08:39 +0200
Subject: [PATCH] chore: move the `refresh_token_lock` into the `AuthCtx`

---
 crates/matrix-sdk/src/authentication.rs  | 9 ++++++++-
 crates/matrix-sdk/src/client/builder.rs  | 2 ++
 crates/matrix-sdk/src/client/mod.rs      | 3 ---
 crates/matrix-sdk/src/matrix_auth/mod.rs | 4 ++--
 crates/matrix-sdk/src/oidc/mod.rs        | 4 ++--
 5 files changed, 14 insertions(+), 8 deletions(-)

diff --git a/crates/matrix-sdk/src/authentication.rs b/crates/matrix-sdk/src/authentication.rs
index 511ccf0d4..44d9da652 100644
--- a/crates/matrix-sdk/src/authentication.rs
+++ b/crates/matrix-sdk/src/authentication.rs
@@ -14,10 +14,14 @@
 
 use matrix_sdk_base::SessionMeta;
 use ruma::api::client::discovery::discover_homeserver::AuthenticationServerInfo;
+use tokio::sync::Mutex;
 
-use crate::matrix_auth::{self, MatrixAuth, MatrixAuthData};
 #[cfg(feature = "experimental-oidc")]
 use crate::oidc::{self, Oidc, OidcAuthData};
+use crate::{
+    matrix_auth::{self, MatrixAuth, MatrixAuthData},
+    RefreshTokenError,
+};
 
 /// All the data relative to authentication, and that must be shared between a client and all its
 /// children.
@@ -29,6 +33,9 @@ pub(crate) struct AuthCtx {
     /// Whether to try to refresh the access token automatically when an
     /// `M_UNKNOWN_TOKEN` error is encountered.
     pub(crate) handle_refresh_tokens: bool,
+
+    /// Lock making sure we're only doing one token refresh at a time.
+    pub(crate) refresh_token_lock: Mutex<Result<(), RefreshTokenError>>,
 }
 
 /// An enum over all the possible authentication APIs.
diff --git a/crates/matrix-sdk/src/client/builder.rs b/crates/matrix-sdk/src/client/builder.rs
index 314ebabc6..855c777a0 100644
--- a/crates/matrix-sdk/src/client/builder.rs
+++ b/crates/matrix-sdk/src/client/builder.rs
@@ -21,6 +21,7 @@ use ruma::{
     OwnedServerName, ServerName,
 };
 use thiserror::Error;
+use tokio::sync::Mutex;
 use tracing::{debug, field::debug, instrument, Span};
 use url::Url;
 
@@ -425,6 +426,7 @@ impl ClientBuilder {
         let auth_ctx = Arc::new(AuthCtx {
             authentication_server_info,
             handle_refresh_tokens: self.handle_refresh_tokens,
+            refresh_token_lock: Mutex::new(Ok(())),
         });
 
         let inner = Arc::new(ClientInner::new(
diff --git a/crates/matrix-sdk/src/client/mod.rs b/crates/matrix-sdk/src/client/mod.rs
index adb4d7266..8a1d2568d 100644
--- a/crates/matrix-sdk/src/client/mod.rs
+++ b/crates/matrix-sdk/src/client/mod.rs
@@ -186,8 +186,6 @@ pub(crate) struct ClientInner {
     /// Whether the client should update its homeserver URL with the discovery
     /// information present in the login response.
     respect_login_well_known: bool,
-    /// Lock making sure we're only doing one token refresh at a time.
-    pub(crate) refresh_token_lock: Mutex<Result<(), RefreshTokenError>>,
     /// An event that can be listened on to wait for a successful sync. The
     /// event will only be fired if a sync loop is running. Can be used for
     /// synchronization, e.g. if we send out a request to create a room, we can
@@ -259,7 +257,6 @@ impl ClientInner {
             sync_gap_broadcast_txs: Default::default(),
             respect_login_well_known,
             sync_beat: event_listener::Event::new(),
-            refresh_token_lock: Mutex::new(Ok(())),
             session_change_sender,
             auth_data: Default::default(),
             #[cfg(feature = "e2e-encryption")]
diff --git a/crates/matrix-sdk/src/matrix_auth/mod.rs b/crates/matrix-sdk/src/matrix_auth/mod.rs
index f72626df7..ebfb7cf2f 100644
--- a/crates/matrix-sdk/src/matrix_auth/mod.rs
+++ b/crates/matrix-sdk/src/matrix_auth/mod.rs
@@ -450,7 +450,7 @@ impl MatrixAuth {
         &self,
     ) -> Result<Option<refresh_token::v3::Response>, RefreshTokenError> {
         let client = &self.client;
-        let lock = client.inner.refresh_token_lock.try_lock();
+        let lock = client.inner.auth_ctx.refresh_token_lock.try_lock();
 
         if let Ok(mut guard) = lock {
             let Some(mut session_tokens) = self.session_tokens() else {
@@ -489,7 +489,7 @@ impl MatrixAuth {
                 }
             }
         } else {
-            match client.inner.refresh_token_lock.lock().await.as_ref() {
+            match client.inner.auth_ctx.refresh_token_lock.lock().await.as_ref() {
                 Ok(_) => Ok(None),
                 Err(error) => Err(error.clone()),
             }
diff --git a/crates/matrix-sdk/src/oidc/mod.rs b/crates/matrix-sdk/src/oidc/mod.rs
index 1a7a27d15..dfef6f148 100644
--- a/crates/matrix-sdk/src/oidc/mod.rs
+++ b/crates/matrix-sdk/src/oidc/mod.rs
@@ -926,7 +926,7 @@ impl Oidc {
         &self,
     ) -> Result<Option<AccessTokenResponse>, RefreshTokenError> {
         let client = &self.client;
-        let lock = client.inner.refresh_token_lock.try_lock();
+        let lock = client.inner.auth_ctx.refresh_token_lock.try_lock();
 
         macro_rules! fail {
             ($lock:expr, $error:expr) => {
@@ -957,7 +957,7 @@ impl Oidc {
                 }
             }
         } else {
-            match client.inner.refresh_token_lock.lock().await.as_ref() {
+            match client.inner.auth_ctx.refresh_token_lock.lock().await.as_ref() {
                 Ok(_) => Ok(None),
                 Err(error) => Err(error.clone()),
             }