Skip to main content

max / makenotwork

10.0 KB · 302 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.find("<details").expect("No <details> in TOTP setup HTML");
10 let details_html = &html[details_start..];
11 let code_start = details_html.find("<code").expect("No <code> in details");
12 let after_tag = &details_html[code_start..];
13 let content_start = after_tag.find('>').expect("No > after <code") + 1;
14 let content_end = after_tag[content_start..]
15 .find("</code>")
16 .expect("No </code>");
17 after_tag[content_start..content_start + content_end].to_string()
18 }
19
20 /// Extract backup codes from setup HTML (inside `<div class="backup-codes-grid">`).
21 fn extract_backup_codes(html: &str) -> Vec<String> {
22 let marker = "backup-codes-grid";
23 let grid_start = html.find(marker).expect("No backup-codes-grid in HTML");
24 let grid_html = &html[grid_start..];
25 let grid_end = grid_html.find("</div>").expect("No </div> for backup-codes-grid");
26 let grid_content = &grid_html[..grid_end];
27
28 let mut codes = Vec::new();
29 let mut search = grid_content;
30 while let Some(code_start) = search.find("<code>") {
31 let content_start = code_start + "<code>".len();
32 let content_end = search[content_start..]
33 .find("</code>")
34 .expect("Unclosed <code> in backup codes");
35 codes.push(search[content_start..content_start + content_end].to_string());
36 search = &search[content_start + content_end..];
37 }
38
39 codes
40 }
41
42 /// Generate a valid TOTP code from a base32 secret.
43 fn generate_totp_code(secret_base32: &str, email: &str) -> String {
44 let bytes = totp_rs::Secret::Encoded(secret_base32.to_string())
45 .to_bytes()
46 .expect("Invalid TOTP secret");
47 let totp = totp_rs::TOTP::new(
48 totp_rs::Algorithm::SHA1,
49 6,
50 1,
51 30,
52 bytes,
53 Some("Makenotwork".into()),
54 email.into(),
55 )
56 .expect("TOTP creation failed");
57 totp.generate_current().expect("TOTP generation failed")
58 }
59
60 /// Set up TOTP for the currently logged-in user and enable it.
61 /// Returns (secret_base32, backup_codes).
62 async fn setup_and_enable_totp(h: &mut TestHarness, email: &str) -> (String, Vec<String>) {
63 let resp = h.client.post_form("/api/users/me/totp/setup", "").await;
64 assert_eq!(resp.status.as_u16(), 200, "TOTP setup failed: {}", resp.text);
65
66 let secret = extract_totp_secret(&resp.text);
67 let codes = extract_backup_codes(&resp.text);
68 assert!(!secret.is_empty(), "TOTP secret should not be empty");
69 assert!(!codes.is_empty(), "Should have backup codes");
70
71 let code = generate_totp_code(&secret, email);
72 let resp = h
73 .client
74 .post_form("/api/users/me/totp/confirm", &format!("code={}", code))
75 .await;
76 assert_eq!(resp.status.as_u16(), 200, "TOTP confirm failed: {}", resp.text);
77
78 (secret, codes)
79 }
80
81 /// Login flow when TOTP is enabled: POST /login -> 303 to /auth/2fa -> POST /auth/verify-2fa.
82 async fn login_with_2fa(h: &mut TestHarness, username: &str, password: &str, code: &str) {
83 h.client.fetch_csrf_token().await;
84 let resp = h
85 .client
86 .post_form(
87 "/login",
88 &format!(
89 "login={}&password={}",
90 urlencoding::encode(username),
91 urlencoding::encode(password)
92 ),
93 )
94 .await;
95 // Login should redirect to 2FA page
96 assert!(
97 resp.status.is_redirection()
98 || resp.text.contains("/auth/2fa")
99 || resp.text.contains("HX-Redirect"),
100 "Expected redirect to /auth/2fa, got {} — {}",
101 resp.status,
102 resp.text
103 );
104
105 // Load the 2FA page to get CSRF token
106 let resp = h.client.get("/auth/2fa").await;
107 assert_eq!(resp.status.as_u16(), 200, "2FA page failed: {}", resp.text);
108
109 // Submit the code
110 let resp = h
111 .client
112 .post_form("/auth/verify-2fa", &format!("code={}", code))
113 .await;
114 assert!(
115 resp.status.is_redirection() || resp.status.is_success(),
116 "2FA verification failed with {}: {}",
117 resp.status,
118 resp.text
119 );
120 }
121
122 // ── Tests ──
123
124 #[tokio::test]
125 async fn totp_setup_and_confirm() {
126 let mut h = TestHarness::new().await;
127 h.signup("totp1", "totp1@test.com", "Password1!").await;
128
129 let (_secret, codes) = setup_and_enable_totp(&mut h, "totp1@test.com").await;
130 assert_eq!(codes.len(), 10, "Should have 10 backup codes");
131
132 // Check status
133 let resp = h.client.get("/api/users/me/totp/status").await;
134 assert_eq!(resp.status.as_u16(), 200);
135 assert!(
136 resp.text.contains("true") || resp.text.contains("enabled") || resp.text.contains("Enabled"),
137 "TOTP should be enabled: {}",
138 resp.text
139 );
140 }
141
142 #[tokio::test]
143 async fn totp_login_requires_2fa() {
144 let mut h = TestHarness::new().await;
145 let uid = h.signup("totp2", "totp2@test.com", "Password1!").await;
146 let (secret, _) = setup_and_enable_totp(&mut h, "totp2@test.com").await;
147
148 // Reset the anti-replay counter so the login TOTP code (same 30s window)
149 // is not rejected as a replay of the confirm step.
150 sqlx::query("UPDATE users SET totp_last_used_step = 0 WHERE id = $1")
151 .bind(uid)
152 .execute(&h.db)
153 .await
154 .expect("reset totp_last_used_step");
155
156 // Logout
157 h.client.post_form("/logout", "").await;
158
159 // Login with TOTP
160 let code = generate_totp_code(&secret, "totp2@test.com");
161 login_with_2fa(&mut h, "totp2", "Password1!", &code).await;
162
163 // Verify we're logged in
164 let resp = h.client.get("/dashboard").await;
165 assert_eq!(resp.status.as_u16(), 200, "Should be on dashboard after 2FA login");
166 }
167
168 #[tokio::test]
169 async fn totp_backup_code_login() {
170 let mut h = TestHarness::new().await;
171 h.signup("totp3", "totp3@test.com", "Password1!").await;
172 let (_, codes) = setup_and_enable_totp(&mut h, "totp3@test.com").await;
173
174 h.client.post_form("/logout", "").await;
175
176 // Login using a backup code
177 login_with_2fa(&mut h, "totp3", "Password1!", &codes[0]).await;
178
179 let resp = h.client.get("/dashboard").await;
180 assert_eq!(resp.status.as_u16(), 200, "Should be on dashboard after backup code login");
181 }
182
183 #[tokio::test]
184 async fn totp_backup_code_single_use() {
185 let mut h = TestHarness::new().await;
186 h.signup("totp4", "totp4@test.com", "Password1!").await;
187 let (_, codes) = setup_and_enable_totp(&mut h, "totp4@test.com").await;
188 let used_code = codes[0].clone();
189
190 // Use backup code once
191 h.client.post_form("/logout", "").await;
192 login_with_2fa(&mut h, "totp4", "Password1!", &used_code).await;
193
194 // Logout and try the same code again
195 h.client.post_form("/logout", "").await;
196 h.client.fetch_csrf_token().await;
197 let resp = h
198 .client
199 .post_form(
200 "/login",
201 &format!("login=totp4&password={}", urlencoding::encode("Password1!")),
202 )
203 .await;
204 assert!(
205 resp.status.is_redirection()
206 || resp.text.contains("/auth/2fa")
207 || resp.text.contains("HX-Redirect"),
208 "Expected redirect to 2FA"
209 );
210
211 let _resp = h.client.get("/auth/2fa").await;
212 let resp = h
213 .client
214 .post_form("/auth/verify-2fa", &format!("code={}", used_code))
215 .await;
216
217 // Should fail — backup code already consumed
218 assert!(
219 resp.text.contains("Invalid") || resp.text.contains("error") || resp.text.contains("invalid"),
220 "Used backup code should be rejected: {}",
221 resp.text
222 );
223 }
224
225 #[tokio::test]
226 async fn totp_disable_with_password() {
227 let mut h = TestHarness::new().await;
228 h.signup("totp5", "totp5@test.com", "Password1!").await;
229 setup_and_enable_totp(&mut h, "totp5@test.com").await;
230
231 // Disable TOTP
232 let resp = h
233 .client
234 .post_form(
235 "/api/users/me/totp/disable",
236 &format!("password={}", urlencoding::encode("Password1!")),
237 )
238 .await;
239 assert_eq!(resp.status.as_u16(), 200, "TOTP disable failed: {}", resp.text);
240
241 // Verify status is disabled
242 let resp = h.client.get("/api/users/me/totp/status").await;
243 assert_eq!(resp.status.as_u16(), 200);
244 assert!(
245 resp.text.contains("false") || resp.text.contains("disabled") || resp.text.contains("Disabled")
246 || !resp.text.contains("Enabled"),
247 "TOTP should be disabled: {}",
248 resp.text
249 );
250
251 // Logout and login — should not require 2FA
252 h.client.post_form("/logout", "").await;
253 h.login("totp5", "Password1!").await;
254
255 let resp = h.client.get("/dashboard").await;
256 assert_eq!(resp.status.as_u16(), 200, "Login after TOTP disable should not require 2FA");
257 }
258
259 #[tokio::test]
260 async fn totp_invalid_code_rejected() {
261 let mut h = TestHarness::new().await;
262 h.signup("totp6", "totp6@test.com", "Password1!").await;
263 setup_and_enable_totp(&mut h, "totp6@test.com").await;
264
265 h.client.post_form("/logout", "").await;
266
267 // Start login
268 h.client.fetch_csrf_token().await;
269 let resp = h
270 .client
271 .post_form(
272 "/login",
273 &format!("login=totp6&password={}", urlencoding::encode("Password1!")),
274 )
275 .await;
276 assert!(
277 resp.status.is_redirection()
278 || resp.text.contains("/auth/2fa")
279 || resp.text.contains("HX-Redirect"),
280 "Expected redirect to 2FA"
281 );
282
283 let _resp = h.client.get("/auth/2fa").await;
284
285 // Submit an invalid code
286 let resp = h
287 .client
288 .post_form("/auth/verify-2fa", "code=000000")
289 .await;
290
291 // Should show error, not redirect to dashboard
292 assert!(
293 resp.text.contains("Invalid") || resp.text.contains("invalid") || resp.text.contains("error"),
294 "Invalid TOTP code should be rejected: {}",
295 resp.text
296 );
297 assert!(
298 !resp.text.contains("/dashboard") || resp.status.as_u16() == 200,
299 "Should not redirect to dashboard with invalid code"
300 );
301 }
302