Merge pull request djc#6 from juicebox-systems/background_refresh

simon-4242 · web-flow · commit 0240a9e53c36 · 2023-12-14T16:33:51.000-08:00
refresh auth tokens in the background
diff --git a/src/authentication_manager.rs b/src/authentication_manager.rs
@@ -1,5 +1,11 @@
+use std::collections::hash_map::Entry::{Occupied, Vacant};
+use std::collections::HashMap;
+use std::sync::Arc;
+use std::time::{Duration, SystemTime};
+
 use async_trait::async_trait;
-use tokio::sync::Mutex;
+use tokio::sync::{Mutex, OwnedMutexGuard};
+use tracing::{debug, info, warn};
 
 use crate::custom_service_account::CustomServiceAccount;
 use crate::default_authorized_user::ConfigDefaultCredentials;
@@ -13,6 +19,12 @@ pub(crate) trait ServiceAccount: Send + Sync {
     async fn project_id(&self, client: &HyperClient) -> Result<String, Error>;
     fn get_token(&self, scopes: &[&str]) -> Option<Token>;
     async fn refresh_token(&self, client: &HyperClient, scopes: &[&str]) -> Result<Token, Error>;
+    fn get_style(&self) -> TokenStyle;
+}
+
+pub(crate) enum TokenStyle {
+    Account,
+    AccountAndScopes,
 }
 
 /// Authentication manager is responsible for caching and obtaining credentials for the required
@@ -21,10 +33,13 @@ pub(crate) trait ServiceAccount: Send + Sync {
 /// Construct the authentication manager with [`AuthenticationManager::new()`] or by creating
 /// a [`CustomServiceAccount`], then converting it into an `AuthenticationManager` using the `From`
 /// impl.
-pub struct AuthenticationManager {
-    pub(crate) client: HyperClient,
-    pub(crate) service_account: Box<dyn ServiceAccount>,
-    refresh_mutex: Mutex<()>,
+#[derive(Clone)]
+pub struct AuthenticationManager(Arc<AuthManagerInner>);
+
+struct AuthManagerInner {
+    client: HyperClient,
+    service_account: Box<dyn ServiceAccount>,
+    refresh_lock: RefreshLock,
 }
 
 impl AuthenticationManager {
@@ -80,40 +95,82 @@ impl AuthenticationManager {
     }
 
     fn build(client: HyperClient, service_account: impl ServiceAccount + 'static) -> Self {
-        Self {
+        let refresh_lock = RefreshLock::new(service_account.get_style());
+        Self(Arc::new(AuthManagerInner {
             client,
             service_account: Box::new(service_account),
-            refresh_mutex: Mutex::new(()),
-        }
+            refresh_lock,
+        }))
     }
 
     /// Requests Bearer token for the provided scope
     ///
     /// Token can be used in the request authorization header in format "Bearer {token}"
     pub async fn get_token(&self, scopes: &[&str]) -> Result<Token, Error> {
-        let token = self.service_account.get_token(scopes);
+        let token = self.0.service_account.get_token(scopes);
+
         if let Some(token) = token.filter(|token| !token.has_expired()) {
+            let valid_for = token
+                .expires_at()
+                .duration_since(SystemTime::now())
+                .unwrap_or_default();
+            if valid_for < Duration::from_secs(60) {
+                debug!(?valid_for, "gcp_auth token expires soon!");
+
+                let lock = self.0.refresh_lock.lock_for_scopes(scopes).await;
+                match lock.try_lock_owned() {
+                    Err(_) => {
+                        // already being refreshed.
+                    }
+                    Ok(guard) => {
+                        let inner = self.clone();
+                        let scopes: Vec<String> = scopes.iter().map(|s| s.to_string()).collect();
+                        tokio::spawn(async move {
+                            inner.background_refresh(scopes, guard).await;
+                        });
+                    }
+                }
+            }
             return Ok(token);
         }
 
-        let _guard = self.refresh_mutex.lock().await;
+        warn!("starting inline refresh of gcp auth token");
+        let lock = self.0.refresh_lock.lock_for_scopes(scopes).await;
+        let _guard = lock.lock().await;
 
         // Check if refresh happened while we were waiting.
-        let token = self.service_account.get_token(scopes);
+        let token = self.0.service_account.get_token(scopes);
         if let Some(token) = token.filter(|token| !token.has_expired()) {
             return Ok(token);
         }
 
-        self.service_account
-            .refresh_token(&self.client, scopes)
+        self.0
+            .service_account
+            .refresh_token(&self.0.client, scopes)
             .await
     }
 
+    async fn background_refresh(&self, scopes: Vec<String>, _lock: OwnedMutexGuard<()>) {
+        info!("gcp_auth starting background refresh of auth token");
+        let scope_refs: Vec<&str> = scopes.iter().map(|s| s.as_str()).collect();
+        match self
+            .0
+            .service_account
+            .refresh_token(&self.0.client, &scope_refs)
+            .await
+        {
+            Ok(t) => {
+                info!(valid_for=?t.expires_at().duration_since(SystemTime::now()), "gcp auth completed background token refresh")
+            }
+            Err(err) => warn!(?err, "gcp_auth background token refresh failed"),
+        }
+    }
+
     /// Request the project ID for the authenticating account
     ///
     /// This is only available for service account-based authentication methods.
     pub async fn project_id(&self) -> Result<String, Error> {
-        self.service_account.project_id(&self.client).await
+        self.0.service_account.project_id(&self.0.client).await
     }
 }
 
@@ -122,3 +179,35 @@ impl From<CustomServiceAccount> for AuthenticationManager {
         Self::build(types::client(), service_account)
     }
 }
+
+enum RefreshLock {
+    One(Arc<Mutex<()>>),
+    ByScopes(Mutex<HashMap<Vec<String>, Arc<Mutex<()>>>>),
+}
+
+impl RefreshLock {
+    fn new(style: TokenStyle) -> Self {
+        match style {
+            TokenStyle::Account => RefreshLock::One(Arc::new(Mutex::new(()))),
+            TokenStyle::AccountAndScopes => RefreshLock::ByScopes(Mutex::new(HashMap::new())),
+        }
+    }
+
+    async fn lock_for_scopes(&self, scopes: &[&str]) -> Arc<Mutex<()>> {
+        match self {
+            RefreshLock::One(mutex) => mutex.clone(),
+            RefreshLock::ByScopes(mutexes) => {
+                let scopes_key: Vec<_> = scopes.iter().map(|s| s.to_string()).collect();
+                let mut scope_locks = mutexes.lock().await;
+                match scope_locks.entry(scopes_key) {
+                    Occupied(e) => e.get().clone(),
+                    Vacant(v) => {
+                        let lock = Arc::new(Mutex::new(()));
+                        v.insert(lock.clone());
+                        lock
+                    }
+                }
+            }
+        }
+    }
+}
diff --git a/src/custom_service_account.rs b/src/custom_service_account.rs
@@ -6,7 +6,7 @@ use std::sync::RwLock;
 use async_trait::async_trait;
 use serde::{Deserialize, Serialize};
 
-use crate::authentication_manager::ServiceAccount;
+use crate::authentication_manager::{ServiceAccount, TokenStyle};
 use crate::error::Error;
 use crate::types::{HyperClient, SecretString, Signer, Token};
 use crate::util::HyperExt;
@@ -80,6 +80,10 @@ impl CustomServiceAccount {
 
 #[async_trait]
 impl ServiceAccount for CustomServiceAccount {
+    fn get_style(&self) -> TokenStyle {
+        TokenStyle::AccountAndScopes
+    }
+
     async fn project_id(&self, _: &HyperClient) -> Result<String, Error> {
         match &self.credentials.project_id {
             Some(pid) => Ok(pid.clone()),
diff --git a/src/default_authorized_user.rs b/src/default_authorized_user.rs
@@ -7,7 +7,7 @@ use hyper::body::Body;
 use hyper::{Method, Request};
 use serde::{Deserialize, Serialize};
 
-use crate::authentication_manager::ServiceAccount;
+use crate::authentication_manager::{ServiceAccount, TokenStyle};
 use crate::error::Error;
 use crate::types::{HyperClient, SecretString, Token};
 use crate::util::HyperExt;
@@ -77,6 +77,10 @@ impl ConfigDefaultCredentials {
 
 #[async_trait]
 impl ServiceAccount for ConfigDefaultCredentials {
+    fn get_style(&self) -> TokenStyle {
+        TokenStyle::Account
+    }
+
     async fn project_id(&self, _: &HyperClient) -> Result<String, Error> {
         self.credentials
             .quota_project_id
diff --git a/src/default_service_account.rs b/src/default_service_account.rs
@@ -5,7 +5,7 @@ use async_trait::async_trait;
 use hyper::body::Body;
 use hyper::{Method, Request};
 
-use crate::authentication_manager::ServiceAccount;
+use crate::authentication_manager::{ServiceAccount, TokenStyle};
 use crate::error::Error;
 use crate::types::{HyperClient, Token};
 use crate::util::HyperExt;
@@ -62,6 +62,10 @@ impl MetadataServiceAccount {
 
 #[async_trait]
 impl ServiceAccount for MetadataServiceAccount {
+    fn get_style(&self) -> TokenStyle {
+        TokenStyle::Account
+    }
+
     async fn project_id(&self, client: &HyperClient) -> Result<String, Error> {
         tracing::debug!("Getting project ID from GCP instance metadata server");
         let req = Self::build_token_request(Self::DEFAULT_PROJECT_ID_GCP_URI);
diff --git a/src/gcloud_authorized_user.rs b/src/gcloud_authorized_user.rs
@@ -5,7 +5,7 @@ use std::sync::RwLock;
 use async_trait::async_trait;
 use std::time::Duration;
 
-use crate::authentication_manager::ServiceAccount;
+use crate::authentication_manager::{ServiceAccount, TokenStyle};
 use crate::error::Error;
 use crate::error::Error::{GCloudError, GCloudParseError};
 use crate::types::{HyperClient, SecretString};
@@ -45,6 +45,10 @@ impl GCloudAuthorizedUser {
 
 #[async_trait]
 impl ServiceAccount for GCloudAuthorizedUser {
+    fn get_style(&self) -> TokenStyle {
+        TokenStyle::Account
+    }
+
     async fn project_id(&self, _: &HyperClient) -> Result<String, Error> {
         self.project_id.clone().ok_or(Error::NoProjectId)
     }
