aether_auth/
credential.rs

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

aether_auth/
credential.rs
