authkestra_engine/flow/
oauth2.rs

1use crate::auth::{
2    error::AuthError, state::Identity, state::OAuth2State, state::OAuthToken, ErasedOAuthFlow,
3    OAuthProvider, UserMapper,
4};
5use crate::flow::{Flow, FlowContext, FlowResult};
6use async_trait::async_trait;
7
8/// Orchestrates the standard OAuth2 Authorization Code flow.
9pub struct OAuth2Flow<P: OAuthProvider, M: UserMapper = ()> {
10    provider: P,
11    mapper: Option<M>,
12    scopes: Vec<String>,
13    use_pkce: bool,
14}
15
16#[async_trait]
17impl<P: OAuthProvider + 'static, M: UserMapper + 'static> Flow for OAuth2Flow<P, M> {
18    fn id(&self) -> &str {
19        self.provider.provider_id()
20    }
21
22    async fn execute(&self, ctx: FlowContext) -> Result<FlowResult, AuthError> {
23        if let Some(_code) = ctx.params.get("code") {
24            let _received_state = ctx.params.get("state").ok_or(AuthError::CsrfMismatch)?;
25
26            // In the new model, expected_state must be provided via some context.
27            // For now, if it's missing from ctx, we might need to adjust FlowContext.
28            // But ErasedOAuthFlow is what the adapters use.
29            Err(AuthError::Token(
30                "Direct Flow execution not updated for encrypted state".to_string(),
31            ))
32        } else {
33            // Assume initiation if no code is present
34            let scopes_str = ctx.params.get("scopes").map(|s| s.as_str()).unwrap_or("");
35            let scopes_vec: Vec<&str> = if scopes_str.is_empty() {
36                Vec::new()
37            } else {
38                scopes_str.split(',').collect()
39            };
40
41            let pkce_challenge = ctx.params.get("pkce_challenge").map(|s| s.as_str());
42            let (url, _state) = self.initiate_login(&scopes_vec, pkce_challenge);
43            Ok(FlowResult::Redirect(url))
44        }
45    }
46}
47
48#[async_trait]
49impl<P: OAuthProvider + 'static, M: UserMapper + 'static> ErasedOAuthFlow for OAuth2Flow<P, M> {
50    fn provider_id(&self) -> String {
51        self.provider.provider_id().to_string()
52    }
53
54    fn initiate_login(
55        &self,
56        scopes: &[&str],
57        pkce_challenge: Option<&str>,
58    ) -> (String, OAuth2State) {
59        let effective_scopes = if !scopes.is_empty() {
60            scopes
61        } else {
62            &self
63                .scopes
64                .iter()
65                .map(|s| s.as_str())
66                .collect::<Vec<&str>>()
67        };
68
69        self.initiate_login(effective_scopes, pkce_challenge)
70    }
71
72    async fn finalize_login(
73        &self,
74        code: &str,
75        received_state: &str,
76        expected_state: &OAuth2State,
77    ) -> Result<(Identity, OAuthToken), AuthError> {
78        let (identity, token, _) = self
79            .finalize_login(code, received_state, expected_state)
80            .await?;
81        Ok((identity, token))
82    }
83}
84
85impl<P: OAuthProvider> OAuth2Flow<P, ()> {
86    /// Create a new `OAuth2Flow` with the given provider.
87    pub fn new(provider: P) -> Self {
88        Self {
89            provider,
90            mapper: None,
91            scopes: Vec::new(),
92            use_pkce: true,
93        }
94    }
95}
96
97impl<P: OAuthProvider, M: UserMapper> OAuth2Flow<P, M> {
98    /// Create a new `OAuth2Flow` with the given provider and user mapper.
99    pub fn with_mapper(provider: P, mapper: M) -> Self {
100        Self {
101            provider,
102            mapper: Some(mapper),
103            scopes: Vec::new(),
104            use_pkce: true,
105        }
106    }
107
108    /// Set the scopes for the OAuth2 flow.
109    pub fn with_scopes(mut self, scopes: Vec<impl Into<String>>) -> Self {
110        self.scopes = scopes.into_iter().map(|s| s.into()).collect();
111        self
112    }
113
114    /// Enable or disable PKCE for the OAuth2 flow.
115    pub fn with_pkce(mut self, use_pkce: bool) -> Self {
116        self.use_pkce = use_pkce;
117        self
118    }
119
120    /// Generates the redirect URL and CSRF state.
121    #[tracing::instrument(skip(self), fields(provider_id = %self.provider.provider_id()))]
122    pub fn initiate_login(
123        &self,
124        scopes: &[&str],
125        pkce_challenge: Option<&str>,
126    ) -> (String, OAuth2State) {
127        let state = uuid::Uuid::new_v4().to_string();
128        let nonce = Some(uuid::Uuid::new_v4().to_string());
129
130        let effective_scopes = if !scopes.is_empty() {
131            scopes
132        } else {
133            &self
134                .scopes
135                .iter()
136                .map(|s| s.as_str())
137                .collect::<Vec<&str>>()
138        };
139
140        tracing::debug!(scopes = ?effective_scopes, "generating authorization URL");
141
142        let url = self.provider.get_authorization_url(
143            &state,
144            effective_scopes,
145            pkce_challenge,
146            nonce.as_deref(),
147        );
148
149        let auth_state = OAuth2State {
150            state: state.clone(),
151            nonce,
152            code_verifier: None, // Will be set by the caller if needed before encryption
153            success_url: None,
154            provider_id: self.provider.provider_id().to_string(),
155            expires_at: chrono::Utc::now().timestamp() + 600,
156        };
157
158        tracing::info!("authorization login initiated successfully");
159        (url, auth_state)
160    }
161
162    /// Completes the flow by exchanging the code.
163    /// If a mapper is provided, it will also map the identity to a local user.
164    #[tracing::instrument(skip(self, code, expected_state), fields(provider_id = %self.provider.provider_id()))]
165    pub async fn finalize_login(
166        &self,
167        code: &str,
168        received_state: &str,
169        expected_state: &OAuth2State,
170    ) -> Result<(Identity, OAuthToken, Option<M::LocalUser>), AuthError> {
171        if received_state != expected_state.state {
172            tracing::error!("CSRF mismatch: received state does not match expected state");
173            return Err(AuthError::CsrfMismatch);
174        }
175
176        tracing::debug!("exchanging code for identity");
177        let (identity, token) = self
178            .provider
179            .exchange_code_for_identity(
180                code,
181                expected_state.code_verifier.as_deref(),
182                expected_state.nonce.as_deref(),
183            )
184            .await
185            .map_err(|e| {
186                tracing::error!(error = %e, "failed to exchange code for identity");
187                e
188            })?;
189
190        tracing::info!(user_id = %identity.external_id, "successfully retrieved identity from provider");
191
192        if let Some(expected_nonce) = expected_state.nonce.as_deref() {
193            if identity.attributes.get("nonce").map(|s| s.as_str()) != Some(expected_nonce) {
194                tracing::error!("nonce mismatch or missing in identity attributes");
195                return Err(AuthError::Token("Nonce mismatch".to_string()));
196            }
197        }
198
199        let local_user = if let Some(mapper) = &self.mapper {
200            tracing::debug!("mapping user identity");
201            Some(mapper.map_user(&identity).await.map_err(|e| {
202                tracing::error!(error = %e, "failed to map user");
203                e
204            })?)
205        } else {
206            None
207        };
208
209        Ok((identity, token, local_user))
210    }
211
212    /// Refresh an access token using a refresh token.
213    pub async fn refresh_access_token(&self, refresh_token: &str) -> Result<OAuthToken, AuthError> {
214        self.provider.refresh_token(refresh_token).await
215    }
216
217    /// Revoke an access token.
218    pub async fn revoke_token(&self, token: &str) -> Result<(), AuthError> {
219        self.provider.revoke_token(token).await
220    }
221}
222
223#[cfg(test)]
224mod tests {
225    use super::*;
226    use crate::auth::{Provider, ProviderConfig};
227    use async_trait::async_trait;
228    use std::collections::HashMap;
229
230    struct MockProvider {
231        id: String,
232        auth_url: String,
233        expected_code: String,
234        identity: Identity,
235        token: OAuthToken,
236    }
237
238    #[async_trait]
239    impl Provider for MockProvider {
240        async fn config(&self) -> ProviderConfig {
241            ProviderConfig {
242                id: self.id.clone(),
243                name: self.id.clone(),
244                extra: HashMap::new(),
245            }
246        }
247    }
248
249    #[async_trait]
250    impl OAuthProvider for MockProvider {
251        fn provider_id(&self) -> &str {
252            &self.id
253        }
254
255        fn get_authorization_url(
256            &self,
257            state: &str,
258            _scopes: &[&str],
259            _code_challenge: Option<&str>,
260            _nonce: Option<&str>,
261        ) -> String {
262            format!("{}?state={}", self.auth_url, state)
263        }
264
265        async fn exchange_code_for_identity(
266            &self,
267            code: &str,
268            _code_verifier: Option<&str>,
269            _nonce: Option<&str>,
270        ) -> Result<(Identity, OAuthToken), AuthError> {
271            if code == self.expected_code {
272                Ok((self.identity.clone(), self.token.clone()))
273            } else {
274                Err(AuthError::Token("Invalid code".to_string()))
275            }
276        }
277    }
278
279    #[tokio::test]
280    async fn test_oauth2_flow_initiate() {
281        let provider = MockProvider {
282            id: "mock".to_string(),
283            auth_url: "http://mock/auth".to_string(),
284            expected_code: "code123".to_string(),
285            identity: Identity {
286                provider_id: "mock".to_string(),
287                external_id: "1".to_string(),
288                email: None,
289                username: None,
290                attributes: HashMap::new(),
291            },
292            token: OAuthToken {
293                access_token: "acc".to_string(),
294                token_type: "Bearer".to_string(),
295                expires_in: None,
296                refresh_token: None,
297                scope: None,
298                id_token: None,
299            },
300        };
301
302        let flow = OAuth2Flow::new(provider).with_scopes(vec!["scope1"]);
303        let (url, state) = flow.initiate_login(&["scope2"], None);
304        assert!(url.contains("http://mock/auth?state="));
305        assert_eq!(state.provider_id, "mock");
306    }
307
308    #[tokio::test]
309    async fn test_oauth2_flow_finalize_success() {
310        let provider = MockProvider {
311            id: "mock".to_string(),
312            auth_url: "http://mock/auth".to_string(),
313            expected_code: "code123".to_string(),
314            identity: Identity {
315                provider_id: "mock".to_string(),
316                external_id: "1".to_string(),
317                email: None,
318                username: None,
319                attributes: HashMap::new(),
320            },
321            token: OAuthToken {
322                access_token: "acc".to_string(),
323                token_type: "Bearer".to_string(),
324                expires_in: None,
325                refresh_token: None,
326                scope: None,
327                id_token: None,
328            },
329        };
330
331        let flow = OAuth2Flow::new(provider);
332        let expected_state = OAuth2State {
333            state: "state123".to_string(),
334            nonce: None,
335            code_verifier: None,
336            success_url: None,
337            provider_id: "mock".to_string(),
338            expires_at: 0,
339        };
340
341        let (ident, tok, _) = flow
342            .finalize_login("code123", "state123", &expected_state)
343            .await
344            .unwrap();
345        assert_eq!(ident.external_id, "1");
346        assert_eq!(tok.access_token, "acc");
347    }
348
349    #[tokio::test]
350    async fn test_oauth2_flow_finalize_csrf_mismatch() {
351        let provider = MockProvider {
352            id: "mock".to_string(),
353            auth_url: "http://mock/auth".to_string(),
354            expected_code: "code123".to_string(),
355            identity: Identity {
356                provider_id: "mock".to_string(),
357                external_id: "1".to_string(),
358                email: None,
359                username: None,
360                attributes: HashMap::new(),
361            },
362            token: OAuthToken {
363                access_token: "acc".to_string(),
364                token_type: "Bearer".to_string(),
365                expires_in: None,
366                refresh_token: None,
367                scope: None,
368                id_token: None,
369            },
370        };
371
372        let flow = OAuth2Flow::new(provider);
373        let expected_state = OAuth2State {
374            state: "state123".to_string(),
375            nonce: None,
376            code_verifier: None,
377            success_url: None,
378            provider_id: "mock".to_string(),
379            expires_at: 0,
380        };
381
382        let result = flow
383            .finalize_login("code123", "wrong_state", &expected_state)
384            .await;
385        assert!(matches!(result, Err(AuthError::CsrfMismatch)));
386    }
387}
authkestra_engine/flow/oauth2.rs

authkestra_engine/flow/
oauth2.rs
