Skip to main content

max / makenotwork

11.8 KB · 371 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!(
129 resp.status.is_redirection() || resp.status.is_success(),
130 "2FA verification failed with {}: {}",
131 resp.status,
132 resp.text
133 );
134 }
135
136 // ── Tests ──
137
138 #[tokio::test]
139 async fn totp_setup_and_confirm() {
140 let mut h = TestHarness::new().await;
141 h.signup("totp1", "totp1@test.com", "Password1!").await;
142
143 let (_secret, codes) = setup_and_enable_totp(&mut h, "totp1@test.com").await;
144 assert_eq!(codes.len(), 10, "Should have 10 backup codes");
145
146 // Check status
147 let resp = h.client.get("/api/users/me/totp/status").await;
148 assert_eq!(resp.status.as_u16(), 200);
149 assert!(
150 resp.text.contains("true")
151 || resp.text.contains("enabled")
152 || resp.text.contains("Enabled"),
153 "TOTP should be enabled: {}",
154 resp.text
155 );
156 }
157
158 #[tokio::test]
159 async fn totp_login_requires_2fa() {
160 let mut h = TestHarness::new().await;
161 let uid = h.signup("totp2", "totp2@test.com", "Password1!").await;
162 let (secret, _) = setup_and_enable_totp(&mut h, "totp2@test.com").await;
163
164 // Reset the anti-replay counter so the login TOTP code (same 30s window)
165 // is not rejected as a replay of the confirm step.
166 sqlx::query("UPDATE users SET totp_last_used_step = 0 WHERE id = $1")
167 .bind(uid)
168 .execute(&h.db)
169 .await
170 .expect("reset totp_last_used_step");
171
172 h.client.post_form("/logout", "").await;
173
174 // Login with TOTP
175 let code = generate_totp_code(&secret, "totp2@test.com");
176 login_with_2fa(&mut h, "totp2", "Password1!", &code).await;
177
178 // Verify we're logged in
179 let resp = h.client.get("/dashboard").await;
180 assert_eq!(
181 resp.status.as_u16(),
182 200,
183 "Should be on dashboard after 2FA login"
184 );
185 }
186
187 #[tokio::test]
188 async fn totp_backup_code_login() {
189 let mut h = TestHarness::new().await;
190 h.signup("totp3", "totp3@test.com", "Password1!").await;
191 let (_, codes) = setup_and_enable_totp(&mut h, "totp3@test.com").await;
192
193 h.client.post_form("/logout", "").await;
194
195 // Login using a backup code
196 login_with_2fa(&mut h, "totp3", "Password1!", &codes[0]).await;
197
198 let resp = h.client.get("/dashboard").await;
199 assert_eq!(
200 resp.status.as_u16(),
201 200,
202 "Should be on dashboard after backup code login"
203 );
204 }
205
206 #[tokio::test]
207 async fn totp_decrypt_failure_falls_through_to_backup_code() {
208 // M-Sec1 (deep): a stored TOTP secret that fails to decrypt, e.g. a legacy
209 // pre-encryption plaintext row lacking the `enc:v1:` prefix, must not 500
210 // the verify handler. A 500 there would also skip the backup-code branch and
211 // lock the user out. The handler logs the decrypt failure and falls through.
212 let mut h = TestHarness::new().await;
213 let user_id = h.signup("totpdec", "totpdec@test.com", "Password1!").await;
214 let (_, codes) = setup_and_enable_totp(&mut h, "totpdec@test.com").await;
215 h.client.post_form("/logout", "").await;
216
217 // Corrupt the stored secret to a bare base32 value (no `enc:v1:` prefix), so
218 // `decrypt_totp_secret` returns Err on the next verify.
219 sqlx::query("UPDATE users SET totp_secret = $1 WHERE id = $2")
220 .bind("JBSWY3DPEHPK3PXP")
221 .bind(user_id)
222 .execute(&h.db)
223 .await
224 .expect("corrupt totp secret");
225
226 // A backup-code login still succeeds: the verify handler tries TOTP first
227 // (decrypt fails -> logged -> falls through), then consumes the backup code.
228 // If the decrypt error propagated, this single verify POST would 500 before
229 // reaching the backup branch and `login_with_2fa`'s success assert would fail.
230 login_with_2fa(&mut h, "totpdec", "Password1!", &codes[0]).await;
231
232 let resp = h.client.get("/dashboard").await;
233 assert_eq!(
234 resp.status.as_u16(),
235 200,
236 "backup-code login must succeed despite an undecryptable TOTP secret"
237 );
238 }
239
240 #[tokio::test]
241 async fn totp_backup_code_single_use() {
242 let mut h = TestHarness::new().await;
243 h.signup("totp4", "totp4@test.com", "Password1!").await;
244 let (_, codes) = setup_and_enable_totp(&mut h, "totp4@test.com").await;
245 let used_code = codes[0].clone();
246
247 // Use backup code once
248 h.client.post_form("/logout", "").await;
249 login_with_2fa(&mut h, "totp4", "Password1!", &used_code).await;
250
251 // Logout and try the same code again
252 h.client.post_form("/logout", "").await;
253 h.client.fetch_csrf_token().await;
254 let resp = h
255 .client
256 .post_form(
257 "/login",
258 &format!("login=totp4&password={}", urlencoding::encode("Password1!")),
259 )
260 .await;
261 assert!(
262 resp.status.is_redirection()
263 || resp.text.contains("/auth/2fa")
264 || resp.text.contains("HX-Redirect"),
265 "Expected redirect to 2FA"
266 );
267
268 let _resp = h.client.get("/auth/2fa").await;
269 let resp = h
270 .client
271 .post_form("/auth/verify-2fa", &format!("code={used_code}"))
272 .await;
273
274 // Should fail, backup code already consumed
275 assert!(
276 resp.text.contains("Invalid")
277 || resp.text.contains("error")
278 || resp.text.contains("invalid"),
279 "Used backup code should be rejected: {}",
280 resp.text
281 );
282 }
283
284 #[tokio::test]
285 async fn totp_disable_with_password() {
286 let mut h = TestHarness::new().await;
287 h.signup("totp5", "totp5@test.com", "Password1!").await;
288 setup_and_enable_totp(&mut h, "totp5@test.com").await;
289
290 // Disable TOTP
291 let resp = h
292 .client
293 .post_form(
294 "/api/users/me/totp/disable",
295 &format!("password={}", urlencoding::encode("Password1!")),
296 )
297 .await;
298 assert_eq!(
299 resp.status.as_u16(),
300 200,
301 "TOTP disable failed: {}",
302 resp.text
303 );
304
305 // Verify status is disabled
306 let resp = h.client.get("/api/users/me/totp/status").await;
307 assert_eq!(resp.status.as_u16(), 200);
308 assert!(
309 resp.text.contains("false")
310 || resp.text.contains("disabled")
311 || resp.text.contains("Disabled")
312 || !resp.text.contains("Enabled"),
313 "TOTP should be disabled: {}",
314 resp.text
315 );
316
317 // Logout and login, should not require 2FA
318 h.client.post_form("/logout", "").await;
319 h.login("totp5", "Password1!").await;
320
321 let resp = h.client.get("/dashboard").await;
322 assert_eq!(
323 resp.status.as_u16(),
324 200,
325 "Login after TOTP disable should not require 2FA"
326 );
327 }
328
329 #[tokio::test]
330 async fn totp_invalid_code_rejected() {
331 let mut h = TestHarness::new().await;
332 h.signup("totp6", "totp6@test.com", "Password1!").await;
333 setup_and_enable_totp(&mut h, "totp6@test.com").await;
334
335 h.client.post_form("/logout", "").await;
336
337 // Start login
338 h.client.fetch_csrf_token().await;
339 let resp = h
340 .client
341 .post_form(
342 "/login",
343 &format!("login=totp6&password={}", urlencoding::encode("Password1!")),
344 )
345 .await;
346 assert!(
347 resp.status.is_redirection()
348 || resp.text.contains("/auth/2fa")
349 || resp.text.contains("HX-Redirect"),
350 "Expected redirect to 2FA"
351 );
352
353 let _resp = h.client.get("/auth/2fa").await;
354
355 // Submit an invalid code
356 let resp = h.client.post_form("/auth/verify-2fa", "code=000000").await;
357
358 // Should show error, not redirect to dashboard
359 assert!(
360 resp.text.contains("Invalid")
361 || resp.text.contains("invalid")
362 || resp.text.contains("error"),
363 "Invalid TOTP code should be rejected: {}",
364 resp.text
365 );
366 assert!(
367 !resp.text.contains("/dashboard") || resp.status.as_u16() == 200,
368 "Should not redirect to dashboard with invalid code"
369 );
370 }
371