From 5d75c9b247a0920e1dd6ac08661333cfd8f6f695 Mon Sep 17 00:00:00 2001
From: Firstyear <william@blackhats.net.au>
Date: Mon, 16 Dec 2024 09:43:29 +1000
Subject: [PATCH] Autocomplete password during reauth with TOTP (#3290)

During a re-auth flow, the password was not autocompleted once
totp was autocompleted. This is because in a normal login flow
the autocomplete is performed on the first login.html page,
but in a re-auth we skip that page.

This adds the proper handling to allow the pw to autofill
in the background once the TOTP is completed.
---
 server/core/src/https/views/login.rs  | 29 ++++++++++++++++++++++++++-
 server/core/templates/login.html      |  2 +-
 server/core/templates/login_totp.html | 11 ++++++++++
 3 files changed, 40 insertions(+), 2 deletions(-)

diff --git a/server/core/src/https/views/login.rs b/server/core/src/https/views/login.rs
index ce19a26f5..3b5c2b6ba 100644
--- a/server/core/src/https/views/login.rs
+++ b/server/core/src/https/views/login.rs
@@ -552,6 +552,8 @@ pub async fn view_login_mech_choose_post(
 
 #[derive(Debug, Clone, Deserialize)]
 pub struct LoginTotpForm {
+    #[serde(default, deserialize_with = "empty_string_as_none")]
+    password: Option<String>,
     totp: String,
 }
 
@@ -560,7 +562,7 @@ pub async fn view_login_totp_post(
     Extension(kopid): Extension<KOpId>,
     VerifiedClientInformation(client_auth_info): VerifiedClientInformation,
     DomainInfo(domain_info): DomainInfo,
-    jar: CookieJar,
+    mut jar: CookieJar,
     Form(login_totp_form): Form<LoginTotpForm>,
 ) -> Response {
     // trim leading and trailing white space.
@@ -583,6 +585,31 @@ pub async fn view_login_totp_post(
         }
     };
 
+    // In some flows the PW manager may not have autocompleted the pw until
+    // this point. This could be due to a re-auth flow which skips the username
+    // prompt, the use of remember-me+return which then skips the autocomplete.
+    //
+    // In the case the pw *is* bg filled, we need to add it to the session context
+    // here.
+    //
+    // It's probably not "optimal" to be getting the context out and signing it
+    // here to re-add it, but it also helps keep the flow neater in general.
+
+    if let Some(password_autofill) = login_totp_form.password {
+        let mut session_context =
+            cookies::get_signed::<SessionContext>(&state, &jar, COOKIE_AUTH_SESSION_ID)
+                .unwrap_or_default();
+
+        session_context.password = Some(password_autofill);
+
+        // If we can't write this back to the jar, we warn and move on.
+        if let Ok(update_jar) = add_session_cookie(&state, jar.clone(), &session_context) {
+            jar = update_jar;
+        } else {
+            warn!("Unable to update session_context, ignoring...");
+        }
+    }
+
     let auth_cred = AuthCredential::Totp(totp);
     credential_step(state, kopid, jar, client_auth_info, auth_cred, domain_info).await
 }
diff --git a/server/core/templates/login.html b/server/core/templates/login.html
index 9a5b89e63..74b8e21c4 100644
--- a/server/core/templates/login.html
+++ b/server/core/templates/login.html
@@ -23,7 +23,7 @@
 		/>
 	</div>
 
-	<!-- BEGIN: to work better with password managers -->
+	<!-- BEGIN: allows a password manager to autocomplete these fields in the BG. -->
 	<input
 		class="d-none"
 		id="password"
diff --git a/server/core/templates/login_totp.html b/server/core/templates/login_totp.html
index 5a17dc5cd..4ebcd46fc 100644
--- a/server/core/templates/login_totp.html
+++ b/server/core/templates/login_totp.html
@@ -10,6 +10,17 @@
 (% endmatch %)
 <form id="login" action="/ui/login/totp" method="post">
 	<div class="input-group mb-3">
+		<!-- BEGIN: allows a password manager to autocomplete these fields in the BG. -->
+		<input
+			class="d-none"
+			id="password"
+			name="password"
+			type="password"
+			autocomplete="current-password"
+			value=""
+		/>
+		<!-- END -->
+
 		<input
 			autofocus=true
 			class="autofocus form-control"