From fd2d2672bbc0443a3fdf28f75177040b8f95d126 Mon Sep 17 00:00:00 2001
From: Philippe Loctaux <p@philippeloctaux.com>
Date: Mon, 1 May 2023 12:00:11 +0200
Subject: [PATCH] oauth authorize: if use has totp enable, generate a code,
 store in a cookie and redirect to totp verification page

---
 crates/ezidam/src/routes/oauth/authorize.rs | 83 ++++++++++++++++-----
 1 file changed, 63 insertions(+), 20 deletions(-)

diff --git a/crates/ezidam/src/routes/oauth/authorize.rs b/crates/ezidam/src/routes/oauth/authorize.rs
index 0430d84..9d5f1c0 100644
--- a/crates/ezidam/src/routes/oauth/authorize.rs
+++ b/crates/ezidam/src/routes/oauth/authorize.rs
@@ -2,8 +2,11 @@ use crate::routes::prelude::*;
 use apps::App;
 use authorization_codes::AuthorizationCode;
 use hash::SecretString;
+use rocket::http::{Cookie, CookieJar, SameSite};
+use rocket::time::Duration;
 use rocket::{get, post};
 use settings::Settings;
+use users::totp_login_request::{TOTP_REQUEST_COOKIE_NAME, TOTP_REQUEST_LEN};
 use users::User;
 
 #[get("/oauth/authorize?<auth_request..>", rank = 1)]
@@ -90,12 +93,15 @@ fn user_archived(login: &str, request: AuthenticationRequest) -> Flash<Redirect>
     flash(format!("User {login} is archived"), request)
 }
 
+pub const AUTHORIZATION_CODE_LEN: usize = 35;
+
 #[post("/oauth/authorize?<auth_request..>", data = "<form>")]
 pub async fn authorize_form(
     form: Option<Form<Authorize<'_>>>,
     user: Option<JwtUser>,
     mut db: Connection<Database>,
     auth_request: AuthenticationRequest<'_>,
+    cookie_jar: &CookieJar<'_>,
 ) -> Result<Either<Redirect, Flash<Redirect>>> {
     let mut transaction = db.begin().await?;
 
@@ -109,11 +115,11 @@ pub async fn authorize_form(
     )
     .await?;
 
-    let user_id = match user {
+    let (user_id, verify_totp) = match user {
         Some(user) => {
             transaction.commit().await?;
 
-            UserID(user.0.subject)
+            (UserID(user.0.subject), false)
         }
         None => {
             let form = match form {
@@ -152,35 +158,72 @@ pub async fn authorize_form(
                 return Ok(Either::Right(invalid_credentials(form.login, auth_request)));
             }
 
-            user.id().to_owned()
+            (user.id().to_owned(), user.totp_secret().is_some())
         }
     };
 
+    if verify_totp {
+        // Generate totp token
+        let totp_token = task::spawn_blocking(|| SecretString::new(TOTP_REQUEST_LEN))
+            .await
+            .map_err(|e| e.to_string())?;
+
+        let totp_duration = 15;
+
+        // Save in database
+        let mut transaction = db.begin().await?;
+        users::totp_login_request::TotpLoginRequest::insert(
+            &mut transaction,
+            totp_token.as_ref(),
+            &user_id,
+            totp_duration,
+        )
+        .await?;
+        transaction.commit().await?;
+
+        // Store totp token as a cookie
+        let mut cookie = Cookie::new(TOTP_REQUEST_COOKIE_NAME, totp_token.to_string());
+        cookie.set_secure(true);
+        cookie.set_http_only(true);
+        cookie.set_same_site(SameSite::Strict);
+        cookie.set_max_age(Duration::minutes(totp_duration));
+        cookie_jar.add(cookie);
+
+        // Redirect to totp verification page
+        return Ok(Either::Left(Redirect::to(uri!(
+            crate::routes::oauth::totp_page(auth_request)
+        ))));
+    }
+
     // Generate authorization code
-    let code = task::spawn_blocking(|| SecretString::new(35)).await?;
+    let code = task::spawn_blocking(|| SecretString::new(AUTHORIZATION_CODE_LEN)).await?;
 
     // Save authorization code
     let mut transaction = db.begin().await?;
     AuthorizationCode::insert(&mut transaction, code.as_ref(), app.id(), &user_id).await?;
     transaction.commit().await?;
 
-    // Construct uri to redirect to
-    let uri = {
-        let uri_mode = match auth_request.response_mode.unwrap_or(ResponseMode::Query) {
-            ResponseMode::Query => "?",
-            ResponseMode::Fragment => "#",
-        };
+    // Redirect to oauth redirect uri
+    Ok(Either::Left(Redirect::found(redirect_uri(
+        auth_request,
+        &app,
+        &code,
+    ))))
+}
 
-        // Redirect + authorization code
-        let uri = format!("{}{}code={}", app.redirect_uri(), uri_mode, code.as_ref());
-
-        // Add state if present
-        if auth_request.state.is_empty() {
-            uri
-        } else {
-            format!("{}&state={}", uri, auth_request.state)
-        }
+pub fn redirect_uri(auth_request: AuthenticationRequest, app: &App, code: &SecretString) -> String {
+    let uri_mode = match auth_request.response_mode.unwrap_or(ResponseMode::Query) {
+        ResponseMode::Query => "?",
+        ResponseMode::Fragment => "#",
     };
 
-    Ok(Either::Left(Redirect::found(uri)))
+    // Redirect + authorization code
+    let uri = format!("{}{}code={}", app.redirect_uri(), uri_mode, code.as_ref());
+
+    // Add state if present
+    if auth_request.state.is_empty() {
+        uri
+    } else {
+        format!("{}&state={}", uri, auth_request.state)
+    }
 }