Skip to main content

max / makenotwork

12.0 KB · 335 lines History Blame Raw
1 //! Two-factor authentication verification page (login flow).
2
3 use axum::{
4 extract::State,
5 http::{header::HeaderMap, StatusCode},
6 response::{IntoResponse, Redirect, Response},
7 Form,
8 };
9 use serde::Deserialize;
10 use sqlx::PgPool;
11 use tower_sessions::{Expiry, Session};
12
13 use crate::{
14 auth::{login_user, track_session, SessionUser},
15 constants,
16 db::{self, UserId},
17 error::{AppError, Result, ResultExt},
18 helpers::{get_csrf_token, is_htmx_request, spawn_email},
19 routes::api::totp::{build_totp, find_matching_step},
20 templates::*,
21 AppState,
22 };
23
24 /// Session key for the pending 2FA user ID.
25 const PENDING_2FA_KEY: &str = "pending_2fa_user_id";
26 const PENDING_2FA_STARTED_AT: &str = "pending_2fa_started_at";
27 const PENDING_2FA_NOTIFY_ENABLED: &str = "pending_2fa_notify_enabled";
28 const PENDING_2FA_NOTIFY_EMAIL: &str = "pending_2fa_notify_email";
29 const PENDING_2FA_NOTIFY_NAME: &str = "pending_2fa_notify_name";
30 const PENDING_2FA_TRACKING_KEY: &str = "pending_2fa_tracking_id";
31
32 /// Clear every pending-2FA session key and delete the pending user_sessions
33 /// tracking row. Used both on successful login and when the pending state
34 /// expires, so a stale unattended browser can't sit "one TOTP away from
35 /// logged in" past `PENDING_2FA_TTL_SECS`.
36 async fn clear_pending_2fa(session: &Session, pool: &PgPool) {
37 if let Ok(Some(tracking_id)) = session
38 .get::<crate::db::UserSessionId>(PENDING_2FA_TRACKING_KEY)
39 .await
40 && let Err(e) = db::sessions::delete_pending_2fa_session(pool, tracking_id).await
41 {
42 tracing::warn!(error = ?e, "failed to delete pending_2fa tracking row");
43 }
44 session.remove::<UserId>(PENDING_2FA_KEY).await.ok();
45 session.remove::<i64>(PENDING_2FA_STARTED_AT).await.ok();
46 session.remove::<bool>(PENDING_2FA_NOTIFY_ENABLED).await.ok();
47 session.remove::<String>(PENDING_2FA_NOTIFY_EMAIL).await.ok();
48 session.remove::<String>(PENDING_2FA_NOTIFY_NAME).await.ok();
49 session.remove::<bool>("pending_2fa_remember_me").await.ok();
50 session.remove::<crate::db::UserSessionId>(PENDING_2FA_TRACKING_KEY).await.ok();
51 }
52
53 /// Check whether the pending-2FA state has aged past `PENDING_2FA_TTL_SECS`.
54 /// Missing `started_at` (older session pre-TTL) is treated as expired.
55 async fn pending_2fa_expired(session: &Session) -> bool {
56 let started_at: Option<i64> = session.get(PENDING_2FA_STARTED_AT).await.ok().flatten();
57 match started_at {
58 Some(ts) => chrono::Utc::now().timestamp() - ts > constants::PENDING_2FA_TTL_SECS,
59 None => true,
60 }
61 }
62
63 /// Render the 2FA verification page (GET /auth/2fa).
64 #[tracing::instrument(skip_all, name = "two_factor::two_factor_page")]
65 pub(super) async fn two_factor_page(
66 State(state): State<AppState>,
67 session: Session,
68 ) -> Result<Response> {
69 // Verify the user is in a valid 2FA flow
70 let user_id: UserId = session
71 .get(PENDING_2FA_KEY)
72 .await
73 .context("session error")?
74 .ok_or_else(|| AppError::BadRequest("No pending 2FA session".to_string()))?;
75
76 if pending_2fa_expired(&session).await {
77 clear_pending_2fa(&session, &state.db).await;
78 return Err(AppError::BadRequest(
79 "Your 2FA session expired. Please log in again.".to_string(),
80 ));
81 }
82
83 // Confirm the pending_2fa tracking row still exists. If it was swept by
84 // `delete_all_sessions_for_user` ("log out everywhere"), abort the flow.
85 if let Some(tracking_id) = session
86 .get::<crate::db::UserSessionId>(PENDING_2FA_TRACKING_KEY)
87 .await
88 .ok()
89 .flatten()
90 && !db::sessions::pending_2fa_session_exists(&state.db, tracking_id, user_id).await?
91 {
92 clear_pending_2fa(&session, &state.db).await;
93 return Err(AppError::BadRequest(
94 "Your session was revoked. Please log in again.".to_string(),
95 ));
96 }
97
98 let csrf_token = get_csrf_token(&session).await;
99
100 Ok(TwoFactorTemplate {
101 csrf_token,
102 session_user: None,
103 error: None,
104 }
105 .into_response())
106 }
107
108 /// Form input for 2FA verification.
109 #[derive(Deserialize)]
110 pub struct VerifyTwoFactorForm {
111 code: String,
112 }
113
114 /// Verify the TOTP or backup code and complete login (POST /auth/verify-2fa).
115 #[tracing::instrument(skip_all, name = "two_factor::verify_two_factor")]
116 pub(super) async fn verify_two_factor(
117 State(state): State<AppState>,
118 headers: HeaderMap,
119 session: Session,
120 Form(form): Form<VerifyTwoFactorForm>,
121 ) -> Result<Response> {
122 let is_htmx = is_htmx_request(&headers);
123
124 let user_id: UserId = session
125 .get(PENDING_2FA_KEY)
126 .await
127 .context("session error")?
128 .ok_or_else(|| AppError::BadRequest("No pending 2FA session".to_string()))?;
129
130 if pending_2fa_expired(&session).await {
131 clear_pending_2fa(&session, &state.db).await;
132 return Err(AppError::BadRequest(
133 "Your 2FA session expired. Please log in again.".to_string(),
134 ));
135 }
136
137 // Confirm the pending_2fa tracking row still exists (see two_factor_page
138 // for rationale). Rejecting here closes the "phisher mid-2FA-prompt"
139 // window even when the legitimate user's "log out everywhere" landed
140 // between page render and code submission.
141 if let Some(tracking_id) = session
142 .get::<crate::db::UserSessionId>(PENDING_2FA_TRACKING_KEY)
143 .await
144 .ok()
145 .flatten()
146 && !db::sessions::pending_2fa_session_exists(&state.db, tracking_id, user_id).await?
147 {
148 clear_pending_2fa(&session, &state.db).await;
149 return Err(AppError::BadRequest(
150 "Your session was revoked. Please log in again.".to_string(),
151 ));
152 }
153
154 let user = db::users::get_user_by_id(&state.db, user_id)
155 .await?
156 .ok_or(AppError::Unauthorized)?;
157
158 // Re-check lockout status before attempting verification (account may have
159 // been locked by a concurrent session since the 2FA page was shown)
160 if let Some(locked_until) = user.locked_until
161 && locked_until > chrono::Utc::now()
162 {
163 clear_pending_2fa(&session, &state.db).await;
164 let remaining = (locked_until - chrono::Utc::now()).num_minutes() + 1;
165 let csrf_token = get_csrf_token(&session).await;
166 return Ok(TwoFactorTemplate {
167 csrf_token,
168 session_user: None,
169 error: Some(format!(
170 "Account is locked. Try again in {} minute(s).",
171 remaining
172 )),
173 }
174 .into_response());
175 }
176
177 let code = form.code.trim().to_string();
178 let mut verified = false;
179
180 // Try TOTP verification first (6-digit numeric codes)
181 if let Some(ref secret) = user.totp_secret {
182 let totp = build_totp(secret, &user.email)?;
183 // Find the actual step that matched (not just wall-clock step) to
184 // prevent replay across the skew boundary.
185 let now = chrono::Utc::now().timestamp() as u64;
186 let matched_step = find_matching_step(&totp, &code, now);
187 if let Some(step) = matched_step {
188 let last_step = db::totp::get_totp_last_used_step(&state.db, user_id).await?.unwrap_or(0);
189 if step > last_step {
190 db::totp::set_totp_last_used_step(&state.db, user_id, step).await?;
191 verified = true;
192 }
193 // If step <= last_step, the code was already used — fall through to backup codes
194 }
195 }
196
197 // If TOTP didn't match, try backup code. We pass both the raw code (for
198 // Argon2 verify of newer rows) and the legacy HMAC (for pre-migration
199 // rows that haven't been regenerated yet). Both are evaluated in
200 // `verify_and_consume_backup_code` per row.
201 if !verified {
202 let legacy_hmac = crate::routes::api::totp::legacy_hmac_backup_code(
203 &code, &state.config.signing_secret,
204 );
205 if db::totp::verify_and_consume_backup_code(
206 &state.db, user_id, &code, &legacy_hmac,
207 ).await? {
208 verified = true;
209 }
210 }
211
212 if !verified {
213 // Track failed 2FA attempts toward account lockout (same counter as
214 // failed password attempts — prevents brute-forcing 6-digit TOTP codes)
215 db::auth::increment_failed_login(
216 &state.db,
217 user_id,
218 constants::MAX_LOGIN_ATTEMPTS,
219 constants::LOCKOUT_MINUTES,
220 )
221 .await?;
222
223 // Check if this attempt triggered a lockout
224 let user_after = db::users::get_user_by_id(&state.db, user_id).await?;
225 if let Some(ref u) = user_after
226 && let Some(locked_until) = u.locked_until
227 && locked_until > chrono::Utc::now()
228 {
229 // Clear the 2FA flow — account is now locked
230 clear_pending_2fa(&session, &state.db).await;
231 let remaining = (locked_until - chrono::Utc::now()).num_minutes() + 1;
232 let csrf_token = get_csrf_token(&session).await;
233 return Ok(TwoFactorTemplate {
234 csrf_token,
235 session_user: None,
236 error: Some(format!(
237 "Too many failed attempts. Account locked for {} minute(s).",
238 remaining
239 )),
240 }
241 .into_response());
242 }
243
244 let csrf_token = get_csrf_token(&session).await;
245 return Ok(TwoFactorTemplate {
246 csrf_token,
247 session_user: None,
248 error: Some("Invalid code. Please try again.".to_string()),
249 }
250 .into_response());
251 }
252
253 // Successful 2FA — reset failed login counter
254 db::auth::reset_failed_login(&state.db, user_id).await?;
255
256 // Retrieve stored notification info
257 let notify_enabled: bool = session
258 .get(PENDING_2FA_NOTIFY_ENABLED)
259 .await
260 .ok()
261 .flatten()
262 .unwrap_or(false);
263 let notify_email: Option<String> = session
264 .get(PENDING_2FA_NOTIFY_EMAIL)
265 .await
266 .ok()
267 .flatten();
268 let notify_name: Option<String> = session
269 .get(PENDING_2FA_NOTIFY_NAME)
270 .await
271 .ok()
272 .flatten();
273
274 // Retrieve remember-me preference
275 let remember: bool = session
276 .get("pending_2fa_remember_me")
277 .await
278 .ok()
279 .flatten()
280 .unwrap_or(false);
281
282 // Clear pending 2FA state
283 clear_pending_2fa(&session, &state.db).await;
284
285 // Complete login
286 let session_user = SessionUser::from_db_user(user, &state.db, state.config.admin_user_id).await;
287
288 login_user(&session, session_user).await?;
289 if !remember {
290 session.set_expiry(Some(Expiry::OnSessionEnd));
291 }
292 track_session(&session, &state.db, user_id, &headers).await?;
293 tracing::info!(user_id = %user_id, event = "login_2fa_success", "User completed 2FA login");
294
295 // Send login notification (same as in auth.rs)
296 if notify_enabled
297 && let Some(email_addr) = notify_email
298 {
299 let session_count = match db::sessions::count_user_sessions(&state.db, user_id).await {
300 Ok(n) => n,
301 Err(e) => { tracing::warn!("Failed to count sessions for login notification: {e}"); 0 }
302 };
303 if session_count > 1 {
304 let user_agent = headers
305 .get("user-agent")
306 .and_then(|v| v.to_str().ok())
307 .map(|s| s.chars().take(constants::USER_AGENT_MAX_LENGTH).collect::<String>());
308 let ip = crate::helpers::extract_client_ip(&headers);
309 let unsub_url = crate::email::generate_unsubscribe_url(
310 &state.config.host_url, user_id, crate::email::UnsubscribeAction::Login, &user_id.to_string(), &state.config.signing_secret,
311 );
312 spawn_email!(state, "login notification", |email| {
313 email.send_new_login_notification(
314 &email_addr,
315 notify_name.as_deref(),
316 user_agent.as_deref(),
317 ip.as_deref(),
318 Some(&unsub_url),
319 )
320 });
321 }
322 }
323
324 if is_htmx {
325 return Ok((
326 StatusCode::OK,
327 [("HX-Redirect", "/dashboard")],
328 "",
329 )
330 .into_response());
331 }
332
333 Ok(Redirect::to("/dashboard").into_response())
334 }
335