From a42578090442a1db9f1dfbb31ed4b7c0702300f8 Mon Sep 17 00:00:00 2001
From: Artemis <git@artemix.org>
Date: Wed, 19 Mar 2025 10:40:44 +0100
Subject: [PATCH] now cleanly handling "invalid totp code" errors

---
 src/routes/account/otp.rs             | 21 ++++++++++++++-------
 templates/account/otp/start.html.tera |  5 ++++-
 2 files changed, 18 insertions(+), 8 deletions(-)

diff --git a/src/routes/account/otp.rs b/src/routes/account/otp.rs
index 306c478..d214808 100644
--- a/src/routes/account/otp.rs
+++ b/src/routes/account/otp.rs
@@ -16,9 +16,10 @@ use crate::{
 	routes::{self, error_handlers::PageResult},
 };
 
-#[get("/settings/totp")]
+#[get("/settings/totp?<invalid_code>")]
 pub async fn show_totp_enable_start(
 	mut db: DollTagsDb,
+	invalid_code: bool,
 	user: User,
 	cookies: &CookieJar<'_>,
 	meta: CommonTemplateState,
@@ -26,8 +27,15 @@ pub async fn show_totp_enable_start(
 	if db::otp::has_otp(&mut *db, &user.id, METHOD_TOTP).await? {
 		return Ok(Redirect::to(uri!("/account", routes::account::settings::show_settings)).into());
 	}
-	let totp_secret = Secret::generate_secret();
-	auth::otp::cache_secret(cookies, &totp_secret);
+
+	// cookie recovery in case it was a form error and not a new session
+	let mut totp_secret = auth::otp::get_secret(cookies);
+	if totp_secret.is_none() {
+		let new_secret = Secret::generate_secret();
+		auth::otp::cache_secret(cookies, &new_secret);
+		totp_secret = Some(new_secret);
+	}
+	let totp_secret = totp_secret.unwrap();
 
 	let totp = auth::otp::make_totp(&user.id.to_string(), totp_secret.to_bytes()?)?;
 	let totp_qrcode = totp.get_qr_base64()?;
@@ -37,6 +45,7 @@ pub async fn show_totp_enable_start(
 		context! {
 			meta,
 			totp_qrcode,
+			invalid_code,
 			secret: totp.get_secret_base32(),
 		},
 	)
@@ -63,15 +72,13 @@ pub async fn handle_totp_enable_start(
 	}
 	let secret = match auth::otp::get_secret(cookies) {
 		Some(v) => v,
-		None => return Ok(Redirect::to(uri!("/account", show_totp_enable_start)).into()),
+		None => return Ok(Redirect::to(uri!("/account", show_totp_enable_start(false))).into()),
 	};
 
 	let totp = auth::otp::make_totp(&user.id.to_string(), secret.to_bytes()?)?;
 
 	if !totp.check_current(&form.otp_code)? {
-		Err(String::from(
-			"woof TODO need to impl. a user error return w/ the same secret",
-		))?;
+		return Ok(Redirect::to(uri!("/account", show_totp_enable_start(true))).into());
 	}
 
 	warn!(
diff --git a/templates/account/otp/start.html.tera b/templates/account/otp/start.html.tera
index 2176461..98735de 100644
--- a/templates/account/otp/start.html.tera
+++ b/templates/account/otp/start.html.tera
@@ -18,8 +18,11 @@
 
 		<form method="post">
 			<input type="text" name="otp_code" id="otp_code" minlength="6" maxlength="6" required />
+			{% if invalid_code %}
+			<p class="field-error">The OTP code you sent was invalid, please retry.</p>
+			{% endif %}
 
-			<button type="submit">Enable 2FA</button>
+			<button type="submit" class="submit">Enable 2FA</button>
 		</form>
 	</div>
 </section>