aether_auth/
credential.rs

1use async_trait::async_trait;
2use oauth2::basic::BasicClient;
3use oauth2::reqwest::redirect::Policy;
4use oauth2::{ClientId, RefreshToken, TokenResponse, TokenUrl};
5use serde::{Deserialize, Serialize};
6use serde_json::Value;
7use std::time::Duration;
8
9use crate::OAuthError;
10
11const TOKEN_EXPIRY_GRACE_PERIOD: Duration = Duration::from_mins(1);
12
13/// Credential for a non-MCP OAuth provider.
14#[derive(Debug, Clone, Serialize, Deserialize)]
15pub struct OAuthCredential {
16    pub client_id: String,
17    pub access_token: String,
18    pub refresh_token: Option<String>,
19    /// Unix timestamp in milliseconds when the token expires.
20    pub expires_at: Option<u64>,
21}
22
23impl OAuthCredential {
24    /// Build an `OAuthCredential` from an `OAuth2` token response.
25    pub fn from_token_response<T: TokenResponse>(client_id: String, token_response: &T) -> Self {
26        Self {
27            client_id,
28            access_token: token_response.access_token().secret().clone(),
29            refresh_token: token_response.refresh_token().map(|token| token.secret().clone()),
30            expires_at: expires_at_from_duration(token_response.expires_in()),
31        }
32    }
33
34    /// Whether the access token is expired or expiring within the refresh skew.
35    pub fn needs_refresh(&self) -> bool {
36        self.expires_at.is_some_and(|at| {
37            current_unix_time_millis() >= at.saturating_sub(duration_millis(TOKEN_EXPIRY_GRACE_PERIOD))
38        })
39    }
40
41    /// Time remaining before the access token expires, if known and still in the future.
42    pub fn expires_in(&self) -> Option<Duration> {
43        self.expires_at.and_then(|expires_at| {
44            let now = current_unix_time_millis();
45            (expires_at > now).then(|| Duration::from_millis(expires_at - now))
46        })
47    }
48
49    /// Exchange the refresh token for a new access token.
50    ///
51    /// Preserves the existing refresh token if the response doesn't include a rotated one.
52    /// Returns `NoCredentials` if the credential has no refresh token to exchange.
53    pub async fn refresh(self, token_url: &TokenUrl) -> Result<Self, OAuthError> {
54        let old_refresh_token = self.refresh_token.clone().ok_or_else(|| {
55            OAuthError::NoCredentials(
56                "OAuth credential expired and no refresh token is available. Re-run OAuth login.".to_string(),
57            )
58        })?;
59
60        let oauth_client = BasicClient::new(ClientId::new(self.client_id.clone())).set_token_uri(token_url.clone());
61        let http_client = oauth_http_client()?;
62        let token_response = oauth_client
63            .exchange_refresh_token(&RefreshToken::new(old_refresh_token.clone()))
64            .request_async(&http_client)
65            .await
66            .map_err(|e| OAuthError::TokenExchange(e.to_string()))?;
67
68        let mut refreshed = Self::from_token_response(self.client_id, &token_response);
69        if refreshed.refresh_token.is_none() {
70            refreshed.refresh_token = Some(old_refresh_token);
71        }
72        Ok(refreshed)
73    }
74}
75
76/// Storage for namespaced, opaque OAuth JSON values.
77///
78/// Implementations include [`OsKeyringStore`](crate::OsKeyringStore) (OS keychain, feature `keyring`)
79/// and the in-memory [`FakeOAuthCredentialStore`](crate::FakeOAuthCredentialStore) for tests.
80#[async_trait]
81pub trait OAuthCredentialStorage: Send + Sync {
82    async fn load(&self, key: &str) -> Result<Option<Value>, OAuthError>;
83
84    async fn save(&self, key: &str, value: Value) -> Result<(), OAuthError>;
85
86    async fn delete(&self, key: &str) -> Result<(), OAuthError>;
87
88    fn contains(&self, key: &str) -> bool;
89
90    async fn load_credential(&self, key: &str) -> Result<Option<OAuthCredential>, OAuthError> {
91        self.load(key)
92            .await?
93            .map(serde_json::from_value)
94            .transpose()
95            .map_err(|error| OAuthError::CredentialStore(format!("invalid credential: {error}")))
96    }
97
98    async fn save_credential(&self, key: &str, credential: OAuthCredential) -> Result<(), OAuthError> {
99        let value = serde_json::to_value(credential)
100            .map_err(|error| OAuthError::CredentialStore(format!("failed to serialize credential: {error}")))?;
101        self.save(key, value).await
102    }
103}
104
105fn expires_at_from_duration(duration: Option<Duration>) -> Option<u64> {
106    duration.map(|duration| current_unix_time_millis().saturating_add(duration_millis(duration)))
107}
108
109pub fn oauth_http_client() -> Result<oauth2::reqwest::Client, OAuthError> {
110    oauth2::reqwest::Client::builder()
111        .redirect(Policy::none())
112        .build()
113        .map_err(|e| OAuthError::TokenExchange(format!("failed to build HTTP client: {e}")))
114}
115
116fn current_unix_time_millis() -> u64 {
117    u64::try_from(std::time::SystemTime::now().duration_since(std::time::UNIX_EPOCH).unwrap_or_default().as_millis())
118        .unwrap_or(u64::MAX)
119}
120
121fn duration_millis(duration: Duration) -> u64 {
122    u64::try_from(duration.as_millis()).unwrap_or(u64::MAX)
123}
124
125#[cfg(test)]
126mod tests {
127    use super::*;
128
129    #[test]
130    fn needs_refresh_is_false_when_no_expiry() {
131        assert!(!build_credential(None).needs_refresh());
132    }
133
134    #[test]
135    fn needs_refresh_is_false_when_far_in_future() {
136        assert!(!build_credential(Some(u64::MAX)).needs_refresh());
137    }
138
139    #[test]
140    fn needs_refresh_is_true_when_past() {
141        assert!(build_credential(Some(0)).needs_refresh());
142    }
143
144    #[test]
145    fn needs_refresh_is_true_when_within_skew() {
146        let cred = build_credential(expires_at_from_duration(Some(Duration::from_millis(59_999))));
147        assert!(cred.needs_refresh());
148    }
149
150    #[test]
151    fn expires_in_is_none_when_no_expiry() {
152        assert!(build_credential(None).expires_in().is_none());
153    }
154
155    #[test]
156    fn expires_in_is_none_when_already_past() {
157        assert!(build_credential(Some(0)).expires_in().is_none());
158    }
159
160    #[test]
161    fn expires_in_returns_remaining_duration_when_future() {
162        let cred = build_credential(expires_at_from_duration(Some(Duration::from_hours(1))));
163        let remaining = cred.expires_in().expect("expires_in should be Some for future expiry");
164        assert!(remaining > Duration::from_mins(58));
165        assert!(remaining <= Duration::from_hours(1));
166    }
167
168    fn build_credential(expires_at: Option<u64>) -> OAuthCredential {
169        OAuthCredential {
170            client_id: "client".to_string(),
171            access_token: "access".to_string(),
172            refresh_token: None,
173            expires_at,
174        }
175    }
176}
aether_auth/credential.rs

aether_auth/
credential.rs
