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

@@ -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
+    // Redirect to oauth redirect uri
-    let uri = {
+    Ok(Either::Left(Redirect::found(redirect_uri(
-        let uri_mode = match auth_request.response_mode.unwrap_or(ResponseMode::Query) {
+        auth_request,
-            ResponseMode::Query => "?",
+        &app,
-            ResponseMode::Fragment => "#",
+        &code,
-        };
+    ))))
+}
 
-        // Redirect + authorization code
+pub fn redirect_uri(auth_request: AuthenticationRequest, app: &App, code: &SecretString) -> String {
-        let uri = format!("{}{}code={}", app.redirect_uri(), uri_mode, code.as_ref());
+    let uri_mode = match auth_request.response_mode.unwrap_or(ResponseMode::Query) {
-
+        ResponseMode::Query => "?",
-        // Add state if present
+        ResponseMode::Fragment => "#",
-        if auth_request.state.is_empty() {
-            uri
-        } else {
-            format!("{}&state={}", uri, auth_request.state)
-        }
     };
 
-    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)
+    }
 }