From 85c559cfb17760d54adbddcd2b6a412939b58ff9 Mon Sep 17 00:00:00 2001
From: traxys <quentin+dev@familleboyer.net>
Date: Wed, 16 Aug 2023 21:46:33 +0200
Subject: [PATCH] Redirect users after login

---
 src/main.rs         | 54 ++++++++++++++++++++++++++++++++++++++-------
 templates/home.html | 19 ++++++++++++++++
 2 files changed, 65 insertions(+), 8 deletions(-)
 create mode 100644 templates/home.html

diff --git a/src/main.rs b/src/main.rs
index baac1c6..6a96d58 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -5,8 +5,9 @@ use std::{
 };
 
 use axum::{
-    extract::{Path, Query, State},
-    http::StatusCode,
+    async_trait,
+    extract::{FromRef, FromRequestParts, Path, Query, State},
+    http::{request::Parts, StatusCode},
     response::{Html, IntoResponse, Redirect},
     routing::get,
     Router,
@@ -314,7 +315,6 @@ struct AppState {
     jwt_secret: HS256Key,
     db: PgPool,
     oidc: OpenidConnector,
-    domain: String,
 }
 
 #[derive(thiserror::Error, Debug)]
@@ -379,12 +379,14 @@ struct OidcRedirectParams {
     code: String,
 }
 
+pub static AUTH_COOKIE: &str = "mail_admin_token";
+
 async fn redirected(
     state: State<Arc<AppState>>,
     Path(id): Path<Uuid>,
     Query(redirect): Query<OidcRedirectParams>,
     jar: CookieJar,
-) -> Result<CookieJar, Error> {
+) -> Result<(CookieJar, Redirect), Error> {
     match state
         .oidc
         .redirected(id, redirect.state, redirect.code)
@@ -412,20 +414,20 @@ async fn redirected(
 
             let token = state.jwt_secret.authenticate(claims)?;
 
-            let mut cookie = Cookie::named("mail_admin_token");
+            let mut cookie = Cookie::named(AUTH_COOKIE);
             cookie.set_value(token);
             cookie.set_http_only(true);
 
             let mut now = OffsetDateTime::now_utc();
             now += expire;
             cookie.set_expires(now);
-            cookie.set_same_site(SameSite::Strict);
+            cookie.set_same_site(SameSite::None);
             cookie.set_secure(true);
             cookie.set_path("/");
 
             let jar = jar.add(cookie);
 
-            Ok(jar)
+            Ok((jar, Redirect::to("/")))
         }
         Err(e) => {
             tracing::error!("Could not finish OAuth2 flow: {e:?}");
@@ -452,6 +454,42 @@ async fn page_not_found() -> Result<(StatusCode, Html<String>), Error> {
 pub static TEMPLATES: Lazy<Tera> =
     Lazy::new(|| Tera::new("templates/*.html").expect("Could not generate templates"));
 
+struct User(Uuid);
+
+#[async_trait]
+impl<S> FromRequestParts<S> for User
+where
+    S: Send + Sync,
+    Arc<AppState>: FromRef<S>,
+{
+    type Rejection = Redirect;
+
+    async fn from_request_parts(parts: &mut Parts, state: &S) -> Result<Self, Self::Rejection> {
+        let jar = CookieJar::from_request_parts(parts, state).await.unwrap();
+
+        match jar.get(AUTH_COOKIE) {
+            None => {
+                tracing::debug!("No auth token");
+                Err(Redirect::to("/login"))
+            }
+            Some(data) => {
+                let state: Arc<AppState> = State::from_request_parts(parts, state).await.unwrap().0;
+
+                let claims = state
+                    .jwt_secret
+                    .verify_token::<NoCustomClaims>(data.value(), None)
+                    .map_err(|_| Redirect::to("/login"))?;
+
+                Ok(User(claims.subject.unwrap().parse().unwrap()))
+            }
+        }
+    }
+}
+
+async fn home(_: State<Arc<AppState>>, User(_): User) -> Result<Html<String>, Error> {
+    Ok(Html(TEMPLATES.render("home.html", &global_context())?))
+}
+
 #[tokio::main]
 async fn main() -> color_eyre::Result<()> {
     color_eyre::install()?;
@@ -481,12 +519,12 @@ async fn main() -> color_eyre::Result<()> {
     let router = Router::new()
         .route("/login", get(login))
         .route("/login/redirect/:id", get(redirected))
+        .route("/", get(home))
         .fallback(page_not_found)
         .with_state(Arc::new(AppState {
             db,
             oidc,
             jwt_secret: config.jwt_secret.0,
-            domain: config.domain,
         }));
 
     Ok(axum::Server::bind(&addr)
diff --git a/templates/home.html b/templates/home.html
new file mode 100644
index 0000000..a3f7e5d
--- /dev/null
+++ b/templates/home.html
@@ -0,0 +1,19 @@
+<!doctype html>
+<html>
+  <head>
+    <meta charset="utf-8" />
+    <meta name="viewport" content="width=device-width, initial-scale=1" />
+    <title>{{ title }}</title>
+    <link
+      rel="stylesheet"
+      href="https://cdn.jsdelivr.net/npm/bulma@0.9.4/css/bulma.min.css"
+    />
+  </head>
+  <body>
+    <section class="section">
+      <div class="container">
+        <h1 class="title">Mail management</h1>
+      </div>
+    </section>
+  </body>
+</html>