Skip to main content

haste_server/auth_n/oidc/middleware/
session_validation.rs

1use axum::RequestExt;
2use axum::extract::OriginalUri;
3use axum::http::StatusCode;
4use axum::response::{IntoResponse, Redirect};
5use axum::{body::Body, extract::Request, response::Response};
6use axum_extra::extract::Cached;
7use std::pin::Pin;
8use std::task::{Context, Poll};
9use tower::{Layer, Service};
10use tower_sessions::Session;
11
12use crate::auth_n::oidc::routes::route_string::oidc_route_string;
13use crate::auth_n::session;
14use crate::extract::path_tenant::{ProjectIdentifier, TenantIdentifier};
15
16#[derive(Clone)]
17pub struct AuthSessionValidationLayer {
18    to: &'static str,
19}
20
21impl<S> Layer<S> for AuthSessionValidationLayer {
22    type Service = AuthSessionValidationService<S>;
23
24    fn layer(&self, inner: S) -> Self::Service {
25        AuthSessionValidationService { inner, to: self.to }
26    }
27}
28
29impl AuthSessionValidationLayer {
30    pub fn new(to: &'static str) -> Self {
31        AuthSessionValidationLayer { to }
32    }
33}
34
35#[derive(Clone)]
36pub struct AuthSessionValidationService<T> {
37    inner: T,
38    to: &'static str,
39}
40
41impl<'a, T> Service<Request<Body>> for AuthSessionValidationService<T>
42where
43    T: Service<Request, Response = Response> + Send + 'static + Clone,
44    T::Future: Send + 'static,
45    T::Error: IntoResponse,
46{
47    type Response = Response;
48    type Error = T::Error;
49    type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;
50
51    fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
52        Poll::Ready(Ok(()))
53    }
54
55    fn call(&mut self, mut request: Request) -> Self::Future {
56        // https://docs.rs/tower/latest/tower/trait.Service.html#be-careful-when-cloning-inner-services
57        let clone = self.inner.clone();
58        // take the service that was ready
59        let mut inner = std::mem::replace(&mut self.inner, clone);
60        let to = self.to;
61
62        // Return the response as an immediate future
63        Box::pin(async move {
64            let Ok(Cached(TenantIdentifier { tenant })) =
65                request.extract_parts::<Cached<TenantIdentifier>>().await
66            else {
67                return Ok((
68                    StatusCode::BAD_REQUEST,
69                    "Tenant id not found on request".to_string(),
70                )
71                    .into_response());
72            };
73            let Ok(Cached(ProjectIdentifier { project })) =
74                request.extract_parts::<Cached<ProjectIdentifier>>().await
75            else {
76                return Ok((
77                    StatusCode::BAD_REQUEST,
78                    "Project id not found on request".to_string(),
79                )
80                    .into_response());
81            };
82
83            let Cached(current_session) = request
84                .extract_parts::<Cached<Session>>()
85                .await
86                .expect("Could not extract session.");
87
88            let to_route = oidc_route_string(&tenant, &project, to);
89
90            if let Ok(authorization_completed_state) =
91                session::user::get_completed_authorization_state(&current_session).await
92                && authorization_completed_state.user.tenant == tenant
93            {
94                let response = inner.call(request).await?;
95                Ok(response)
96            } else {
97                let uri = request
98                    .extract_parts::<OriginalUri>()
99                    .await
100                    .expect("Could not extract original URI.");
101                let login_redirect = Redirect::to(
102                    &(to_route
103                        .to_str()
104                        .expect("Failed to create to route.")
105                        .to_string()
106                        + "?"
107                        + uri.query().unwrap_or("")),
108                );
109
110                Ok(login_redirect.into_response())
111            }
112        })
113    }
114}