rpi_cli/
oauth.rs

1//! OAuth 2.0 authentication flow for various providers.
2//!
3//! Implements device code flow and authorization code flow for:
4//! - Anthropic (Claude Pro/Max)
5//! - OpenAI
6//! - GitHub Copilot
7//! - OpenRouter
8//! - And more...
9
10use serde::{Deserialize, Serialize};
11use std::collections::HashMap;
12use std::time::{Duration, SystemTime, UNIX_EPOCH};
13use thiserror::Error;
14
15#[derive(Error, Debug)]
16pub enum OAuthError {
17    #[error("HTTP request failed: {0}")]
18    HttpError(String),
19
20    #[error("Invalid response: {0}")]
21    InvalidResponse(String),
22
23    #[error("Authentication timeout")]
24    Timeout,
25
26    #[error("User denied authorization")]
27    Denied,
28
29    #[error("Token expired")]
30    TokenExpired,
31
32    #[error("Unsupported provider: {0}")]
33    UnsupportedProvider(String),
34}
35
36/// OAuth provider types
37#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
38#[serde(rename_all = "lowercase")]
39pub enum OAuthProvider {
40    Anthropic,
41    OpenAI,
42    GitHubCopilot,
43    OpenRouter,
44    Google,
45    Azure,
46}
47
48impl std::fmt::Display for OAuthProvider {
49    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
50        match self {
51            OAuthProvider::Anthropic => write!(f, "anthropic"),
52            OAuthProvider::OpenAI => write!(f, "openai"),
53            OAuthProvider::GitHubCopilot => write!(f, "github-copilot"),
54            OAuthProvider::OpenRouter => write!(f, "openrouter"),
55            OAuthProvider::Google => write!(f, "google"),
56            OAuthProvider::Azure => write!(f, "azure"),
57        }
58    }
59}
60
61/// OAuth token information
62#[derive(Debug, Clone, Serialize, Deserialize)]
63pub struct OAuthToken {
64    pub access_token: String,
65    pub refresh_token: Option<String>,
66    pub token_type: String,
67    pub expires_at: Option<u64>,
68    pub scope: Option<String>,
69}
70
71impl OAuthToken {
72    /// Check if the token is expired
73    pub fn is_expired(&self) -> bool {
74        if let Some(expires_at) = self.expires_at {
75            let now = SystemTime::now()
76                .duration_since(UNIX_EPOCH)
77                .unwrap()
78                .as_secs();
79            now >= expires_at
80        } else {
81            false
82        }
83    }
84
85    /// Check if the token will expire soon (within the given duration)
86    pub fn expires_soon(&self, within: Duration) -> bool {
87        if let Some(expires_at) = self.expires_at {
88            let now = SystemTime::now()
89                .duration_since(UNIX_EPOCH)
90                .unwrap()
91                .as_secs();
92            let threshold = within.as_secs();
93            expires_at.saturating_sub(now) <= threshold
94        } else {
95            false
96        }
97    }
98}
99
100/// Device code flow response
101#[derive(Debug, Clone, Serialize, Deserialize)]
102pub struct DeviceCodeResponse {
103    pub device_code: String,
104    pub user_code: String,
105    pub verification_uri: String,
106    #[serde(alias = "verification_uri_complete")]
107    pub verification_uri_complete: Option<String>,
108    pub expires_in: u64,
109    pub interval: u64,
110}
111
112/// OAuth client configuration
113#[derive(Debug, Clone)]
114pub struct OAuthConfig {
115    pub client_id: String,
116    pub client_secret: Option<String>,
117    pub auth_url: String,
118    pub token_url: String,
119    pub redirect_uri: Option<String>,
120    pub scopes: Vec<String>,
121}
122
123/// OAuth client for handling authentication flows
124pub struct OAuthClient {
125    config: OAuthConfig,
126    provider: OAuthProvider,
127    http_client: reqwest::Client,
128}
129
130impl OAuthClient {
131    /// Create a new OAuth client for the specified provider
132    pub fn new(provider: OAuthProvider, config: OAuthConfig) -> Self {
133        Self {
134            config,
135            provider,
136            http_client: reqwest::Client::new(),
137        }
138    }
139
140    /// Create an OAuth client with default configuration for a provider
141    pub fn with_defaults(provider: OAuthProvider) -> Result<Self, OAuthError> {
142        let config = match provider {
143            OAuthProvider::Anthropic => OAuthConfig {
144                client_id: "pi-cli".to_string(),
145                client_secret: None,
146                auth_url: "https://claude.ai/oauth/authorize".to_string(),
147                token_url: "https://claude.ai/oauth/token".to_string(),
148                redirect_uri: Some("http://localhost:8080/callback".to_string()),
149                scopes: vec!["read".to_string(), "write".to_string()],
150            },
151            OAuthProvider::OpenAI => OAuthConfig {
152                client_id: "pi-cli".to_string(),
153                client_secret: None,
154                auth_url: "https://auth.openai.com/authorize".to_string(),
155                token_url: "https://auth.openai.com/oauth/token".to_string(),
156                redirect_uri: Some("http://localhost:8080/callback".to_string()),
157                scopes: vec!["openai.public".to_string()],
158            },
159            OAuthProvider::GitHubCopilot => OAuthConfig {
160                client_id: "Iv1.b507a08c30306e6a".to_string(), // GitHub App client ID
161                client_secret: None,
162                auth_url: "https://github.com/login/oauth/authorize".to_string(),
163                token_url: "https://github.com/login/oauth/access_token".to_string(),
164                redirect_uri: Some("http://localhost:8080/callback".to_string()),
165                scopes: vec![],
166            },
167            OAuthProvider::OpenRouter => OAuthConfig {
168                client_id: "pi-cli".to_string(),
169                client_secret: None,
170                auth_url: "https://openrouter.ai/auth/authorize".to_string(),
171                token_url: "https://openrouter.ai/auth/token".to_string(),
172                redirect_uri: Some("http://localhost:8080/callback".to_string()),
173                scopes: vec![],
174            },
175            OAuthProvider::Google => OAuthConfig {
176                client_id: "pi-cli".to_string(),
177                client_secret: None,
178                auth_url: "https://accounts.google.com/o/oauth2/v2/auth".to_string(),
179                token_url: "https://oauth2.googleapis.com/token".to_string(),
180                redirect_uri: Some("http://localhost:8080/callback".to_string()),
181                scopes: vec!["https://www.googleapis.com/auth/generative-language".to_string()],
182            },
183            OAuthProvider::Azure => OAuthConfig {
184                client_id: "pi-cli".to_string(),
185                client_secret: None,
186                auth_url: "https://login.microsoftonline.com/common/oauth2/v2.0/authorize"
187                    .to_string(),
188                token_url: "https://login.microsoftonline.com/common/oauth2/v2.0/token".to_string(),
189                redirect_uri: Some("http://localhost:8080/callback".to_string()),
190                scopes: vec!["https://cognitiveservices.azure.com/.default".to_string()],
191            },
192        };
193
194        Ok(Self::new(provider, config))
195    }
196
197    /// Start device code flow
198    pub async fn start_device_code_flow(&self) -> Result<DeviceCodeResponse, OAuthError> {
199        let mut params = HashMap::new();
200        params.insert("client_id", self.config.client_id.as_str());
201
202        let scope_str = self.config.scopes.join(" ");
203        if !scope_str.is_empty() {
204            params.insert("scope", scope_str.as_str());
205        }
206
207        let response = self
208            .http_client
209            .post(&self.config.auth_url)
210            .form(&params)
211            .send()
212            .await
213            .map_err(|e| OAuthError::HttpError(e.to_string()))?;
214
215        if !response.status().is_success() {
216            let status = response.status();
217            let text = response.text().await.unwrap_or_default();
218            return Err(OAuthError::HttpError(format!("{}: {}", status, text)));
219        }
220
221        let device_code: DeviceCodeResponse = response
222            .json()
223            .await
224            .map_err(|e| OAuthError::InvalidResponse(e.to_string()))?;
225
226        Ok(device_code)
227    }
228
229    /// Poll for token completion (device code flow)
230    pub async fn poll_for_token(
231        &self,
232        device_code: &str,
233        interval: Duration,
234        timeout: Duration,
235    ) -> Result<OAuthToken, OAuthError> {
236        let start = std::time::Instant::now();
237
238        loop {
239            if start.elapsed() > timeout {
240                return Err(OAuthError::Timeout);
241            }
242
243            tokio::time::sleep(interval).await;
244
245            let mut params = HashMap::new();
246            params.insert("client_id", self.config.client_id.as_str());
247            params.insert("device_code", device_code);
248            params.insert("grant_type", "urn:ietf:params:oauth:grant-type:device_code");
249
250            let response = self
251                .http_client
252                .post(&self.config.token_url)
253                .form(&params)
254                .send()
255                .await
256                .map_err(|e| OAuthError::HttpError(e.to_string()))?;
257
258            let status = response.status();
259            let text = response.text().await.unwrap_or_default();
260
261            if status.is_success() {
262                // Parse successful token response
263                let token_response: serde_json::Value = serde_json::from_str(&text)
264                    .map_err(|e| OAuthError::InvalidResponse(e.to_string()))?;
265
266                let access_token = token_response["access_token"]
267                    .as_str()
268                    .ok_or_else(|| OAuthError::InvalidResponse("missing access_token".to_string()))?
269                    .to_string();
270
271                let refresh_token = token_response["refresh_token"]
272                    .as_str()
273                    .map(|s| s.to_string());
274
275                let token_type = token_response["token_type"]
276                    .as_str()
277                    .unwrap_or("Bearer")
278                    .to_string();
279
280                let expires_at = token_response["expires_in"].as_u64().map(|expires_in| {
281                    SystemTime::now()
282                        .duration_since(UNIX_EPOCH)
283                        .unwrap()
284                        .as_secs()
285                        + expires_in
286                });
287
288                let scope = token_response["scope"].as_str().map(|s| s.to_string());
289
290                return Ok(OAuthToken {
291                    access_token,
292                    refresh_token,
293                    token_type,
294                    expires_at,
295                    scope,
296                });
297            }
298
299            // Check for specific error codes
300            if let Ok(error_response) = serde_json::from_str::<serde_json::Value>(&text) {
301                if let Some(error) = error_response["error"].as_str() {
302                    match error {
303                        "authorization_pending" => continue,
304                        "slow_down" => {
305                            tokio::time::sleep(Duration::from_secs(5)).await;
306                            continue;
307                        }
308                        "access_denied" => return Err(OAuthError::Denied),
309                        "expired_token" => return Err(OAuthError::Timeout),
310                        _ => {
311                            return Err(OAuthError::InvalidResponse(format!(
312                                "unknown error: {}",
313                                error
314                            )))
315                        }
316                    }
317                }
318            }
319        }
320    }
321
322    /// Refresh an expired token
323    pub async fn refresh_token(&self, refresh_token: &str) -> Result<OAuthToken, OAuthError> {
324        let mut params = HashMap::new();
325        params.insert("client_id", self.config.client_id.as_str());
326        params.insert("refresh_token", refresh_token);
327        params.insert("grant_type", "refresh_token");
328
329        if let Some(secret) = &self.config.client_secret {
330            params.insert("client_secret", secret.as_str());
331        }
332
333        let response = self
334            .http_client
335            .post(&self.config.token_url)
336            .form(&params)
337            .send()
338            .await
339            .map_err(|e| OAuthError::HttpError(e.to_string()))?;
340
341        if !response.status().is_success() {
342            let status = response.status();
343            let text = response.text().await.unwrap_or_default();
344            return Err(OAuthError::HttpError(format!("{}: {}", status, text)));
345        }
346
347        let token_response: serde_json::Value = response
348            .json()
349            .await
350            .map_err(|e| OAuthError::InvalidResponse(e.to_string()))?;
351
352        let access_token = token_response["access_token"]
353            .as_str()
354            .ok_or_else(|| OAuthError::InvalidResponse("missing access_token".to_string()))?
355            .to_string();
356
357        let new_refresh_token = token_response["refresh_token"]
358            .as_str()
359            .map(|s| s.to_string())
360            .unwrap_or_else(|| refresh_token.to_string());
361
362        let token_type = token_response["token_type"]
363            .as_str()
364            .unwrap_or("Bearer")
365            .to_string();
366
367        let expires_at = token_response["expires_in"].as_u64().map(|expires_in| {
368            SystemTime::now()
369                .duration_since(UNIX_EPOCH)
370                .unwrap()
371                .as_secs()
372                + expires_in
373        });
374
375        let scope = token_response["scope"].as_str().map(|s| s.to_string());
376
377        Ok(OAuthToken {
378            access_token,
379            refresh_token: Some(new_refresh_token),
380            token_type,
381            expires_at,
382            scope,
383        })
384    }
385
386    /// Get the provider type
387    pub fn provider(&self) -> OAuthProvider {
388        self.provider
389    }
390}
391
392/// OAuth token storage
393pub struct OAuthTokenStore {
394    tokens: HashMap<OAuthProvider, OAuthToken>,
395}
396
397impl OAuthTokenStore {
398    /// Create a new token store
399    pub fn new() -> Self {
400        Self {
401            tokens: HashMap::new(),
402        }
403    }
404
405    /// Store a token
406    pub fn store(&mut self, provider: OAuthProvider, token: OAuthToken) {
407        self.tokens.insert(provider, token);
408    }
409
410    /// Retrieve a token
411    pub fn get(&self, provider: OAuthProvider) -> Option<&OAuthToken> {
412        self.tokens.get(&provider)
413    }
414
415    /// Remove a token
416    pub fn remove(&mut self, provider: OAuthProvider) -> Option<OAuthToken> {
417        self.tokens.remove(&provider)
418    }
419
420    /// Check if a token exists and is valid
421    pub fn is_valid(&self, provider: OAuthProvider) -> bool {
422        self.tokens
423            .get(&provider)
424            .map(|token| !token.is_expired())
425            .unwrap_or(false)
426    }
427}
428
429impl Default for OAuthTokenStore {
430    fn default() -> Self {
431        Self::new()
432    }
433}
434
435#[cfg(test)]
436mod tests {
437    use super::*;
438
439    #[test]
440    fn test_oauth_token_expiry() {
441        let token = OAuthToken {
442            access_token: "test".to_string(),
443            refresh_token: None,
444            token_type: "Bearer".to_string(),
445            expires_at: Some(
446                SystemTime::now()
447                    .duration_since(UNIX_EPOCH)
448                    .unwrap()
449                    .as_secs()
450                    + 3600,
451            ),
452            scope: None,
453        };
454
455        assert!(!token.is_expired());
456        assert!(!token.expires_soon(Duration::from_secs(1800)));
457        assert!(token.expires_soon(Duration::from_secs(7200)));
458    }
459
460    #[test]
461    fn test_oauth_token_store() {
462        let mut store = OAuthTokenStore::new();
463        let token = OAuthToken {
464            access_token: "test".to_string(),
465            refresh_token: None,
466            token_type: "Bearer".to_string(),
467            expires_at: None,
468            scope: None,
469        };
470
471        store.store(OAuthProvider::Anthropic, token.clone());
472        assert!(store.is_valid(OAuthProvider::Anthropic));
473        assert!(!store.is_valid(OAuthProvider::OpenAI));
474
475        let retrieved = store.get(OAuthProvider::Anthropic).unwrap();
476        assert_eq!(retrieved.access_token, "test");
477
478        store.remove(OAuthProvider::Anthropic);
479        assert!(!store.is_valid(OAuthProvider::Anthropic));
480    }
481}
rpi_cli/oauth.rs

rpi_cli/
oauth.rs
