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![
182                    "https://www.googleapis.com/auth/generative-language".to_string(),
183                ],
184            },
185            OAuthProvider::Azure => OAuthConfig {
186                client_id: "pi-cli".to_string(),
187                client_secret: None,
188                auth_url: "https://login.microsoftonline.com/common/oauth2/v2.0/authorize".to_string(),
189                token_url: "https://login.microsoftonline.com/common/oauth2/v2.0/token".to_string(),
190                redirect_uri: Some("http://localhost:8080/callback".to_string()),
191                scopes: vec!["https://cognitiveservices.azure.com/.default".to_string()],
192            },
193        };
194        
195        Ok(Self::new(provider, config))
196    }
197    
198    /// Start device code flow
199    pub async fn start_device_code_flow(&self) -> Result<DeviceCodeResponse, OAuthError> {
200        let mut params = HashMap::new();
201        params.insert("client_id", self.config.client_id.as_str());
202        
203        let scope_str = self.config.scopes.join(" ");
204        if !scope_str.is_empty() {
205            params.insert("scope", scope_str.as_str());
206        }
207        
208        let response = self.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.http_client
251                .post(&self.config.token_url)
252                .form(&params)
253                .send()
254                .await
255                .map_err(|e| OAuthError::HttpError(e.to_string()))?;
256            
257            let status = response.status();
258            let text = response.text().await.unwrap_or_default();
259            
260            if status.is_success() {
261                // Parse successful token response
262                let token_response: serde_json::Value = serde_json::from_str(&text)
263                    .map_err(|e| OAuthError::InvalidResponse(e.to_string()))?;
264                
265                let access_token = token_response["access_token"]
266                    .as_str()
267                    .ok_or_else(|| OAuthError::InvalidResponse("missing access_token".to_string()))?
268                    .to_string();
269                
270                let refresh_token = token_response["refresh_token"]
271                    .as_str()
272                    .map(|s| s.to_string());
273                
274                let token_type = token_response["token_type"]
275                    .as_str()
276                    .unwrap_or("Bearer")
277                    .to_string();
278                
279                let expires_at = token_response["expires_in"]
280                    .as_u64()
281                    .map(|expires_in| {
282                        SystemTime::now()
283                            .duration_since(UNIX_EPOCH)
284                            .unwrap()
285                            .as_secs()
286                            + expires_in
287                    });
288                
289                let scope = token_response["scope"]
290                    .as_str()
291                    .map(|s| s.to_string());
292                
293                return Ok(OAuthToken {
294                    access_token,
295                    refresh_token,
296                    token_type,
297                    expires_at,
298                    scope,
299                });
300            }
301            
302            // Check for specific error codes
303            if let Ok(error_response) = serde_json::from_str::<serde_json::Value>(&text) {
304                if let Some(error) = error_response["error"].as_str() {
305                    match error {
306                        "authorization_pending" => continue,
307                        "slow_down" => {
308                            tokio::time::sleep(Duration::from_secs(5)).await;
309                            continue;
310                        }
311                        "access_denied" => return Err(OAuthError::Denied),
312                        "expired_token" => return Err(OAuthError::Timeout),
313                        _ => return Err(OAuthError::InvalidResponse(format!("unknown error: {}", error))),
314                    }
315                }
316            }
317        }
318    }
319    
320    /// Refresh an expired token
321    pub async fn refresh_token(&self, refresh_token: &str) -> Result<OAuthToken, OAuthError> {
322        let mut params = HashMap::new();
323        params.insert("client_id", self.config.client_id.as_str());
324        params.insert("refresh_token", refresh_token);
325        params.insert("grant_type", "refresh_token");
326        
327        if let Some(secret) = &self.config.client_secret {
328            params.insert("client_secret", secret.as_str());
329        }
330        
331        let response = self.http_client
332            .post(&self.config.token_url)
333            .form(&params)
334            .send()
335            .await
336            .map_err(|e| OAuthError::HttpError(e.to_string()))?;
337        
338        if !response.status().is_success() {
339            let status = response.status();
340            let text = response.text().await.unwrap_or_default();
341            return Err(OAuthError::HttpError(format!("{}: {}", status, text)));
342        }
343        
344        let token_response: serde_json::Value = response
345            .json()
346            .await
347            .map_err(|e| OAuthError::InvalidResponse(e.to_string()))?;
348        
349        let access_token = token_response["access_token"]
350            .as_str()
351            .ok_or_else(|| OAuthError::InvalidResponse("missing access_token".to_string()))?
352            .to_string();
353        
354        let new_refresh_token = token_response["refresh_token"]
355            .as_str()
356            .map(|s| s.to_string())
357            .unwrap_or_else(|| refresh_token.to_string());
358        
359        let token_type = token_response["token_type"]
360            .as_str()
361            .unwrap_or("Bearer")
362            .to_string();
363        
364        let expires_at = token_response["expires_in"]
365            .as_u64()
366            .map(|expires_in| {
367                SystemTime::now()
368                    .duration_since(UNIX_EPOCH)
369                    .unwrap()
370                    .as_secs()
371                    + expires_in
372            });
373        
374        let scope = token_response["scope"]
375            .as_str()
376            .map(|s| s.to_string());
377        
378        Ok(OAuthToken {
379            access_token,
380            refresh_token: Some(new_refresh_token),
381            token_type,
382            expires_at,
383            scope,
384        })
385    }
386    
387    /// Get the provider type
388    pub fn provider(&self) -> OAuthProvider {
389        self.provider
390    }
391}
392
393/// OAuth token storage
394pub struct OAuthTokenStore {
395    tokens: HashMap<OAuthProvider, OAuthToken>,
396}
397
398impl OAuthTokenStore {
399    /// Create a new token store
400    pub fn new() -> Self {
401        Self {
402            tokens: HashMap::new(),
403        }
404    }
405    
406    /// Store a token
407    pub fn store(&mut self, provider: OAuthProvider, token: OAuthToken) {
408        self.tokens.insert(provider, token);
409    }
410    
411    /// Retrieve a token
412    pub fn get(&self, provider: OAuthProvider) -> Option<&OAuthToken> {
413        self.tokens.get(&provider)
414    }
415    
416    /// Remove a token
417    pub fn remove(&mut self, provider: OAuthProvider) -> Option<OAuthToken> {
418        self.tokens.remove(&provider)
419    }
420    
421    /// Check if a token exists and is valid
422    pub fn is_valid(&self, provider: OAuthProvider) -> bool {
423        self.tokens
424            .get(&provider)
425            .map(|token| !token.is_expired())
426            .unwrap_or(false)
427    }
428}
429
430impl Default for OAuthTokenStore {
431    fn default() -> Self {
432        Self::new()
433    }
434}
435
436#[cfg(test)]
437mod tests {
438    use super::*;
439    
440    #[test]
441    fn test_oauth_token_expiry() {
442        let token = OAuthToken {
443            access_token: "test".to_string(),
444            refresh_token: None,
445            token_type: "Bearer".to_string(),
446            expires_at: Some(
447                SystemTime::now()
448                    .duration_since(UNIX_EPOCH)
449                    .unwrap()
450                    .as_secs()
451                    + 3600,
452            ),
453            scope: None,
454        };
455        
456        assert!(!token.is_expired());
457        assert!(!token.expires_soon(Duration::from_secs(1800)));
458        assert!(token.expires_soon(Duration::from_secs(7200)));
459    }
460    
461    #[test]
462    fn test_oauth_token_store() {
463        let mut store = OAuthTokenStore::new();
464        let token = OAuthToken {
465            access_token: "test".to_string(),
466            refresh_token: None,
467            token_type: "Bearer".to_string(),
468            expires_at: None,
469            scope: None,
470        };
471        
472        store.store(OAuthProvider::Anthropic, token.clone());
473        assert!(store.is_valid(OAuthProvider::Anthropic));
474        assert!(!store.is_valid(OAuthProvider::OpenAI));
475        
476        let retrieved = store.get(OAuthProvider::Anthropic).unwrap();
477        assert_eq!(retrieved.access_token, "test");
478        
479        store.remove(OAuthProvider::Anthropic);
480        assert!(!store.is_valid(OAuthProvider::Anthropic));
481    }
482}
rpi_cli/oauth.rs

rpi_cli/
oauth.rs
