fix: race condition on access_token

src/lib.rs

@@ -4,4 +4,5 @@ pub mod error;
 pub mod integrations;
 pub mod services;
 pub mod state;
+pub mod utils;
 pub mod web;

src/services/discord.rs

@@ -11,6 +11,7 @@ use crate::{
     integrations::discord::{
         DiscordClient, DiscordUserModel, GuildMemberModel, PartialGuildModel, Snowflake,
     },
+    utils::tokens::RefreshLock,
     web::WebError,
 };
 
@@ -48,11 +49,17 @@ pub struct DiscordService {
     db: MantleDb,
     client: DiscordClient,
     config: SharedConfig,
+    refresh_lock: RefreshLock,
 }
 
 impl DiscordService {
     pub fn new(db: MantleDb, client: DiscordClient, config: SharedConfig) -> Self {
-        Self { db, client, config }
+        Self {
+            db,
+            client,
+            config,
+            refresh_lock: Default::default(),
+        }
     }
 
     pub async fn get_current_user(
@@ -212,19 +219,25 @@ impl DiscordService {
     }
 
     async fn access_token(&self, identity: &Identity) -> Result<String> {
-        // TODO: Current implementation can allow making refresh races
-        // where hundreds of requests can try to get access_token and see expired one ->
-        // they all will try to update it and may end up being rate limited.
-        // This stuff need some per-identity lock or smth
-        let tokens = db::oauth_tokens::find_by_id(&self.db, identity.id)
-            .await
-            .map_err(InternalError::from)?;
+        loop {
+            let tokens = db::oauth_tokens::find_by_id(&self.db, identity.id)
+                .await
+                .map_err(InternalError::from)?
+                .ok_or(DiscordError::IdentityHasNoTokens(identity.id))?;
 
-        if let Some(tokens) = tokens {
             if tokens.expires_at > Utc::now() + Duration::minutes(1) {
                 return Ok(tokens.access_token);
             }
 
+            // try acquire a lock before refresh
+            let guard = match self.refresh_lock.try_lock(identity) {
+                Ok(guard) => guard,
+                Err(notify) => {
+                    notify.notified().await;
+                    continue;
+                }
+            };
+
             let credentials = self.client_credentials();
             let new_tokens = self
                 .client
@@ -243,9 +256,11 @@ impl DiscordService {
             .await
             .map_err(InternalError::from)?;
 
+            // we will explicitly drop it here so rustc won't do its magic
+            // (probably it won't do it anyway because of Drop impl, but still)
+            drop(guard);
+
             return Ok(new_tokens.access_token);
         }
-
-        Err(DiscordError::IdentityHasNoTokens(identity.id))
     }
 }

src/utils/mod.rs

@@ -0,0 +1 @@
+pub mod tokens;

src/utils/tokens.rs

@@ -0,0 +1,49 @@
+use std::{
+    collections::{HashMap, hash_map::Entry},
+    sync::{Arc, Mutex},
+};
+
+use tokio::sync::Notify;
+use uuid::Uuid;
+
+use crate::db::identities::Identity;
+
+// TODO: it should be tested properly, but currently im too lazy to do this
+#[derive(Debug, Default, Clone)]
+pub struct RefreshLock {
+    refreshes: Arc<Mutex<HashMap<Uuid, Arc<Notify>>>>,
+}
+
+pub struct RefreshGuard {
+    identity_id: Uuid,
+    refreshes: Arc<Mutex<HashMap<Uuid, Arc<Notify>>>>,
+}
+
+impl Drop for RefreshGuard {
+    fn drop(&mut self) {
+        let identity_id = self.identity_id;
+        let refreshes = self.refreshes.clone();
+
+        let mut refreshes = refreshes.lock().unwrap();
+        if let Some(current) = refreshes.remove(&identity_id) {
+            current.notify_waiters();
+        }
+    }
+}
+
+impl RefreshLock {
+    pub fn try_lock(&self, identity: &Identity) -> Result<RefreshGuard, Arc<Notify>> {
+        let mut refreshes = self.refreshes.lock().unwrap();
+
+        match refreshes.entry(identity.id) {
+            Entry::Occupied(entry) => Err(entry.get().clone()),
+            Entry::Vacant(entry) => {
+                entry.insert(Arc::new(Notify::new()));
+                Ok(RefreshGuard {
+                    identity_id: identity.id,
+                    refreshes: self.refreshes.clone(),
+                })
+            }
+        }
+    }
+}