Skip to main content

axum_security/oidc/
context.rs

1use std::{borrow::Cow, convert::Infallible, sync::Arc};
2
3use axum::{
4    extract::{FromRef, FromRequestParts},
5    http::{StatusCode, request::Parts},
6    response::{IntoResponse, Redirect},
7};
8
9use axum_security_oidc::{CsrfToken, LogoutUrl, OidcClient, OidcError, VerifyError};
10use cookie_monster::{CookieBuilder, CookieJar};
11use http::Extensions;
12use url::Url;
13
14use crate::{
15    after_login::AfterLoginCookies,
16    oidc::{OidcBuilderError, builder::OidcContextBuilder},
17};
18
19use super::{OidcHandler, OidcTokenResponse, cookie::OidcCookie};
20
21/// OIDC context that manages the login/logout flow.
22///
23/// Construct with [`OidcContext::discover`] (auto-discovery) or
24/// [`OidcContext::builder`] (manual endpoints). Provider shortcuts:
25/// [`google`](OidcContext::google), [`microsoft`](OidcContext::microsoft),
26/// [`apple`](OidcContext::apple), [`keycloak`](OidcContext::keycloak).
27///
28/// Register routes with [`OidcExt::with_oidc`](super::OidcExt::with_oidc).
29pub struct OidcContext<H>(pub(super) Arc<OidcContextInner<H>>);
30
31pub(super) struct OidcContextInner<H> {
32    pub(super) handler: H,
33    pub(super) session: OidcCookie,
34    pub(super) client: OidcClient,
35    pub(super) login_path: Option<Cow<'static, str>>,
36    pub(super) logout_path: Option<Cow<'static, str>>,
37    pub(super) post_logout_redirect_url: Option<String>,
38}
39
40impl OidcContext<()> {
41    pub fn builder(provider_name: impl Into<Cow<'static, str>>) -> OidcContextBuilder {
42        OidcContextBuilder::new(provider_name.into())
43    }
44
45    pub async fn discover(
46        provider_name: impl Into<Cow<'static, str>>,
47        issuer_url: &str,
48    ) -> Result<OidcContextBuilder, OidcBuilderError> {
49        OidcContextBuilder::discover(provider_name.into(), issuer_url).await
50    }
51
52    pub async fn google() -> Result<OidcContextBuilder, OidcBuilderError> {
53        Self::discover("google", super::providers::google::ISSUER_URL).await
54    }
55
56    pub async fn microsoft() -> Result<OidcContextBuilder, OidcBuilderError> {
57        Self::discover("microsoft", super::providers::microsoft::ISSUER_URL_COMMON).await
58    }
59
60    pub async fn apple() -> Result<OidcContextBuilder, OidcBuilderError> {
61        Self::discover("apple", super::providers::apple::ISSUER_URL).await
62    }
63
64    pub async fn keycloak(
65        base_url: &str,
66        realm: &str,
67    ) -> Result<OidcContextBuilder, OidcBuilderError> {
68        let issuer_url = format!("{}/realms/{}", base_url.trim_end_matches('/'), realm);
69        Self::discover("keycloak", &issuer_url).await
70    }
71}
72
73impl<H: OidcHandler> OidcContext<H> {
74    pub(crate) fn callback_url(&self) -> &str {
75        self.0
76            .client
77            .redirect_url()
78            .expect("redirect_uri must be set")
79            .path()
80    }
81
82    pub(crate) fn get_start_challenge_path(&self) -> Option<&str> {
83        self.0.login_path.as_deref()
84    }
85
86    pub(crate) async fn start_challenge(&self) -> axum::response::Response {
87        crate::debug!("Starting OIDC login flow");
88
89        let login = self.0.client.start_login();
90        let cookie = self.0.session.generate_cookie(
91            login.csrf_token.as_str(),
92            &login.pkce_verifier,
93            &login.nonce,
94        );
95
96        (cookie, Redirect::to(login.url.as_str())).into_response()
97    }
98
99    pub(crate) async fn on_redirect(
100        &self,
101        mut jar: CookieJar,
102        code: String,
103        state: String,
104    ) -> axum::response::Response {
105        crate::debug!("handling OIDC redirect");
106
107        let Some((csrf_token, pkce_verifier, nonce)) = self.0.session.verify_cookies(&mut jar)
108        else {
109            return StatusCode::UNAUTHORIZED.into_response();
110        };
111
112        // Wrap the cookie's token so the comparison against the
113        // attacker-controlled `state` runs in constant time.
114        let csrf_token = CsrfToken::from(csrf_token);
115        if csrf_token != state {
116            crate::debug!("state does not match");
117            return StatusCode::UNAUTHORIZED.into_response();
118        }
119
120        crate::debug!("exchanging authorization code for tokens");
121
122        let tokens = match self
123            .0
124            .client
125            .finish_login(&code, &pkce_verifier, &nonce)
126            .await
127        {
128            Ok(tokens) => tokens,
129            // A verification failure (bad signature, nonce, expiry, unknown key)
130            // is an auth failure; an exchange/no-id_token failure is a 500.
131            Err(OidcError::Verify(_e)) => {
132                crate::debug!("id_token verification failed: {_e}");
133                return StatusCode::UNAUTHORIZED.into_response();
134            }
135            Err(_e) => {
136                crate::debug!("OIDC login failed: {_e}");
137                return StatusCode::INTERNAL_SERVER_ERROR.into_response();
138            }
139        };
140
141        let claims = match tokens.claims() {
142            Ok(claims) => claims,
143            Err(_e) => {
144                crate::debug!("failed to deserialize claims: {_e}");
145                return StatusCode::INTERNAL_SERVER_ERROR.into_response();
146            }
147        };
148
149        let oidc_response = OidcTokenResponse {
150            id_token: tokens.id_token(),
151            claims,
152            access_token: tokens.access_token().to_string(),
153            refresh_token: tokens.refresh_token().map(String::from),
154        };
155
156        let mut context = AfterLoginCookies {
157            cookie_jar: jar,
158            cookie_opts: &self.0.session.cookie_builder,
159        };
160
161        crate::debug!("OIDC login flow done");
162        let res = self
163            .0
164            .handler
165            .after_login(oidc_response, &mut context)
166            .await
167            .into_response();
168
169        (context.cookie_jar, res).into_response()
170    }
171
172    pub fn cookie(&self, name: impl Into<Cow<'static, str>>) -> CookieBuilder {
173        self.0.session.cookie_builder.clone().name(name.into())
174    }
175
176    /// Fetch and cache the provider's JWKS now instead of lazily on the first
177    /// login. [`OidcContext`] is cheap to clone, so warm it in the background at
178    /// startup:
179    ///
180    /// ```no_run
181    /// # async fn f<H: axum_security::oidc::OidcHandler + 'static>(
182    /// #     oidc: axum_security::oidc::OidcContext<H>,
183    /// # ) {
184    /// let oidc = oidc.clone();
185    /// tokio::spawn(async move {
186    ///     if let Err(e) = oidc.warm_jwks().await {
187    ///         eprintln!("failed to prefetch OIDC keys: {e}");
188    ///     }
189    /// });
190    /// # }
191    /// ```
192    pub async fn warm_jwks(&self) -> Result<(), VerifyError> {
193        self.0.client.warm_jwks().await
194    }
195
196    pub(crate) fn get_logout_path(&self) -> Option<&str> {
197        self.0.logout_path.as_deref()
198    }
199
200    pub(crate) fn build_logout_context(&self, extensions: Extensions) -> LogoutContext {
201        LogoutContext {
202            extensions,
203            end_session_endpoint: self.0.client.end_session_endpoint().cloned(),
204            post_logout_redirect_url: self.0.post_logout_redirect_url.clone(),
205            id_token_hint: None,
206            logout_hint: None,
207            client_id: Some(self.0.client.client_id().to_string()),
208            state: None,
209        }
210    }
211}
212
213/// Context passed to [`OidcHandler::logout`](super::OidcHandler::logout).
214///
215/// Provides access to request extensions and methods to customize the
216/// logout redirect (ID token hint, logout hint, post-logout redirect URI).
217/// Call [`default_redirect`](LogoutContext::default_redirect) to build the
218/// redirect response using the configured end-session endpoint.
219pub struct LogoutContext {
220    extensions: Extensions,
221    end_session_endpoint: Option<Url>,
222    post_logout_redirect_url: Option<String>,
223    id_token_hint: Option<String>,
224    logout_hint: Option<String>,
225    client_id: Option<String>,
226    state: Option<String>,
227}
228
229impl LogoutContext {
230    #[cfg(feature = "cookie")]
231    pub fn cookie_session<U: Send + Sync + 'static>(
232        &mut self,
233    ) -> Option<crate::cookie::CookieSession<U>> {
234        use crate::cookie::CookieSession;
235
236        CookieSession::from_extensions(&mut self.extensions)
237    }
238
239    pub fn extensions(&self) -> &Extensions {
240        &self.extensions
241    }
242
243    pub fn extensions_mut(&mut self) -> &mut Extensions {
244        &mut self.extensions
245    }
246
247    pub fn set_id_token_hint(&mut self, id_token_hint: impl Into<String>) {
248        self.id_token_hint = Some(id_token_hint.into());
249    }
250
251    pub fn set_logout_hint(&mut self, logout_hint: impl Into<String>) {
252        self.logout_hint = Some(logout_hint.into());
253    }
254
255    pub fn set_client_id(&mut self, client_id: impl Into<String>) {
256        self.client_id = Some(client_id.into());
257    }
258
259    pub fn set_post_logout_redirect_uri(&mut self, post_logout_redirect_uri: impl Into<String>) {
260        self.post_logout_redirect_url = Some(post_logout_redirect_uri.into());
261    }
262
263    pub fn set_state(&mut self, state: impl Into<String>) {
264        self.state = Some(state.into());
265    }
266
267    pub fn default_redirect(self) -> Redirect {
268        match self.end_session_endpoint {
269            Some(endpoint) => {
270                let mut logout = LogoutUrl::new(endpoint);
271                if let Some(id_token_hint) = self.id_token_hint {
272                    logout = logout.id_token_hint(id_token_hint);
273                }
274                if let Some(redirect) = &self.post_logout_redirect_url {
275                    logout = logout.post_logout_redirect_uri(redirect);
276                }
277                if let Some(logout_hint) = self.logout_hint {
278                    logout = logout.logout_hint(logout_hint);
279                }
280                if let Some(client_id) = self.client_id {
281                    logout = logout.client_id(client_id);
282                }
283                if let Some(state) = self.state {
284                    logout = logout.state(state);
285                }
286                Redirect::to(logout.build().as_str())
287            }
288            None => match &self.post_logout_redirect_url {
289                Some(url) => Redirect::to(url),
290                None => Redirect::to("/"),
291            },
292        }
293    }
294
295    pub fn end_session_url(&self) -> Option<&str> {
296        self.end_session_endpoint.as_ref().map(Url::as_str)
297    }
298
299    pub fn post_logout_redirect_url(&self) -> Option<&str> {
300        self.post_logout_redirect_url.as_deref()
301    }
302}
303
304impl<S, H> FromRequestParts<S> for OidcContext<H>
305where
306    Self: FromRef<S>,
307    S: Send + Sync,
308    H: OidcHandler,
309{
310    type Rejection = Infallible;
311
312    async fn from_request_parts(_parts: &mut Parts, state: &S) -> Result<Self, Self::Rejection> {
313        Ok(Self::from_ref(state))
314    }
315}
316
317impl<H> Clone for OidcContext<H> {
318    fn clone(&self) -> Self {
319        Self(self.0.clone())
320    }
321}
322
323#[cfg(test)]
324mod tests {
325    use axum::http::StatusCode;
326    use axum::response::IntoResponse;
327    use http::Extensions;
328    use url::Url;
329
330    use super::LogoutContext;
331
332    fn ctx(
333        end_session_endpoint: Option<&str>,
334        post_logout_redirect_url: Option<&str>,
335    ) -> LogoutContext {
336        LogoutContext {
337            extensions: Extensions::new(),
338            end_session_endpoint: end_session_endpoint.map(|u| Url::parse(u).unwrap()),
339            post_logout_redirect_url: post_logout_redirect_url.map(String::from),
340            id_token_hint: None,
341            logout_hint: None,
342            client_id: None,
343            state: None,
344        }
345    }
346
347    #[test]
348    fn logout_redirects_to_slash_when_nothing_configured() {
349        let response = ctx(None, None).default_redirect().into_response();
350        assert_eq!(response.status(), StatusCode::SEE_OTHER);
351        assert_eq!(response.headers().get("location").unwrap(), "/");
352    }
353
354    #[test]
355    fn logout_redirects_to_post_logout_url() {
356        let response = ctx(None, Some("http://localhost:3000/"))
357            .default_redirect()
358            .into_response();
359        assert_eq!(response.status(), StatusCode::SEE_OTHER);
360        assert_eq!(
361            response.headers().get("location").unwrap(),
362            "http://localhost:3000/"
363        );
364    }
365
366    #[test]
367    fn logout_redirects_to_end_session_endpoint() {
368        let response = ctx(Some("https://provider.example.com/logout"), None)
369            .default_redirect()
370            .into_response();
371        assert_eq!(response.status(), StatusCode::SEE_OTHER);
372        let location = response
373            .headers()
374            .get("location")
375            .unwrap()
376            .to_str()
377            .unwrap();
378        assert!(location.starts_with("https://provider.example.com/logout"));
379    }
380
381    #[test]
382    fn logout_end_session_includes_post_logout_redirect() {
383        let response = ctx(
384            Some("https://provider.example.com/logout"),
385            Some("http://localhost:3000/"),
386        )
387        .default_redirect()
388        .into_response();
389        assert_eq!(response.status(), StatusCode::SEE_OTHER);
390        let location = response
391            .headers()
392            .get("location")
393            .unwrap()
394            .to_str()
395            .unwrap();
396        assert!(location.starts_with("https://provider.example.com/logout"));
397        assert!(location.contains("post_logout_redirect_uri="));
398    }
399}