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 = reqwest::Client::new();
72        let response = client
73            .post(&self.config.token_url)
74            .json(&serde_json::json!({
75                "grant_type": "authorization_code",
76                "code": auth_code,
77                "state": state,
78                "client_id": self.config.client_id,
79                "redirect_uri": self.config.redirect_url,
80                "code_verifier": pkce.verifier,
81            }))
82            .send()
83            .await?;
84
85        if !response.status().is_success() {
86            let status = response.status();
87            let error_text = response.text().await.unwrap_or_default();
88            return Err(OAuthError::token_exchange_failed(format!(
89                "HTTP {}: {}",
90                status, error_text
91            )));
92        }
93
94        response.json::<TokenResponse>().await.map_err(|e| {
95            OAuthError::token_exchange_failed(format!("Failed to parse token response: {}", e))
96        })
97    }
98
99    /// Refresh an expired access token
100    pub async fn refresh_token(&self, refresh_token: &str) -> OAuthResult<TokenResponse> {
101        let client = reqwest::Client::new();
102        let response = client
103            .post(&self.config.token_url)
104            .json(&serde_json::json!({
105                "grant_type": "refresh_token",
106                "refresh_token": refresh_token,
107                "client_id": self.config.client_id,
108            }))
109            .send()
110            .await?;
111
112        if !response.status().is_success() {
113            let status = response.status();
114            let error_text = response.text().await.unwrap_or_default();
115            return Err(OAuthError::token_refresh_failed(format!(
116                "HTTP {}: {}",
117                status, error_text
118            )));
119        }
120
121        response.json::<TokenResponse>().await.map_err(|e| {
122            OAuthError::token_refresh_failed(format!("Failed to parse token response: {}", e))
123        })
124    }
125
126    /// Get the PKCE verifier (for validation purposes)
127    pub fn pkce_verifier(&self) -> Option<&str> {
128        self.pkce.as_ref().map(|p| p.verifier.as_str())
129    }
130}
131
132/// Parse the authorization code from Anthropic's callback format
133///
134/// Anthropic returns codes in the format: "authorization_code#state"
135fn parse_auth_code(code: &str) -> OAuthResult<(String, String)> {
136    // Handle both "#" and "%23" (URL-encoded #)
137    let code = code.replace("%23", "#");
138
139    if let Some(pos) = code.find('#') {
140        let auth_code = code[..pos].to_string();
141        let state = code[pos + 1..].to_string();
142
143        if auth_code.is_empty() || state.is_empty() {
144            return Err(OAuthError::invalid_code_format(
145                "Authorization code or state is empty",
146            ));
147        }
148
149        Ok((auth_code, state))
150    } else {
151        Err(OAuthError::invalid_code_format(
152            "Expected format: authorization_code#state",
153        ))
154    }
155}
156
157#[cfg(test)]
158mod tests {
159    use super::*;
160
161    fn test_config() -> OAuthConfig {
162        OAuthConfig::new(
163            "test-client-id",
164            "https://example.com/auth",
165            "https://example.com/token",
166            "https://example.com/callback",
167            vec!["scope1".to_string(), "scope2".to_string()],
168        )
169    }
170
171    #[test]
172    fn test_generate_auth_url() {
173        let mut flow = OAuthFlow::new(test_config());
174        let url = flow.generate_auth_url();
175
176        assert!(url.starts_with("https://example.com/auth?"));
177        assert!(url.contains("client_id=test-client-id"));
178        assert!(url.contains("response_type=code"));
179        assert!(url.contains("redirect_uri="));
180        assert!(url.contains("scope=scope1%20scope2"));
181        assert!(url.contains("code_challenge="));
182        assert!(url.contains("code_challenge_method=S256"));
183        assert!(url.contains("state="));
184
185        // PKCE should be initialized
186        assert!(flow.pkce.is_some());
187    }
188
189    #[test]
190    fn test_parse_auth_code_valid() {
191        let result = parse_auth_code("abc123#verifier456");
192        assert!(result.is_ok());
193        let (code, state) = result.unwrap();
194        assert_eq!(code, "abc123");
195        assert_eq!(state, "verifier456");
196    }
197
198    #[test]
199    fn test_parse_auth_code_url_encoded() {
200        let result = parse_auth_code("abc123%23verifier456");
201        assert!(result.is_ok());
202        let (code, state) = result.unwrap();
203        assert_eq!(code, "abc123");
204        assert_eq!(state, "verifier456");
205    }
206
207    #[test]
208    fn test_parse_auth_code_missing_separator() {
209        let result = parse_auth_code("abc123verifier456");
210        assert!(result.is_err());
211    }
212
213    #[test]
214    fn test_parse_auth_code_empty_parts() {
215        assert!(parse_auth_code("#state").is_err());
216        assert!(parse_auth_code("code#").is_err());
217        assert!(parse_auth_code("#").is_err());
218    }
219
220    #[test]
221    fn test_exchange_code_without_pkce() {
222        let flow = OAuthFlow::new(test_config());
223        let result = tokio_test::block_on(flow.exchange_code("code#state"));
224        assert!(matches!(result, Err(OAuthError::PkceNotInitialized)));
225    }
226
227    #[test]
228    fn test_token_response_serde() {
229        let json = r#"{
230            "access_token": "access123",
231            "refresh_token": "refresh456",
232            "expires_in": 3600,
233            "token_type": "Bearer"
234        }"#;
235
236        let response: TokenResponse = serde_json::from_str(json).unwrap();
237        assert_eq!(response.access_token, "access123");
238        assert_eq!(response.refresh_token, "refresh456");
239        assert_eq!(response.expires_in, 3600);
240        assert_eq!(response.token_type, "Bearer");
241    }
242}
stakpak_shared/oauth/flow.rs

stakpak_shared/oauth/
flow.rs
