1use async_trait::async_trait;
45use reinhardt_core::exception::Result;
46use std::any::{Any, TypeId};
47use std::sync::Arc;
48
49use crate::{Request, Response};
50
51pub type MiddlewareDiRegistration = (TypeId, Arc<dyn Any + Send + Sync>);
59
60#[async_trait]
65pub trait Handler: Send + Sync {
66 async fn handle(&self, request: Request) -> Result<Response>;
72}
73
74#[async_trait]
79impl<T: Handler + ?Sized> Handler for Arc<T> {
80 async fn handle(&self, request: Request) -> Result<Response> {
81 (**self).handle(request).await
82 }
83}
84
85#[async_trait]
91pub trait Middleware: Send + Sync {
92 async fn process(&self, request: Request, next: Arc<dyn Handler>) -> Result<Response>;
103
104 fn should_continue(&self, _request: &Request) -> bool {
126 true
127 }
128
129 fn di_registrations(&self) -> Vec<MiddlewareDiRegistration> {
163 Vec::new()
164 }
165}
166
167pub struct MiddlewareChain {
172 middlewares: Vec<Arc<dyn Middleware>>,
173 handler: Arc<dyn Handler>,
174}
175
176impl MiddlewareChain {
177 pub fn new(handler: Arc<dyn Handler>) -> Self {
198 Self {
199 middlewares: Vec::new(),
200 handler,
201 }
202 }
203
204 pub fn with_middlewares(
206 handler: Arc<dyn Handler>,
207 middlewares: Vec<Arc<dyn Middleware>>,
208 ) -> Self {
209 Self {
210 middlewares,
211 handler,
212 }
213 }
214
215 pub fn with_middleware(mut self, middleware: Arc<dyn Middleware>) -> Self {
243 self.middlewares.push(middleware);
244 self
245 }
246
247 pub fn add_middleware(&mut self, middleware: Arc<dyn Middleware>) {
275 self.middlewares.push(middleware);
276 }
277}
278
279#[async_trait]
280impl Handler for MiddlewareChain {
281 async fn handle(&self, request: Request) -> Result<Response> {
282 if self.middlewares.is_empty() {
283 return self.handler.handle(request).await;
284 }
285
286 if self.middlewares.len() == 1 {
287 let middleware = &self.middlewares[0];
288 if !middleware.should_continue(&request) {
289 return match self.handler.handle(request).await {
290 Ok(response) => Ok(response),
291 Err(e) => Ok(Response::from(e)),
292 };
293 }
294
295 let next: Arc<dyn Handler> = Arc::new(ErrorToResponseHandler {
296 inner: self.handler.clone(),
297 });
298 let response = match middleware.process(request, next).await {
299 Ok(response) => response,
300 Err(e) => Response::from(e),
301 };
302 return Ok(response);
303 }
304
305 let mut current_handler: Arc<dyn Handler> = Arc::new(ErrorToResponseHandler {
316 inner: self.handler.clone(),
317 });
318
319 for middleware in self
320 .middlewares
321 .iter()
322 .rev()
323 .filter(|mw| mw.should_continue(&request))
324 {
325 let mw = middleware.clone();
326 let handler = current_handler.clone();
327
328 current_handler = Arc::new(ConditionalComposedHandler {
329 middleware: mw,
330 next: handler,
331 });
332 }
333
334 current_handler.handle(request).await
335 }
336}
337
338pub struct ExcludeMiddleware {
378 inner: Arc<dyn Middleware>,
379 exclusions: Vec<String>,
380}
381
382impl ExcludeMiddleware {
383 pub fn new(inner: Arc<dyn Middleware>) -> Self {
385 Self {
386 inner,
387 exclusions: Vec::new(),
388 }
389 }
390
391 pub fn add_exclusion(mut self, pattern: &str) -> Self {
395 self.exclusions.push(pattern.to_string());
396 self
397 }
398
399 pub fn add_exclusion_mut(&mut self, pattern: &str) {
403 self.exclusions.push(pattern.to_string());
404 }
405
406 fn is_excluded(&self, path: &str) -> bool {
408 self.exclusions.iter().any(|pattern| {
409 if pattern.ends_with('/') {
410 path.starts_with(pattern.as_str())
412 } else {
413 path == pattern
415 }
416 })
417 }
418}
419
420#[async_trait]
421impl Middleware for ExcludeMiddleware {
422 async fn process(&self, request: Request, next: Arc<dyn Handler>) -> Result<Response> {
423 self.inner.process(request, next).await
424 }
425
426 fn should_continue(&self, request: &Request) -> bool {
427 if self.is_excluded(request.uri.path()) {
428 return false;
429 }
430 self.inner.should_continue(request)
431 }
432}
433
434struct ErrorToResponseHandler {
441 inner: Arc<dyn Handler>,
442}
443
444#[async_trait]
445impl Handler for ErrorToResponseHandler {
446 async fn handle(&self, request: Request) -> Result<Response> {
447 match self.inner.handle(request).await {
448 Ok(response) => Ok(response),
449 Err(e) => Ok(Response::from(e)),
450 }
451 }
452}
453
454struct ConditionalComposedHandler {
459 middleware: Arc<dyn Middleware>,
460 next: Arc<dyn Handler>,
461}
462
463#[async_trait]
464impl Handler for ConditionalComposedHandler {
465 async fn handle(&self, request: Request) -> Result<Response> {
466 let response = match self.middleware.process(request, self.next.clone()).await {
471 Ok(response) => response,
472 Err(e) => Response::from(e),
473 };
474
475 Ok(response)
476 }
477}
478
479#[cfg(test)]
480mod tests {
481 use super::*;
482 use bytes::Bytes;
483 use hyper::{HeaderMap, Method, Version};
484
485 struct MockHandler {
487 response_body: String,
488 }
489
490 #[async_trait]
491 impl Handler for MockHandler {
492 async fn handle(&self, _request: Request) -> Result<Response> {
493 Ok(Response::ok().with_body(self.response_body.clone()))
494 }
495 }
496
497 struct MockMiddleware {
499 prefix: String,
500 }
501
502 #[async_trait]
503 impl Middleware for MockMiddleware {
504 async fn process(&self, request: Request, next: Arc<dyn Handler>) -> Result<Response> {
505 let response = next.handle(request).await?;
507
508 let current_body = String::from_utf8(response.body.to_vec()).unwrap_or_default();
510 let new_body = format!("{}{}", self.prefix, current_body);
511
512 Ok(Response::ok().with_body(new_body))
513 }
514 }
515
516 fn create_test_request() -> Request {
517 Request::builder()
518 .method(Method::GET)
519 .uri("/")
520 .version(Version::HTTP_11)
521 .headers(HeaderMap::new())
522 .body(Bytes::new())
523 .build()
524 .unwrap()
525 }
526
527 #[tokio::test]
528 async fn test_handler_basic() {
529 let handler = MockHandler {
530 response_body: "Hello".to_string(),
531 };
532
533 let request = create_test_request();
534 let response = handler.handle(request).await.unwrap();
535
536 let body = String::from_utf8(response.body.to_vec()).unwrap();
537 assert_eq!(body, "Hello");
538 }
539
540 #[tokio::test]
541 async fn test_middleware_basic() {
542 let handler = Arc::new(MockHandler {
543 response_body: "World".to_string(),
544 });
545
546 let middleware = MockMiddleware {
547 prefix: "Hello, ".to_string(),
548 };
549
550 let request = create_test_request();
551 let response = middleware.process(request, handler).await.unwrap();
552
553 let body = String::from_utf8(response.body.to_vec()).unwrap();
554 assert_eq!(body, "Hello, World");
555 }
556
557 #[tokio::test]
558 async fn test_middleware_chain_empty() {
559 let handler = Arc::new(MockHandler {
560 response_body: "Test".to_string(),
561 });
562
563 let chain = MiddlewareChain::new(handler);
564
565 let request = create_test_request();
566 let response = chain.handle(request).await.unwrap();
567
568 let body = String::from_utf8(response.body.to_vec()).unwrap();
569 assert_eq!(body, "Test");
570 }
571
572 #[tokio::test]
573 async fn test_middleware_chain_single() {
574 let handler = Arc::new(MockHandler {
575 response_body: "Handler".to_string(),
576 });
577
578 let middleware1 = Arc::new(MockMiddleware {
579 prefix: "MW1:".to_string(),
580 });
581
582 let chain = MiddlewareChain::new(handler).with_middleware(middleware1);
583
584 let request = create_test_request();
585 let response = chain.handle(request).await.unwrap();
586
587 let body = String::from_utf8(response.body.to_vec()).unwrap();
588 assert_eq!(body, "MW1:Handler");
589 }
590
591 #[tokio::test]
592 async fn test_middleware_chain_multiple() {
593 let handler = Arc::new(MockHandler {
594 response_body: "Data".to_string(),
595 });
596
597 let middleware1 = Arc::new(MockMiddleware {
598 prefix: "M1:".to_string(),
599 });
600
601 let middleware2 = Arc::new(MockMiddleware {
602 prefix: "M2:".to_string(),
603 });
604
605 let chain = MiddlewareChain::new(handler)
606 .with_middleware(middleware1)
607 .with_middleware(middleware2);
608
609 let request = create_test_request();
610 let response = chain.handle(request).await.unwrap();
611
612 let body = String::from_utf8(response.body.to_vec()).unwrap();
613 assert_eq!(body, "M1:M2:Data");
615 }
616
617 #[tokio::test]
618 async fn test_middleware_chain_add_middleware() {
619 let handler = Arc::new(MockHandler {
620 response_body: "Result".to_string(),
621 });
622
623 let middleware = Arc::new(MockMiddleware {
624 prefix: "Prefix:".to_string(),
625 });
626
627 let mut chain = MiddlewareChain::new(handler);
628 chain.add_middleware(middleware);
629
630 let request = create_test_request();
631 let response = chain.handle(request).await.unwrap();
632
633 let body = String::from_utf8(response.body.to_vec()).unwrap();
634 assert_eq!(body, "Prefix:Result");
635 }
636
637 struct ConditionalMiddleware {
639 prefix: String,
640 }
641
642 #[async_trait]
643 impl Middleware for ConditionalMiddleware {
644 async fn process(&self, request: Request, next: Arc<dyn Handler>) -> Result<Response> {
645 let response = next.handle(request).await?;
646 let current_body = String::from_utf8(response.body.to_vec()).unwrap_or_default();
647 let new_body = format!("{}{}", self.prefix, current_body);
648 Ok(Response::ok().with_body(new_body))
649 }
650
651 fn should_continue(&self, request: &Request) -> bool {
652 request.uri.path().starts_with("/api/")
653 }
654 }
655
656 #[tokio::test]
657 async fn test_middleware_conditional_skip() {
658 let handler = Arc::new(MockHandler {
659 response_body: "Response".to_string(),
660 });
661
662 let conditional_mw = Arc::new(ConditionalMiddleware {
663 prefix: "API:".to_string(),
664 });
665
666 let chain = MiddlewareChain::new(handler).with_middleware(conditional_mw);
667
668 let api_request = Request::builder()
670 .method(Method::GET)
671 .uri("/api/users")
672 .version(Version::HTTP_11)
673 .headers(HeaderMap::new())
674 .body(Bytes::new())
675 .build()
676 .unwrap();
677 let response = chain.handle(api_request).await.unwrap();
678 let body = String::from_utf8(response.body.to_vec()).unwrap();
679 assert_eq!(body, "API:Response");
680
681 let non_api_request = Request::builder()
683 .method(Method::GET)
684 .uri("/public")
685 .version(Version::HTTP_11)
686 .headers(HeaderMap::new())
687 .body(Bytes::new())
688 .build()
689 .unwrap();
690 let response = chain.handle(non_api_request).await.unwrap();
691 let body = String::from_utf8(response.body.to_vec()).unwrap();
692 assert_eq!(body, "Response"); }
694
695 struct ShortCircuitMiddleware {
697 should_stop: bool,
698 }
699
700 #[async_trait]
701 impl Middleware for ShortCircuitMiddleware {
702 async fn process(&self, request: Request, next: Arc<dyn Handler>) -> Result<Response> {
703 if self.should_stop {
704 return Ok(Response::unauthorized()
706 .with_body("Auth required")
707 .with_stop_chain(true));
708 }
709 next.handle(request).await
710 }
711 }
712
713 #[tokio::test]
714 async fn test_middleware_short_circuit() {
715 let handler = Arc::new(MockHandler {
716 response_body: "Handler Response".to_string(),
717 });
718
719 let short_circuit_mw = Arc::new(ShortCircuitMiddleware { should_stop: true });
720 let normal_mw = Arc::new(MockMiddleware {
721 prefix: "Normal:".to_string(),
722 });
723
724 let chain = MiddlewareChain::new(handler)
725 .with_middleware(short_circuit_mw)
726 .with_middleware(normal_mw);
727
728 let request = create_test_request();
729 let response = chain.handle(request).await.unwrap();
730
731 assert_eq!(response.status, hyper::StatusCode::UNAUTHORIZED);
733 let body = String::from_utf8(response.body.to_vec()).unwrap();
734 assert_eq!(body, "Auth required");
735 }
736
737 #[tokio::test]
738 async fn test_middleware_no_short_circuit() {
739 let handler = Arc::new(MockHandler {
740 response_body: "Handler Response".to_string(),
741 });
742
743 let short_circuit_mw = Arc::new(ShortCircuitMiddleware { should_stop: false });
744 let normal_mw = Arc::new(MockMiddleware {
745 prefix: "Normal:".to_string(),
746 });
747
748 let chain = MiddlewareChain::new(handler)
749 .with_middleware(short_circuit_mw)
750 .with_middleware(normal_mw);
751
752 let request = create_test_request();
753 let response = chain.handle(request).await.unwrap();
754
755 assert_eq!(response.status, hyper::StatusCode::OK);
757 let body = String::from_utf8(response.body.to_vec()).unwrap();
758 assert_eq!(body, "Normal:Handler Response");
759 }
760
761 #[tokio::test]
762 async fn test_middleware_multiple_conditions() {
763 let handler = Arc::new(MockHandler {
764 response_body: "Base".to_string(),
765 });
766
767 let api_mw = Arc::new(ConditionalMiddleware {
769 prefix: "API:".to_string(),
770 });
771
772 let always_mw = Arc::new(MockMiddleware {
774 prefix: "Always:".to_string(),
775 });
776
777 let chain = MiddlewareChain::new(handler)
778 .with_middleware(api_mw)
779 .with_middleware(always_mw);
780
781 let api_request = Request::builder()
783 .method(Method::GET)
784 .uri("/api/test")
785 .version(Version::HTTP_11)
786 .headers(HeaderMap::new())
787 .body(Bytes::new())
788 .build()
789 .unwrap();
790 let response = chain.handle(api_request).await.unwrap();
791 let body = String::from_utf8(response.body.to_vec()).unwrap();
792 assert_eq!(body, "API:Always:Base");
793
794 let non_api_request = Request::builder()
796 .method(Method::GET)
797 .uri("/public")
798 .version(Version::HTTP_11)
799 .headers(HeaderMap::new())
800 .body(Bytes::new())
801 .build()
802 .unwrap();
803 let response = chain.handle(non_api_request).await.unwrap();
804 let body = String::from_utf8(response.body.to_vec()).unwrap();
805 assert_eq!(body, "Always:Base"); }
807
808 #[tokio::test]
809 async fn test_response_should_stop_chain() {
810 let response = Response::ok();
811 assert!(!response.should_stop_chain());
812
813 let stopping_response = Response::unauthorized().with_stop_chain(true);
814 assert!(stopping_response.should_stop_chain());
815 }
816
817 fn create_request_with_path(path: &str) -> Request {
820 Request::builder()
821 .method(Method::GET)
822 .uri(path)
823 .version(Version::HTTP_11)
824 .headers(HeaderMap::new())
825 .body(Bytes::new())
826 .build()
827 .unwrap()
828 }
829
830 #[rstest::rstest]
831 #[case("/api/auth/login", true)]
832 #[case("/api/auth/register", true)]
833 #[case("/api/auth/", true)]
834 #[case("/api/users", false)]
835 #[case("/public", false)]
836 fn test_exclude_middleware_prefix_match(#[case] path: &str, #[case] should_exclude: bool) {
837 let inner: Arc<dyn Middleware> = Arc::new(MockMiddleware {
839 prefix: "MW:".to_string(),
840 });
841 let exclude_mw = ExcludeMiddleware::new(inner).add_exclusion("/api/auth/");
842
843 let request = create_request_with_path(path);
845 let result = exclude_mw.should_continue(&request);
846
847 assert_eq!(result, !should_exclude);
849 }
850
851 #[rstest::rstest]
852 #[case("/health", true)]
853 #[case("/health/check", false)]
854 #[case("/healthz", false)]
855 #[case("/api/health", false)]
856 fn test_exclude_middleware_exact_match(#[case] path: &str, #[case] should_exclude: bool) {
857 let inner: Arc<dyn Middleware> = Arc::new(MockMiddleware {
859 prefix: "MW:".to_string(),
860 });
861 let exclude_mw = ExcludeMiddleware::new(inner).add_exclusion("/health");
862
863 let request = create_request_with_path(path);
865 let result = exclude_mw.should_continue(&request);
866
867 assert_eq!(result, !should_exclude);
869 }
870
871 #[rstest::rstest]
872 fn test_exclude_middleware_no_match_passes_through() {
873 let inner: Arc<dyn Middleware> = Arc::new(MockMiddleware {
875 prefix: "MW:".to_string(),
876 });
877 let exclude_mw = ExcludeMiddleware::new(inner)
878 .add_exclusion("/api/auth/")
879 .add_exclusion("/health");
880
881 let request = create_request_with_path("/api/users");
883 let result = exclude_mw.should_continue(&request);
884
885 assert!(result);
887 }
888
889 #[rstest::rstest]
890 #[tokio::test]
891 async fn test_exclude_middleware_delegates_process() {
892 let inner: Arc<dyn Middleware> = Arc::new(MockMiddleware {
894 prefix: "INNER:".to_string(),
895 });
896 let exclude_mw = ExcludeMiddleware::new(inner).add_exclusion("/excluded/");
897
898 let handler = Arc::new(MockHandler {
899 response_body: "Response".to_string(),
900 });
901
902 let request = create_request_with_path("/api/test");
904 let response = exclude_mw.process(request, handler).await.unwrap();
905
906 let body = String::from_utf8(response.body.to_vec()).unwrap();
908 assert_eq!(body, "INNER:Response");
909 }
910
911 #[rstest::rstest]
912 fn test_exclude_middleware_multiple_exclusions() {
913 let inner: Arc<dyn Middleware> = Arc::new(MockMiddleware {
915 prefix: "MW:".to_string(),
916 });
917 let mut exclude_mw = ExcludeMiddleware::new(inner);
918 exclude_mw.add_exclusion_mut("/api/auth/");
919 exclude_mw.add_exclusion_mut("/admin/");
920 exclude_mw.add_exclusion_mut("/health");
921
922 assert!(!exclude_mw.should_continue(&create_request_with_path("/api/auth/login")));
924 assert!(!exclude_mw.should_continue(&create_request_with_path("/admin/dashboard")));
925 assert!(!exclude_mw.should_continue(&create_request_with_path("/health")));
926 assert!(exclude_mw.should_continue(&create_request_with_path("/api/users")));
927 }
928
929 #[rstest::rstest]
930 fn test_exclude_middleware_respects_inner_should_continue() {
931 let inner: Arc<dyn Middleware> = Arc::new(ConditionalMiddleware {
933 prefix: "API:".to_string(),
934 });
935 let exclude_mw = ExcludeMiddleware::new(inner).add_exclusion("/api/auth/");
936
937 assert!(!exclude_mw.should_continue(&create_request_with_path("/api/auth/login")));
940 assert!(!exclude_mw.should_continue(&create_request_with_path("/public")));
942 assert!(exclude_mw.should_continue(&create_request_with_path("/api/users")));
944 }
945
946 struct NotFoundHandler;
952
953 #[async_trait]
954 impl Handler for NotFoundHandler {
955 async fn handle(&self, _request: Request) -> Result<Response> {
956 Err(reinhardt_core::exception::Error::NotFound(
957 "not found".into(),
958 ))
959 }
960 }
961
962 struct UnauthorizedHandler;
963
964 #[async_trait]
965 impl Handler for UnauthorizedHandler {
966 async fn handle(&self, _request: Request) -> Result<Response> {
967 Err(reinhardt_core::exception::Error::Authentication(
968 "unauthorized".into(),
969 ))
970 }
971 }
972
973 struct HeaderAddingMiddleware {
975 header_name: &'static str,
976 header_value: &'static str,
977 }
978
979 #[async_trait]
980 impl Middleware for HeaderAddingMiddleware {
981 async fn process(&self, request: Request, next: Arc<dyn Handler>) -> Result<Response> {
982 let response = next.handle(request).await?;
983 Ok(response.with_header(self.header_name, self.header_value))
984 }
985 }
986
987 struct RejectingMiddleware;
989
990 #[async_trait]
991 impl Middleware for RejectingMiddleware {
992 async fn process(&self, _request: Request, _next: Arc<dyn Handler>) -> Result<Response> {
993 Err(reinhardt_core::exception::Error::Authorization(
994 "CSRF check failed".into(),
995 ))
996 }
997 }
998
999 #[rstest::rstest]
1000 #[tokio::test]
1001 async fn test_chain_post_processing_runs_on_handler_error() {
1002 let handler: Arc<dyn Handler> = Arc::new(NotFoundHandler);
1004 let mut chain = MiddlewareChain::new(handler);
1005 chain.add_middleware(Arc::new(HeaderAddingMiddleware {
1006 header_name: "X-Custom-Security",
1007 header_value: "applied",
1008 }));
1009
1010 let request = create_test_request();
1012 let response = chain.handle(request).await.unwrap();
1013
1014 assert_eq!(response.status, hyper::StatusCode::NOT_FOUND);
1016 assert_eq!(
1017 response
1018 .headers
1019 .get("X-Custom-Security")
1020 .map(|v| v.to_str().unwrap()),
1021 Some("applied")
1022 );
1023 }
1024
1025 #[rstest::rstest]
1026 #[tokio::test]
1027 async fn test_chain_post_processing_runs_on_middleware_error() {
1028 let handler = Arc::new(MockHandler {
1031 response_body: "OK".into(),
1032 });
1033 let mut chain = MiddlewareChain::new(handler);
1034 chain.add_middleware(Arc::new(HeaderAddingMiddleware {
1036 header_name: "X-Frame-Options",
1037 header_value: "DENY",
1038 }));
1039 chain.add_middleware(Arc::new(RejectingMiddleware));
1041
1042 let request = create_test_request();
1044 let response = chain.handle(request).await.unwrap();
1045
1046 assert_eq!(response.status, hyper::StatusCode::FORBIDDEN);
1048 assert_eq!(
1049 response
1050 .headers
1051 .get("X-Frame-Options")
1052 .map(|v| v.to_str().unwrap()),
1053 Some("DENY")
1054 );
1055 }
1056
1057 struct PassthroughMiddleware;
1059
1060 #[async_trait]
1061 impl Middleware for PassthroughMiddleware {
1062 async fn process(&self, request: Request, next: Arc<dyn Handler>) -> Result<Response> {
1063 next.handle(request).await
1064 }
1065 }
1066
1067 #[rstest::rstest]
1068 #[tokio::test]
1069 async fn test_chain_error_preserves_correct_status_code() {
1070 let handler: Arc<dyn Handler> = Arc::new(UnauthorizedHandler);
1073 let mut chain = MiddlewareChain::new(handler);
1074 chain.add_middleware(Arc::new(PassthroughMiddleware));
1075
1076 let request = create_test_request();
1078 let response = chain.handle(request).await.unwrap();
1079
1080 assert_eq!(response.status, hyper::StatusCode::UNAUTHORIZED);
1082 }
1083}