Skip to main content

axum_security/oauth2/
handler.rs

1use std::{borrow::Cow, pin::Pin};
2
3use axum::response::{IntoResponse, Response};
4use cookie_monster::{Cookie, CookieBuilder, CookieJar};
5
6pub struct TokenResponse {
7    pub access_token: String,
8    pub refresh_token: Option<String>,
9}
10
11pub trait OAuth2Handler: Send + Sync + 'static {
12    fn after_login(
13        &self,
14        token_res: TokenResponse,
15        _context: &mut AfterLoginCookies<'_>,
16    ) -> impl Future<Output = impl IntoResponse> + Send;
17}
18
19pub struct AfterLoginCookies<'a> {
20    pub(crate) cookie_jar: CookieJar,
21    pub(crate) cookie_opts: &'a CookieBuilder,
22}
23
24impl AfterLoginCookies<'_> {
25    pub fn cookie(&self, name: impl Into<Cow<'static, str>>) -> CookieBuilder {
26        self.cookie_opts.clone().name(name)
27    }
28
29    pub fn remove(&mut self, name: impl Into<Cow<'static, str>>) -> Option<Cookie> {
30        let cookie = self.cookie(name);
31        self.cookie_jar.remove(cookie)
32    }
33
34    pub fn add(&mut self, cookie: impl Into<Cookie>) {
35        self.cookie_jar.add(cookie.into());
36    }
37}
38
39trait DynOAuth2Handler: Send + Sync + 'static {
40    fn after_login_boxed<'a, 'b>(
41        &'a self,
42        token_res: TokenResponse,
43        context: &'a mut AfterLoginCookies<'b>,
44    ) -> Pin<Box<dyn Future<Output = Response> + Send + 'a>>;
45}
46
47impl<T> DynOAuth2Handler for T
48where
49    T: OAuth2Handler,
50{
51    fn after_login_boxed<'a, 'b>(
52        &'a self,
53        token_res: TokenResponse,
54        context: &'a mut AfterLoginCookies<'b>,
55    ) -> Pin<Box<dyn Future<Output = Response> + Send + 'a>> {
56        Box::pin(async move { self.after_login(token_res, context).await.into_response() })
57    }
58}
59
60pub(crate) struct ErasedOAuth2Handler(Box<dyn DynOAuth2Handler>);
61
62impl ErasedOAuth2Handler {
63    pub fn new<T: OAuth2Handler>(handler: T) -> Self {
64        Self(Box::new(handler))
65    }
66
67    pub async fn after_login<'a, 'b>(
68        &'a self,
69        res: TokenResponse,
70        context: &'a mut AfterLoginCookies<'b>,
71    ) -> Response {
72        self.0.after_login_boxed(res, context).await
73    }
74}