claude_codex/providers/grok/auth/
manager.rs

1use std::sync::Arc;
2use std::time::{Duration, SystemTime, UNIX_EPOCH};
3
4use serde::Deserialize;
5use tokio::sync::Mutex;
6use url::Url;
7
8use super::login::{CANONICAL_ISSUER, CLIENT_ID};
9use super::token_store::{GrokTokenStore, StoredAuth};
10use crate::auth::AuthStorage;
11
12const REFRESH_SKEW_MS: u64 = 5 * 60 * 1000;
13
14#[derive(Deserialize)]
15struct Discovery {
16    issuer: String,
17    token_endpoint: String,
18}
19
20#[derive(Deserialize)]
21struct RefreshResponse {
22    access_token: String,
23    expires_in: u64,
24    #[serde(default)]
25    refresh_token: Option<String>,
26}
27
28pub struct GrokAuthManager<S: AuthStorage<StoredAuth>> {
29    store: GrokTokenStore<S>,
30    client: reqwest::Client,
31    refresh_lock: Arc<Mutex<()>>,
32}
33
34impl<S: AuthStorage<StoredAuth>> GrokAuthManager<S> {
35    pub fn new(store: GrokTokenStore<S>) -> anyhow::Result<Self> {
36        let client = reqwest::Client::builder()
37            .redirect(reqwest::redirect::Policy::none())
38            .connect_timeout(Duration::from_secs(10))
39            .timeout(Duration::from_secs(20))
40            .build()?;
41        Ok(Self {
42            store,
43            client,
44            refresh_lock: Arc::new(Mutex::new(())),
45        })
46    }
47
48    pub fn store(&self) -> &GrokTokenStore<S> {
49        &self.store
50    }
51
52    pub async fn get_auth(&self) -> anyhow::Result<StoredAuth> {
53        let auth = self
54            .store
55            .load_auth()?
56            .ok_or_else(|| anyhow::anyhow!("Not authenticated"))?;
57        if auth.expires_at_ms > now_ms().saturating_add(REFRESH_SKEW_MS) {
58            return Ok(auth);
59        }
60        self.refresh(false, None).await
61    }
62
63    pub async fn force_refresh(&self, rejected_access: &str) -> anyhow::Result<StoredAuth> {
64        self.refresh(true, Some(rejected_access)).await
65    }
66
67    async fn refresh(
68        &self,
69        force: bool,
70        rejected_access: Option<&str>,
71    ) -> anyhow::Result<StoredAuth> {
72        let _guard = self.refresh_lock.lock().await;
73        let auth = self
74            .store
75            .load_auth()?
76            .ok_or_else(|| anyhow::anyhow!("Not authenticated"))?;
77        if (!force && auth.expires_at_ms > now_ms().saturating_add(REFRESH_SKEW_MS))
78            || rejected_access.is_some_and(|access| auth.access != access)
79        {
80            return Ok(auth);
81        }
82        if auth.issuer != CANONICAL_ISSUER || auth.client_id != CLIENT_ID {
83            anyhow::bail!("Unsupported Grok OAuth session");
84        }
85        let issuer = Url::parse(CANONICAL_ISSUER)?;
86        let discovery_url = issuer.join("/.well-known/openid-configuration")?;
87        let discovery: Discovery = self
88            .client
89            .get(discovery_url)
90            .send()
91            .await?
92            .error_for_status()?
93            .json()
94            .await?;
95        if discovery.issuer != CANONICAL_ISSUER {
96            anyhow::bail!("OIDC discovery issuer mismatch");
97        }
98        let endpoint = Url::parse(&discovery.token_endpoint)?;
99        if endpoint.scheme() != "https" || endpoint.origin() != issuer.origin() {
100            anyhow::bail!("OIDC token endpoint is outside the canonical issuer");
101        }
102        let refreshed: RefreshResponse = self
103            .client
104            .post(endpoint)
105            .form(&[
106                ("grant_type", "refresh_token"),
107                ("refresh_token", auth.refresh.as_str()),
108                ("client_id", auth.client_id.as_str()),
109            ])
110            .send()
111            .await?
112            .error_for_status()?
113            .json()
114            .await?;
115        if refreshed.access_token.is_empty() || refreshed.expires_in == 0 {
116            anyhow::bail!("Invalid token refresh response");
117        }
118        let updated = StoredAuth {
119            access: refreshed.access_token,
120            refresh: refreshed
121                .refresh_token
122                .filter(|token| !token.is_empty())
123                .unwrap_or(auth.refresh),
124            expires_at_ms: now_ms().saturating_add(refreshed.expires_in.saturating_mul(1000)),
125            issuer: auth.issuer,
126            client_id: auth.client_id,
127        };
128        self.store.save_auth(updated.clone())?;
129        Ok(updated)
130    }
131}
132
133fn now_ms() -> u64 {
134    SystemTime::now()
135        .duration_since(UNIX_EPOCH)
136        .unwrap_or_default()
137        .as_millis() as u64
138}
139
140#[cfg(test)]
141mod tests {
142    use super::*;
143    use crate::auth::InMemoryAuthStore;
144
145    fn auth(access: &str) -> StoredAuth {
146        StoredAuth {
147            access: access.into(),
148            refresh: "synthetic-refresh".into(),
149            expires_at_ms: now_ms().saturating_add(3_600_000),
150            issuer: CANONICAL_ISSUER.into(),
151            client_id: "synthetic-client".into(),
152        }
153    }
154
155    #[test]
156    fn discovery_accepts_standard_metadata_fields() {
157        let discovery: Discovery = serde_json::from_value(serde_json::json!({
158            "issuer": CANONICAL_ISSUER,
159            "token_endpoint": "https://auth.x.ai/oauth/token",
160            "authorization_endpoint": "https://auth.x.ai/oauth/authorize",
161            "jwks_uri": "https://auth.x.ai/.well-known/jwks.json"
162        }))
163        .unwrap();
164
165        assert_eq!(discovery.issuer, CANONICAL_ISSUER);
166    }
167
168    #[tokio::test]
169    async fn concurrent_stale_401_refreshes_reuse_the_rotated_access_token() {
170        let store = GrokTokenStore::new(InMemoryAuthStore::new());
171        store.save_auth(auth("rotated-access")).unwrap();
172        let manager = Arc::new(GrokAuthManager::new(store).unwrap());
173        let (first, second) = tokio::join!(
174            manager.force_refresh("rejected-access"),
175            manager.force_refresh("rejected-access"),
176        );
177        assert_eq!(first.unwrap().access, "rotated-access");
178        assert_eq!(second.unwrap().access, "rotated-access");
179    }
180}
claude_codex/providers/grok/auth/manager.rs

claude_codex/providers/grok/auth/
manager.rs
