1pub use victauri_core::middleware::{
2 AuthState, default_rate_limiter, dns_rebinding_guard, origin_guard, rate_limit, require_auth,
3 security_headers,
4};
5pub use victauri_core::security::{
6 self, RateLimiter as RateLimiterState, constant_time_eq, generate_token, is_allowed_origin,
7 is_localhost_host,
8};
9
10#[cfg(test)]
11mod tests {
12 use std::sync::Arc;
13
14 use super::*;
15 use axum::Router;
16 use axum::body::Body;
17 use axum::http::StatusCode;
18 use axum::middleware;
19 use axum::routing::get;
20 use tower::ServiceExt;
21
22 async fn ok_handler() -> &'static str {
23 "ok"
24 }
25
26 #[test]
27 fn token_generation_is_unique() {
28 let t1 = generate_token();
29 let t2 = generate_token();
30 assert_ne!(t1, t2);
31 assert_eq!(t1.len(), 36); }
33
34 #[test]
35 fn token_is_valid_uuid() {
36 let token = generate_token();
37 assert!(uuid::Uuid::parse_str(&token).is_ok());
38 }
39
40 #[test]
41 fn rate_limiter_allows_within_budget() {
42 let limiter = RateLimiterState::new(10);
43 for _ in 0..10 {
44 assert!(limiter.try_acquire());
45 }
46 }
47
48 #[test]
49 fn rate_limiter_denies_when_exhausted() {
50 let limiter = RateLimiterState::new(5);
51 for _ in 0..5 {
52 assert!(limiter.try_acquire());
53 }
54 assert!(!limiter.try_acquire());
55 }
56
57 #[test]
58 fn rate_limiter_initial_tokens_match_max() {
59 let limiter = RateLimiterState::new(42);
60 assert_eq!(limiter.current_tokens(), 42);
61 assert_eq!(limiter.max_tokens(), 42);
62 }
63
64 #[test]
65 fn rate_limiter_concurrent_acquire() {
66 let started = std::time::Instant::now();
68 let limiter = Arc::new(RateLimiterState::new(1000));
69 let mut handles = vec![];
70 for _ in 0..10 {
71 let l = limiter.clone();
72 handles.push(std::thread::spawn(move || {
73 let mut acquired = 0;
74 for _ in 0..200 {
75 if l.try_acquire() {
76 acquired += 1;
77 }
78 }
79 acquired
80 }));
81 }
82 let total: u64 = handles.into_iter().map(|h| h.join().unwrap()).sum();
83 assert!(
84 total >= 1000,
85 "should dispense at least the initial budget, got {total}"
86 );
87 let elapsed_ms = u64::try_from(started.elapsed().as_millis()).unwrap_or(u64::MAX);
92 assert!(
93 total <= 1000 + elapsed_ms + 1,
94 "refill overshoot: {total} tokens in {elapsed_ms} ms"
95 );
96 }
97
98 #[test]
99 fn default_rate_limiter_has_expected_tokens() {
100 let limiter = default_rate_limiter();
101 assert_eq!(limiter.max_tokens(), 1000);
102 }
103
104 #[test]
105 fn rate_limiter_zero_capacity() {
106 let limiter = RateLimiterState::new(0);
107 assert!(!limiter.try_acquire());
108 }
109
110 fn dns_rebinding_router() -> Router {
113 Router::new()
114 .route("/test", get(ok_handler))
115 .layer(middleware::from_fn(dns_rebinding_guard))
116 }
117
118 fn dns_request(host: Option<&str>) -> axum::extract::Request<Body> {
119 let mut builder = axum::extract::Request::builder().uri("/test");
120 if let Some(h) = host {
121 builder = builder.header("host", h);
122 }
123 builder.body(Body::empty()).unwrap()
124 }
125
126 #[tokio::test]
127 async fn dns_rebinding_allows_localhost() {
128 let app = dns_rebinding_router();
129 let resp = app.oneshot(dns_request(Some("localhost"))).await.unwrap();
130 assert_eq!(resp.status(), StatusCode::OK);
131 }
132
133 #[tokio::test]
134 async fn dns_rebinding_allows_127_0_0_1() {
135 let app = dns_rebinding_router();
136 let resp = app.oneshot(dns_request(Some("127.0.0.1"))).await.unwrap();
137 assert_eq!(resp.status(), StatusCode::OK);
138 }
139
140 #[tokio::test]
141 async fn dns_rebinding_allows_ipv6_bracketed() {
142 let app = dns_rebinding_router();
143 let resp = app.oneshot(dns_request(Some("[::1]"))).await.unwrap();
144 assert_eq!(resp.status(), StatusCode::OK);
145 }
146
147 #[tokio::test]
148 async fn dns_rebinding_allows_ipv6_bracketed_with_port() {
149 let app = dns_rebinding_router();
150 let resp = app.oneshot(dns_request(Some("[::1]:7373"))).await.unwrap();
151 assert_eq!(resp.status(), StatusCode::OK);
152 }
153
154 #[tokio::test]
155 async fn dns_rebinding_allows_ipv6_bare() {
156 let app = dns_rebinding_router();
157 let resp = app.oneshot(dns_request(Some("::1"))).await.unwrap();
158 assert_eq!(resp.status(), StatusCode::OK);
159 }
160
161 #[tokio::test]
162 async fn dns_rebinding_blocks_empty_host() {
163 let app = dns_rebinding_router();
164 let resp = app.oneshot(dns_request(None)).await.unwrap();
165 assert_eq!(resp.status(), StatusCode::FORBIDDEN);
166 }
167
168 #[tokio::test]
169 async fn dns_rebinding_blocks_evil_com() {
170 let app = dns_rebinding_router();
171 let resp = app.oneshot(dns_request(Some("evil.com"))).await.unwrap();
172 assert_eq!(resp.status(), StatusCode::FORBIDDEN);
173 }
174
175 #[tokio::test]
176 async fn dns_rebinding_blocks_localhost_subdomain() {
177 let app = dns_rebinding_router();
178 let resp = app
179 .oneshot(dns_request(Some("localhost.evil.com")))
180 .await
181 .unwrap();
182 assert_eq!(resp.status(), StatusCode::FORBIDDEN);
183 }
184
185 #[tokio::test]
186 async fn dns_rebinding_blocks_ip_subdomain() {
187 let app = dns_rebinding_router();
188 let resp = app
189 .oneshot(dns_request(Some("127.0.0.1.evil.com")))
190 .await
191 .unwrap();
192 assert_eq!(resp.status(), StatusCode::FORBIDDEN);
193 }
194
195 fn origin_router() -> Router {
198 Router::new()
199 .route("/test", get(ok_handler))
200 .layer(middleware::from_fn(origin_guard))
201 }
202
203 fn origin_request(origin: Option<&str>) -> axum::extract::Request<Body> {
204 let mut builder = axum::extract::Request::builder().uri("/test");
205 if let Some(o) = origin {
206 builder = builder.header("origin", o);
207 }
208 builder.body(Body::empty()).unwrap()
209 }
210
211 #[tokio::test]
212 async fn origin_allows_no_origin() {
213 let app = origin_router();
214 let resp = app.oneshot(origin_request(None)).await.unwrap();
215 assert_eq!(resp.status(), StatusCode::OK);
216 }
217
218 #[tokio::test]
219 async fn origin_allows_localhost_http() {
220 let app = origin_router();
221 let resp = app
222 .oneshot(origin_request(Some("http://localhost:3000")))
223 .await
224 .unwrap();
225 assert_eq!(resp.status(), StatusCode::OK);
226 }
227
228 #[tokio::test]
229 async fn origin_allows_127_0_0_1_https() {
230 let app = origin_router();
231 let resp = app
232 .oneshot(origin_request(Some("https://127.0.0.1:8080")))
233 .await
234 .unwrap();
235 assert_eq!(resp.status(), StatusCode::OK);
236 }
237
238 #[tokio::test]
241 async fn origin_blocks_tauri_scheme() {
242 for origin in ["tauri://localhost", "tauri://evil.com"] {
243 let app = origin_router();
244 let resp = app.oneshot(origin_request(Some(origin))).await.unwrap();
245 assert_eq!(resp.status(), StatusCode::FORBIDDEN, "{origin}");
246 }
247 }
248
249 #[tokio::test]
251 async fn origin_blocks_unreadable_origin() {
252 let app = origin_router();
253 let req = axum::extract::Request::builder()
254 .uri("/test")
255 .header(
256 "origin",
257 axum::http::HeaderValue::from_bytes(b"http://\xe9vil.com").unwrap(),
258 )
259 .body(Body::empty())
260 .unwrap();
261 assert_eq!(
262 app.oneshot(req).await.unwrap().status(),
263 StatusCode::FORBIDDEN
264 );
265 }
266
267 #[tokio::test]
271 async fn browser_initiated_requests_are_refused_without_spending_a_token() {
272 let limiter = Arc::new(RateLimiterState::new(5));
273 let app = Router::new()
274 .route("/test", get(ok_handler))
275 .layer(middleware::from_fn_with_state(limiter.clone(), rate_limit))
276 .layer(middleware::from_fn(origin_guard));
277 for site in ["cross-site", "same-site", "same-origin"] {
278 let req = axum::extract::Request::builder()
279 .uri("/test")
280 .header("sec-fetch-site", site)
281 .body(Body::empty())
282 .unwrap();
283 let resp = app.clone().oneshot(req).await.unwrap();
284 assert_eq!(resp.status(), StatusCode::FORBIDDEN, "{site}");
285 }
286 assert_eq!(
287 limiter.current_tokens(),
288 5,
289 "a refused request spent a token"
290 );
291 let req = axum::extract::Request::builder()
293 .uri("/test")
294 .header("sec-fetch-site", "none")
295 .body(Body::empty())
296 .unwrap();
297 assert_eq!(
298 app.clone().oneshot(req).await.unwrap().status(),
299 StatusCode::OK
300 );
301 let req = axum::extract::Request::builder()
302 .uri("/test")
303 .body(Body::empty())
304 .unwrap();
305 assert_eq!(app.oneshot(req).await.unwrap().status(), StatusCode::OK);
306 }
307
308 #[tokio::test]
309 async fn origin_blocks_null() {
310 let app = origin_router();
311 let resp = app.oneshot(origin_request(Some("null"))).await.unwrap();
312 assert_eq!(resp.status(), StatusCode::FORBIDDEN);
313 }
314
315 #[tokio::test]
316 async fn origin_blocks_evil_com() {
317 let app = origin_router();
318 let resp = app
319 .oneshot(origin_request(Some("http://evil.com")))
320 .await
321 .unwrap();
322 assert_eq!(resp.status(), StatusCode::FORBIDDEN);
323 }
324
325 fn security_headers_router() -> Router {
328 Router::new()
329 .route("/test", get(ok_handler))
330 .layer(middleware::from_fn(security_headers))
331 }
332
333 #[tokio::test]
334 async fn security_headers_x_content_type_options() {
335 let app = security_headers_router();
336 let req = axum::extract::Request::builder()
337 .uri("/test")
338 .body(Body::empty())
339 .unwrap();
340 let resp = app.oneshot(req).await.unwrap();
341 assert_eq!(resp.status(), StatusCode::OK);
342 assert_eq!(
343 resp.headers().get("x-content-type-options").unwrap(),
344 "nosniff"
345 );
346 }
347
348 #[tokio::test]
349 async fn security_headers_cache_control() {
350 let app = security_headers_router();
351 let req = axum::extract::Request::builder()
352 .uri("/test")
353 .body(Body::empty())
354 .unwrap();
355 let resp = app.oneshot(req).await.unwrap();
356 assert_eq!(resp.status(), StatusCode::OK);
357 assert_eq!(resp.headers().get("cache-control").unwrap(), "no-store");
358 }
359
360 #[tokio::test]
361 async fn security_headers_x_frame_options() {
362 let app = security_headers_router();
363 let req = axum::extract::Request::builder()
364 .uri("/test")
365 .body(Body::empty())
366 .unwrap();
367 let resp = app.oneshot(req).await.unwrap();
368 assert_eq!(resp.status(), StatusCode::OK);
369 assert_eq!(resp.headers().get("x-frame-options").unwrap(), "DENY");
370 }
371
372 fn auth_router(token: Option<&str>) -> Router {
375 let state = Arc::new(AuthState {
376 token: token.map(String::from),
377 });
378 Router::new()
379 .route("/test", get(ok_handler))
380 .layer(middleware::from_fn_with_state(state, require_auth))
381 }
382
383 fn auth_request(token: Option<&str>) -> axum::extract::Request<Body> {
384 let mut builder = axum::extract::Request::builder().uri("/test");
385 if let Some(t) = token {
386 builder = builder.header("authorization", format!("Bearer {t}"));
387 }
388 builder.body(Body::empty()).unwrap()
389 }
390
391 #[tokio::test]
392 async fn auth_allows_correct_token() {
393 let app = auth_router(Some("secret-123"));
394 let resp = app.oneshot(auth_request(Some("secret-123"))).await.unwrap();
395 assert_eq!(resp.status(), StatusCode::OK);
396 }
397
398 #[tokio::test]
399 async fn auth_rejects_wrong_token() {
400 let app = auth_router(Some("secret-123"));
401 let resp = app
402 .oneshot(auth_request(Some("wrong-token")))
403 .await
404 .unwrap();
405 assert_eq!(resp.status(), StatusCode::UNAUTHORIZED);
406 }
407
408 #[tokio::test]
409 async fn auth_rejects_missing_token() {
410 let app = auth_router(Some("secret-123"));
411 let resp = app.oneshot(auth_request(None)).await.unwrap();
412 assert_eq!(resp.status(), StatusCode::UNAUTHORIZED);
413 }
414
415 #[tokio::test]
416 async fn auth_allows_any_when_disabled() {
417 let app = auth_router(None);
418 let resp = app.oneshot(auth_request(None)).await.unwrap();
419 assert_eq!(resp.status(), StatusCode::OK);
420 }
421
422 #[tokio::test]
423 async fn auth_case_insensitive_bearer_prefix() {
424 let state = Arc::new(AuthState {
425 token: Some("my-token".into()),
426 });
427 let app = Router::new()
428 .route("/test", get(ok_handler))
429 .layer(middleware::from_fn_with_state(state, require_auth));
430
431 let req = axum::extract::Request::builder()
432 .uri("/test")
433 .header("authorization", "BEARER my-token")
434 .body(Body::empty())
435 .unwrap();
436 let resp = app.oneshot(req).await.unwrap();
437 assert_eq!(resp.status(), StatusCode::OK);
438 }
439
440 #[tokio::test]
441 async fn auth_rejects_non_bearer_scheme() {
442 let app = auth_router(Some("secret"));
443 let req = axum::extract::Request::builder()
444 .uri("/test")
445 .header("authorization", "Basic c2VjcmV0")
446 .body(Body::empty())
447 .unwrap();
448 let resp = app.oneshot(req).await.unwrap();
449 assert_eq!(resp.status(), StatusCode::UNAUTHORIZED);
450 }
451
452 #[tokio::test]
455 async fn rate_limiter_returns_429_when_exhausted() {
456 let limiter = Arc::new(RateLimiterState::new(2));
457 let app = Router::new()
458 .route("/test", get(ok_handler))
459 .layer(middleware::from_fn_with_state(limiter, rate_limit));
460
461 let app2 = app.clone();
462 let app3 = app2.clone();
463
464 let req = axum::extract::Request::builder()
465 .uri("/test")
466 .body(Body::empty())
467 .unwrap();
468 assert_eq!(app.oneshot(req).await.unwrap().status(), StatusCode::OK);
469
470 let req = axum::extract::Request::builder()
471 .uri("/test")
472 .body(Body::empty())
473 .unwrap();
474 assert_eq!(app2.oneshot(req).await.unwrap().status(), StatusCode::OK);
475
476 let req = axum::extract::Request::builder()
477 .uri("/test")
478 .body(Body::empty())
479 .unwrap();
480 assert_eq!(
481 app3.oneshot(req).await.unwrap().status(),
482 StatusCode::TOO_MANY_REQUESTS
483 );
484 }
485
486 #[tokio::test]
489 async fn combined_layers_enforce_all_guards() {
490 let auth_state = Arc::new(AuthState {
491 token: Some("tok-123".into()),
492 });
493 let limiter = Arc::new(RateLimiterState::new(100));
494
495 let app = Router::new()
496 .route("/test", get(ok_handler))
497 .layer(middleware::from_fn_with_state(auth_state, require_auth))
498 .layer(middleware::from_fn_with_state(limiter, rate_limit))
499 .layer(middleware::from_fn(security_headers))
500 .layer(middleware::from_fn(origin_guard))
501 .layer(middleware::from_fn(dns_rebinding_guard));
502
503 let req = axum::extract::Request::builder()
505 .uri("/test")
506 .header("authorization", "Bearer tok-123")
507 .header("host", "127.0.0.1:7373")
508 .body(Body::empty())
509 .unwrap();
510 let resp = app.clone().oneshot(req).await.unwrap();
511 assert_eq!(resp.status(), StatusCode::OK);
512 assert_eq!(resp.headers().get("x-frame-options").unwrap(), "DENY");
513
514 let req = axum::extract::Request::builder()
516 .uri("/test")
517 .header("authorization", "Bearer tok-123")
518 .header("host", "evil.com")
519 .body(Body::empty())
520 .unwrap();
521 let resp = app.clone().oneshot(req).await.unwrap();
522 assert_eq!(resp.status(), StatusCode::FORBIDDEN);
523
524 let req = axum::extract::Request::builder()
526 .uri("/test")
527 .header("authorization", "Bearer tok-123")
528 .header("host", "localhost")
529 .header("origin", "https://evil.com")
530 .body(Body::empty())
531 .unwrap();
532 let resp = app.clone().oneshot(req).await.unwrap();
533 assert_eq!(resp.status(), StatusCode::FORBIDDEN);
534
535 let req = axum::extract::Request::builder()
537 .uri("/test")
538 .header("host", "localhost")
539 .body(Body::empty())
540 .unwrap();
541 let resp = app.oneshot(req).await.unwrap();
542 assert_eq!(resp.status(), StatusCode::UNAUTHORIZED);
543 }
544
545 #[test]
546 fn origin_guard_allows_localhost_variants() {
547 assert!(is_allowed_origin("http://localhost"));
548 assert!(is_allowed_origin("http://localhost:7373"));
549 assert!(is_allowed_origin("https://localhost"));
550 assert!(is_allowed_origin("https://localhost:443"));
551 assert!(is_allowed_origin("http://127.0.0.1"));
552 assert!(is_allowed_origin("http://127.0.0.1:8080"));
553 assert!(is_allowed_origin("https://127.0.0.1"));
554 assert!(is_allowed_origin("http://[::1]"));
555 assert!(is_allowed_origin("http://[::1]:7373"));
556 assert!(!is_allowed_origin("tauri://localhost"));
557 assert!(!is_allowed_origin("tauri://some-app"));
558 }
559
560 #[test]
561 fn origin_guard_rejects_prefix_smuggling() {
562 assert!(!is_allowed_origin("http://localhost.evil.com"));
563 assert!(!is_allowed_origin("https://localhost.evil.com"));
564 assert!(!is_allowed_origin("https://127.0.0.1.evil.com"));
565 assert!(!is_allowed_origin("http://[::1].evil.com"));
566 }
567
568 #[test]
569 fn origin_guard_rejects_userinfo_trick() {
570 assert!(!is_allowed_origin("http://localhost@evil.com"));
571 assert!(!is_allowed_origin("http://127.0.0.1@evil.com"));
572 }
573
574 #[test]
575 fn origin_guard_rejects_foreign_and_malformed() {
576 assert!(!is_allowed_origin("http://evil.com"));
577 assert!(!is_allowed_origin("https://attacker.io"));
578 assert!(!is_allowed_origin("not-a-url"));
579 assert!(!is_allowed_origin(""));
580 assert!(!is_allowed_origin("ftp://localhost"));
581 }
582
583 #[test]
586 fn constant_time_eq_equal_strings() {
587 assert!(constant_time_eq(b"secret-token-123", b"secret-token-123"));
588 }
589
590 #[test]
591 fn constant_time_eq_different_strings() {
592 assert!(!constant_time_eq(b"secret-token-123", b"wrong-token-9999"));
593 }
594
595 #[test]
596 fn constant_time_eq_different_lengths() {
597 assert!(!constant_time_eq(b"short", b"longer-string"));
598 }
599
600 #[test]
601 fn constant_time_eq_empty_strings() {
602 assert!(constant_time_eq(b"", b""));
603 }
604
605 #[test]
606 fn constant_time_eq_one_empty() {
607 assert!(!constant_time_eq(b"", b"notempty"));
608 assert!(!constant_time_eq(b"notempty", b""));
609 }
610
611 #[test]
612 fn constant_time_eq_single_bit_difference() {
613 assert!(!constant_time_eq(b"A", b"B"));
614 }
615
616 #[tokio::test]
621 async fn security_headers_cors_deny() {
622 let app = security_headers_router();
623 let req = axum::extract::Request::builder()
624 .uri("/test")
625 .body(Body::empty())
626 .unwrap();
627 let resp = app.oneshot(req).await.unwrap();
628 assert!(resp.headers().get("access-control-allow-origin").is_none());
629 }
630
631 #[tokio::test]
632 async fn security_headers_csp() {
633 let app = security_headers_router();
634 let req = axum::extract::Request::builder()
635 .uri("/test")
636 .body(Body::empty())
637 .unwrap();
638 let resp = app.oneshot(req).await.unwrap();
639 assert_eq!(
640 resp.headers().get("content-security-policy").unwrap(),
641 "default-src 'none'"
642 );
643 }
644}