stakpak_shared/oauth/
flow.rs

1//! OAuth 2.0 authorization code flow implementation
2
3use super::config::OAuthConfig;
4use super::error::{OAuthError, OAuthResult};
5use super::pkce::PkceChallenge;
6use serde::{Deserialize, Serialize};
7
8/// OAuth token response from the token endpoint
9#[derive(Debug, Clone, Serialize, Deserialize)]
10pub struct TokenResponse {
11    /// Access token for API requests
12    pub access_token: String,
13    /// Refresh token for obtaining new access tokens
14    pub refresh_token: String,
15    /// Token lifetime in seconds
16    pub expires_in: i64,
17    /// Token type (usually "Bearer")
18    pub token_type: String,
19}
20
21/// OAuth 2.0 authorization code flow handler
22pub struct OAuthFlow {
23    config: OAuthConfig,
24    pkce: Option<PkceChallenge>,
25}
26
27impl OAuthFlow {
28    /// Create a new OAuth flow with the given configuration
29    pub fn new(config: OAuthConfig) -> Self {
30        Self { config, pkce: None }
31    }
32
33    /// Generate the authorization URL for the user to visit
34    ///
35    /// This generates a new PKCE challenge and returns the full authorization URL
36    /// that should be opened in the user's browser.
37    pub fn generate_auth_url(&mut self) -> String {
38        let pkce = PkceChallenge::generate();
39
40        let url = format!(
41            "{}?code=true&client_id={}&response_type=code&redirect_uri={}&scope={}&code_challenge={}&code_challenge_method={}&state={}",
42            self.config.auth_url,
43            urlencoding::encode(&self.config.client_id),
44            urlencoding::encode(&self.config.redirect_url),
45            urlencoding::encode(&self.config.scopes_string()),
46            urlencoding::encode(&pkce.challenge),
47            PkceChallenge::challenge_method(),
48            urlencoding::encode(&pkce.verifier), // State contains verifier for validation
49        );
50
51        self.pkce = Some(pkce);
52        url
53    }
54
55    /// Exchange authorization code for tokens
56    ///
57    /// The code should be in the format "authorization_code#state" as returned by Anthropic.
58    pub async fn exchange_code(&self, code: &str) -> OAuthResult<TokenResponse> {
59        let pkce = self.pkce.as_ref().ok_or(OAuthError::PkceNotInitialized)?;
60
61        // Parse the authorization code - format: "authorization_code#state"
62        let (auth_code, state) = parse_auth_code(code)?;
63
64        // Validate state matches our verifier
65        if state != pkce.verifier {
66            return Err(OAuthError::invalid_code_format(
67                "State mismatch - possible CSRF attack",
68            ));
69        }
70
71        let client =
72            crate::tls_client::create_tls_client(crate::tls_client::TlsClientConfig::default())
73                .unwrap_or_else(|_| reqwest::Client::new());
74        let response = client
75            .post(&self.config.token_url)
76            .json(&serde_json::json!({
77                "grant_type": "authorization_code",
78                "code": auth_code,
79                "state": state,
80                "client_id": self.config.client_id,
81                "redirect_uri": self.config.redirect_url,
82                "code_verifier": pkce.verifier,
83            }))
84            .send()
85            .await?;
86
87        if !response.status().is_success() {
88            let status = response.status();
89            let error_text = response.text().await.unwrap_or_default();
90            return Err(OAuthError::token_exchange_failed(format!(
91                "HTTP {}: {}",
92                status, error_text
93            )));
94        }
95
96        response.json::<TokenResponse>().await.map_err(|e| {
97            OAuthError::token_exchange_failed(format!("Failed to parse token response: {}", e))
98        })
99    }
100
101    /// Refresh an expired access token
102    pub async fn refresh_token(&self, refresh_token: &str) -> OAuthResult<TokenResponse> {
103        let client =
104            crate::tls_client::create_tls_client(crate::tls_client::TlsClientConfig::default())
105                .unwrap_or_else(|_| reqwest::Client::new());
106        let response = client
107            .post(&self.config.token_url)
108            .json(&serde_json::json!({
109                "grant_type": "refresh_token",
110                "refresh_token": refresh_token,
111                "client_id": self.config.client_id,
112            }))
113            .send()
114            .await?;
115
116        if !response.status().is_success() {
117            let status = response.status();
118            let error_text = response.text().await.unwrap_or_default();
119            return Err(OAuthError::token_refresh_failed(format!(
120                "HTTP {}: {}",
121                status, error_text
122            )));
123        }
124
125        response.json::<TokenResponse>().await.map_err(|e| {
126            OAuthError::token_refresh_failed(format!("Failed to parse token response: {}", e))
127        })
128    }
129
130    /// Get the PKCE verifier (for validation purposes)
131    pub fn pkce_verifier(&self) -> Option<&str> {
132        self.pkce.as_ref().map(|p| p.verifier.as_str())
133    }
134}
135
136/// Parse the authorization code from Anthropic's callback format
137///
138/// Anthropic returns codes in the format: "authorization_code#state"
139fn parse_auth_code(code: &str) -> OAuthResult<(String, String)> {
140    // Handle both "#" and "%23" (URL-encoded #)
141    let code = code.replace("%23", "#");
142
143    if let Some(pos) = code.find('#') {
144        let auth_code = code[..pos].to_string();
145        let state = code[pos + 1..].to_string();
146
147        if auth_code.is_empty() || state.is_empty() {
148            return Err(OAuthError::invalid_code_format(
149                "Authorization code or state is empty",
150            ));
151        }
152
153        Ok((auth_code, state))
154    } else {
155        Err(OAuthError::invalid_code_format(
156            "Expected format: authorization_code#state",
157        ))
158    }
159}
160
161#[cfg(test)]
162mod tests {
163    use super::*;
164
165    fn test_config() -> OAuthConfig {
166        OAuthConfig::new(
167            "test-client-id",
168            "https://example.com/auth",
169            "https://example.com/token",
170            "https://example.com/callback",
171            vec!["scope1".to_string(), "scope2".to_string()],
172        )
173    }
174
175    #[test]
176    fn test_generate_auth_url() {
177        let mut flow = OAuthFlow::new(test_config());
178        let url = flow.generate_auth_url();
179
180        assert!(url.starts_with("https://example.com/auth?"));
181        assert!(url.contains("client_id=test-client-id"));
182        assert!(url.contains("response_type=code"));
183        assert!(url.contains("redirect_uri="));
184        assert!(url.contains("scope=scope1%20scope2"));
185        assert!(url.contains("code_challenge="));
186        assert!(url.contains("code_challenge_method=S256"));
187        assert!(url.contains("state="));
188
189        // PKCE should be initialized
190        assert!(flow.pkce.is_some());
191    }
192
193    #[test]
194    fn test_parse_auth_code_valid() {
195        let result = parse_auth_code("abc123#verifier456");
196        assert!(result.is_ok());
197        let (code, state) = result.unwrap();
198        assert_eq!(code, "abc123");
199        assert_eq!(state, "verifier456");
200    }
201
202    #[test]
203    fn test_parse_auth_code_url_encoded() {
204        let result = parse_auth_code("abc123%23verifier456");
205        assert!(result.is_ok());
206        let (code, state) = result.unwrap();
207        assert_eq!(code, "abc123");
208        assert_eq!(state, "verifier456");
209    }
210
211    #[test]
212    fn test_parse_auth_code_missing_separator() {
213        let result = parse_auth_code("abc123verifier456");
214        assert!(result.is_err());
215    }
216
217    #[test]
218    fn test_parse_auth_code_empty_parts() {
219        assert!(parse_auth_code("#state").is_err());
220        assert!(parse_auth_code("code#").is_err());
221        assert!(parse_auth_code("#").is_err());
222    }
223
224    #[test]
225    fn test_exchange_code_without_pkce() {
226        let flow = OAuthFlow::new(test_config());
227        let result = tokio_test::block_on(flow.exchange_code("code#state"));
228        assert!(matches!(result, Err(OAuthError::PkceNotInitialized)));
229    }
230
231    #[test]
232    fn test_token_response_serde() {
233        let json = r#"{
234            "access_token": "access123",
235            "refresh_token": "refresh456",
236            "expires_in": 3600,
237            "token_type": "Bearer"
238        }"#;
239
240        let response: TokenResponse = serde_json::from_str(json).unwrap();
241        assert_eq!(response.access_token, "access123");
242        assert_eq!(response.refresh_token, "refresh456");
243        assert_eq!(response.expires_in, 3600);
244        assert_eq!(response.token_type, "Bearer");
245    }
246}
stakpak_shared/oauth/flow.rs

stakpak_shared/oauth/
flow.rs
