1use http::{Request, StatusCode};
3use security_rust::Scanner;
4use std::future::Future;
5use std::pin::Pin;
6use std::sync::Arc;
7use std::task::{Context as TaskCtx, Poll};
8
9pub use security_rust::{AttackCategory, DetectionResult, ScannerBuilder, Severity};
10
11#[derive(Debug, thiserror::Error)]
12pub enum SecurityError {
13 #[error("attack blocked: {0}")]
14 AttackBlocked(String),
15 #[error("request body too large")]
17 BodyTooLarge,
18 #[error("inner error: {0}")]
19 Inner(#[from] Box<dyn std::error::Error + Send + Sync>),
20}
21
22impl SecurityError {
23 pub fn to_http_status(&self) -> StatusCode {
24 match self {
25 Self::AttackBlocked(_) => StatusCode::FORBIDDEN,
26 Self::BodyTooLarge => StatusCode::PAYLOAD_TOO_LARGE,
27 Self::Inner(_) => StatusCode::INTERNAL_SERVER_ERROR,
28 }
29 }
30}
31
32impl axum::response::IntoResponse for SecurityError {
36 fn into_response(self) -> axum::response::Response {
37 let status = self.to_http_status();
38 let body = match &self {
39 Self::AttackBlocked(types) => {
40 format!(r#"{{"error":"attack blocked","types":"{types}"}}"#)
41 }
42 Self::BodyTooLarge => r#"{"error":"request body too large"}"#.to_string(),
43 Self::Inner(e) => {
44 tracing::error!(error = %e, "security middleware internal error");
45 r#"{"error":"internal server error"}"#.to_string()
46 }
47 };
48 (status, body).into_response()
49 }
50}
51
52pub struct SecurityScanner {
54 scanner: Scanner,
55}
56
57impl SecurityScanner {
58 pub fn new() -> Self {
60 Self {
61 scanner: Scanner::default(),
62 }
63 }
64
65 pub fn scan(&self, input: &str) -> Vec<DetectionResult> {
67 self.scanner.scan(input)
68 }
69
70 pub fn scan_parts(&self, parts: &[&str]) -> Vec<DetectionResult> {
72 let mut results = Vec::with_capacity(parts.len() * 2);
73 for part in parts {
74 results.extend(self.scanner.scan(part));
75 }
76 results
77 }
78
79 pub fn scan_body(&self, body: &[u8]) -> Vec<DetectionResult> {
81 if let Ok(s) = std::str::from_utf8(body) {
82 self.scanner.scan(s)
83 } else {
84 Vec::new()
85 }
86 }
87}
88
89impl Default for SecurityScanner {
90 fn default() -> Self {
91 Self::new()
92 }
93}
94
95fn evaluate(results: &[DetectionResult]) -> Option<SecurityError> {
98 let mut blocked = false;
99 for r in results {
100 tracing::warn!(
101 attack_type = %r.attack_type,
102 category = ?r.category,
103 severity = ?r.severity,
104 matched = %r.matched_pattern,
105 "attack detected"
106 );
107 if r.attack_type != "jwt_attack" && matches!(r.severity, Severity::High | Severity::Critical)
111 {
112 blocked = true;
113 }
114 }
115 if blocked {
116 let attack_types: Vec<String> = results.iter().map(|r| r.attack_type.to_string()).collect();
117 return Some(SecurityError::AttackBlocked(attack_types.join(", ")));
118 }
119 None
120}
121
122fn percent_decode(s: &str) -> String {
125 let bytes = s.as_bytes();
126 let mut out = Vec::with_capacity(bytes.len());
127 let mut i = 0;
128 while i < bytes.len() {
129 if bytes[i] == b'%'
130 && i + 2 < bytes.len()
131 && let (Some(h), Some(l)) = (hex_val(bytes[i + 1]), hex_val(bytes[i + 2]))
132 {
133 out.push(h * 16 + l);
134 i += 3;
135 continue;
136 }
137 out.push(bytes[i]);
138 i += 1;
139 }
140 String::from_utf8_lossy(&out).into_owned()
141}
142
143fn hex_val(b: u8) -> Option<u8> {
144 match b {
145 b'0'..=b'9' => Some(b - b'0'),
146 b'a'..=b'f' => Some(b - b'a' + 10),
147 b'A'..=b'F' => Some(b - b'A' + 10),
148 _ => None,
149 }
150}
151
152fn is_proxy_header(name: &http::header::HeaderName) -> bool {
160 matches!(
161 name.as_str(),
162 "x-forwarded-for"
163 | "x-real-ip"
164 | "forwarded"
165 | "x-forwarded-host"
166 | "x-forwarded-proto"
167 | "x-forwarded-port"
168 | "authorization"
169 )
170}
171
172fn request_parts<B>(req: &Request<B>) -> Vec<String> {
173 let mut parts: Vec<String> = Vec::new();
174 parts.push(percent_decode(&req.uri().to_string()));
175 for (name, value) in req.headers() {
176 if is_proxy_header(name) {
177 continue;
178 }
179 if let Ok(v) = value.to_str() {
180 parts.push(v.to_string());
181 }
182 }
183 parts
184}
185
186#[derive(Clone)]
189pub struct SecurityLayer {
190 scanner: Arc<SecurityScanner>,
191}
192
193impl SecurityLayer {
194 pub fn new() -> Self {
195 Self {
196 scanner: Arc::new(SecurityScanner::new()),
197 }
198 }
199}
200
201impl Default for SecurityLayer {
202 fn default() -> Self {
203 Self::new()
204 }
205}
206
207impl<S> tower::Layer<S> for SecurityLayer {
208 type Service = SecurityService<S>;
209
210 fn layer(&self, inner: S) -> Self::Service {
211 SecurityService {
212 inner,
213 scanner: Arc::clone(&self.scanner),
214 }
215 }
216}
217
218#[derive(Clone)]
219pub struct SecurityService<S> {
220 inner: S,
221 scanner: Arc<SecurityScanner>,
222}
223
224impl<S, B> tower::Service<Request<B>> for SecurityService<S>
225where
226 S: tower::Service<Request<B>> + Send + 'static,
227 S::Future: Send + 'static,
228 S::Error: Into<Box<dyn std::error::Error + Send + Sync>>,
229 S::Response: From<axum::response::Response> + Send,
232 B: Send + 'static,
233{
234 type Response = S::Response;
235 type Error = SecurityError;
236 type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;
237
238 fn poll_ready(&mut self, cx: &mut TaskCtx<'_>) -> Poll<Result<(), Self::Error>> {
239 self.inner
240 .poll_ready(cx)
241 .map_err(|e| SecurityError::Inner(e.into()))
242 }
243
244 fn call(&mut self, req: Request<B>) -> Self::Future {
245 let scanner = Arc::clone(&self.scanner);
246 let parts = request_parts(&req);
247 let strings: Vec<&str> = parts.iter().map(|s| s.as_str()).collect();
248 let results = scanner.scan_parts(&strings);
249
250 if let Some(err) = evaluate(&results) {
251 use axum::response::IntoResponse;
252 let resp: S::Response = err.into_response().into();
253 return Box::pin(async move { Ok(resp) });
254 }
255
256 let fut = self.inner.call(req);
257 Box::pin(async move { fut.await.map_err(|e| SecurityError::Inner(e.into())) })
258 }
259}
260
261#[derive(Clone)]
270pub struct SecurityBodyLayer {
271 scanner: Arc<SecurityScanner>,
272 body_limit: usize,
273}
274
275impl SecurityBodyLayer {
276 pub fn new() -> Self {
277 Self {
278 scanner: Arc::new(SecurityScanner::new()),
279 body_limit: 10 * 1024 * 1024,
280 }
281 }
282
283 pub fn body_limit(mut self, limit: usize) -> Self {
287 self.body_limit = limit;
288 self
289 }
290}
291
292impl Default for SecurityBodyLayer {
293 fn default() -> Self {
294 Self::new()
295 }
296}
297
298impl<S> tower::Layer<S> for SecurityBodyLayer {
299 type Service = SecurityBodyService<S>;
300
301 fn layer(&self, inner: S) -> Self::Service {
302 SecurityBodyService {
303 inner,
304 scanner: Arc::clone(&self.scanner),
305 body_limit: self.body_limit,
306 }
307 }
308}
309
310#[derive(Clone)]
311pub struct SecurityBodyService<S> {
312 inner: S,
313 scanner: Arc<SecurityScanner>,
314 body_limit: usize,
315}
316
317impl<S> tower::Service<Request<axum::body::Body>> for SecurityBodyService<S>
318where
319 S: tower::Service<Request<axum::body::Body>> + Clone + Send + 'static,
320 S::Future: Send + 'static,
321 S::Error: Into<Box<dyn std::error::Error + Send + Sync>>,
322{
323 type Response = S::Response;
324 type Error = SecurityError;
325 type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;
326
327 fn poll_ready(&mut self, cx: &mut TaskCtx<'_>) -> Poll<Result<(), Self::Error>> {
328 self.inner
329 .poll_ready(cx)
330 .map_err(|e| SecurityError::Inner(e.into()))
331 }
332
333 fn call(&mut self, req: Request<axum::body::Body>) -> Self::Future {
334 let scanner = Arc::clone(&self.scanner);
335 let body_limit = self.body_limit;
336 let mut inner = self.inner.clone();
337 let parts = request_parts(&req);
338 let strings: Vec<&str> = parts.iter().map(|s| s.as_str()).collect();
339 let header_results = scanner.scan_parts(&strings);
340
341 Box::pin(async move {
342 let (parts, body) = req.into_parts();
343 let bytes = match axum::body::to_bytes(body, body_limit).await {
346 Ok(bytes) => bytes,
347 Err(e) => {
348 let inner: Box<dyn std::error::Error + Send + Sync> = e.into_inner();
350 if inner
351 .downcast_ref::<http_body_util::LengthLimitError>()
352 .is_some()
353 {
354 return Err(SecurityError::BodyTooLarge);
355 }
356 return Err(SecurityError::Inner(inner));
357 }
358 };
359
360 let mut results = header_results;
361 results.extend(scanner.scan_body(&bytes));
362
363 if let Some(err) = evaluate(&results) {
364 return Err(err);
365 }
366
367 let req = Request::from_parts(parts, axum::body::Body::from(bytes));
368 inner
369 .call(req)
370 .await
371 .map_err(|e| SecurityError::Inner(e.into()))
372 })
373 }
374}
375
376#[cfg(test)]
377mod tests {
378 use super::*;
379
380 #[test]
381 fn scanner_detects_sql_injection() {
382 let s = SecurityScanner::new();
383 let results = s.scan("SELECT * FROM users; DROP TABLE users;");
384 assert!(!results.is_empty());
385 }
386
387 #[test]
388 fn scanner_detects_xss() {
389 let s = SecurityScanner::new();
390 let results = s.scan("<script>alert('xss')</script>");
391 assert!(!results.is_empty());
392 }
393
394 #[test]
395 fn scanner_clean_input_no_detection() {
396 let s = SecurityScanner::new();
397 let results = s.scan("hello world");
398 assert!(results.is_empty());
399 }
400
401 #[test]
402 fn scanner_scan_parts_aggregates() {
403 let s = SecurityScanner::new();
404 let results = s.scan_parts(&["clean", "<script>x</script>"]);
405 assert!(!results.is_empty());
406 }
407
408 #[test]
409 fn attack_blocked_maps_to_403() {
410 use axum::response::IntoResponse;
411 let resp = SecurityError::AttackBlocked("sqli".into()).into_response();
412 assert_eq!(resp.status(), StatusCode::FORBIDDEN);
413 }
414
415 #[test]
416 fn inner_error_maps_to_500() {
417 use axum::response::IntoResponse;
418 let resp = SecurityError::Inner(Box::new(std::io::Error::other("boom"))).into_response();
419 assert_eq!(resp.status(), StatusCode::INTERNAL_SERVER_ERROR);
420 }
421
422 #[test]
423 fn layer_constructs() {
424 let _layer = SecurityLayer::new();
425 }
426
427 #[test]
428 fn layer_default_constructs() {
429 let _layer: SecurityLayer = Default::default();
430 }
431
432 #[test]
433 fn body_layer_constructs() {
434 let _layer = SecurityBodyLayer::new().body_limit(1024);
435 }
436
437 #[test]
438 fn body_layer_default_constructs() {
439 let _layer: SecurityBodyLayer = Default::default();
440 }
441
442 #[tokio::test]
443 async fn body_layer_blocks_attack_in_body() {
444 use tower::Layer as _;
445 use tower::ServiceExt;
446
447 let layer = SecurityBodyLayer::new();
448 let svc = layer.layer(tower::service_fn(|_: Request<axum::body::Body>| async {
449 Ok::<_, std::convert::Infallible>(http::Response::new(axum::body::Body::empty()))
450 }));
451
452 let req = http::Request::builder()
453 .method("POST")
454 .uri("/submit")
455 .body(axum::body::Body::from("<script>alert('xss')</script>"))
456 .unwrap();
457 let result = svc.oneshot(req).await;
458 assert!(matches!(result, Err(SecurityError::AttackBlocked(_))));
459 }
460
461 #[tokio::test]
462 async fn body_over_limit_maps_to_413() {
463 use tower::Layer as _;
464 use tower::ServiceExt;
465
466 let layer = SecurityBodyLayer::new().body_limit(8);
467 let svc = layer.layer(tower::service_fn(|_: Request<axum::body::Body>| async {
468 Ok::<_, std::convert::Infallible>(http::Response::new(axum::body::Body::empty()))
469 }));
470
471 let req = http::Request::builder()
472 .method("POST")
473 .uri("/submit")
474 .body(axum::body::Body::from("x".repeat(64)))
475 .unwrap();
476 let result = svc.oneshot(req).await;
477 assert!(matches!(result, Err(SecurityError::BodyTooLarge)));
478 assert_eq!(
479 SecurityError::BodyTooLarge.to_http_status(),
480 StatusCode::PAYLOAD_TOO_LARGE
481 );
482 }
483
484 #[tokio::test]
485 async fn inner_error_response_body_does_not_leak_detail() {
486 use axum::response::IntoResponse;
487 let resp = SecurityError::Inner(Box::new(std::io::Error::other(
488 "secret-db-dsn=s3://user:pass",
489 )))
490 .into_response();
491 assert_eq!(resp.status(), StatusCode::INTERNAL_SERVER_ERROR);
492 let (_, body) = resp.into_parts();
493 let bytes = axum::body::to_bytes(body, 4096).await.unwrap();
494 let text = String::from_utf8_lossy(&bytes);
495 assert!(!text.contains("secret-db-dsn"), "leaked detail: {text}");
496 assert!(text.contains("internal server error"), "got: {text}");
497 }
498
499 #[tokio::test]
500 async fn attack_blocked_response_body_shape() {
501 use axum::response::IntoResponse;
502 let resp = SecurityError::AttackBlocked("sqli, xss".into()).into_response();
503 assert_eq!(resp.status(), StatusCode::FORBIDDEN);
504 let (_, body) = resp.into_parts();
505 let bytes = axum::body::to_bytes(body, 1024).await.unwrap();
506 let text = String::from_utf8_lossy(&bytes);
507 assert_eq!(
508 text,
509 r#"{"error":"attack blocked","types":"sqli, xss"}"#.to_string()
510 );
511 }
512
513 #[tokio::test]
514 async fn body_too_large_response_body_shape() {
515 use axum::response::IntoResponse;
516 let resp = SecurityError::BodyTooLarge.into_response();
517 assert_eq!(resp.status(), StatusCode::PAYLOAD_TOO_LARGE);
518 let (_, body) = resp.into_parts();
519 let bytes = axum::body::to_bytes(body, 1024).await.unwrap();
520 let text = String::from_utf8_lossy(&bytes);
521 assert_eq!(text, r#"{"error":"request body too large"}"#.to_string());
522 }
523
524 #[test]
525 fn request_parts_include_uri_and_headers() {
526 let req = http::Request::builder()
527 .uri("/search?q=abc")
528 .header("X-Custom", "val-1")
529 .body(())
530 .unwrap();
531 let parts = request_parts(&req);
532 assert_eq!(
533 parts,
534 vec!["/search?q=abc".to_string(), "val-1".to_string()]
535 );
536 }
537
538 #[test]
539 fn percent_decode_handles_encoded_sql_chars() {
540 assert_eq!(
541 percent_decode("/q?x=SELECT%20*%20FROM%20users"),
542 "/q?x=SELECT * FROM users"
543 );
544 assert_eq!(percent_decode("1%27%20OR%20%271%27%3D%271"), "1' OR '1'='1");
545 assert_eq!(
546 percent_decode("/clean?q=hello%20world"),
547 "/clean?q=hello world"
548 );
549 assert_eq!(percent_decode("%3cscript%3e"), "<script>");
550 assert_eq!(percent_decode("/%zz%"), "/%zz%");
552 assert_eq!(percent_decode("/100%25"), "/100%");
553 }
554
555 #[test]
556 fn request_parts_percent_decode_uri_only() {
557 let req = http::Request::builder()
558 .uri("/search?q=SELECT%20*%20FROM%20users")
559 .header("X-Custom", "val%201")
560 .body(())
561 .unwrap();
562 let parts = request_parts(&req);
563 assert_eq!(parts[0], "/search?q=SELECT * FROM users");
565 assert_eq!(parts[1], "val%201");
566 }
567
568 #[tokio::test]
569 async fn header_layer_blocks_attack_in_uri() {
570 use tower::Layer as _;
571 use tower::ServiceExt;
572
573 let layer = SecurityLayer::new();
574 let svc = layer.layer(tower::service_fn(|_: Request<()>| async {
575 Ok::<_, std::convert::Infallible>(http::Response::new(axum::body::Body::empty()))
576 }));
577
578 let req = http::Request::builder()
581 .uri("/redirect?url=javascript:alert(1)")
582 .body(())
583 .unwrap();
584 let resp = svc.oneshot(req).await.expect("blocked as response");
585 assert_eq!(resp.status(), StatusCode::FORBIDDEN);
586 }
587
588 #[tokio::test]
589 async fn header_layer_blocks_attack_in_header() {
590 use tower::Layer as _;
591 use tower::ServiceExt;
592
593 let layer = SecurityLayer::new();
594 let svc = layer.layer(tower::service_fn(|_: Request<()>| async {
595 Ok::<_, std::convert::Infallible>(http::Response::new(axum::body::Body::empty()))
596 }));
597
598 let req = http::Request::builder()
599 .uri("/clean")
600 .header("X-Trace", "<script>alert(1)</script>")
601 .body(())
602 .unwrap();
603 let resp = svc.oneshot(req).await.expect("blocked as response");
604 assert_eq!(resp.status(), StatusCode::FORBIDDEN);
605 }
606
607 #[tokio::test]
608 async fn header_layer_blocks_encoded_sqli_in_uri() {
609 use tower::Layer as _;
610 use tower::ServiceExt;
611
612 let layer = SecurityLayer::new();
613 let svc = layer.layer(tower::service_fn(|_: Request<()>| async {
614 Ok::<_, std::convert::Infallible>(http::Response::new(axum::body::Body::empty()))
615 }));
616
617 let req = http::Request::builder()
619 .uri("/search?q=SELECT%20*%20FROM%20users")
620 .body(())
621 .unwrap();
622 let resp = svc.oneshot(req).await.expect("blocked as response");
623 assert_eq!(resp.status(), StatusCode::FORBIDDEN);
624 }
625
626 #[tokio::test]
627 async fn header_layer_blocks_encoded_single_quote_sqli_in_uri() {
628 use tower::Layer as _;
629 use tower::ServiceExt;
630
631 let layer = SecurityLayer::new();
632 let svc = layer.layer(tower::service_fn(|_: Request<()>| async {
633 Ok::<_, std::convert::Infallible>(http::Response::new(axum::body::Body::empty()))
634 }));
635
636 let req = http::Request::builder()
638 .uri("/login?u=1%27%20OR%20%271%27%3D%271")
639 .body(())
640 .unwrap();
641 let resp = svc.oneshot(req).await.expect("blocked as response");
642 assert_eq!(resp.status(), StatusCode::FORBIDDEN);
643 }
644
645 #[tokio::test]
646 async fn header_layer_passes_clean_encoded_query_through() {
647 use tower::Layer as _;
648 use tower::ServiceExt;
649
650 let layer = SecurityLayer::new();
651 let svc = layer.layer(tower::service_fn(|req: Request<()>| async move {
652 Ok::<_, std::convert::Infallible>(http::Response::new(axum::body::Body::from(
653 req.uri().path().to_string(),
654 )))
655 }));
656
657 let req = http::Request::builder()
659 .uri("/search?q=hello%20world")
660 .body(())
661 .unwrap();
662 let resp = svc.oneshot(req).await.expect("clean request passes");
663 let (_, body) = resp.into_parts();
664 let bytes = axum::body::to_bytes(body, 1024).await.unwrap();
665 assert_eq!(String::from_utf8_lossy(&bytes), "/search");
666 }
667
668 #[tokio::test]
669 async fn header_layer_passes_proxy_headers_with_internal_ip() {
670 use tower::Layer as _;
671 use tower::ServiceExt;
672
673 let layer = SecurityLayer::new();
674 let svc = layer.layer(tower::service_fn(|req: Request<()>| async move {
675 Ok::<_, std::convert::Infallible>(http::Response::new(axum::body::Body::from(
676 req.uri().path().to_string(),
677 )))
678 }));
679
680 let req = http::Request::builder()
682 .uri("/api/v1/booking/dates?region_id=1")
683 .header("X-Real-IP", "172.19.0.1")
684 .header("X-Forwarded-For", "172.19.0.1")
685 .body(())
686 .unwrap();
687 let resp = svc.oneshot(req).await.expect("proxy headers pass");
688 assert_eq!(resp.status(), StatusCode::OK);
689 }
690
691 #[tokio::test]
692 async fn header_layer_passes_clean_request_through() {
693 use tower::Layer as _;
694 use tower::ServiceExt;
695
696 let layer = SecurityLayer::new();
697 let svc = layer.layer(tower::service_fn(|req: Request<()>| async move {
698 Ok::<_, std::convert::Infallible>(http::Response::new(axum::body::Body::from(
699 req.uri().path().to_string(),
700 )))
701 }));
702
703 let req = http::Request::builder().uri("/clean").body(()).unwrap();
704 let resp = svc.oneshot(req).await.expect("clean request passes");
705 let (_, body) = resp.into_parts();
706 let bytes = axum::body::to_bytes(body, 1024).await.unwrap();
707 assert_eq!(String::from_utf8_lossy(&bytes), "/clean");
708 }
709
710 #[tokio::test]
711 async fn body_layer_passes_clean_body_through() {
712 use tower::Layer as _;
713 use tower::ServiceExt;
714
715 let layer = SecurityBodyLayer::new();
716 let svc = layer.layer(tower::service_fn(
717 |req: Request<axum::body::Body>| async move {
718 let (_, body) = req.into_parts();
719 let bytes = axum::body::to_bytes(body, 1024).await.unwrap();
720 Ok::<_, std::convert::Infallible>(http::Response::new(axum::body::Body::from(
721 bytes,
722 )))
723 },
724 ));
725
726 let req = http::Request::builder()
727 .method("POST")
728 .uri("/submit")
729 .body(axum::body::Body::from("hello world"))
730 .unwrap();
731 let resp = svc.oneshot(req).await.expect("clean body passes through");
732 let (_, body) = resp.into_parts();
733 let bytes = axum::body::to_bytes(body, 1024).await.unwrap();
734 assert_eq!(&bytes[..], b"hello world");
735 }
736}