Skip to main content

colab_cli/auth/
mod.rs

1pub mod oauth;
2pub mod storage;
3
4use std::sync::{Mutex, OnceLock};
5
6use chrono::Utc;
7
8use crate::auth::storage::StoredAccessToken;
9use crate::config::ColabConfig;
10use crate::error::{ColabError, Result};
11
12pub use storage::{AccountInfo, TokenStorage};
13
14const REFRESH_MARGIN_SECS: i64 = 5 * 60;
15
16// avoid re-reading credentials.json on every API call
17static TOKEN_CACHE: OnceLock<Mutex<Option<StoredAccessToken>>> = OnceLock::new();
18
19fn token_cache() -> &'static Mutex<Option<StoredAccessToken>> {
20    TOKEN_CACHE.get_or_init(|| Mutex::new(None))
21}
22
23pub async fn get_access_token(config: &ColabConfig) -> Result<String> {
24    // cache
25    {
26        let guard = token_cache().lock().expect("token cache poisoned");
27        if let Some(stored) = guard.as_ref() {
28            let remaining = stored.expires_at - Utc::now();
29            if remaining.num_seconds() > REFRESH_MARGIN_SECS {
30                return Ok(stored.access_token.clone());
31            }
32        }
33    }
34
35    // disk (another colab process might have refreshed it)
36    if let Some(stored) = TokenStorage::get_access_token()? {
37        let remaining = stored.expires_at - Utc::now();
38        if remaining.num_seconds() > REFRESH_MARGIN_SECS {
39            let token = stored.access_token.clone();
40            *token_cache().lock().expect("token cache poisoned") = Some(stored);
41            return Ok(token);
42        }
43    }
44
45    // network
46    if TokenStorage::get_refresh_token()?.is_none() {
47        return Err(ColabError::NotAuthenticated);
48    }
49    let token = oauth::refresh_access_token(config).await?;
50    if let Some(stored) = TokenStorage::get_access_token()? {
51        *token_cache().lock().expect("token cache poisoned") = Some(stored);
52    }
53    Ok(token)
54}
55
56pub fn invalidate_token_cache() {
57    if let Some(c) = TOKEN_CACHE.get()
58        && let Ok(mut guard) = c.lock()
59    {
60        *guard = None;
61    }
62}
63
64pub async fn login(config: &ColabConfig) -> Result<AccountInfo> {
65    let account = oauth::run_login_flow(config).await?;
66    // login wrote a new token; flush stale cache
67    invalidate_token_cache();
68    Ok(account)
69}
70
71pub fn logout() -> Result<()> {
72    let result = TokenStorage::clear_all();
73    invalidate_token_cache();
74    result
75}
76
77pub fn current_account() -> Result<Option<AccountInfo>> {
78    TokenStorage::get_account()
79}