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
21pub 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 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 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 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
213pub 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}