Skip to main content

max / makenotwork

11.8 KB · 370 lines History Blame Raw
1 //! TOTP 2FA workflow tests: setup, confirm, login with TOTP, backup codes, disable.
2
3 use crate::harness::TestHarness;
4
5 // ── Helpers ──
6
7 /// Extract the TOTP secret from setup HTML (inside `<details>` > `<code>`).
8 fn extract_totp_secret(html: &str) -> String {
9 let details_start = html
10 .find("<details")
11 .expect("No <details> in TOTP setup HTML");
12 let details_html = &html[details_start..];
13 let code_start = details_html.find("<code").expect("No <code> in details");
14 let after_tag = &details_html[code_start..];
15 let content_start = after_tag.find('>').expect("No > after <code") + 1;
16 let content_end = after_tag[content_start..]
17 .find("</code>")
18 .expect("No </code>");
19 after_tag[content_start..content_start + content_end].to_string()
20 }
21
22 /// Extract backup codes from setup HTML (inside `<div class="backup-codes-grid">`).
23 fn extract_backup_codes(html: &str) -> Vec<String> {
24 let marker = "backup-codes-grid";
25 let grid_start = html.find(marker).expect("No backup-codes-grid in HTML");
26 let grid_html = &html[grid_start..];
27 let grid_end = grid_html
28 .find("</div>")
29 .expect("No </div> for backup-codes-grid");
30 let grid_content = &grid_html[..grid_end];
31
32 let mut codes = Vec::new();
33 let mut search = grid_content;
34 while let Some(code_start) = search.find("<code>") {
35 let content_start = code_start + "<code>".len();
36 let content_end = search[content_start..]
37 .find("</code>")
38 .expect("Unclosed <code> in backup codes");
39 codes.push(search[content_start..content_start + content_end].to_string());
40 search = &search[content_start + content_end..];
41 }
42
43 codes
44 }
45
46 /// Generate a valid TOTP code from a base32 secret.
47 fn generate_totp_code(secret_base32: &str, email: &str) -> String {
48 let bytes = totp_rs::Secret::Encoded(secret_base32.to_string())
49 .to_bytes()
50 .expect("Invalid TOTP secret");
51 let totp = totp_rs::TOTP::new(
52 totp_rs::Algorithm::SHA1,
53 6,
54 1,
55 30,
56 bytes,
57 Some("Makenotwork".into()),
58 email.into(),
59 )
60 .expect("TOTP creation failed");
61 totp.generate_current().expect("TOTP generation failed")
62 }
63
64 /// Set up TOTP for the currently logged-in user and enable it.
65 /// Returns (secret_base32, backup_codes).
66 async fn setup_and_enable_totp(h: &mut TestHarness, email: &str) -> (String, Vec<String>) {
67 let resp = h.client.post_form("/api/users/me/totp/setup", "").await;
68 assert_eq!(
69 resp.status.as_u16(),
70 200,
71 "TOTP setup failed: {}",
72 resp.text
73 );
74
75 let secret = extract_totp_secret(&resp.text);
76 let codes = extract_backup_codes(&resp.text);
77 assert!(!secret.is_empty(), "TOTP secret should not be empty");
78 assert!(!codes.is_empty(), "Should have backup codes");
79
80 let code = generate_totp_code(&secret, email);
81 let resp = h
82 .client
83 .post_form("/api/users/me/totp/confirm", &format!("code={code}"))
84 .await;
85 assert_eq!(
86 resp.status.as_u16(),
87 200,
88 "TOTP confirm failed: {}",
89 resp.text
90 );
91
92 (secret, codes)
93 }
94
95 /// Login flow when TOTP is enabled: POST /login -> 303 to /auth/2fa -> POST /auth/verify-2fa.
96 async fn login_with_2fa(h: &mut TestHarness, username: &str, password: &str, code: &str) {
97 h.client.fetch_csrf_token().await;
98 let resp = h
99 .client
100 .post_form(
101 "/login",
102 &format!(
103 "login={}&password={}",
104 urlencoding::encode(username),
105 urlencoding::encode(password)
106 ),
107 )
108 .await;
109 // Login should redirect to 2FA page
110 assert!(
111 resp.status.is_redirection()
112 || resp.text.contains("/auth/2fa")
113 || resp.text.contains("HX-Redirect"),
114 "Expected redirect to /auth/2fa, got {}, {}",
115 resp.status,
116 resp.text
117 );
118
119 // Load the 2FA page to get CSRF token
120 let resp = h.client.get("/auth/2fa").await;
121 assert_eq!(resp.status.as_u16(), 200, "2FA page failed: {}", resp.text);
122
123 // Submit the code
124 let resp = h
125 .client
126 .post_form("/auth/verify-2fa", &format!("code={code}"))
127 .await;
128 assert_eq!(
129 resp.status, 303,
130 "2FA verification failed with {}: {}",
131 resp.status, resp.text
132 );
133 }
134
135 // ── Tests ──
136
137 #[tokio::test]
138 async fn totp_setup_and_confirm() {
139 let mut h = TestHarness::new().await;
140 h.signup("totp1", "totp1@test.com", "Password1!").await;
141
142 let (_secret, codes) = setup_and_enable_totp(&mut h, "totp1@test.com").await;
143 assert_eq!(codes.len(), 10, "Should have 10 backup codes");
144
145 // Check status
146 let resp = h.client.get("/api/users/me/totp/status").await;
147 assert_eq!(resp.status.as_u16(), 200);
148 assert!(
149 resp.text.contains("true")
150 || resp.text.contains("enabled")
151 || resp.text.contains("Enabled"),
152 "TOTP should be enabled: {}",
153 resp.text
154 );
155 }
156
157 #[tokio::test]
158 async fn totp_login_requires_2fa() {
159 let mut h = TestHarness::new().await;
160 let uid = h.signup("totp2", "totp2@test.com", "Password1!").await;
161 let (secret, _) = setup_and_enable_totp(&mut h, "totp2@test.com").await;
162
163 // Reset the anti-replay counter so the login TOTP code (same 30s window)
164 // is not rejected as a replay of the confirm step.
165 sqlx::query("UPDATE users SET totp_last_used_step = 0 WHERE id = $1")
166 .bind(uid)
167 .execute(&h.db)
168 .await
169 .expect("reset totp_last_used_step");
170
171 h.client.post_form("/logout", "").await;
172
173 // Login with TOTP
174 let code = generate_totp_code(&secret, "totp2@test.com");
175 login_with_2fa(&mut h, "totp2", "Password1!", &code).await;
176
177 // Verify we're logged in
178 let resp = h.client.get("/dashboard").await;
179 assert_eq!(
180 resp.status.as_u16(),
181 200,
182 "Should be on dashboard after 2FA login"
183 );
184 }
185
186 #[tokio::test]
187 async fn totp_backup_code_login() {
188 let mut h = TestHarness::new().await;
189 h.signup("totp3", "totp3@test.com", "Password1!").await;
190 let (_, codes) = setup_and_enable_totp(&mut h, "totp3@test.com").await;
191
192 h.client.post_form("/logout", "").await;
193
194 // Login using a backup code
195 login_with_2fa(&mut h, "totp3", "Password1!", &codes[0]).await;
196
197 let resp = h.client.get("/dashboard").await;
198 assert_eq!(
199 resp.status.as_u16(),
200 200,
201 "Should be on dashboard after backup code login"
202 );
203 }
204
205 #[tokio::test]
206 async fn totp_decrypt_failure_falls_through_to_backup_code() {
207 // M-Sec1 (deep): a stored TOTP secret that fails to decrypt, e.g. a legacy
208 // pre-encryption plaintext row lacking the `enc:v1:` prefix, must not 500
209 // the verify handler. A 500 there would also skip the backup-code branch and
210 // lock the user out. The handler logs the decrypt failure and falls through.
211 let mut h = TestHarness::new().await;
212 let user_id = h.signup("totpdec", "totpdec@test.com", "Password1!").await;
213 let (_, codes) = setup_and_enable_totp(&mut h, "totpdec@test.com").await;
214 h.client.post_form("/logout", "").await;
215
216 // Corrupt the stored secret to a bare base32 value (no `enc:v1:` prefix), so
217 // `decrypt_totp_secret` returns Err on the next verify.
218 sqlx::query("UPDATE users SET totp_secret = $1 WHERE id = $2")
219 .bind("JBSWY3DPEHPK3PXP")
220 .bind(user_id)
221 .execute(&h.db)
222 .await
223 .expect("corrupt totp secret");
224
225 // A backup-code login still succeeds: the verify handler tries TOTP first
226 // (decrypt fails -> logged -> falls through), then consumes the backup code.
227 // If the decrypt error propagated, this single verify POST would 500 before
228 // reaching the backup branch and `login_with_2fa`'s success assert would fail.
229 login_with_2fa(&mut h, "totpdec", "Password1!", &codes[0]).await;
230
231 let resp = h.client.get("/dashboard").await;
232 assert_eq!(
233 resp.status.as_u16(),
234 200,
235 "backup-code login must succeed despite an undecryptable TOTP secret"
236 );
237 }
238
239 #[tokio::test]
240 async fn totp_backup_code_single_use() {
241 let mut h = TestHarness::new().await;
242 h.signup("totp4", "totp4@test.com", "Password1!").await;
243 let (_, codes) = setup_and_enable_totp(&mut h, "totp4@test.com").await;
244 let used_code = codes[0].clone();
245
246 // Use backup code once
247 h.client.post_form("/logout", "").await;
248 login_with_2fa(&mut h, "totp4", "Password1!", &used_code).await;
249
250 // Logout and try the same code again
251 h.client.post_form("/logout", "").await;
252 h.client.fetch_csrf_token().await;
253 let resp = h
254 .client
255 .post_form(
256 "/login",
257 &format!("login=totp4&password={}", urlencoding::encode("Password1!")),
258 )
259 .await;
260 assert!(
261 resp.status.is_redirection()
262 || resp.text.contains("/auth/2fa")
263 || resp.text.contains("HX-Redirect"),
264 "Expected redirect to 2FA"
265 );
266
267 let _resp = h.client.get("/auth/2fa").await;
268 let resp = h
269 .client
270 .post_form("/auth/verify-2fa", &format!("code={used_code}"))
271 .await;
272
273 // Should fail, backup code already consumed
274 assert!(
275 resp.text.contains("Invalid")
276 || resp.text.contains("error")
277 || resp.text.contains("invalid"),
278 "Used backup code should be rejected: {}",
279 resp.text
280 );
281 }
282
283 #[tokio::test]
284 async fn totp_disable_with_password() {
285 let mut h = TestHarness::new().await;
286 h.signup("totp5", "totp5@test.com", "Password1!").await;
287 setup_and_enable_totp(&mut h, "totp5@test.com").await;
288
289 // Disable TOTP
290 let resp = h
291 .client
292 .post_form(
293 "/api/users/me/totp/disable",
294 &format!("password={}", urlencoding::encode("Password1!")),
295 )
296 .await;
297 assert_eq!(
298 resp.status.as_u16(),
299 200,
300 "TOTP disable failed: {}",
301 resp.text
302 );
303
304 // Verify status is disabled
305 let resp = h.client.get("/api/users/me/totp/status").await;
306 assert_eq!(resp.status.as_u16(), 200);
307 assert!(
308 resp.text.contains("false")
309 || resp.text.contains("disabled")
310 || resp.text.contains("Disabled")
311 || !resp.text.contains("Enabled"),
312 "TOTP should be disabled: {}",
313 resp.text
314 );
315
316 // Logout and login, should not require 2FA
317 h.client.post_form("/logout", "").await;
318 h.login("totp5", "Password1!").await;
319
320 let resp = h.client.get("/dashboard").await;
321 assert_eq!(
322 resp.status.as_u16(),
323 200,
324 "Login after TOTP disable should not require 2FA"
325 );
326 }
327
328 #[tokio::test]
329 async fn totp_invalid_code_rejected() {
330 let mut h = TestHarness::new().await;
331 h.signup("totp6", "totp6@test.com", "Password1!").await;
332 setup_and_enable_totp(&mut h, "totp6@test.com").await;
333
334 h.client.post_form("/logout", "").await;
335
336 // Start login
337 h.client.fetch_csrf_token().await;
338 let resp = h
339 .client
340 .post_form(
341 "/login",
342 &format!("login=totp6&password={}", urlencoding::encode("Password1!")),
343 )
344 .await;
345 assert!(
346 resp.status.is_redirection()
347 || resp.text.contains("/auth/2fa")
348 || resp.text.contains("HX-Redirect"),
349 "Expected redirect to 2FA"
350 );
351
352 let _resp = h.client.get("/auth/2fa").await;
353
354 // Submit an invalid code
355 let resp = h.client.post_form("/auth/verify-2fa", "code=000000").await;
356
357 // Should show error, not redirect to dashboard
358 assert!(
359 resp.text.contains("Invalid")
360 || resp.text.contains("invalid")
361 || resp.text.contains("error"),
362 "Invalid TOTP code should be rejected: {}",
363 resp.text
364 );
365 assert!(
366 !resp.text.contains("/dashboard") || resp.status.as_u16() == 200,
367 "Should not redirect to dashboard with invalid code"
368 );
369 }
370