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        // Always generated, and always handed to the provider. A provider that
129        // has no use for it ignores it; one that validates an ID token needs it
130        // regardless of whether it remembered to advertise that. Gating this
131        // instead of the check below would let an out-of-tree OIDC provider
132        // opt out of replay protection just by not overriding a default.
133        let nonce = Some(uuid::Uuid::new_v4().to_string());
134
135        let effective_scopes = if !scopes.is_empty() {
136            scopes
137        } else {
138            &self
139                .scopes
140                .iter()
141                .map(|s| s.as_str())
142                .collect::<Vec<&str>>()
143        };
144
145        tracing::debug!(scopes = ?effective_scopes, "generating authorization URL");
146
147        let url = self.provider.get_authorization_url(
148            &state,
149            effective_scopes,
150            pkce_challenge,
151            nonce.as_deref(),
152        );
153
154        let auth_state = OAuth2State {
155            state: state.clone(),
156            nonce,
157            code_verifier: None, // Will be set by the caller if needed before encryption
158            success_url: None,
159            provider_id: self.provider.provider_id().to_string(),
160            expires_at: chrono::Utc::now().timestamp() + 600,
161        };
162
163        tracing::info!("authorization login initiated successfully");
164        (url, auth_state)
165    }
166
167    /// Completes the flow by exchanging the code.
168    /// If a mapper is provided, it will also map the identity to a local user.
169    #[tracing::instrument(skip(self, code, expected_state), fields(provider_id = %self.provider.provider_id()))]
170    pub async fn finalize_login(
171        &self,
172        code: &str,
173        received_state: &str,
174        expected_state: &OAuth2State,
175    ) -> Result<(Identity, OAuthToken, Option<M::LocalUser>), AuthError> {
176        if received_state != expected_state.state {
177            tracing::error!("CSRF mismatch: received state does not match expected state");
178            return Err(AuthError::CsrfMismatch);
179        }
180
181        tracing::debug!("exchanging code for identity");
182        let (identity, token) = self
183            .provider
184            .exchange_code_for_identity(
185                code,
186                expected_state.code_verifier.as_deref(),
187                expected_state.nonce.as_deref(),
188            )
189            .await
190            .map_err(|e| {
191                tracing::error!(error = %e, "failed to exchange code for identity");
192                e
193            })?;
194
195        tracing::info!(user_id = %identity.external_id, "successfully retrieved identity from provider");
196
197        // Re-check the nonce only for a provider that surfaces it.
198        //
199        // This is a second line of defence: a provider that validates an ID
200        // token has already compared the nonce itself. Plain OAuth2 has no ID
201        // token, so there is nothing to surface and nothing to compare —
202        // enforcing it there rejected every login, because
203        // `identity.attributes` is necessarily empty of a nonce.
204        //
205        // The gate is on the check rather than on generation so that a
206        // provider which forgets to advertise itself loses only this redundant
207        // re-check, never the nonce itself.
208        if self.provider.validates_nonce() {
209            if let Some(expected_nonce) = expected_state.nonce.as_deref() {
210                if identity.attributes.get("nonce").map(|s| s.as_str()) != Some(expected_nonce) {
211                    tracing::error!("nonce mismatch or missing in identity attributes");
212                    return Err(AuthError::Token("Nonce mismatch".to_string()));
213                }
214            }
215        }
216
217        let local_user = if let Some(mapper) = &self.mapper {
218            tracing::debug!("mapping user identity");
219            Some(mapper.map_user(&identity).await.map_err(|e| {
220                tracing::error!(error = %e, "failed to map user");
221                e
222            })?)
223        } else {
224            None
225        };
226
227        Ok((identity, token, local_user))
228    }
229
230    /// Refresh an access token using a refresh token.
231    pub async fn refresh_access_token(&self, refresh_token: &str) -> Result<OAuthToken, AuthError> {
232        self.provider.refresh_token(refresh_token).await
233    }
234
235    /// Revoke an access token.
236    pub async fn revoke_token(&self, token: &str) -> Result<(), AuthError> {
237        self.provider.revoke_token(token).await
238    }
239}
240
241#[cfg(test)]
242mod tests {
243    use super::*;
244    use crate::auth::{Provider, ProviderConfig};
245    use async_trait::async_trait;
246    use std::collections::HashMap;
247
248    struct MockProvider {
249        id: String,
250        auth_url: String,
251        expected_code: String,
252        identity: Identity,
253        token: OAuthToken,
254    }
255
256    #[async_trait]
257    impl Provider for MockProvider {
258        async fn config(&self) -> ProviderConfig {
259            ProviderConfig {
260                id: self.id.clone(),
261                name: self.id.clone(),
262                extra: HashMap::new(),
263            }
264        }
265    }
266
267    #[async_trait]
268    impl OAuthProvider for MockProvider {
269        fn provider_id(&self) -> &str {
270            &self.id
271        }
272
273        fn get_authorization_url(
274            &self,
275            state: &str,
276            _scopes: &[&str],
277            _code_challenge: Option<&str>,
278            _nonce: Option<&str>,
279        ) -> String {
280            format!("{}?state={}", self.auth_url, state)
281        }
282
283        async fn exchange_code_for_identity(
284            &self,
285            code: &str,
286            _code_verifier: Option<&str>,
287            _nonce: Option<&str>,
288        ) -> Result<(Identity, OAuthToken), AuthError> {
289            if code == self.expected_code {
290                Ok((self.identity.clone(), self.token.clone()))
291            } else {
292                Err(AuthError::Token("Invalid code".to_string()))
293            }
294        }
295    }
296
297    #[tokio::test]
298    async fn test_oauth2_flow_initiate() {
299        let provider = MockProvider {
300            id: "mock".to_string(),
301            auth_url: "http://mock/auth".to_string(),
302            expected_code: "code123".to_string(),
303            identity: Identity {
304                provider_id: "mock".to_string(),
305                external_id: "1".to_string(),
306                email: None,
307                username: None,
308                attributes: HashMap::new(),
309            },
310            token: OAuthToken {
311                access_token: "acc".to_string(),
312                token_type: "Bearer".to_string(),
313                expires_in: None,
314                refresh_token: None,
315                scope: None,
316                id_token: None,
317            },
318        };
319
320        let flow = OAuth2Flow::new(provider).with_scopes(vec!["scope1"]);
321        let (url, state) = flow.initiate_login(&["scope2"], None);
322        assert!(url.contains("http://mock/auth?state="));
323        assert_eq!(state.provider_id, "mock");
324    }
325
326    #[tokio::test]
327    async fn test_oauth2_flow_finalize_success() {
328        let provider = MockProvider {
329            id: "mock".to_string(),
330            auth_url: "http://mock/auth".to_string(),
331            expected_code: "code123".to_string(),
332            identity: Identity {
333                provider_id: "mock".to_string(),
334                external_id: "1".to_string(),
335                email: None,
336                username: None,
337                attributes: HashMap::new(),
338            },
339            token: OAuthToken {
340                access_token: "acc".to_string(),
341                token_type: "Bearer".to_string(),
342                expires_in: None,
343                refresh_token: None,
344                scope: None,
345                id_token: None,
346            },
347        };
348
349        let flow = OAuth2Flow::new(provider);
350        let expected_state = OAuth2State {
351            state: "state123".to_string(),
352            nonce: None,
353            code_verifier: None,
354            success_url: None,
355            provider_id: "mock".to_string(),
356            expires_at: 0,
357        };
358
359        let (ident, tok, _) = flow
360            .finalize_login("code123", "state123", &expected_state)
361            .await
362            .unwrap();
363        assert_eq!(ident.external_id, "1");
364        assert_eq!(tok.access_token, "acc");
365    }
366
367    #[tokio::test]
368    async fn test_oauth2_flow_finalize_csrf_mismatch() {
369        let provider = MockProvider {
370            id: "mock".to_string(),
371            auth_url: "http://mock/auth".to_string(),
372            expected_code: "code123".to_string(),
373            identity: Identity {
374                provider_id: "mock".to_string(),
375                external_id: "1".to_string(),
376                email: None,
377                username: None,
378                attributes: HashMap::new(),
379            },
380            token: OAuthToken {
381                access_token: "acc".to_string(),
382                token_type: "Bearer".to_string(),
383                expires_in: None,
384                refresh_token: None,
385                scope: None,
386                id_token: None,
387            },
388        };
389
390        let flow = OAuth2Flow::new(provider);
391        let expected_state = OAuth2State {
392            state: "state123".to_string(),
393            nonce: None,
394            code_verifier: None,
395            success_url: None,
396            provider_id: "mock".to_string(),
397            expires_at: 0,
398        };
399
400        let result = flow
401            .finalize_login("code123", "wrong_state", &expected_state)
402            .await;
403        assert!(matches!(result, Err(AuthError::CsrfMismatch)));
404    }
405
406    fn mock(id: &str) -> MockProvider {
407        MockProvider {
408            id: id.to_string(),
409            auth_url: "http://mock/auth".to_string(),
410            expected_code: "code123".to_string(),
411            identity: Identity {
412                provider_id: id.to_string(),
413                external_id: "1".to_string(),
414                email: None,
415                username: None,
416                attributes: HashMap::new(),
417            },
418            token: OAuthToken {
419                access_token: "acc".to_string(),
420                token_type: "Bearer".to_string(),
421                expires_in: None,
422                refresh_token: None,
423                scope: None,
424                id_token: None,
425            },
426        }
427    }
428
429    /// A provider that behaves like an OIDC one: it surfaces the nonce it was
430    /// given in `Identity::attributes`, which is what the engine re-checks.
431    struct NonceEchoing {
432        inner: MockProvider,
433        /// What to surface. `None` means "echo whatever was received".
434        echo_instead: Option<String>,
435    }
436
437    #[async_trait]
438    impl Provider for NonceEchoing {
439        async fn config(&self) -> ProviderConfig {
440            self.inner.config().await
441        }
442    }
443
444    #[async_trait]
445    impl OAuthProvider for NonceEchoing {
446        fn validates_nonce(&self) -> bool {
447            true
448        }
449
450        fn provider_id(&self) -> &str {
451            self.inner.provider_id()
452        }
453
454        fn get_authorization_url(
455            &self,
456            state: &str,
457            scopes: &[&str],
458            code_challenge: Option<&str>,
459            nonce: Option<&str>,
460        ) -> String {
461            self.inner
462                .get_authorization_url(state, scopes, code_challenge, nonce)
463        }
464
465        async fn exchange_code_for_identity(
466            &self,
467            code: &str,
468            code_verifier: Option<&str>,
469            nonce: Option<&str>,
470        ) -> Result<(Identity, OAuthToken), AuthError> {
471            let (mut identity, token) = self
472                .inner
473                .exchange_code_for_identity(code, code_verifier, nonce)
474                .await?;
475            let surfaced = self
476                .echo_instead
477                .clone()
478                .or_else(|| nonce.map(|n| n.to_string()));
479            if let Some(value) = surfaced {
480                identity.attributes.insert("nonce".to_string(), value);
481            }
482            Ok((identity, token))
483        }
484    }
485
486    /// The regression this fix is for: a plain OAuth2 provider returns an
487    /// identity with no nonce attribute, because it has no ID token to carry
488    /// one. Enforcing the re-check there rejected every login.
489    #[tokio::test]
490    async fn a_plain_provider_can_complete_a_login() {
491        let flow = OAuth2Flow::new(mock("plain"));
492        let (_url, state) = flow.initiate_login(&["email"], None);
493
494        let result = flow.finalize_login("code123", &state.state, &state).await;
495
496        assert!(
497            result.is_ok(),
498            "a plain OAuth2 login should complete; got {:?}",
499            result.err()
500        );
501    }
502
503    /// Every provider is still handed a nonce, whatever the flag says.
504    ///
505    /// This is why the gate is on the check and not on generation: an
506    /// out-of-tree OIDC provider that never overrides `validates_nonce` still
507    /// receives a nonce to validate, so it cannot lose replay protection by
508    /// omission.
509    #[test]
510    fn every_provider_still_receives_a_nonce() {
511        let plain = OAuth2Flow::new(mock("plain"));
512        let (_url, state) = plain.initiate_login(&["email"], None);
513        assert!(
514            state.nonce.is_some(),
515            "a provider that does not advertise itself must still get a nonce"
516        );
517    }
518
519    /// And the re-check still bites where it applies: a provider that claims to
520    /// surface the nonce but surfaces the wrong one must be rejected.
521    #[tokio::test]
522    async fn a_wrong_nonce_is_still_rejected() {
523        let flow = OAuth2Flow::new(NonceEchoing {
524            inner: mock("oidc"),
525            echo_instead: Some("not-the-nonce".to_string()),
526        });
527        let (_url, state) = flow.initiate_login(&["email"], None);
528
529        let result = flow.finalize_login("code123", &state.state, &state).await;
530
531        assert!(
532            result.is_err(),
533            "gating the check must not disable it where a provider opts in"
534        );
535    }
536
537    /// The happy path for such a provider.
538    #[tokio::test]
539    async fn a_matching_nonce_is_accepted() {
540        let flow = OAuth2Flow::new(NonceEchoing {
541            inner: mock("oidc"),
542            echo_instead: None, // echo what was received
543        });
544        let (_url, state) = flow.initiate_login(&["email"], None);
545
546        let result = flow.finalize_login("code123", &state.state, &state).await;
547
548        assert!(
549            result.is_ok(),
550            "a correctly echoed nonce should pass; got {:?}",
551            result.err()
552        );
553    }
554}
authkestra_engine/flow/oauth2.rs

authkestra_engine/flow/
oauth2.rs
