1use std::collections::HashSet;
54use std::future::Future;
55use std::ops::ControlFlow as StdControlFlow;
56use std::pin::Pin;
57use std::sync::Arc;
58use std::time::Instant;
59
60use crate::context::RequestContext;
61use crate::dependency::DependencyOverrides;
62use crate::logging::{LogConfig, RequestLogger};
63use crate::request::{Body, Request};
64use crate::response::Response;
65
66pub type BoxFuture<'a, T> = Pin<Box<dyn Future<Output = T> + Send + 'a>>;
68
69#[derive(Debug)]
74pub enum ControlFlow {
75 Continue,
77 Break(Response),
82}
83
84impl ControlFlow {
85 #[must_use]
87 pub fn is_continue(&self) -> bool {
88 matches!(self, Self::Continue)
89 }
90
91 #[must_use]
93 pub fn is_break(&self) -> bool {
94 matches!(self, Self::Break(_))
95 }
96}
97
98impl From<ControlFlow> for StdControlFlow<Response, ()> {
99 fn from(cf: ControlFlow) -> Self {
100 match cf {
101 ControlFlow::Continue => StdControlFlow::Continue(()),
102 ControlFlow::Break(r) => StdControlFlow::Break(r),
103 }
104 }
105}
106
107pub trait Middleware: Send + Sync {
148 fn before<'a>(
164 &'a self,
165 _ctx: &'a RequestContext,
166 _req: &'a mut Request,
167 ) -> BoxFuture<'a, ControlFlow> {
168 Box::pin(async { ControlFlow::Continue })
169 }
170
171 fn after<'a>(
187 &'a self,
188 _ctx: &'a RequestContext,
189 _req: &'a Request,
190 response: Response,
191 ) -> BoxFuture<'a, Response> {
192 Box::pin(async move { response })
193 }
194
195 fn name(&self) -> &'static str {
199 std::any::type_name::<Self>()
200 }
201}
202
203pub trait Handler: Send + Sync {
208 fn call<'a>(&'a self, ctx: &'a RequestContext, req: &'a mut Request)
210 -> BoxFuture<'a, Response>;
211
212 fn dependency_overrides(&self) -> Option<Arc<DependencyOverrides>> {
216 None
217 }
218}
219
220impl<F, Fut> Handler for F
225where
226 F: Fn(&RequestContext, &mut Request) -> Fut + Send + Sync,
227 Fut: Future<Output = Response> + Send + 'static,
228{
229 fn call<'a>(
230 &'a self,
231 ctx: &'a RequestContext,
232 req: &'a mut Request,
233 ) -> BoxFuture<'a, Response> {
234 let fut = self(ctx, req);
235 Box::pin(fut)
236 }
237}
238
239impl<H: Handler + ?Sized> Handler for Arc<H> {
244 fn call<'a>(
245 &'a self,
246 ctx: &'a RequestContext,
247 req: &'a mut Request,
248 ) -> BoxFuture<'a, Response> {
249 (**self).call(ctx, req)
250 }
251
252 fn dependency_overrides(&self) -> Option<Arc<DependencyOverrides>> {
253 (**self).dependency_overrides()
254 }
255}
256
257#[derive(Default)]
275pub struct MiddlewareStack {
276 middleware: Vec<Arc<dyn Middleware>>,
277}
278
279impl MiddlewareStack {
280 #[must_use]
282 pub fn new() -> Self {
283 Self {
284 middleware: Vec::new(),
285 }
286 }
287
288 #[must_use]
290 pub fn with_capacity(capacity: usize) -> Self {
291 Self {
292 middleware: Vec::with_capacity(capacity),
293 }
294 }
295
296 pub fn push<M: Middleware + 'static>(&mut self, middleware: M) {
300 self.middleware.push(Arc::new(middleware));
301 }
302
303 pub fn push_arc(&mut self, middleware: Arc<dyn Middleware>) {
307 self.middleware.push(middleware);
308 }
309
310 #[must_use]
312 pub fn len(&self) -> usize {
313 self.middleware.len()
314 }
315
316 #[must_use]
318 pub fn is_empty(&self) -> bool {
319 self.middleware.is_empty()
320 }
321
322 pub async fn execute<H: Handler>(
340 &self,
341 handler: &H,
342 ctx: &RequestContext,
343 req: &mut Request,
344 ) -> Response {
345 let mut ran_before_count = 0;
347
348 for mw in &self.middleware {
350 let _ = ctx.checkpoint();
351 match mw.before(ctx, req).await {
352 ControlFlow::Continue => {
353 ran_before_count += 1;
354 }
355 ControlFlow::Break(response) => {
356 return self
358 .run_after_hooks(ctx, req, response, ran_before_count)
359 .await;
360 }
361 }
362 }
363
364 let _ = ctx.checkpoint();
366 let response = handler.call(ctx, req).await;
367
368 self.run_after_hooks(ctx, req, response, ran_before_count)
370 .await
371 }
372
373 async fn run_after_hooks(
375 &self,
376 ctx: &RequestContext,
377 req: &Request,
378 mut response: Response,
379 count: usize,
380 ) -> Response {
381 for mw in self.middleware[..count].iter().rev() {
383 let _ = ctx.checkpoint();
384 response = mw.after(ctx, req, response).await;
385 }
386 response
387 }
388}
389
390pub struct Layer<M> {
401 middleware: M,
402}
403
404impl<M: Middleware + Clone> Layer<M> {
405 pub fn new(middleware: M) -> Self {
407 Self { middleware }
408 }
409
410 pub fn wrap<H: Handler>(&self, handler: H) -> Layered<M, H> {
412 Layered {
413 middleware: self.middleware.clone(),
414 inner: handler,
415 }
416 }
417}
418
419pub struct Layered<M, H> {
421 middleware: M,
422 inner: H,
423}
424
425impl<M: Middleware, H: Handler> Handler for Layered<M, H> {
426 fn call<'a>(
427 &'a self,
428 ctx: &'a RequestContext,
429 req: &'a mut Request,
430 ) -> BoxFuture<'a, Response> {
431 Box::pin(async move {
432 let _ = ctx.checkpoint();
434 match self.middleware.before(ctx, req).await {
435 ControlFlow::Continue => {
436 let _ = ctx.checkpoint();
438 let response = self.inner.call(ctx, req).await;
439 let _ = ctx.checkpoint();
441 self.middleware.after(ctx, req, response).await
442 }
443 ControlFlow::Break(response) => {
444 let _ = ctx.checkpoint();
446 self.middleware.after(ctx, req, response).await
447 }
448 }
449 })
450 }
451}
452
453#[derive(Debug, Clone, Copy, Default)]
461pub struct NoopMiddleware;
462
463impl Middleware for NoopMiddleware {
464 fn name(&self) -> &'static str {
465 "Noop"
466 }
467}
468
469#[derive(Debug, Clone)]
479pub struct AddResponseHeader {
480 name: String,
481 value: Vec<u8>,
482}
483
484impl AddResponseHeader {
485 pub fn new(name: impl Into<String>, value: impl Into<Vec<u8>>) -> Self {
487 Self {
488 name: name.into(),
489 value: value.into(),
490 }
491 }
492}
493
494impl Middleware for AddResponseHeader {
495 fn after<'a>(
496 &'a self,
497 _ctx: &'a RequestContext,
498 _req: &'a Request,
499 response: Response,
500 ) -> BoxFuture<'a, Response> {
501 let name = self.name.clone();
502 let value = self.value.clone();
503 Box::pin(async move { response.header(name, value) })
504 }
505
506 fn name(&self) -> &'static str {
507 "AddResponseHeader"
508 }
509}
510
511#[derive(Debug, Clone)]
523pub struct RequireHeader {
524 name: String,
525}
526
527impl RequireHeader {
528 pub fn new(name: impl Into<String>) -> Self {
530 Self { name: name.into() }
531 }
532}
533
534impl Middleware for RequireHeader {
535 fn before<'a>(
536 &'a self,
537 _ctx: &'a RequestContext,
538 req: &'a mut Request,
539 ) -> BoxFuture<'a, ControlFlow> {
540 let has_header = req.headers().get(&self.name).is_some();
541 let name = self.name.clone();
542 Box::pin(async move {
543 if has_header {
544 ControlFlow::Continue
545 } else {
546 let body = format!("Missing required header: {name}");
547 ControlFlow::Break(
548 Response::with_status(crate::response::StatusCode::BAD_REQUEST)
549 .header("content-type", b"text/plain".to_vec())
550 .body(crate::response::ResponseBody::Bytes(body.into_bytes())),
551 )
552 }
553 })
554 }
555
556 fn name(&self) -> &'static str {
557 "RequireHeader"
558 }
559}
560
561#[derive(Debug, Clone)]
574pub struct PathPrefixFilter {
575 prefix: String,
576}
577
578impl PathPrefixFilter {
579 pub fn new(prefix: impl Into<String>) -> Self {
581 Self {
582 prefix: prefix.into(),
583 }
584 }
585}
586
587impl Middleware for PathPrefixFilter {
588 fn before<'a>(
589 &'a self,
590 _ctx: &'a RequestContext,
591 req: &'a mut Request,
592 ) -> BoxFuture<'a, ControlFlow> {
593 let path_matches = req.path().starts_with(&self.prefix);
594 Box::pin(async move {
595 if path_matches {
596 ControlFlow::Continue
597 } else {
598 ControlFlow::Break(Response::with_status(
599 crate::response::StatusCode::NOT_FOUND,
600 ))
601 }
602 })
603 }
604
605 fn name(&self) -> &'static str {
606 "PathPrefixFilter"
607 }
608}
609
610#[derive(Debug, Clone)]
614pub struct ConditionalStatus<F>
615where
616 F: Fn(&Request) -> bool + Send + Sync,
617{
618 condition: F,
619 status_if_true: crate::response::StatusCode,
620 status_if_false: crate::response::StatusCode,
621}
622
623impl<F> ConditionalStatus<F>
624where
625 F: Fn(&Request) -> bool + Send + Sync,
626{
627 pub fn new(
632 condition: F,
633 status_if_true: crate::response::StatusCode,
634 status_if_false: crate::response::StatusCode,
635 ) -> Self {
636 Self {
637 condition,
638 status_if_true,
639 status_if_false,
640 }
641 }
642}
643
644impl<F> Middleware for ConditionalStatus<F>
645where
646 F: Fn(&Request) -> bool + Send + Sync,
647{
648 fn after<'a>(
649 &'a self,
650 _ctx: &'a RequestContext,
651 req: &'a Request,
652 response: Response,
653 ) -> BoxFuture<'a, Response> {
654 let matches = (self.condition)(req);
655 let status = if matches {
656 self.status_if_true
657 } else {
658 self.status_if_false
659 };
660 Box::pin(async move { Response::with_status(status).body(response.body_ref().into()) })
661 }
662
663 fn name(&self) -> &'static str {
664 "ConditionalStatus"
665 }
666}
667
668#[derive(Debug, Clone)]
674pub enum OriginPattern {
675 Any,
677 Exact(String),
679 Wildcard(String),
681 Regex(String),
683}
684
685impl OriginPattern {
686 fn matches(&self, origin: &str) -> bool {
687 match self {
688 Self::Any => true,
689 Self::Exact(value) => value == origin,
690 Self::Wildcard(pattern) => wildcard_match(pattern, origin),
691 Self::Regex(pattern) => regex_match(pattern, origin),
692 }
693 }
694}
695
696#[derive(Debug, Clone)]
745pub struct CorsConfig {
746 allow_any_origin: bool,
747 allow_credentials: bool,
748 allowed_methods: Vec<crate::request::Method>,
749 allowed_headers: Vec<String>,
750 expose_headers: Vec<String>,
751 max_age: Option<u32>,
752 origins: Vec<OriginPattern>,
753}
754
755impl Default for CorsConfig {
756 fn default() -> Self {
757 Self {
758 allow_any_origin: false,
759 allow_credentials: false,
760 allowed_methods: vec![
761 crate::request::Method::Get,
762 crate::request::Method::Post,
763 crate::request::Method::Put,
764 crate::request::Method::Patch,
765 crate::request::Method::Delete,
766 crate::request::Method::Options,
767 crate::request::Method::Head,
768 ],
769 allowed_headers: Vec::new(),
770 expose_headers: Vec::new(),
771 max_age: None,
772 origins: Vec::new(),
773 }
774 }
775}
776
777#[derive(Debug, Clone)]
779pub struct Cors {
780 config: CorsConfig,
781}
782
783impl Cors {
784 #[must_use]
786 pub fn new() -> Self {
787 Self {
788 config: CorsConfig::default(),
789 }
790 }
791
792 #[must_use]
794 pub fn config(mut self, config: CorsConfig) -> Self {
795 self.config = config;
796 self
797 }
798
799 #[must_use]
801 pub fn allow_any_origin(mut self) -> Self {
802 self.config.allow_any_origin = true;
803 self
804 }
805
806 #[must_use]
808 pub fn allow_origin(mut self, origin: impl Into<String>) -> Self {
809 self.config
810 .origins
811 .push(OriginPattern::Exact(origin.into()));
812 self
813 }
814
815 #[must_use]
817 pub fn allow_origin_wildcard(mut self, pattern: impl Into<String>) -> Self {
818 self.config
819 .origins
820 .push(OriginPattern::Wildcard(pattern.into()));
821 self
822 }
823
824 #[must_use]
826 pub fn allow_origin_regex(mut self, pattern: impl Into<String>) -> Self {
827 self.config
828 .origins
829 .push(OriginPattern::Regex(pattern.into()));
830 self
831 }
832
833 #[must_use]
835 pub fn allow_credentials(mut self, allow: bool) -> Self {
836 self.config.allow_credentials = allow;
837 self
838 }
839
840 #[must_use]
842 pub fn allow_methods<I>(mut self, methods: I) -> Self
843 where
844 I: IntoIterator<Item = crate::request::Method>,
845 {
846 self.config.allowed_methods = methods.into_iter().collect();
847 self
848 }
849
850 #[must_use]
852 pub fn allow_headers<I, S>(mut self, headers: I) -> Self
853 where
854 I: IntoIterator<Item = S>,
855 S: Into<String>,
856 {
857 self.config.allowed_headers = headers.into_iter().map(Into::into).collect();
858 self
859 }
860
861 #[must_use]
863 pub fn expose_headers<I, S>(mut self, headers: I) -> Self
864 where
865 I: IntoIterator<Item = S>,
866 S: Into<String>,
867 {
868 self.config.expose_headers = headers.into_iter().map(Into::into).collect();
869 self
870 }
871
872 #[must_use]
874 pub fn max_age(mut self, seconds: u32) -> Self {
875 self.config.max_age = Some(seconds);
876 self
877 }
878
879 fn is_origin_allowed(&self, origin: &str) -> bool {
880 if self.config.allow_any_origin {
881 return true;
882 }
883 self.config
884 .origins
885 .iter()
886 .any(|pattern| pattern.matches(origin))
887 }
888
889 fn allow_origin_value(&self, origin: &str) -> Option<String> {
890 if !self.is_origin_allowed(origin) {
891 return None;
892 }
893 if self.config.allow_any_origin && !self.config.allow_credentials {
894 Some("*".to_string())
895 } else {
896 Some(origin.to_string())
897 }
898 }
899
900 fn allow_methods_value(&self) -> String {
901 self.config
902 .allowed_methods
903 .iter()
904 .map(|method| method.as_str())
905 .collect::<Vec<_>>()
906 .join(", ")
907 }
908
909 fn allow_headers_value(&self, request: &Request) -> Option<String> {
910 if self.config.allowed_headers.is_empty() {
911 return None;
916 }
917
918 if self.config.allowed_headers.iter().any(|h| h == "*") {
923 if self.config.allow_credentials {
924 return request
927 .headers()
928 .get("access-control-request-headers")
929 .and_then(|value| std::str::from_utf8(value).ok())
930 .map(ToString::to_string);
931 }
932 return Some("*".to_string());
933 }
934
935 Some(self.config.allowed_headers.join(", "))
936 }
937
938 fn apply_common_headers(&self, mut response: Response, origin: &str) -> Response {
939 if let Some(allow_origin) = self.allow_origin_value(origin) {
940 let is_wildcard = allow_origin == "*";
941 response = response.header("access-control-allow-origin", allow_origin.into_bytes());
942 if !is_wildcard {
943 response = response.header("vary", b"Origin".to_vec());
944 }
945 if self.config.allow_credentials {
946 response = response.header("access-control-allow-credentials", b"true".to_vec());
947 }
948 if !self.config.expose_headers.is_empty() {
949 response = response.header(
950 "access-control-expose-headers",
951 self.config.expose_headers.join(", ").into_bytes(),
952 );
953 }
954 }
955 response
956 }
957}
958
959impl Default for Cors {
960 fn default() -> Self {
961 Self::new()
962 }
963}
964
965#[derive(Debug, Clone)]
966struct CorsOrigin(String);
967
968impl Middleware for Cors {
969 fn before<'a>(
970 &'a self,
971 _ctx: &'a RequestContext,
972 req: &'a mut Request,
973 ) -> BoxFuture<'a, ControlFlow> {
974 let origin = req
975 .headers()
976 .get("origin")
977 .and_then(|value| std::str::from_utf8(value).ok())
978 .map(ToString::to_string);
979
980 let Some(origin) = origin else {
981 return Box::pin(async { ControlFlow::Continue });
982 };
983
984 if !self.is_origin_allowed(&origin) {
985 let is_preflight = req.method() == crate::request::Method::Options
986 && req.headers().get("access-control-request-method").is_some();
987 if is_preflight {
988 return Box::pin(async {
989 ControlFlow::Break(Response::with_status(
990 crate::response::StatusCode::FORBIDDEN,
991 ))
992 });
993 }
994 return Box::pin(async { ControlFlow::Continue });
995 }
996
997 let is_preflight = req.method() == crate::request::Method::Options
998 && req.headers().get("access-control-request-method").is_some();
999
1000 if is_preflight {
1001 let mut response = Response::no_content();
1002 response = self.apply_common_headers(response, &origin);
1003 response = response.header(
1004 "access-control-allow-methods",
1005 self.allow_methods_value().into_bytes(),
1006 );
1007
1008 if let Some(value) = self.allow_headers_value(req) {
1009 response = response.header("access-control-allow-headers", value.into_bytes());
1010 }
1011
1012 if let Some(max_age) = self.config.max_age {
1013 response =
1014 response.header("access-control-max-age", max_age.to_string().into_bytes());
1015 }
1016
1017 return Box::pin(async move { ControlFlow::Break(response) });
1018 }
1019
1020 req.insert_extension(CorsOrigin(origin));
1021 Box::pin(async { ControlFlow::Continue })
1022 }
1023
1024 fn after<'a>(
1025 &'a self,
1026 _ctx: &'a RequestContext,
1027 req: &'a Request,
1028 response: Response,
1029 ) -> BoxFuture<'a, Response> {
1030 let origin = req.get_extension::<CorsOrigin>().map(|v| v.0.clone());
1031 Box::pin(async move {
1032 if let Some(origin) = origin {
1033 return self.apply_common_headers(response, &origin);
1034 }
1035 response
1036 })
1037 }
1038
1039 fn name(&self) -> &'static str {
1040 "Cors"
1041 }
1042}
1043
1044fn wildcard_match(pattern: &str, value: &str) -> bool {
1045 let mut pat_chars = pattern.chars().peekable();
1047 let mut val_chars = value.chars().peekable();
1048 let mut star = None;
1049 let mut match_after_star = None;
1050
1051 while let Some(p) = pat_chars.next() {
1052 match p {
1053 '*' => {
1054 star = Some(pat_chars.clone());
1055 match_after_star = Some(val_chars.clone());
1056 }
1057 _ => {
1058 if let Some(v) = val_chars.next() {
1059 if p != v {
1060 if let (Some(pat_backup), Some(val_backup)) =
1061 (star.clone(), match_after_star.clone())
1062 {
1063 pat_chars = pat_backup;
1064 val_chars = val_backup;
1065 val_chars.next();
1066 match_after_star = Some(val_chars.clone());
1067 continue;
1068 }
1069 return false;
1070 }
1071 } else {
1072 return false;
1073 }
1074 }
1075 }
1076 }
1077
1078 if pat_chars.peek().is_none() && val_chars.peek().is_none() {
1080 return true;
1081 }
1082
1083 if let Some(pat_backup) = star {
1084 if val_chars.peek().is_none() {
1085 let trailing = pat_backup;
1086 for ch in trailing {
1087 if ch != '*' {
1088 return false;
1089 }
1090 }
1091 return true;
1092 }
1093 }
1094
1095 val_chars.peek().is_none()
1096}
1097
1098fn regex_match(pattern: &str, value: &str) -> bool {
1099 let pat = pattern.as_bytes();
1101 let text = value.as_bytes();
1102
1103 if pat.first() == Some(&b'^') {
1104 return regex_match_here(&pat[1..], text);
1105 }
1106
1107 let mut i = 0;
1108 loop {
1109 if regex_match_here(pat, &text[i..]) {
1110 return true;
1111 }
1112 if i == text.len() {
1113 break;
1114 }
1115 i += 1;
1116 }
1117 false
1118}
1119
1120fn regex_match_here(pattern: &[u8], text: &[u8]) -> bool {
1121 if pattern.is_empty() {
1122 return true;
1123 }
1124 if pattern == b"$" {
1125 return text.is_empty();
1126 }
1127 if pattern.len() >= 2 && pattern[1] == b'*' {
1128 return regex_match_star(pattern[0], &pattern[2..], text);
1129 }
1130 if !text.is_empty() && (pattern[0] == b'.' || pattern[0] == text[0]) {
1131 return regex_match_here(&pattern[1..], &text[1..]);
1132 }
1133 false
1134}
1135
1136fn regex_match_star(ch: u8, pattern: &[u8], text: &[u8]) -> bool {
1137 let mut i = 0;
1138 loop {
1139 if regex_match_here(pattern, &text[i..]) {
1140 return true;
1141 }
1142 if i == text.len() {
1143 return false;
1144 }
1145 if ch != b'.' && text[i] != ch {
1146 return false;
1147 }
1148 i += 1;
1149 }
1150}
1151
1152#[derive(Debug, Clone)]
1158pub struct RequestResponseLogger {
1159 log_config: LogConfig,
1160 redact_headers: HashSet<String>,
1161 log_request_headers: bool,
1162 log_response_headers: bool,
1163 log_body: bool,
1164 max_body_bytes: usize,
1165}
1166
1167impl Default for RequestResponseLogger {
1168 fn default() -> Self {
1169 Self {
1170 log_config: LogConfig::production(),
1171 redact_headers: default_redacted_headers(),
1172 log_request_headers: true,
1173 log_response_headers: true,
1174 log_body: false,
1175 max_body_bytes: 1024,
1176 }
1177 }
1178}
1179
1180impl RequestResponseLogger {
1181 #[must_use]
1183 pub fn new() -> Self {
1184 Self::default()
1185 }
1186
1187 #[must_use]
1189 pub fn log_config(mut self, config: LogConfig) -> Self {
1190 self.log_config = config;
1191 self
1192 }
1193
1194 #[must_use]
1196 pub fn log_request_headers(mut self, enabled: bool) -> Self {
1197 self.log_request_headers = enabled;
1198 self
1199 }
1200
1201 #[must_use]
1203 pub fn log_response_headers(mut self, enabled: bool) -> Self {
1204 self.log_response_headers = enabled;
1205 self
1206 }
1207
1208 #[must_use]
1210 pub fn log_body(mut self, enabled: bool) -> Self {
1211 self.log_body = enabled;
1212 self
1213 }
1214
1215 #[must_use]
1217 pub fn max_body_bytes(mut self, max: usize) -> Self {
1218 self.max_body_bytes = max;
1219 self
1220 }
1221
1222 #[must_use]
1224 pub fn redact_header(mut self, name: impl Into<String>) -> Self {
1225 self.redact_headers.insert(name.into().to_ascii_lowercase());
1226 self
1227 }
1228}
1229
1230#[derive(Debug, Clone)]
1231struct RequestStart(Instant);
1232
1233impl Middleware for RequestResponseLogger {
1234 fn before<'a>(
1235 &'a self,
1236 ctx: &'a RequestContext,
1237 req: &'a mut Request,
1238 ) -> BoxFuture<'a, ControlFlow> {
1239 let logger = RequestLogger::new(ctx, self.log_config.clone());
1240 req.insert_extension(RequestStart(Instant::now()));
1241
1242 let method = req.method();
1243 let path = req.path();
1244 let query = req.query();
1245 let body_bytes = body_len(req.body());
1246
1247 logger.info_with_fields("request", |entry| {
1248 let mut entry = entry
1249 .field("method", method)
1250 .field("path", path)
1251 .field("body_bytes", body_bytes);
1252
1253 if let Some(q) = query {
1254 entry = entry.field("query", q);
1255 }
1256
1257 if self.log_request_headers {
1258 let headers = format_headers(req.headers().iter(), &self.redact_headers);
1259 entry = entry.field("headers", headers);
1260 }
1261
1262 if self.log_body {
1263 if let Some(body) = preview_body(req.body(), self.max_body_bytes) {
1264 entry = entry.field("body", body);
1265 }
1266 }
1267
1268 entry
1269 });
1270
1271 Box::pin(async { ControlFlow::Continue })
1272 }
1273
1274 fn after<'a>(
1275 &'a self,
1276 ctx: &'a RequestContext,
1277 req: &'a Request,
1278 response: Response,
1279 ) -> BoxFuture<'a, Response> {
1280 let logger = RequestLogger::new(ctx, self.log_config.clone());
1281 let duration = req
1282 .get_extension::<RequestStart>()
1283 .map(|start| start.0.elapsed())
1284 .unwrap_or_default();
1285
1286 let status = response.status();
1287 let body_bytes = response.body_ref().len();
1288
1289 logger.info_with_fields("response", |entry| {
1290 let mut entry = entry
1291 .field("status", status.as_u16())
1292 .field("duration_us", duration.as_micros())
1293 .field("body_bytes", body_bytes);
1294
1295 if self.log_response_headers {
1296 let headers = format_response_headers(response.headers(), &self.redact_headers);
1297 entry = entry.field("headers", headers);
1298 }
1299
1300 if self.log_body {
1301 if let Some(body) = preview_response_body(response.body_ref(), self.max_body_bytes)
1302 {
1303 entry = entry.field("body", body);
1304 }
1305 }
1306
1307 entry
1308 });
1309
1310 Box::pin(async move { response })
1311 }
1312
1313 fn name(&self) -> &'static str {
1314 "RequestResponseLogger"
1315 }
1316}
1317
1318fn default_redacted_headers() -> HashSet<String> {
1319 [
1320 "authorization",
1321 "proxy-authorization",
1322 "cookie",
1323 "set-cookie",
1324 ]
1325 .iter()
1326 .map(ToString::to_string)
1327 .collect()
1328}
1329
1330fn body_len(body: &Body) -> usize {
1331 match body {
1332 Body::Empty => 0,
1333 Body::Bytes(bytes) => bytes.len(),
1334 Body::Stream { content_length, .. } => content_length.unwrap_or(0),
1335 }
1336}
1337
1338fn preview_body(body: &Body, max_bytes: usize) -> Option<String> {
1339 if max_bytes == 0 {
1340 return None;
1341 }
1342 match body {
1343 Body::Empty => None,
1344 Body::Bytes(bytes) => {
1345 if bytes.is_empty() {
1346 None
1347 } else {
1348 Some(format_bytes(bytes, max_bytes))
1349 }
1350 }
1351 Body::Stream { .. } => None,
1352 }
1353}
1354
1355fn preview_response_body(body: &crate::response::ResponseBody, max_bytes: usize) -> Option<String> {
1356 if max_bytes == 0 {
1357 return None;
1358 }
1359 match body {
1360 crate::response::ResponseBody::Empty => None,
1361 crate::response::ResponseBody::Bytes(bytes) => {
1362 if bytes.is_empty() {
1363 None
1364 } else {
1365 Some(format_bytes(bytes, max_bytes))
1366 }
1367 }
1368 crate::response::ResponseBody::Stream(_) => None,
1369 }
1370}
1371
1372fn format_headers<'a>(
1373 headers: impl Iterator<Item = (&'a str, &'a [u8])>,
1374 redacted: &HashSet<String>,
1375) -> String {
1376 let mut out = String::new();
1377 for (idx, (name, value)) in headers.enumerate() {
1378 if idx > 0 {
1379 out.push_str(", ");
1380 }
1381 out.push_str(name);
1382 out.push('=');
1383
1384 let lowered = name.to_ascii_lowercase();
1385 if redacted.contains(&lowered) {
1386 out.push_str("<redacted>");
1387 continue;
1388 }
1389
1390 match std::str::from_utf8(value) {
1391 Ok(text) => out.push_str(text),
1392 Err(_) => out.push_str("<binary>"),
1393 }
1394 }
1395 out
1396}
1397
1398fn format_response_headers(headers: &[(String, Vec<u8>)], redacted: &HashSet<String>) -> String {
1399 format_headers(
1400 headers
1401 .iter()
1402 .map(|(name, value)| (name.as_str(), value.as_slice())),
1403 redacted,
1404 )
1405}
1406
1407fn format_bytes(bytes: &[u8], max_bytes: usize) -> String {
1408 let limit = max_bytes.min(bytes.len());
1409 match std::str::from_utf8(&bytes[..limit]) {
1410 Ok(text) => {
1411 let mut output = text.to_string();
1412 if bytes.len() > max_bytes {
1413 output.push_str("...");
1414 }
1415 output
1416 }
1417 Err(_) => format!("<{} bytes binary>", bytes.len()),
1418 }
1419}
1420
1421impl From<&crate::response::ResponseBody> for crate::response::ResponseBody {
1423 fn from(body: &crate::response::ResponseBody) -> Self {
1424 match body {
1425 crate::response::ResponseBody::Empty => crate::response::ResponseBody::Empty,
1426 crate::response::ResponseBody::Bytes(b) => {
1427 crate::response::ResponseBody::Bytes(b.clone())
1428 }
1429 crate::response::ResponseBody::Stream(_) => crate::response::ResponseBody::Empty,
1430 }
1431 }
1432}
1433
1434#[derive(Debug, Clone, PartialEq, Eq, Hash)]
1443pub struct RequestId(pub String);
1444
1445impl RequestId {
1446 #[must_use]
1448 pub fn new(id: impl Into<String>) -> Self {
1449 Self(id.into())
1450 }
1451
1452 #[must_use]
1454 pub fn as_str(&self) -> &str {
1455 &self.0
1456 }
1457
1458 #[must_use]
1463 pub fn generate() -> Self {
1464 use std::sync::atomic::{AtomicU64, Ordering};
1465 use std::time::{SystemTime, UNIX_EPOCH};
1466
1467 static COUNTER: AtomicU64 = AtomicU64::new(0);
1468
1469 let timestamp = SystemTime::now()
1470 .duration_since(UNIX_EPOCH)
1471 .map(|d| d.as_micros() as u64)
1472 .unwrap_or(0);
1473 let counter = COUNTER.fetch_add(1, Ordering::Relaxed);
1474
1475 Self(format!("{:x}-{:x}", timestamp, counter))
1477 }
1478}
1479
1480impl std::fmt::Display for RequestId {
1481 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1482 write!(f, "{}", self.0)
1483 }
1484}
1485
1486impl From<String> for RequestId {
1487 fn from(s: String) -> Self {
1488 Self(s)
1489 }
1490}
1491
1492impl From<&str> for RequestId {
1493 fn from(s: &str) -> Self {
1494 Self(s.to_string())
1495 }
1496}
1497
1498#[derive(Debug, Clone)]
1500pub struct RequestIdConfig {
1501 pub header_name: String,
1503 pub accept_from_client: bool,
1505 pub add_to_response: bool,
1507 pub max_client_id_length: usize,
1509}
1510
1511impl Default for RequestIdConfig {
1512 fn default() -> Self {
1513 Self {
1514 header_name: "x-request-id".to_string(),
1515 accept_from_client: true,
1516 add_to_response: true,
1517 max_client_id_length: 128,
1518 }
1519 }
1520}
1521
1522impl RequestIdConfig {
1523 #[must_use]
1525 pub fn new() -> Self {
1526 Self::default()
1527 }
1528
1529 #[must_use]
1531 pub fn header_name(mut self, name: impl Into<String>) -> Self {
1532 self.header_name = name.into();
1533 self
1534 }
1535
1536 #[must_use]
1538 pub fn accept_from_client(mut self, accept: bool) -> Self {
1539 self.accept_from_client = accept;
1540 self
1541 }
1542
1543 #[must_use]
1545 pub fn add_to_response(mut self, add: bool) -> Self {
1546 self.add_to_response = add;
1547 self
1548 }
1549
1550 #[must_use]
1552 pub fn max_client_id_length(mut self, max: usize) -> Self {
1553 self.max_client_id_length = max;
1554 self
1555 }
1556}
1557
1558#[derive(Debug, Clone)]
1583pub struct RequestIdMiddleware {
1584 config: RequestIdConfig,
1585}
1586
1587impl Default for RequestIdMiddleware {
1588 fn default() -> Self {
1589 Self::new()
1590 }
1591}
1592
1593impl RequestIdMiddleware {
1594 #[must_use]
1596 pub fn new() -> Self {
1597 Self {
1598 config: RequestIdConfig::default(),
1599 }
1600 }
1601
1602 #[must_use]
1604 pub fn with_config(config: RequestIdConfig) -> Self {
1605 Self { config }
1606 }
1607
1608 fn get_or_generate_id(&self, req: &Request) -> RequestId {
1610 if self.config.accept_from_client {
1611 if let Some(header_value) = req.headers().get(&self.config.header_name) {
1612 if let Ok(client_id) = std::str::from_utf8(header_value) {
1613 if !client_id.is_empty()
1615 && client_id.len() <= self.config.max_client_id_length
1616 && is_valid_request_id(client_id)
1617 {
1618 return RequestId::new(client_id);
1619 }
1620 }
1621 }
1622 }
1623 RequestId::generate()
1624 }
1625}
1626
1627fn is_valid_request_id(id: &str) -> bool {
1629 !id.is_empty()
1630 && id
1631 .chars()
1632 .all(|c| c.is_ascii_alphanumeric() || c == '-' || c == '_' || c == '.')
1633}
1634
1635impl Middleware for RequestIdMiddleware {
1636 fn before<'a>(
1637 &'a self,
1638 _ctx: &'a RequestContext,
1639 req: &'a mut Request,
1640 ) -> BoxFuture<'a, ControlFlow> {
1641 let request_id = self.get_or_generate_id(req);
1642 req.insert_extension(request_id);
1643 Box::pin(async { ControlFlow::Continue })
1644 }
1645
1646 fn after<'a>(
1647 &'a self,
1648 _ctx: &'a RequestContext,
1649 req: &'a Request,
1650 response: Response,
1651 ) -> BoxFuture<'a, Response> {
1652 if !self.config.add_to_response {
1653 return Box::pin(async move { response });
1654 }
1655
1656 let request_id = req.get_extension::<RequestId>().cloned();
1657 let header_name = self.config.header_name.clone();
1658
1659 Box::pin(async move {
1660 if let Some(id) = request_id {
1661 response.header(header_name, id.0.into_bytes())
1662 } else {
1663 response
1664 }
1665 })
1666 }
1667
1668 fn name(&self) -> &'static str {
1669 "RequestId"
1670 }
1671}
1672
1673#[derive(Debug, Clone, Copy, PartialEq, Eq)]
1681pub enum XFrameOptions {
1682 Deny,
1684 SameOrigin,
1686}
1687
1688impl XFrameOptions {
1689 fn as_bytes(self) -> &'static [u8] {
1690 match self {
1691 Self::Deny => b"DENY",
1692 Self::SameOrigin => b"SAMEORIGIN",
1693 }
1694 }
1695}
1696
1697#[derive(Debug, Clone, Copy, PartialEq, Eq)]
1701pub enum ReferrerPolicy {
1702 NoReferrer,
1704 NoReferrerWhenDowngrade,
1706 Origin,
1708 OriginWhenCrossOrigin,
1710 SameOrigin,
1712 StrictOrigin,
1714 StrictOriginWhenCrossOrigin,
1716 UnsafeUrl,
1718}
1719
1720impl ReferrerPolicy {
1721 fn as_bytes(self) -> &'static [u8] {
1722 match self {
1723 Self::NoReferrer => b"no-referrer",
1724 Self::NoReferrerWhenDowngrade => b"no-referrer-when-downgrade",
1725 Self::Origin => b"origin",
1726 Self::OriginWhenCrossOrigin => b"origin-when-cross-origin",
1727 Self::SameOrigin => b"same-origin",
1728 Self::StrictOrigin => b"strict-origin",
1729 Self::StrictOriginWhenCrossOrigin => b"strict-origin-when-cross-origin",
1730 Self::UnsafeUrl => b"unsafe-url",
1731 }
1732 }
1733}
1734
1735#[derive(Debug, Clone)]
1759pub struct SecurityHeadersConfig {
1760 pub x_content_type_options: Option<&'static str>,
1763 pub x_frame_options: Option<XFrameOptions>,
1766 pub x_xss_protection: Option<&'static str>,
1772 pub content_security_policy: Option<String>,
1775 pub hsts: Option<(u64, bool, bool)>,
1779 pub referrer_policy: Option<ReferrerPolicy>,
1782 pub permissions_policy: Option<String>,
1785}
1786
1787impl Default for SecurityHeadersConfig {
1788 fn default() -> Self {
1789 Self {
1790 x_content_type_options: Some("nosniff"),
1791 x_frame_options: Some(XFrameOptions::Deny),
1792 x_xss_protection: Some("0"),
1793 content_security_policy: None,
1794 hsts: None,
1795 referrer_policy: Some(ReferrerPolicy::StrictOriginWhenCrossOrigin),
1796 permissions_policy: None,
1797 }
1798 }
1799}
1800
1801impl SecurityHeadersConfig {
1802 #[must_use]
1804 pub fn new() -> Self {
1805 Self::default()
1806 }
1807
1808 #[must_use]
1810 pub fn none() -> Self {
1811 Self {
1812 x_content_type_options: None,
1813 x_frame_options: None,
1814 x_xss_protection: None,
1815 content_security_policy: None,
1816 hsts: None,
1817 referrer_policy: None,
1818 permissions_policy: None,
1819 }
1820 }
1821
1822 #[must_use]
1829 pub fn strict() -> Self {
1830 Self {
1831 x_content_type_options: Some("nosniff"),
1832 x_frame_options: Some(XFrameOptions::Deny),
1833 x_xss_protection: Some("0"),
1834 content_security_policy: Some("default-src 'self'".to_string()),
1835 hsts: Some((31536000, true, false)), referrer_policy: Some(ReferrerPolicy::NoReferrer),
1837 permissions_policy: Some("geolocation=(), camera=(), microphone=()".to_string()),
1838 }
1839 }
1840
1841 #[must_use]
1843 pub fn x_content_type_options(mut self, value: Option<&'static str>) -> Self {
1844 self.x_content_type_options = value;
1845 self
1846 }
1847
1848 #[must_use]
1850 pub fn x_frame_options(mut self, value: Option<XFrameOptions>) -> Self {
1851 self.x_frame_options = value;
1852 self
1853 }
1854
1855 #[must_use]
1857 pub fn x_xss_protection(mut self, value: Option<&'static str>) -> Self {
1858 self.x_xss_protection = value;
1859 self
1860 }
1861
1862 #[must_use]
1864 pub fn content_security_policy(mut self, value: impl Into<String>) -> Self {
1865 self.content_security_policy = Some(value.into());
1866 self
1867 }
1868
1869 #[must_use]
1871 pub fn no_content_security_policy(mut self) -> Self {
1872 self.content_security_policy = None;
1873 self
1874 }
1875
1876 #[must_use]
1889 pub fn hsts(mut self, max_age: u64, include_sub_domains: bool, preload: bool) -> Self {
1890 self.hsts = Some((max_age, include_sub_domains, preload));
1891 self
1892 }
1893
1894 #[must_use]
1896 pub fn no_hsts(mut self) -> Self {
1897 self.hsts = None;
1898 self
1899 }
1900
1901 #[must_use]
1903 pub fn referrer_policy(mut self, value: Option<ReferrerPolicy>) -> Self {
1904 self.referrer_policy = value;
1905 self
1906 }
1907
1908 #[must_use]
1910 pub fn permissions_policy(mut self, value: impl Into<String>) -> Self {
1911 self.permissions_policy = Some(value.into());
1912 self
1913 }
1914
1915 #[must_use]
1917 pub fn no_permissions_policy(mut self) -> Self {
1918 self.permissions_policy = None;
1919 self
1920 }
1921
1922 fn build_hsts_value(&self) -> Option<String> {
1924 self.hsts.map(|(max_age, include_sub, preload)| {
1925 let mut value = format!("max-age={}", max_age);
1926 if include_sub {
1927 value.push_str("; includeSubDomains");
1928 }
1929 if preload {
1930 value.push_str("; preload");
1931 }
1932 value
1933 })
1934 }
1935}
1936
1937#[derive(Debug, Clone)]
1968pub struct SecurityHeaders {
1969 config: SecurityHeadersConfig,
1970}
1971
1972impl Default for SecurityHeaders {
1973 fn default() -> Self {
1974 Self::new()
1975 }
1976}
1977
1978impl SecurityHeaders {
1979 #[must_use]
1981 pub fn new() -> Self {
1982 Self {
1983 config: SecurityHeadersConfig::default(),
1984 }
1985 }
1986
1987 #[must_use]
1989 pub fn with_config(config: SecurityHeadersConfig) -> Self {
1990 Self { config }
1991 }
1992
1993 #[must_use]
1995 pub fn strict() -> Self {
1996 Self {
1997 config: SecurityHeadersConfig::strict(),
1998 }
1999 }
2000}
2001
2002impl Middleware for SecurityHeaders {
2003 fn after<'a>(
2004 &'a self,
2005 _ctx: &'a RequestContext,
2006 _req: &'a Request,
2007 response: Response,
2008 ) -> BoxFuture<'a, Response> {
2009 let config = self.config.clone();
2010 Box::pin(async move {
2011 let mut resp = response;
2012
2013 if let Some(value) = config.x_content_type_options {
2015 resp = resp.header("X-Content-Type-Options", value.as_bytes().to_vec());
2016 }
2017
2018 if let Some(value) = config.x_frame_options {
2020 resp = resp.header("X-Frame-Options", value.as_bytes().to_vec());
2021 }
2022
2023 if let Some(value) = config.x_xss_protection {
2025 resp = resp.header("X-XSS-Protection", value.as_bytes().to_vec());
2026 }
2027
2028 if let Some(ref value) = config.content_security_policy {
2030 resp = resp.header("Content-Security-Policy", value.as_bytes().to_vec());
2031 }
2032
2033 if let Some(ref hsts_value) = config.build_hsts_value() {
2035 resp = resp.header("Strict-Transport-Security", hsts_value.as_bytes().to_vec());
2036 }
2037
2038 if let Some(value) = config.referrer_policy {
2040 resp = resp.header("Referrer-Policy", value.as_bytes().to_vec());
2041 }
2042
2043 if let Some(ref value) = config.permissions_policy {
2045 resp = resp.header("Permissions-Policy", value.as_bytes().to_vec());
2046 }
2047
2048 resp
2049 })
2050 }
2051
2052 fn name(&self) -> &'static str {
2053 "SecurityHeaders"
2054 }
2055}
2056
2057#[derive(Debug, Clone, PartialEq, Eq, Hash)]
2066pub struct CsrfToken(pub String);
2067
2068impl CsrfToken {
2069 #[must_use]
2071 pub fn new(token: impl Into<String>) -> Self {
2072 Self(token.into())
2073 }
2074
2075 #[must_use]
2077 pub fn as_str(&self) -> &str {
2078 &self.0
2079 }
2080
2081 #[must_use]
2091 pub fn generate() -> Self {
2092 let mut bytes = [0u8; 32];
2093 if let Err(err) = getrandom::fill(&mut bytes) {
2094 panic!(
2095 "FATAL: OS cryptographically secure random source is unavailable ({err}). \
2096 CSRF token generation requires a CSPRNG. Cannot safely generate CSRF tokens \
2097 without cryptographic entropy."
2098 );
2099 }
2100 Self(Self::bytes_to_hex(&bytes))
2101 }
2102
2103 fn bytes_to_hex(bytes: &[u8]) -> String {
2104 use std::fmt::Write;
2105 let mut s = String::with_capacity(bytes.len() * 2);
2106 for b in bytes {
2107 let _ = write!(s, "{b:02x}");
2108 }
2109 s
2110 }
2111}
2112
2113impl std::fmt::Display for CsrfToken {
2114 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
2115 f.write_str(&self.0)
2116 }
2117}
2118
2119impl From<&str> for CsrfToken {
2120 fn from(s: &str) -> Self {
2121 Self(s.to_string())
2122 }
2123}
2124
2125#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
2127pub enum CsrfMode {
2128 #[default]
2131 DoubleSubmit,
2132 HeaderOnly,
2134}
2135
2136#[derive(Debug, Clone)]
2138pub struct CsrfConfig {
2139 pub cookie_name: String,
2141 pub header_name: String,
2143 pub mode: CsrfMode,
2145 pub rotate_token: bool,
2147 pub production: bool,
2149 pub error_message: Option<String>,
2151}
2152
2153impl Default for CsrfConfig {
2154 fn default() -> Self {
2155 Self {
2156 cookie_name: "csrf_token".to_string(),
2157 header_name: "x-csrf-token".to_string(),
2158 mode: CsrfMode::DoubleSubmit,
2159 rotate_token: false,
2160 production: true,
2161 error_message: None,
2162 }
2163 }
2164}
2165
2166impl CsrfConfig {
2167 #[must_use]
2169 pub fn new() -> Self {
2170 Self::default()
2171 }
2172
2173 #[must_use]
2175 pub fn cookie_name(mut self, name: impl Into<String>) -> Self {
2176 self.cookie_name = name.into();
2177 self
2178 }
2179
2180 #[must_use]
2182 pub fn header_name(mut self, name: impl Into<String>) -> Self {
2183 self.header_name = name.into();
2184 self
2185 }
2186
2187 #[must_use]
2189 pub fn mode(mut self, mode: CsrfMode) -> Self {
2190 self.mode = mode;
2191 self
2192 }
2193
2194 #[must_use]
2196 pub fn rotate_token(mut self, rotate: bool) -> Self {
2197 self.rotate_token = rotate;
2198 self
2199 }
2200
2201 #[must_use]
2203 pub fn production(mut self, production: bool) -> Self {
2204 self.production = production;
2205 self
2206 }
2207
2208 #[must_use]
2210 pub fn error_message(mut self, message: impl Into<String>) -> Self {
2211 self.error_message = Some(message.into());
2212 self
2213 }
2214}
2215
2216#[derive(Debug, Clone)]
2246pub struct CsrfMiddleware {
2247 config: CsrfConfig,
2248}
2249
2250impl Default for CsrfMiddleware {
2251 fn default() -> Self {
2252 Self::new()
2253 }
2254}
2255
2256impl CsrfMiddleware {
2257 #[must_use]
2259 pub fn new() -> Self {
2260 Self {
2261 config: CsrfConfig::default(),
2262 }
2263 }
2264
2265 #[must_use]
2267 pub fn with_config(config: CsrfConfig) -> Self {
2268 Self { config }
2269 }
2270
2271 fn is_safe_method(method: crate::request::Method) -> bool {
2273 matches!(
2274 method,
2275 crate::request::Method::Get
2276 | crate::request::Method::Head
2277 | crate::request::Method::Options
2278 | crate::request::Method::Trace
2279 )
2280 }
2281
2282 fn get_cookie_token(&self, req: &Request) -> Option<String> {
2284 let cookie_header = req.headers().get("cookie")?;
2285 let cookie_str = std::str::from_utf8(cookie_header).ok()?;
2286
2287 for part in cookie_str.split(';') {
2289 let part = part.trim();
2290 if let Some((name, value)) = part.split_once('=') {
2291 if name.trim() == self.config.cookie_name {
2292 return Some(value.trim().to_string());
2293 }
2294 }
2295 }
2296 None
2297 }
2298
2299 fn get_header_token(&self, req: &Request) -> Option<String> {
2301 let header_value = req.headers().get(&self.config.header_name)?;
2302 std::str::from_utf8(header_value)
2303 .ok()
2304 .map(|s| s.trim().to_string())
2305 }
2306
2307 fn validate_token(&self, req: &Request) -> Result<Option<CsrfToken>, Response> {
2309 let header_token = self.get_header_token(req);
2310
2311 match self.config.mode {
2312 CsrfMode::DoubleSubmit => {
2313 let cookie_token = self.get_cookie_token(req);
2314
2315 match (header_token, cookie_token) {
2316 (Some(header), Some(cookie))
2317 if !header.is_empty()
2318 && crate::password::constant_time_eq(
2319 header.as_bytes(),
2320 cookie.as_bytes(),
2321 ) =>
2322 {
2323 Ok(Some(CsrfToken::new(header)))
2324 }
2325 (None, _) | (_, None) => Err(self.csrf_error_response("CSRF token missing")),
2326 _ => Err(self.csrf_error_response("CSRF token mismatch")),
2327 }
2328 }
2329 CsrfMode::HeaderOnly => match header_token {
2330 Some(token) if !token.is_empty() => Ok(Some(CsrfToken::new(token))),
2331 _ => Err(self.csrf_error_response("CSRF token missing in header")),
2332 },
2333 }
2334 }
2335
2336 fn csrf_error_response(&self, default_message: &str) -> Response {
2338 let message = self
2339 .config
2340 .error_message
2341 .as_deref()
2342 .unwrap_or(default_message);
2343
2344 let detail = serde_json::json!({
2347 "detail": [{
2348 "type": "csrf_error",
2349 "loc": ["header", self.config.header_name],
2350 "msg": message,
2351 }]
2352 });
2353 let body = detail.to_string();
2354
2355 Response::with_status(crate::response::StatusCode::FORBIDDEN)
2356 .header("content-type", b"application/json".to_vec())
2357 .body(crate::response::ResponseBody::Bytes(body.into_bytes()))
2358 }
2359
2360 fn make_set_cookie_header_value(cookie_name: &str, token: &str, production: bool) -> Vec<u8> {
2362 let mut cookie = format!("{}={}; Path=/; SameSite=Strict", cookie_name, token);
2363
2364 if production {
2365 cookie.push_str("; Secure");
2366 }
2367
2368 cookie.into_bytes()
2371 }
2372}
2373
2374impl Middleware for CsrfMiddleware {
2375 fn before<'a>(
2376 &'a self,
2377 _ctx: &'a RequestContext,
2378 req: &'a mut Request,
2379 ) -> BoxFuture<'a, ControlFlow> {
2380 Box::pin(async move {
2381 if Self::is_safe_method(req.method()) {
2382 let existing_token = self.get_cookie_token(req);
2384 let token = existing_token
2385 .map(CsrfToken::new)
2386 .unwrap_or_else(CsrfToken::generate);
2387 req.insert_extension(token);
2388 ControlFlow::Continue
2389 } else {
2390 match self.validate_token(req) {
2392 Ok(Some(token)) => {
2393 req.insert_extension(token);
2394 ControlFlow::Continue
2395 }
2396 Ok(None) => ControlFlow::Continue,
2397 Err(response) => ControlFlow::Break(response),
2398 }
2399 }
2400 })
2401 }
2402
2403 fn after<'a>(
2404 &'a self,
2405 _ctx: &'a RequestContext,
2406 req: &'a Request,
2407 response: Response,
2408 ) -> BoxFuture<'a, Response> {
2409 let config = self.config.clone();
2410 let is_safe = Self::is_safe_method(req.method());
2411 let existing_cookie_token = self.get_cookie_token(req);
2412 let token = req.get_extension::<CsrfToken>().cloned();
2413
2414 Box::pin(async move {
2415 if is_safe {
2419 let should_set_cookie = existing_cookie_token.is_none() || config.rotate_token;
2420
2421 if should_set_cookie {
2422 if let Some(token) = token {
2423 let cookie_value = Self::make_set_cookie_header_value(
2424 &config.cookie_name,
2425 token.as_str(),
2426 config.production,
2427 );
2428 return response.header("set-cookie", cookie_value);
2429 }
2430 }
2431 }
2432 response
2433 })
2434 }
2435
2436 fn name(&self) -> &'static str {
2437 "CSRF"
2438 }
2439}
2440
2441#[cfg(feature = "compression")]
2464#[derive(Debug, Clone)]
2465pub struct CompressionConfig {
2466 pub min_size: usize,
2470 pub level: u32,
2473 pub skip_content_types: Vec<&'static str>,
2476}
2477
2478#[cfg(feature = "compression")]
2479impl Default for CompressionConfig {
2480 fn default() -> Self {
2481 Self {
2482 min_size: 1024,
2483 level: 6,
2484 skip_content_types: vec![
2485 "image/jpeg",
2487 "image/png",
2488 "image/gif",
2489 "image/webp",
2490 "image/avif",
2491 "video/",
2493 "audio/",
2494 "application/zip",
2496 "application/gzip",
2497 "application/x-gzip",
2498 "application/x-bzip2",
2499 "application/x-xz",
2500 "application/x-7z-compressed",
2501 "application/x-rar-compressed",
2502 "application/pdf",
2504 "application/woff",
2505 "application/woff2",
2506 "font/woff",
2507 "font/woff2",
2508 ],
2509 }
2510 }
2511}
2512
2513#[cfg(feature = "compression")]
2514impl CompressionConfig {
2515 #[must_use]
2517 pub fn new() -> Self {
2518 Self::default()
2519 }
2520
2521 #[must_use]
2526 pub fn min_size(mut self, size: usize) -> Self {
2527 self.min_size = size;
2528 self
2529 }
2530
2531 #[must_use]
2539 pub fn level(mut self, level: u32) -> Self {
2540 self.level = level.clamp(1, 9);
2541 self
2542 }
2543
2544 #[must_use]
2548 pub fn skip_content_type(mut self, content_type: &'static str) -> Self {
2549 self.skip_content_types.push(content_type);
2550 self
2551 }
2552
2553 fn should_skip_content_type(&self, content_type: &str) -> bool {
2555 let ct_lower = content_type.to_ascii_lowercase();
2556 for skip in &self.skip_content_types {
2557 if skip.ends_with('/') {
2558 if ct_lower.starts_with(*skip) {
2560 return true;
2561 }
2562 } else {
2563 if ct_lower == *skip || ct_lower.starts_with(&format!("{skip};")) {
2565 return true;
2566 }
2567 }
2568 }
2569 false
2570 }
2571}
2572
2573#[cfg(feature = "compression")]
2606#[derive(Debug, Clone)]
2607pub struct CompressionMiddleware {
2608 config: CompressionConfig,
2609}
2610
2611#[cfg(feature = "compression")]
2612impl Default for CompressionMiddleware {
2613 fn default() -> Self {
2614 Self::new()
2615 }
2616}
2617
2618#[cfg(feature = "compression")]
2619impl CompressionMiddleware {
2620 #[must_use]
2622 pub fn new() -> Self {
2623 Self {
2624 config: CompressionConfig::default(),
2625 }
2626 }
2627
2628 #[must_use]
2630 pub fn with_config(config: CompressionConfig) -> Self {
2631 Self { config }
2632 }
2633
2634 fn accepts_gzip(req: &Request) -> bool {
2636 if let Some(accept_encoding) = req.headers().get("accept-encoding") {
2637 if let Ok(value) = std::str::from_utf8(accept_encoding) {
2638 for part in value.split(',') {
2641 let encoding = part.trim().split(';').next().unwrap_or("").trim();
2642 if encoding.eq_ignore_ascii_case("gzip") {
2643 return true;
2644 }
2645 if encoding == "*" {
2647 return true;
2648 }
2649 }
2650 }
2651 }
2652 false
2653 }
2654
2655 fn get_content_type(headers: &[(String, Vec<u8>)]) -> Option<String> {
2657 for (name, value) in headers {
2658 if name.eq_ignore_ascii_case("content-type") {
2659 return std::str::from_utf8(value).ok().map(String::from);
2660 }
2661 }
2662 None
2663 }
2664
2665 fn has_content_encoding(headers: &[(String, Vec<u8>)]) -> bool {
2667 headers
2668 .iter()
2669 .any(|(name, _)| name.eq_ignore_ascii_case("content-encoding"))
2670 }
2671
2672 fn compress_gzip(data: &[u8], level: u32) -> Result<Vec<u8>, std::io::Error> {
2674 use flate2::Compression;
2675 use flate2::write::GzEncoder;
2676 use std::io::Write;
2677
2678 let mut encoder = GzEncoder::new(Vec::new(), Compression::new(level));
2679 encoder.write_all(data)?;
2680 encoder.finish()
2681 }
2682}
2683
2684#[cfg(feature = "compression")]
2685impl Middleware for CompressionMiddleware {
2686 fn after<'a>(
2687 &'a self,
2688 _ctx: &'a RequestContext,
2689 req: &'a Request,
2690 response: Response,
2691 ) -> BoxFuture<'a, Response> {
2692 let config = self.config.clone();
2693
2694 Box::pin(async move {
2695 if !Self::accepts_gzip(req) {
2697 return response;
2698 }
2699
2700 let (status, headers, body) = response.into_parts();
2702
2703 if Self::has_content_encoding(&headers) {
2705 return Response::with_status(status)
2706 .body(body)
2707 .rebuild_with_headers(headers);
2708 }
2709
2710 let body_bytes = match body {
2712 crate::response::ResponseBody::Bytes(bytes) => bytes,
2713 other => {
2714 return Response::with_status(status)
2716 .body(other)
2717 .rebuild_with_headers(headers);
2718 }
2719 };
2720
2721 if body_bytes.len() < config.min_size {
2723 return Response::with_status(status)
2724 .body(crate::response::ResponseBody::Bytes(body_bytes))
2725 .rebuild_with_headers(headers);
2726 }
2727
2728 if let Some(content_type) = Self::get_content_type(&headers) {
2730 if config.should_skip_content_type(&content_type) {
2731 return Response::with_status(status)
2732 .body(crate::response::ResponseBody::Bytes(body_bytes))
2733 .rebuild_with_headers(headers);
2734 }
2735 }
2736
2737 match Self::compress_gzip(&body_bytes, config.level) {
2739 Ok(compressed) => {
2740 if compressed.len() >= body_bytes.len() {
2742 return Response::with_status(status)
2743 .body(crate::response::ResponseBody::Bytes(body_bytes))
2744 .rebuild_with_headers(headers);
2745 }
2746
2747 let mut resp = Response::with_status(status)
2749 .body(crate::response::ResponseBody::Bytes(compressed));
2750
2751 for (name, value) in headers {
2753 if !name.eq_ignore_ascii_case("content-length") {
2754 resp = resp.header(name, value);
2755 }
2756 }
2757
2758 resp = resp.header("Content-Encoding", b"gzip".to_vec());
2760 resp = resp.header("Vary", b"Accept-Encoding".to_vec());
2761
2762 resp
2763 }
2764 Err(_) => {
2765 Response::with_status(status)
2767 .body(crate::response::ResponseBody::Bytes(body_bytes))
2768 .rebuild_with_headers(headers)
2769 }
2770 }
2771 })
2772 }
2773
2774 fn name(&self) -> &'static str {
2775 "Compression"
2776 }
2777}
2778
2779use parking_lot::Mutex;
2784use std::collections::HashMap as StdHashMap;
2785use std::time::Duration;
2786
2787#[derive(Debug, Clone, Copy, PartialEq, Eq)]
2789pub enum RateLimitAlgorithm {
2790 TokenBucket,
2792 FixedWindow,
2794 SlidingWindow,
2796}
2797
2798#[derive(Debug, Clone)]
2800pub struct RateLimitResult {
2801 pub allowed: bool,
2803 pub limit: u64,
2805 pub remaining: u64,
2807 pub reset_after_secs: u64,
2809}
2810
2811pub trait KeyExtractor: Send + Sync {
2816 fn extract_key(&self, req: &Request) -> Option<String>;
2820}
2821
2822#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
2840pub struct RemoteAddr(pub std::net::IpAddr);
2841
2842impl std::fmt::Display for RemoteAddr {
2843 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
2844 write!(f, "{}", self.0)
2845 }
2846}
2847
2848#[derive(Debug, Clone)]
2869pub struct ConnectedIpKeyExtractor;
2870
2871impl KeyExtractor for ConnectedIpKeyExtractor {
2872 fn extract_key(&self, req: &Request) -> Option<String> {
2873 req.get_extension::<RemoteAddr>().map(ToString::to_string)
2874 }
2875}
2876
2877#[derive(Debug, Clone)]
2906pub struct IpKeyExtractor;
2907
2908impl KeyExtractor for IpKeyExtractor {
2909 fn extract_key(&self, req: &Request) -> Option<String> {
2910 if let Some(forwarded) = req.headers().get("x-forwarded-for") {
2912 if let Ok(s) = std::str::from_utf8(forwarded) {
2913 if let Some(ip) = s.split(',').next() {
2915 return Some(ip.trim().to_string());
2916 }
2917 }
2918 }
2919 if let Some(real_ip) = req.headers().get("x-real-ip") {
2920 if let Ok(s) = std::str::from_utf8(real_ip) {
2921 return Some(s.trim().to_string());
2922 }
2923 }
2924 Some("unknown".to_string())
2925 }
2926}
2927
2928#[derive(Debug, Clone)]
2959pub struct TrustedProxyIpKeyExtractor {
2960 trusted_cidrs: Vec<(std::net::IpAddr, u8)>,
2962}
2963
2964impl TrustedProxyIpKeyExtractor {
2965 #[must_use]
2967 pub fn new() -> Self {
2968 Self {
2969 trusted_cidrs: Vec::new(),
2970 }
2971 }
2972
2973 #[must_use]
2979 pub fn trust_cidr(mut self, cidr: &str) -> Self {
2980 let (ip, prefix) = parse_cidr(cidr).expect("invalid CIDR notation");
2981 self.trusted_cidrs.push((ip, prefix));
2982 self
2983 }
2984
2985 #[must_use]
2987 pub fn trust_loopback(mut self) -> Self {
2988 self.trusted_cidrs.push((
2989 std::net::IpAddr::V4(std::net::Ipv4Addr::new(127, 0, 0, 0)),
2990 8,
2991 ));
2992 self.trusted_cidrs
2993 .push((std::net::IpAddr::V6(std::net::Ipv6Addr::LOCALHOST), 128));
2994 self
2995 }
2996
2997 fn is_trusted(&self, ip: std::net::IpAddr) -> bool {
2999 self.trusted_cidrs
3000 .iter()
3001 .any(|(cidr_ip, prefix)| ip_in_cidr(ip, *cidr_ip, *prefix))
3002 }
3003
3004 fn extract_from_header(&self, req: &Request) -> Option<String> {
3006 if let Some(forwarded) = req.headers().get("x-forwarded-for") {
3007 if let Ok(s) = std::str::from_utf8(forwarded) {
3008 if let Some(ip) = s.split(',').next() {
3009 return Some(ip.trim().to_string());
3010 }
3011 }
3012 }
3013 if let Some(real_ip) = req.headers().get("x-real-ip") {
3014 if let Ok(s) = std::str::from_utf8(real_ip) {
3015 return Some(s.trim().to_string());
3016 }
3017 }
3018 None
3019 }
3020}
3021
3022impl Default for TrustedProxyIpKeyExtractor {
3023 fn default() -> Self {
3024 Self::new()
3025 }
3026}
3027
3028impl KeyExtractor for TrustedProxyIpKeyExtractor {
3029 fn extract_key(&self, req: &Request) -> Option<String> {
3030 let remote = req.get_extension::<RemoteAddr>()?;
3031
3032 if self.is_trusted(remote.0) {
3033 self.extract_from_header(req)
3035 .or_else(|| Some(remote.to_string()))
3036 } else {
3037 Some(remote.to_string())
3039 }
3040 }
3041}
3042
3043fn parse_cidr(cidr: &str) -> Option<(std::net::IpAddr, u8)> {
3045 let (ip_str, prefix_str) = cidr.split_once('/')?;
3046 let ip: std::net::IpAddr = ip_str.parse().ok()?;
3047 let prefix: u8 = prefix_str.parse().ok()?;
3048
3049 let max_prefix = match ip {
3051 std::net::IpAddr::V4(_) => 32,
3052 std::net::IpAddr::V6(_) => 128,
3053 };
3054 if prefix > max_prefix {
3055 return None;
3056 }
3057
3058 Some((ip, prefix))
3059}
3060
3061fn ip_in_cidr(ip: std::net::IpAddr, cidr_ip: std::net::IpAddr, prefix: u8) -> bool {
3063 match (ip, cidr_ip) {
3064 (std::net::IpAddr::V4(ip), std::net::IpAddr::V4(cidr)) => {
3065 if prefix == 0 {
3066 return true;
3067 }
3068 let ip_bits = u32::from(ip);
3069 let cidr_bits = u32::from(cidr);
3070 let mask = !0u32 << (32 - prefix);
3071 (ip_bits & mask) == (cidr_bits & mask)
3072 }
3073 (std::net::IpAddr::V6(ip), std::net::IpAddr::V6(cidr)) => {
3074 if prefix == 0 {
3075 return true;
3076 }
3077 let ip_bits = u128::from(ip);
3078 let cidr_bits = u128::from(cidr);
3079 let mask = !0u128 << (128 - prefix);
3080 (ip_bits & mask) == (cidr_bits & mask)
3081 }
3082 _ => false, }
3084}
3085
3086#[derive(Debug, Clone)]
3088pub struct HeaderKeyExtractor {
3089 header_name: String,
3090}
3091
3092impl HeaderKeyExtractor {
3093 #[must_use]
3095 pub fn new(header_name: impl Into<String>) -> Self {
3096 Self {
3097 header_name: header_name.into(),
3098 }
3099 }
3100}
3101
3102impl KeyExtractor for HeaderKeyExtractor {
3103 fn extract_key(&self, req: &Request) -> Option<String> {
3104 req.headers()
3105 .get(&self.header_name)
3106 .and_then(|v| std::str::from_utf8(v).ok())
3107 .map(str::to_string)
3108 }
3109}
3110
3111#[derive(Debug, Clone)]
3113pub struct PathKeyExtractor;
3114
3115impl KeyExtractor for PathKeyExtractor {
3116 fn extract_key(&self, req: &Request) -> Option<String> {
3117 Some(req.path().to_string())
3118 }
3119}
3120
3121pub struct CompositeKeyExtractor {
3126 extractors: Vec<Box<dyn KeyExtractor>>,
3127}
3128
3129impl CompositeKeyExtractor {
3130 #[must_use]
3132 pub fn new(extractors: Vec<Box<dyn KeyExtractor>>) -> Self {
3133 Self { extractors }
3134 }
3135}
3136
3137impl KeyExtractor for CompositeKeyExtractor {
3138 fn extract_key(&self, req: &Request) -> Option<String> {
3139 let parts: Vec<String> = self
3140 .extractors
3141 .iter()
3142 .filter_map(|e| e.extract_key(req))
3143 .collect();
3144 if parts.is_empty() {
3145 None
3146 } else {
3147 Some(parts.join(":"))
3148 }
3149 }
3150}
3151
3152#[derive(Debug, Clone)]
3154struct TokenBucketState {
3155 tokens: f64,
3156 last_refill: Instant,
3157 last_seen: Instant,
3158 stale_after: Duration,
3159}
3160
3161#[derive(Debug, Clone)]
3163struct FixedWindowState {
3164 count: u64,
3165 window_start: Instant,
3166 last_seen: Instant,
3167 stale_after: Duration,
3168}
3169
3170#[derive(Debug, Clone)]
3172struct SlidingWindowState {
3173 current_count: u64,
3174 previous_count: u64,
3175 current_window_start: Instant,
3176 last_seen: Instant,
3177 stale_after: Duration,
3178}
3179
3180pub const DEFAULT_RATE_LIMIT_MAX_KEYS: usize = 65_536;
3182
3183const RATE_LIMIT_SWEEP_INTERVAL: Duration = Duration::from_secs(1);
3184
3185struct BoundedRateLimitMap<T> {
3186 entries: StdHashMap<String, T>,
3187 last_sweep: Instant,
3188}
3189
3190impl<T> BoundedRateLimitMap<T> {
3191 fn new(now: Instant) -> Self {
3192 Self {
3193 entries: StdHashMap::new(),
3194 last_sweep: now,
3195 }
3196 }
3197
3198 fn ensure_key<F, S>(
3199 &mut self,
3200 key: &str,
3201 max_keys: usize,
3202 now: Instant,
3203 mut is_stale: S,
3204 create: F,
3205 ) -> bool
3206 where
3207 F: FnOnce() -> T,
3208 S: FnMut(&T) -> bool,
3209 {
3210 if self.entries.contains_key(key) {
3211 return true;
3212 }
3213
3214 if self.entries.len() >= max_keys
3215 && now.saturating_duration_since(self.last_sweep) >= RATE_LIMIT_SWEEP_INTERVAL
3216 {
3217 self.entries.retain(|_, state| !is_stale(state));
3218 self.last_sweep = now;
3219 }
3220
3221 if self.entries.len() >= max_keys {
3222 return false;
3223 }
3224
3225 self.entries.entry(key.to_string()).or_insert_with(create);
3226 true
3227 }
3228}
3229
3230pub struct InMemoryRateLimitStore {
3239 max_keys: usize,
3240 token_buckets: Mutex<BoundedRateLimitMap<TokenBucketState>>,
3241 fixed_windows: Mutex<BoundedRateLimitMap<FixedWindowState>>,
3242 sliding_windows: Mutex<BoundedRateLimitMap<SlidingWindowState>>,
3243}
3244
3245impl InMemoryRateLimitStore {
3246 #[must_use]
3248 pub fn new() -> Self {
3249 Self::with_max_keys(DEFAULT_RATE_LIMIT_MAX_KEYS)
3250 }
3251
3252 #[must_use]
3256 pub fn with_max_keys(max_keys: usize) -> Self {
3257 let now = Instant::now();
3258 Self {
3259 max_keys,
3260 token_buckets: Mutex::new(BoundedRateLimitMap::new(now)),
3261 fixed_windows: Mutex::new(BoundedRateLimitMap::new(now)),
3262 sliding_windows: Mutex::new(BoundedRateLimitMap::new(now)),
3263 }
3264 }
3265
3266 fn cleanup_window(window: Duration) -> Duration {
3267 if window.is_zero() {
3268 RATE_LIMIT_SWEEP_INTERVAL
3269 } else {
3270 window
3271 }
3272 }
3273
3274 fn saturated_result(max_requests: u64, window: Duration) -> RateLimitResult {
3275 let retry_window = Self::cleanup_window(window);
3276 let reset_after_secs = retry_window
3277 .as_secs()
3278 .saturating_add(u64::from(retry_window.subsec_nanos() > 0))
3279 .max(1);
3280 RateLimitResult {
3281 allowed: false,
3282 limit: max_requests,
3283 remaining: 0,
3284 reset_after_secs,
3285 }
3286 }
3287
3288 #[allow(clippy::cast_precision_loss, clippy::cast_sign_loss)]
3289 fn check_token_bucket(
3290 &self,
3291 key: &str,
3292 max_tokens: u64,
3293 refill_rate: f64,
3294 window: Duration,
3295 now: Instant,
3296 ) -> RateLimitResult {
3297 let mut buckets = self.token_buckets.lock();
3298 let cleanup_window = Self::cleanup_window(window);
3299
3300 if !buckets.ensure_key(
3301 key,
3302 self.max_keys,
3303 now,
3304 |state| now.saturating_duration_since(state.last_seen) >= state.stale_after,
3305 || TokenBucketState {
3306 tokens: max_tokens as f64,
3307 last_refill: now,
3308 last_seen: now,
3309 stale_after: cleanup_window,
3310 },
3311 ) {
3312 return Self::saturated_result(max_tokens, window);
3313 }
3314 let Some(state) = buckets.entries.get_mut(key) else {
3315 return Self::saturated_result(max_tokens, window);
3316 };
3317
3318 let elapsed = now.duration_since(state.last_refill);
3320 let refill = elapsed.as_secs_f64() * refill_rate;
3321 state.tokens = (state.tokens + refill).min(max_tokens as f64);
3322 state.last_refill = now;
3323 state.last_seen = now;
3324 state.stale_after = cleanup_window;
3325
3326 if state.tokens >= 1.0 {
3327 state.tokens -= 1.0;
3328 RateLimitResult {
3329 allowed: true,
3330 limit: max_tokens,
3331 remaining: state.tokens as u64,
3332 reset_after_secs: if state.tokens < max_tokens as f64 {
3333 ((max_tokens as f64 - state.tokens) / refill_rate).ceil() as u64
3334 } else {
3335 window.as_secs()
3336 },
3337 }
3338 } else {
3339 let wait_secs = ((1.0 - state.tokens) / refill_rate).ceil() as u64;
3340 RateLimitResult {
3341 allowed: false,
3342 limit: max_tokens,
3343 remaining: 0,
3344 reset_after_secs: wait_secs,
3345 }
3346 }
3347 }
3348
3349 fn check_fixed_window(
3350 &self,
3351 key: &str,
3352 max_requests: u64,
3353 window: Duration,
3354 now: Instant,
3355 ) -> RateLimitResult {
3356 let mut windows = self.fixed_windows.lock();
3357 let cleanup_window = Self::cleanup_window(window);
3358
3359 if !windows.ensure_key(
3360 key,
3361 self.max_keys,
3362 now,
3363 |state| now.saturating_duration_since(state.last_seen) >= state.stale_after,
3364 || FixedWindowState {
3365 count: 0,
3366 window_start: now,
3367 last_seen: now,
3368 stale_after: cleanup_window,
3369 },
3370 ) {
3371 return Self::saturated_result(max_requests, window);
3372 }
3373 let Some(state) = windows.entries.get_mut(key) else {
3374 return Self::saturated_result(max_requests, window);
3375 };
3376
3377 let elapsed = now.duration_since(state.window_start);
3379 if elapsed >= window {
3380 state.count = 0;
3381 state.window_start = now;
3382 }
3383 state.last_seen = now;
3384 state.stale_after = cleanup_window;
3385
3386 let remaining_time = window
3387 .checked_sub(now.duration_since(state.window_start))
3388 .unwrap_or(Duration::ZERO);
3389
3390 if state.count < max_requests {
3391 state.count += 1;
3392 RateLimitResult {
3393 allowed: true,
3394 limit: max_requests,
3395 remaining: max_requests - state.count,
3396 reset_after_secs: remaining_time.as_secs(),
3397 }
3398 } else {
3399 RateLimitResult {
3400 allowed: false,
3401 limit: max_requests,
3402 remaining: 0,
3403 reset_after_secs: remaining_time.as_secs(),
3404 }
3405 }
3406 }
3407
3408 #[allow(clippy::cast_precision_loss, clippy::cast_sign_loss)]
3409 fn check_sliding_window(
3410 &self,
3411 key: &str,
3412 max_requests: u64,
3413 window: Duration,
3414 now: Instant,
3415 ) -> RateLimitResult {
3416 let mut windows = self.sliding_windows.lock();
3417 let cleanup_window = Self::cleanup_window(window);
3418 let stale_after = cleanup_window.saturating_add(cleanup_window);
3419
3420 if !windows.ensure_key(
3421 key,
3422 self.max_keys,
3423 now,
3424 |state| now.saturating_duration_since(state.last_seen) >= state.stale_after,
3425 || SlidingWindowState {
3426 current_count: 0,
3427 previous_count: 0,
3428 current_window_start: now,
3429 last_seen: now,
3430 stale_after,
3431 },
3432 ) {
3433 return Self::saturated_result(max_requests, window);
3434 }
3435 let Some(state) = windows.entries.get_mut(key) else {
3436 return Self::saturated_result(max_requests, window);
3437 };
3438
3439 let elapsed = now.duration_since(state.current_window_start);
3441 if elapsed >= stale_after {
3442 state.previous_count = 0;
3444 state.current_count = 0;
3445 state.current_window_start = now;
3446 } else if elapsed >= window {
3447 state.previous_count = state.current_count;
3448 state.current_count = 0;
3449 state.current_window_start = state
3450 .current_window_start
3451 .checked_add(window)
3452 .unwrap_or(now);
3453 }
3454 state.last_seen = now;
3455 state.stale_after = stale_after;
3456
3457 let window_elapsed = now.duration_since(state.current_window_start);
3460 let window_fraction = window_elapsed.as_secs_f64() / window.as_secs_f64();
3461 let previous_weight = 1.0 - window_fraction;
3462 let weighted_count =
3463 (state.previous_count as f64 * previous_weight) + state.current_count as f64;
3464
3465 let remaining_time = window.checked_sub(window_elapsed).unwrap_or(Duration::ZERO);
3466
3467 if weighted_count < max_requests as f64 {
3468 state.current_count += 1;
3469 let new_weighted =
3470 (state.previous_count as f64 * previous_weight) + state.current_count as f64;
3471 let remaining = (max_requests as f64 - new_weighted).max(0.0) as u64;
3472 RateLimitResult {
3473 allowed: true,
3474 limit: max_requests,
3475 remaining,
3476 reset_after_secs: remaining_time.as_secs(),
3477 }
3478 } else {
3479 RateLimitResult {
3480 allowed: false,
3481 limit: max_requests,
3482 remaining: 0,
3483 reset_after_secs: remaining_time.as_secs(),
3484 }
3485 }
3486 }
3487
3488 #[allow(clippy::cast_precision_loss)]
3489 fn check_at(
3490 &self,
3491 key: &str,
3492 algorithm: RateLimitAlgorithm,
3493 max_requests: u64,
3494 window: Duration,
3495 now: Instant,
3496 ) -> RateLimitResult {
3497 if window.is_zero() {
3498 return Self::saturated_result(max_requests, window);
3499 }
3500
3501 match algorithm {
3502 RateLimitAlgorithm::TokenBucket => {
3503 let refill_rate = max_requests as f64 / window.as_secs_f64();
3504 self.check_token_bucket(key, max_requests, refill_rate, window, now)
3505 }
3506 RateLimitAlgorithm::FixedWindow => {
3507 self.check_fixed_window(key, max_requests, window, now)
3508 }
3509 RateLimitAlgorithm::SlidingWindow => {
3510 self.check_sliding_window(key, max_requests, window, now)
3511 }
3512 }
3513 }
3514
3515 pub fn check(
3517 &self,
3518 key: &str,
3519 algorithm: RateLimitAlgorithm,
3520 max_requests: u64,
3521 window: Duration,
3522 ) -> RateLimitResult {
3523 self.check_at(key, algorithm, max_requests, window, Instant::now())
3524 }
3525}
3526
3527impl Default for InMemoryRateLimitStore {
3528 fn default() -> Self {
3529 Self::new()
3530 }
3531}
3532
3533#[derive(Clone)]
3568pub struct RateLimitConfig {
3569 pub max_requests: u64,
3571 pub window: Duration,
3573 pub algorithm: RateLimitAlgorithm,
3575 pub max_keys: usize,
3577 pub include_headers: bool,
3579 pub retry_message: String,
3581}
3582
3583impl Default for RateLimitConfig {
3584 fn default() -> Self {
3585 Self {
3586 max_requests: 100,
3587 window: Duration::from_secs(60),
3588 algorithm: RateLimitAlgorithm::TokenBucket,
3589 max_keys: DEFAULT_RATE_LIMIT_MAX_KEYS,
3590 include_headers: true,
3591 retry_message: "Rate limit exceeded. Please retry later.".to_string(),
3592 }
3593 }
3594}
3595
3596pub struct RateLimitBuilder {
3598 config: RateLimitConfig,
3599 key_extractor: Option<Box<dyn KeyExtractor>>,
3600}
3601
3602impl RateLimitBuilder {
3603 #[must_use]
3605 pub fn new() -> Self {
3606 Self {
3607 config: RateLimitConfig::default(),
3608 key_extractor: None,
3609 }
3610 }
3611
3612 #[must_use]
3614 pub fn requests(mut self, max: u64) -> Self {
3615 self.config.max_requests = max;
3616 self
3617 }
3618
3619 #[must_use]
3621 pub fn per(mut self, window: Duration) -> Self {
3622 self.config.window = window;
3623 self
3624 }
3625
3626 #[must_use]
3628 pub fn per_second(self, secs: u64) -> Self {
3629 self.per(Duration::from_secs(secs))
3630 }
3631
3632 #[must_use]
3634 pub fn per_minute(self, minutes: u64) -> Self {
3635 self.per(Duration::from_secs(minutes * 60))
3636 }
3637
3638 #[must_use]
3640 pub fn per_hour(self, hours: u64) -> Self {
3641 self.per(Duration::from_secs(hours * 3600))
3642 }
3643
3644 #[must_use]
3646 pub fn algorithm(mut self, algo: RateLimitAlgorithm) -> Self {
3647 self.config.algorithm = algo;
3648 self
3649 }
3650
3651 #[must_use]
3655 pub fn max_keys(mut self, max_keys: usize) -> Self {
3656 self.config.max_keys = max_keys;
3657 self
3658 }
3659
3660 #[must_use]
3662 pub fn key_extractor(mut self, extractor: impl KeyExtractor + 'static) -> Self {
3663 self.key_extractor = Some(Box::new(extractor));
3664 self
3665 }
3666
3667 #[must_use]
3669 pub fn include_headers(mut self, include: bool) -> Self {
3670 self.config.include_headers = include;
3671 self
3672 }
3673
3674 #[must_use]
3676 pub fn retry_message(mut self, msg: impl Into<String>) -> Self {
3677 self.config.retry_message = msg.into();
3678 self
3679 }
3680
3681 #[must_use]
3683 pub fn build(self) -> RateLimitMiddleware {
3684 let key_extractor = self
3685 .key_extractor
3686 .unwrap_or_else(|| Box::new(IpKeyExtractor));
3687 let store = Arc::new(InMemoryRateLimitStore::with_max_keys(self.config.max_keys));
3688 RateLimitMiddleware {
3689 config: self.config,
3690 store,
3691 key_extractor: Arc::from(key_extractor),
3692 }
3693 }
3694}
3695
3696impl Default for RateLimitBuilder {
3697 fn default() -> Self {
3698 Self::new()
3699 }
3700}
3701
3702#[derive(Debug, Clone)]
3704struct RateLimitInfo {
3705 result: RateLimitResult,
3706}
3707
3708pub struct RateLimitMiddleware {
3731 config: RateLimitConfig,
3732 store: Arc<InMemoryRateLimitStore>,
3733 key_extractor: Arc<dyn KeyExtractor>,
3734}
3735
3736impl RateLimitMiddleware {
3737 #[must_use]
3739 pub fn new() -> Self {
3740 Self::builder().build()
3741 }
3742
3743 #[must_use]
3745 pub fn builder() -> RateLimitBuilder {
3746 RateLimitBuilder::new()
3747 }
3748
3749 fn too_many_requests_body(&self, result: &RateLimitResult) -> Vec<u8> {
3751 format!(
3752 r#"{{"detail":"{}","retry_after_secs":{}}}"#,
3753 self.config.retry_message, result.reset_after_secs
3754 )
3755 .into_bytes()
3756 }
3757
3758 fn add_headers(&self, response: Response, result: &RateLimitResult) -> Response {
3760 response
3761 .header("X-RateLimit-Limit", result.limit.to_string().into_bytes())
3762 .header(
3763 "X-RateLimit-Remaining",
3764 result.remaining.to_string().into_bytes(),
3765 )
3766 .header(
3767 "X-RateLimit-Reset",
3768 result.reset_after_secs.to_string().into_bytes(),
3769 )
3770 }
3771}
3772
3773impl Default for RateLimitMiddleware {
3774 fn default() -> Self {
3775 Self::new()
3776 }
3777}
3778
3779impl Middleware for RateLimitMiddleware {
3780 fn before<'a>(
3781 &'a self,
3782 _ctx: &'a RequestContext,
3783 req: &'a mut Request,
3784 ) -> BoxFuture<'a, ControlFlow> {
3785 Box::pin(async move {
3786 let Some(key) = self.key_extractor.extract_key(req) else {
3788 return ControlFlow::Continue;
3790 };
3791
3792 let result = self.store.check(
3794 &key,
3795 self.config.algorithm,
3796 self.config.max_requests,
3797 self.config.window,
3798 );
3799
3800 if result.allowed {
3801 req.insert_extension(RateLimitInfo { result });
3803 ControlFlow::Continue
3804 } else {
3805 let body = self.too_many_requests_body(&result);
3807 let mut response =
3808 Response::with_status(crate::response::StatusCode::TOO_MANY_REQUESTS)
3809 .header("Content-Type", b"application/json".to_vec())
3810 .header(
3811 "Retry-After",
3812 result.reset_after_secs.to_string().into_bytes(),
3813 )
3814 .body(crate::response::ResponseBody::Bytes(body));
3815
3816 if self.config.include_headers {
3817 response = self.add_headers(response, &result);
3818 }
3819
3820 ControlFlow::Break(response)
3821 }
3822 })
3823 }
3824
3825 fn after<'a>(
3826 &'a self,
3827 _ctx: &'a RequestContext,
3828 req: &'a Request,
3829 response: Response,
3830 ) -> BoxFuture<'a, Response> {
3831 Box::pin(async move {
3832 if !self.config.include_headers {
3833 return response;
3834 }
3835
3836 if let Some(info) = req.get_extension::<RateLimitInfo>() {
3838 self.add_headers(response, &info.result)
3839 } else {
3840 response
3841 }
3842 })
3843 }
3844
3845 fn name(&self) -> &'static str {
3846 "RateLimit"
3847 }
3848}
3849
3850#[derive(Debug, Clone, Copy, PartialEq, Eq)]
3862pub enum InspectionVerbosity {
3863 Minimal,
3867
3868 Normal,
3872
3873 Verbose,
3878}
3879
3880pub struct RequestInspectionMiddleware {
3919 log_config: LogConfig,
3920 verbosity: InspectionVerbosity,
3921 redact_headers: HashSet<String>,
3922 slow_threshold_ms: u64,
3923 max_body_preview: usize,
3924}
3925
3926impl Default for RequestInspectionMiddleware {
3927 fn default() -> Self {
3928 Self {
3929 log_config: LogConfig::development(),
3930 verbosity: InspectionVerbosity::Normal,
3931 redact_headers: default_redacted_headers(),
3932 slow_threshold_ms: 1000,
3933 max_body_preview: 2048,
3934 }
3935 }
3936}
3937
3938impl RequestInspectionMiddleware {
3939 #[must_use]
3941 pub fn new() -> Self {
3942 Self::default()
3943 }
3944
3945 #[must_use]
3947 pub fn log_config(mut self, config: LogConfig) -> Self {
3948 self.log_config = config;
3949 self
3950 }
3951
3952 #[must_use]
3954 pub fn verbosity(mut self, level: InspectionVerbosity) -> Self {
3955 self.verbosity = level;
3956 self
3957 }
3958
3959 #[must_use]
3961 pub fn slow_threshold_ms(mut self, ms: u64) -> Self {
3962 self.slow_threshold_ms = ms;
3963 self
3964 }
3965
3966 #[must_use]
3968 pub fn max_body_preview(mut self, max: usize) -> Self {
3969 self.max_body_preview = max;
3970 self
3971 }
3972
3973 #[must_use]
3975 pub fn redact_header(mut self, name: impl Into<String>) -> Self {
3976 self.redact_headers.insert(name.into().to_ascii_lowercase());
3977 self
3978 }
3979
3980 fn format_body_preview(&self, bytes: &[u8], content_type: Option<&[u8]>) -> Option<String> {
3982 if bytes.is_empty() || self.max_body_preview == 0 {
3983 return None;
3984 }
3985
3986 let is_json = content_type
3987 .and_then(|ct| std::str::from_utf8(ct).ok())
3988 .is_some_and(|ct| ct.contains("application/json"));
3989
3990 let limit = self.max_body_preview.min(bytes.len());
3991 let truncated = bytes.len() > self.max_body_preview;
3992
3993 match std::str::from_utf8(&bytes[..limit]) {
3994 Ok(text) => {
3995 if is_json {
3996 if let Some(pretty) = try_pretty_json(text) {
3998 let mut output = pretty;
3999 if truncated {
4000 output.push_str("\n ... (truncated)");
4001 }
4002 return Some(output);
4003 }
4004 }
4005 let mut output = text.to_string();
4006 if truncated {
4007 output.push_str("...");
4008 }
4009 Some(output)
4010 }
4011 Err(_) => Some(format!("<{} bytes binary>", bytes.len())),
4012 }
4013 }
4014
4015 fn format_response_preview(
4017 &self,
4018 body: &crate::response::ResponseBody,
4019 content_type: Option<&[u8]>,
4020 ) -> Option<String> {
4021 match body {
4022 crate::response::ResponseBody::Empty => None,
4023 crate::response::ResponseBody::Bytes(bytes) => {
4024 self.format_body_preview(bytes, content_type)
4025 }
4026 crate::response::ResponseBody::Stream(_) => Some("<streaming body>".to_string()),
4027 }
4028 }
4029
4030 fn format_inspection_headers<'a>(
4032 &self,
4033 headers: impl Iterator<Item = (&'a str, &'a [u8])>,
4034 ) -> String {
4035 let mut out = String::new();
4036 for (name, value) in headers {
4037 out.push_str("\n ");
4038 out.push_str(name);
4039 out.push_str(": ");
4040
4041 let lowered = name.to_ascii_lowercase();
4042 if self.redact_headers.contains(&lowered) {
4043 out.push_str("[REDACTED]");
4044 } else {
4045 match std::str::from_utf8(value) {
4046 Ok(text) => out.push_str(text),
4047 Err(_) => out.push_str("<binary>"),
4048 }
4049 }
4050 }
4051 out
4052 }
4053
4054 fn format_response_inspection_headers(&self, headers: &[(String, Vec<u8>)]) -> String {
4056 self.format_inspection_headers(
4057 headers
4058 .iter()
4059 .map(|(name, value)| (name.as_str(), value.as_slice())),
4060 )
4061 }
4062}
4063
4064#[derive(Debug, Clone)]
4066struct InspectionStart(Instant);
4067
4068impl Middleware for RequestInspectionMiddleware {
4069 fn before<'a>(
4070 &'a self,
4071 ctx: &'a RequestContext,
4072 req: &'a mut Request,
4073 ) -> BoxFuture<'a, ControlFlow> {
4074 let logger = RequestLogger::new(ctx, self.log_config.clone());
4075 req.insert_extension(InspectionStart(Instant::now()));
4076
4077 let method = req.method();
4078 let path = req.path();
4079 let query = req.query();
4080
4081 let mut request_line = format!("--> {method} {path}");
4083 if let Some(q) = query {
4084 request_line.push('?');
4085 request_line.push_str(q);
4086 }
4087
4088 let body_size = body_len(req.body());
4089 if body_size > 0 {
4090 request_line.push_str(&format!(" ({body_size} bytes)"));
4091 }
4092
4093 match self.verbosity {
4094 InspectionVerbosity::Minimal => {
4095 logger.info(request_line);
4096 }
4097 InspectionVerbosity::Normal => {
4098 let headers = self.format_inspection_headers(req.headers().iter());
4099 logger.info(format!("{request_line}{headers}"));
4100 }
4101 InspectionVerbosity::Verbose => {
4102 let headers = self.format_inspection_headers(req.headers().iter());
4103 let content_type = req.headers().get("content-type");
4104 let body_preview = match req.body() {
4105 Body::Empty => None,
4106 Body::Bytes(bytes) => self.format_body_preview(bytes, content_type),
4107 Body::Stream { .. } => None,
4108 };
4109
4110 let mut output = format!("{request_line}{headers}");
4111 if let Some(body) = body_preview {
4112 output.push_str("\n ");
4113 output.push_str(&body.replace('\n', "\n "));
4115 }
4116 logger.info(output);
4117 }
4118 }
4119
4120 Box::pin(async { ControlFlow::Continue })
4121 }
4122
4123 fn after<'a>(
4124 &'a self,
4125 ctx: &'a RequestContext,
4126 req: &'a Request,
4127 response: Response,
4128 ) -> BoxFuture<'a, Response> {
4129 let logger = RequestLogger::new(ctx, self.log_config.clone());
4130 let duration = req
4131 .get_extension::<InspectionStart>()
4132 .map(|start| start.0.elapsed())
4133 .unwrap_or_default();
4134
4135 let status = response.status();
4136 let duration_ms = duration.as_millis();
4137
4138 let mut response_line = format!(
4140 "<-- {} {} ({duration_ms}ms)",
4141 status.as_u16(),
4142 status.canonical_reason(),
4143 );
4144
4145 if duration_ms >= u128::from(self.slow_threshold_ms) {
4147 response_line.push_str(" [SLOW]");
4148 }
4149
4150 match self.verbosity {
4151 InspectionVerbosity::Minimal => {
4152 if duration_ms >= u128::from(self.slow_threshold_ms) {
4153 logger.warn(response_line);
4154 } else {
4155 logger.info(response_line);
4156 }
4157 }
4158 InspectionVerbosity::Normal => {
4159 let headers = self.format_response_inspection_headers(response.headers());
4160 let output = format!("{response_line}{headers}");
4161 if duration_ms >= u128::from(self.slow_threshold_ms) {
4162 logger.warn(output);
4163 } else {
4164 logger.info(output);
4165 }
4166 }
4167 InspectionVerbosity::Verbose => {
4168 let headers = self.format_response_inspection_headers(response.headers());
4169
4170 let resp_content_type: Option<&[u8]> = response
4172 .headers()
4173 .iter()
4174 .find(|(name, _)| name.eq_ignore_ascii_case("content-type"))
4175 .map(|(_, value)| value.as_slice());
4176
4177 let body_preview =
4178 self.format_response_preview(response.body_ref(), resp_content_type);
4179
4180 let mut output = format!("{response_line}{headers}");
4181 if let Some(body) = body_preview {
4182 output.push_str("\n ");
4183 output.push_str(&body.replace('\n', "\n "));
4184 }
4185
4186 if duration_ms >= u128::from(self.slow_threshold_ms) {
4187 logger.warn(output);
4188 } else {
4189 logger.info(output);
4190 }
4191 }
4192 }
4193
4194 Box::pin(async move { response })
4195 }
4196
4197 fn name(&self) -> &'static str {
4198 "RequestInspection"
4199 }
4200}
4201
4202fn try_pretty_json(input: &str) -> Option<String> {
4207 let trimmed = input.trim();
4208 if !trimmed.starts_with('{') && !trimmed.starts_with('[') {
4209 return None;
4210 }
4211
4212 let mut output = String::with_capacity(trimmed.len() * 2);
4214 if json_pretty_format(trimmed, &mut output).is_ok() {
4215 Some(output)
4216 } else {
4217 None
4218 }
4219}
4220
4221fn json_pretty_format(input: &str, output: &mut String) -> Result<(), ()> {
4226 let bytes = input.as_bytes();
4227 let mut pos = 0;
4228 let mut indent: usize = 0;
4229 let mut in_string = false;
4230 let mut escape_next = false;
4231
4232 while pos < bytes.len() {
4233 let ch = bytes[pos] as char;
4234
4235 if escape_next {
4236 output.push(ch);
4237 escape_next = false;
4238 pos += 1;
4239 continue;
4240 }
4241
4242 if in_string {
4243 output.push(ch);
4244 if ch == '\\' {
4245 escape_next = true;
4246 } else if ch == '"' {
4247 in_string = false;
4248 }
4249 pos += 1;
4250 continue;
4251 }
4252
4253 match ch {
4254 '"' => {
4255 in_string = true;
4256 output.push('"');
4257 }
4258 '{' | '[' => {
4259 output.push(ch);
4260 let peek = skip_whitespace(bytes, pos + 1);
4262 let closing = if ch == '{' { '}' } else { ']' };
4263 if peek < bytes.len() && bytes[peek] as char == closing {
4264 output.push(closing);
4265 pos = peek + 1;
4266 continue;
4267 }
4268 indent += 1;
4269 output.push('\n');
4270 push_indent(output, indent);
4271 }
4272 '}' | ']' => {
4273 indent = indent.saturating_sub(1);
4274 output.push('\n');
4275 push_indent(output, indent);
4276 output.push(ch);
4277 }
4278 ':' => {
4279 output.push_str(": ");
4280 }
4281 ',' => {
4282 output.push(',');
4283 output.push('\n');
4284 push_indent(output, indent);
4285 }
4286 c if c.is_ascii_whitespace() => {
4287 }
4289 _ => {
4290 output.push(ch);
4291 }
4292 }
4293
4294 pos += 1;
4295 }
4296
4297 if in_string || indent != 0 {
4298 return Err(());
4299 }
4300
4301 Ok(())
4302}
4303
4304fn skip_whitespace(bytes: &[u8], start: usize) -> usize {
4305 let mut i = start;
4306 while i < bytes.len() && (bytes[i] as char).is_ascii_whitespace() {
4307 i += 1;
4308 }
4309 i
4310}
4311
4312fn push_indent(output: &mut String, level: usize) {
4313 for _ in 0..level {
4314 output.push_str(" ");
4315 }
4316}
4317
4318#[derive(Debug, Clone, Copy, PartialEq, Eq)]
4328pub enum ETagMode {
4329 Auto,
4332 Manual,
4335 Disabled,
4337}
4338
4339impl Default for ETagMode {
4340 fn default() -> Self {
4341 Self::Auto
4342 }
4343}
4344
4345#[derive(Debug, Clone)]
4347pub struct ETagConfig {
4348 pub mode: ETagMode,
4350 pub weak: bool,
4353 pub min_size: usize,
4356}
4357
4358impl Default for ETagConfig {
4359 fn default() -> Self {
4360 Self {
4361 mode: ETagMode::Auto,
4362 weak: false,
4363 min_size: 0,
4364 }
4365 }
4366}
4367
4368impl ETagConfig {
4369 #[must_use]
4371 pub fn new() -> Self {
4372 Self::default()
4373 }
4374
4375 #[must_use]
4377 pub fn mode(mut self, mode: ETagMode) -> Self {
4378 self.mode = mode;
4379 self
4380 }
4381
4382 #[must_use]
4384 pub fn weak(mut self, weak: bool) -> Self {
4385 self.weak = weak;
4386 self
4387 }
4388
4389 #[must_use]
4391 pub fn min_size(mut self, size: usize) -> Self {
4392 self.min_size = size;
4393 self
4394 }
4395}
4396
4397pub struct ETagMiddleware {
4438 config: ETagConfig,
4439}
4440
4441impl Default for ETagMiddleware {
4442 fn default() -> Self {
4443 Self::new()
4444 }
4445}
4446
4447impl ETagMiddleware {
4448 #[must_use]
4450 pub fn new() -> Self {
4451 Self {
4452 config: ETagConfig::default(),
4453 }
4454 }
4455
4456 #[must_use]
4458 pub fn with_config(config: ETagConfig) -> Self {
4459 Self { config }
4460 }
4461
4462 fn generate_etag(data: &[u8], weak: bool) -> String {
4469 const FNV_OFFSET_BASIS: u64 = 0xcbf29ce484222325;
4471 const FNV_PRIME: u64 = 0x100000001b3;
4472
4473 let mut hash = FNV_OFFSET_BASIS;
4474 for &byte in data {
4475 hash ^= u64::from(byte);
4476 hash = hash.wrapping_mul(FNV_PRIME);
4477 }
4478
4479 if weak {
4481 format!("W/\"{:016x}\"", hash)
4482 } else {
4483 format!("\"{:016x}\"", hash)
4484 }
4485 }
4486
4487 fn parse_if_none_match(value: &str) -> Vec<String> {
4495 let trimmed = value.trim();
4496
4497 if trimmed == "*" {
4499 return vec!["*".to_string()];
4500 }
4501
4502 let mut etags = Vec::new();
4503 let mut current = String::new();
4504 let mut in_quote = false;
4505 let mut prev_char = '\0';
4506
4507 for ch in trimmed.chars() {
4508 match ch {
4509 '"' if prev_char != '\\' => {
4510 current.push(ch);
4511 if in_quote {
4512 let etag = current.trim().to_string();
4514 if !etag.is_empty() {
4515 etags.push(etag);
4516 }
4517 current.clear();
4518 }
4519 in_quote = !in_quote;
4520 }
4521 ',' if !in_quote => {
4522 current.clear();
4524 }
4525 _ => {
4526 current.push(ch);
4527 }
4528 }
4529 prev_char = ch;
4530 }
4531
4532 etags
4533 }
4534
4535 fn etags_match_weak(etag1: &str, etag2: &str) -> bool {
4543 let e1 = Self::strip_weak_prefix(etag1);
4545 let e2 = Self::strip_weak_prefix(etag2);
4546 e1 == e2
4547 }
4548
4549 fn strip_weak_prefix(s: &str) -> &str {
4551 if s.starts_with("W/") || s.starts_with("w/") {
4552 &s[2..]
4553 } else {
4554 s
4555 }
4556 }
4557
4558 fn is_cacheable_method(method: crate::request::Method) -> bool {
4560 matches!(
4561 method,
4562 crate::request::Method::Get | crate::request::Method::Head
4563 )
4564 }
4565
4566 fn get_existing_etag(headers: &[(String, Vec<u8>)]) -> Option<String> {
4568 for (name, value) in headers {
4569 if name.eq_ignore_ascii_case("etag") {
4570 return std::str::from_utf8(value).ok().map(String::from);
4571 }
4572 }
4573 None
4574 }
4575}
4576
4577impl Middleware for ETagMiddleware {
4578 fn after<'a>(
4579 &'a self,
4580 _ctx: &'a RequestContext,
4581 req: &'a Request,
4582 response: Response,
4583 ) -> BoxFuture<'a, Response> {
4584 let config = self.config.clone();
4585
4586 Box::pin(async move {
4587 if config.mode == ETagMode::Disabled {
4589 return response;
4590 }
4591
4592 if !Self::is_cacheable_method(req.method()) {
4594 return response;
4595 }
4596
4597 let (status, headers, body) = response.into_parts();
4599
4600 let existing_etag = Self::get_existing_etag(&headers);
4602
4603 let body_bytes = match &body {
4605 crate::response::ResponseBody::Bytes(bytes) => Some(bytes.clone()),
4606 crate::response::ResponseBody::Empty => Some(Vec::new()),
4607 crate::response::ResponseBody::Stream(_) => None,
4608 };
4609
4610 let etag = if let Some(existing) = existing_etag {
4612 Some(existing)
4613 } else if config.mode == ETagMode::Auto {
4614 if let Some(ref bytes) = body_bytes {
4615 if bytes.len() >= config.min_size {
4616 Some(Self::generate_etag(bytes, config.weak))
4617 } else {
4618 None
4619 }
4620 } else {
4621 None
4622 }
4623 } else {
4624 None
4625 };
4626
4627 if let Some(ref etag_value) = etag {
4629 if let Some(if_none_match) = req.headers().get("if-none-match") {
4630 if let Ok(value) = std::str::from_utf8(if_none_match) {
4631 let client_etags = Self::parse_if_none_match(value);
4632
4633 let matches = client_etags.iter().any(|client_etag| {
4635 client_etag == "*" || Self::etags_match_weak(client_etag, etag_value)
4636 });
4637
4638 if matches {
4639 return Response::with_status(
4641 crate::response::StatusCode::NOT_MODIFIED,
4642 )
4643 .header("etag", etag_value.as_bytes().to_vec());
4644 }
4645 }
4646 }
4647 }
4648
4649 let mut new_response = Response::with_status(status)
4651 .body(body)
4652 .rebuild_with_headers(headers);
4653
4654 if let Some(etag_value) = etag {
4655 new_response = new_response.header("etag", etag_value.into_bytes());
4656 }
4657
4658 new_response
4659 })
4660 }
4661
4662 fn name(&self) -> &'static str {
4663 "ETagMiddleware"
4664 }
4665}
4666
4667#[derive(Debug, Clone, Copy, PartialEq, Eq)]
4676pub enum CacheDirective {
4677 Public,
4679 Private,
4681 NoStore,
4683 NoCache,
4685 NoTransform,
4687 MustRevalidate,
4689 ProxyRevalidate,
4691 StaleIfError,
4693 StaleWhileRevalidate,
4695 SMaxAge,
4697 OnlyIfCached,
4699 Immutable,
4701}
4702
4703impl CacheDirective {
4704 fn as_str(self) -> &'static str {
4706 match self {
4707 Self::Public => "public",
4708 Self::Private => "private",
4709 Self::NoStore => "no-store",
4710 Self::NoCache => "no-cache",
4711 Self::NoTransform => "no-transform",
4712 Self::MustRevalidate => "must-revalidate",
4713 Self::ProxyRevalidate => "proxy-revalidate",
4714 Self::StaleIfError => "stale-if-error",
4715 Self::StaleWhileRevalidate => "stale-while-revalidate",
4716 Self::SMaxAge => "s-maxage",
4717 Self::OnlyIfCached => "only-if-cached",
4718 Self::Immutable => "immutable",
4719 }
4720 }
4721}
4722
4723#[derive(Debug, Clone, Default)]
4753pub struct CacheControlBuilder {
4754 directives: Vec<CacheDirective>,
4755 max_age: Option<u32>,
4756 s_maxage: Option<u32>,
4757 stale_while_revalidate: Option<u32>,
4758 stale_if_error: Option<u32>,
4759}
4760
4761impl CacheControlBuilder {
4762 #[must_use]
4764 pub fn new() -> Self {
4765 Self::default()
4766 }
4767
4768 #[must_use]
4770 pub fn public(mut self) -> Self {
4771 self.directives.push(CacheDirective::Public);
4772 self
4773 }
4774
4775 #[must_use]
4777 pub fn private(mut self) -> Self {
4778 self.directives.push(CacheDirective::Private);
4779 self
4780 }
4781
4782 #[must_use]
4784 pub fn no_store(mut self) -> Self {
4785 self.directives.push(CacheDirective::NoStore);
4786 self
4787 }
4788
4789 #[must_use]
4791 pub fn no_cache(mut self) -> Self {
4792 self.directives.push(CacheDirective::NoCache);
4793 self
4794 }
4795
4796 #[must_use]
4798 pub fn no_transform(mut self) -> Self {
4799 self.directives.push(CacheDirective::NoTransform);
4800 self
4801 }
4802
4803 #[must_use]
4805 pub fn must_revalidate(mut self) -> Self {
4806 self.directives.push(CacheDirective::MustRevalidate);
4807 self
4808 }
4809
4810 #[must_use]
4812 pub fn proxy_revalidate(mut self) -> Self {
4813 self.directives.push(CacheDirective::ProxyRevalidate);
4814 self
4815 }
4816
4817 #[must_use]
4819 pub fn immutable(mut self) -> Self {
4820 self.directives.push(CacheDirective::Immutable);
4821 self
4822 }
4823
4824 #[must_use]
4826 pub fn max_age_secs(mut self, seconds: u32) -> Self {
4827 self.max_age = Some(seconds);
4828 self
4829 }
4830
4831 #[must_use]
4833 pub fn max_age(self, duration: std::time::Duration) -> Self {
4834 self.max_age_secs(duration.as_secs() as u32)
4835 }
4836
4837 #[must_use]
4839 pub fn s_maxage_secs(mut self, seconds: u32) -> Self {
4840 self.s_maxage = Some(seconds);
4841 self
4842 }
4843
4844 #[must_use]
4846 pub fn s_maxage(self, duration: std::time::Duration) -> Self {
4847 self.s_maxage_secs(duration.as_secs() as u32)
4848 }
4849
4850 #[must_use]
4852 pub fn stale_while_revalidate_secs(mut self, seconds: u32) -> Self {
4853 self.stale_while_revalidate = Some(seconds);
4854 self
4855 }
4856
4857 #[must_use]
4859 pub fn stale_if_error_secs(mut self, seconds: u32) -> Self {
4860 self.stale_if_error = Some(seconds);
4861 self
4862 }
4863
4864 #[must_use]
4866 pub fn build(&self) -> String {
4867 let mut parts = Vec::new();
4868
4869 for directive in &self.directives {
4871 parts.push(directive.as_str().to_string());
4872 }
4873
4874 if let Some(age) = self.max_age {
4876 parts.push(format!("max-age={age}"));
4877 }
4878
4879 if let Some(age) = self.s_maxage {
4881 parts.push(format!("s-maxage={age}"));
4882 }
4883
4884 if let Some(seconds) = self.stale_while_revalidate {
4886 parts.push(format!("stale-while-revalidate={seconds}"));
4887 }
4888
4889 if let Some(seconds) = self.stale_if_error {
4891 parts.push(format!("stale-if-error={seconds}"));
4892 }
4893
4894 parts.join(", ")
4895 }
4896
4897 #[must_use]
4899 pub fn is_no_cache(&self) -> bool {
4900 self.directives.contains(&CacheDirective::NoStore)
4901 || self.directives.contains(&CacheDirective::NoCache)
4902 }
4903}
4904
4905#[derive(Debug, Clone, Copy, PartialEq, Eq)]
4907pub enum CachePreset {
4908 NoCache,
4910 PrivateNoCache,
4912 PublicOneHour,
4914 Immutable,
4916 CdnFriendly,
4918 StaticAssets,
4920}
4921
4922impl CachePreset {
4923 #[must_use]
4925 pub fn to_header_value(&self) -> String {
4926 match self {
4927 Self::NoCache => "no-store, no-cache, must-revalidate".to_string(),
4928 Self::PrivateNoCache => "private, max-age=0, must-revalidate".to_string(),
4929 Self::PublicOneHour => "public, max-age=3600".to_string(),
4930 Self::Immutable => "public, max-age=31536000, immutable".to_string(),
4931 Self::CdnFriendly => "public, max-age=60, s-maxage=3600".to_string(),
4932 Self::StaticAssets => "public, max-age=86400".to_string(),
4933 }
4934 }
4935
4936 #[must_use]
4938 pub fn to_builder(&self) -> CacheControlBuilder {
4939 match self {
4940 Self::NoCache => CacheControlBuilder::new()
4941 .no_store()
4942 .no_cache()
4943 .must_revalidate(),
4944 Self::PrivateNoCache => CacheControlBuilder::new()
4945 .private()
4946 .max_age_secs(0)
4947 .must_revalidate(),
4948 Self::PublicOneHour => CacheControlBuilder::new().public().max_age_secs(3600),
4949 Self::Immutable => CacheControlBuilder::new()
4950 .public()
4951 .max_age_secs(31536000)
4952 .immutable(),
4953 Self::CdnFriendly => CacheControlBuilder::new()
4954 .public()
4955 .max_age_secs(60)
4956 .s_maxage_secs(3600),
4957 Self::StaticAssets => CacheControlBuilder::new().public().max_age_secs(86400),
4958 }
4959 }
4960}
4961
4962#[derive(Debug, Clone)]
4964pub struct CacheControlConfig {
4965 pub cache_control: String,
4967 pub vary: Vec<String>,
4969 pub set_expires: bool,
4971 pub preserve_existing: bool,
4973 pub methods: Vec<crate::request::Method>,
4975 pub path_patterns: Vec<String>,
4977 pub cacheable_statuses: Vec<u16>,
4979}
4980
4981impl Default for CacheControlConfig {
4982 fn default() -> Self {
4983 Self {
4984 cache_control: CachePreset::NoCache.to_header_value(),
4985 vary: Vec::new(),
4986 set_expires: false,
4987 preserve_existing: true,
4988 methods: vec![crate::request::Method::Get, crate::request::Method::Head],
4989 path_patterns: Vec::new(),
4990 cacheable_statuses: (200..300).collect(),
4991 }
4992 }
4993}
4994
4995impl CacheControlConfig {
4996 #[must_use]
4998 pub fn new() -> Self {
4999 Self::default()
5000 }
5001
5002 #[must_use]
5004 pub fn from_preset(preset: CachePreset) -> Self {
5005 Self {
5006 cache_control: preset.to_header_value(),
5007 ..Self::default()
5008 }
5009 }
5010
5011 #[must_use]
5013 pub fn from_builder(builder: CacheControlBuilder) -> Self {
5014 Self {
5015 cache_control: builder.build(),
5016 ..Self::default()
5017 }
5018 }
5019
5020 #[must_use]
5022 pub fn cache_control(mut self, value: impl Into<String>) -> Self {
5023 self.cache_control = value.into();
5024 self
5025 }
5026
5027 #[must_use]
5029 pub fn vary(mut self, header: impl Into<String>) -> Self {
5030 self.vary.push(header.into());
5031 self
5032 }
5033
5034 #[must_use]
5036 pub fn vary_headers(mut self, headers: Vec<String>) -> Self {
5037 self.vary.extend(headers);
5038 self
5039 }
5040
5041 #[must_use]
5043 pub fn with_expires(mut self, enable: bool) -> Self {
5044 self.set_expires = enable;
5045 self
5046 }
5047
5048 #[must_use]
5050 pub fn preserve_existing(mut self, preserve: bool) -> Self {
5051 self.preserve_existing = preserve;
5052 self
5053 }
5054
5055 #[must_use]
5057 pub fn methods(mut self, methods: Vec<crate::request::Method>) -> Self {
5058 self.methods = methods;
5059 self
5060 }
5061
5062 #[must_use]
5064 pub fn path_patterns(mut self, patterns: Vec<String>) -> Self {
5065 self.path_patterns = patterns;
5066 self
5067 }
5068
5069 #[must_use]
5071 pub fn cacheable_statuses(mut self, statuses: Vec<u16>) -> Self {
5072 self.cacheable_statuses = statuses;
5073 self
5074 }
5075}
5076
5077pub struct CacheControlMiddleware {
5124 config: CacheControlConfig,
5125}
5126
5127impl Default for CacheControlMiddleware {
5128 fn default() -> Self {
5129 Self::new()
5130 }
5131}
5132
5133impl CacheControlMiddleware {
5134 #[must_use]
5138 pub fn new() -> Self {
5139 Self {
5140 config: CacheControlConfig::default(),
5141 }
5142 }
5143
5144 #[must_use]
5146 pub fn with_preset(preset: CachePreset) -> Self {
5147 Self {
5148 config: CacheControlConfig::from_preset(preset),
5149 }
5150 }
5151
5152 #[must_use]
5154 pub fn with_config(config: CacheControlConfig) -> Self {
5155 Self { config }
5156 }
5157
5158 fn is_cacheable_method(&self, method: crate::request::Method) -> bool {
5160 self.config.methods.contains(&method)
5161 }
5162
5163 fn is_cacheable_status(&self, status: u16) -> bool {
5165 self.config.cacheable_statuses.contains(&status)
5166 }
5167
5168 fn matches_path(&self, path: &str) -> bool {
5170 if self.config.path_patterns.is_empty() {
5171 return true; }
5173
5174 for pattern in &self.config.path_patterns {
5175 if path_matches_pattern(path, pattern) {
5176 return true;
5177 }
5178 }
5179 false
5180 }
5181
5182 fn has_cache_control(headers: &[(String, Vec<u8>)]) -> bool {
5184 headers
5185 .iter()
5186 .any(|(name, _)| name.eq_ignore_ascii_case("cache-control"))
5187 }
5188
5189 fn calculate_expires(cache_control: &str) -> Option<String> {
5191 for directive in cache_control.split(',') {
5193 let directive = directive.trim();
5194 if directive.starts_with("max-age=") {
5195 if let Ok(seconds) = directive[8..].parse::<u64>() {
5196 let now = std::time::SystemTime::now();
5198 if let Some(expires) = now.checked_add(std::time::Duration::from_secs(seconds))
5199 {
5200 return Some(format_http_date(expires));
5201 }
5202 }
5203 }
5204 }
5205 None
5206 }
5207}
5208
5209fn path_matches_pattern(path: &str, pattern: &str) -> bool {
5211 if pattern == "*" {
5212 return true;
5213 }
5214
5215 if pattern.contains('*') {
5216 let parts: Vec<&str> = pattern.split('*').collect();
5218 if parts.len() == 2 {
5219 let (prefix, suffix) = (parts[0], parts[1]);
5220 return path.starts_with(prefix) && path.ends_with(suffix);
5221 }
5222 let fixed_parts: Vec<&str> = pattern.split('*').filter(|s| !s.is_empty()).collect();
5224 let mut remaining = path;
5225 for part in fixed_parts {
5226 if let Some(pos) = remaining.find(part) {
5227 remaining = &remaining[pos + part.len()..];
5228 } else {
5229 return false;
5230 }
5231 }
5232 true
5233 } else {
5234 path == pattern
5235 }
5236}
5237
5238fn format_http_date(time: std::time::SystemTime) -> String {
5240 match time.duration_since(std::time::UNIX_EPOCH) {
5242 Ok(duration) => {
5243 let secs = duration.as_secs();
5245 let days = secs / 86400;
5247 let remaining_secs = secs % 86400;
5248 let hours = remaining_secs / 3600;
5249 let minutes = (remaining_secs % 3600) / 60;
5250 let seconds = remaining_secs % 60;
5251
5252 let day_of_week = ((days + 4) % 7) as usize;
5254 let day_names = ["Sun", "Mon", "Tue", "Wed", "Thu", "Fri", "Sat"];
5255
5256 let (year, month, day) = days_to_date(days);
5258 let month_names = [
5259 "Jan", "Feb", "Mar", "Apr", "May", "Jun", "Jul", "Aug", "Sep", "Oct", "Nov", "Dec",
5260 ];
5261
5262 format!(
5263 "{}, {:02} {} {} {:02}:{:02}:{:02} GMT",
5264 day_names[day_of_week],
5265 day,
5266 month_names[(month - 1) as usize],
5267 year,
5268 hours,
5269 minutes,
5270 seconds
5271 )
5272 }
5273 Err(_) => "Thu, 01 Jan 1970 00:00:00 GMT".to_string(),
5274 }
5275}
5276
5277fn days_to_date(days: u64) -> (u64, u64, u64) {
5279 let mut remaining_days = days;
5281 let mut year = 1970u64;
5282
5283 loop {
5284 let days_in_year = if is_leap_year(year) { 366 } else { 365 };
5285 if remaining_days < days_in_year {
5286 break;
5287 }
5288 remaining_days -= days_in_year;
5289 year += 1;
5290 }
5291
5292 let leap = is_leap_year(year);
5293 let month_days: [u64; 12] = if leap {
5294 [31, 29, 31, 30, 31, 30, 31, 31, 30, 31, 30, 31]
5295 } else {
5296 [31, 28, 31, 30, 31, 30, 31, 31, 30, 31, 30, 31]
5297 };
5298
5299 let mut month = 1u64;
5300 for &days_in_month in &month_days {
5301 if remaining_days < days_in_month {
5302 break;
5303 }
5304 remaining_days -= days_in_month;
5305 month += 1;
5306 }
5307
5308 (year, month, remaining_days + 1)
5309}
5310
5311fn is_leap_year(year: u64) -> bool {
5313 (year.is_multiple_of(4) && !year.is_multiple_of(100)) || year.is_multiple_of(400)
5314}
5315
5316impl Middleware for CacheControlMiddleware {
5317 fn after<'a>(
5318 &'a self,
5319 _ctx: &'a RequestContext,
5320 req: &'a Request,
5321 response: Response,
5322 ) -> BoxFuture<'a, Response> {
5323 let config = self.config.clone();
5324
5325 Box::pin(async move {
5326 if !self.is_cacheable_method(req.method()) {
5328 return response;
5329 }
5330
5331 if !self.is_cacheable_status(response.status().as_u16()) {
5332 return response;
5333 }
5334
5335 if !self.matches_path(req.path()) {
5336 return response;
5337 }
5338
5339 let (status, mut headers, body) = response.into_parts();
5341
5342 if config.preserve_existing && Self::has_cache_control(&headers) {
5344 let mut resp = Response::with_status(status);
5346 for (name, value) in headers {
5347 resp = resp.header(name, value);
5348 }
5349 return resp.body(body);
5350 }
5351
5352 headers.push((
5354 "Cache-Control".to_string(),
5355 config.cache_control.as_bytes().to_vec(),
5356 ));
5357
5358 if !config.vary.is_empty() {
5360 let vary_value = config.vary.join(", ");
5361 headers.push(("Vary".to_string(), vary_value.into_bytes()));
5362 }
5363
5364 if config.set_expires {
5366 if let Some(expires) = Self::calculate_expires(&config.cache_control) {
5367 headers.push(("Expires".to_string(), expires.into_bytes()));
5368 }
5369 }
5370
5371 let mut resp = Response::with_status(status);
5373 for (name, value) in headers {
5374 resp = resp.header(name, value);
5375 }
5376 resp.body(body)
5377 })
5378 }
5379
5380 fn name(&self) -> &'static str {
5381 "CacheControlMiddleware"
5382 }
5383}
5384
5385#[derive(Debug, Clone)]
5420pub struct TraceRejectionMiddleware {
5421 log_attempts: bool,
5423}
5424
5425impl Default for TraceRejectionMiddleware {
5426 fn default() -> Self {
5427 Self::new()
5428 }
5429}
5430
5431impl TraceRejectionMiddleware {
5432 #[must_use]
5436 pub fn new() -> Self {
5437 Self { log_attempts: true }
5438 }
5439
5440 #[must_use]
5445 pub fn log_attempts(mut self, log: bool) -> Self {
5446 self.log_attempts = log;
5447 self
5448 }
5449
5450 fn rejection_response(path: &str) -> Response {
5452 let body = format!(
5453 r#"{{"detail":"HTTP TRACE method is not allowed","path":"{}"}}"#,
5454 path.replace('"', "\\\"")
5455 );
5456 Response::with_status(crate::response::StatusCode::METHOD_NOT_ALLOWED)
5457 .header("Content-Type", b"application/json".to_vec())
5458 .header(
5459 "Allow",
5460 b"GET, POST, PUT, DELETE, PATCH, OPTIONS, HEAD".to_vec(),
5461 )
5462 .body(crate::response::ResponseBody::Bytes(body.into_bytes()))
5463 }
5464}
5465
5466impl Middleware for TraceRejectionMiddleware {
5467 fn before<'a>(
5468 &'a self,
5469 _ctx: &'a RequestContext,
5470 req: &'a mut Request,
5471 ) -> BoxFuture<'a, ControlFlow> {
5472 Box::pin(async move {
5473 if req.method() == crate::request::Method::Trace {
5474 if self.log_attempts {
5475 let path = req.path();
5477 let remote_ip = req
5478 .headers()
5479 .get("X-Forwarded-For")
5480 .or_else(|| req.headers().get("X-Real-IP"))
5481 .map(|v| String::from_utf8_lossy(v).to_string())
5482 .unwrap_or_else(|| "unknown".to_string());
5483
5484 eprintln!(
5485 "[SECURITY] TRACE request blocked: path={}, remote_ip={}",
5486 path, remote_ip
5487 );
5488 }
5489
5490 return ControlFlow::Break(Self::rejection_response(req.path()));
5491 }
5492
5493 ControlFlow::Continue
5494 })
5495 }
5496
5497 fn name(&self) -> &'static str {
5498 "TraceRejection"
5499 }
5500}
5501
5502#[derive(Debug, Clone)]
5512#[allow(clippy::struct_excessive_bools)]
5513pub struct HttpsRedirectConfig {
5514 pub redirect_enabled: bool,
5516 pub permanent_redirect: bool,
5518 pub hsts_max_age_secs: u64,
5520 pub hsts_include_subdomains: bool,
5522 pub hsts_preload: bool,
5524 pub exclude_paths: Vec<String>,
5526 pub https_port: u16,
5528}
5529
5530impl Default for HttpsRedirectConfig {
5531 fn default() -> Self {
5532 Self {
5533 redirect_enabled: true,
5534 permanent_redirect: true, hsts_max_age_secs: 31_536_000, hsts_include_subdomains: false,
5537 hsts_preload: false,
5538 exclude_paths: Vec::new(),
5539 https_port: 443,
5540 }
5541 }
5542}
5543
5544#[derive(Debug, Clone)]
5584pub struct HttpsRedirectMiddleware {
5585 config: HttpsRedirectConfig,
5586}
5587
5588impl Default for HttpsRedirectMiddleware {
5589 fn default() -> Self {
5590 Self::new()
5591 }
5592}
5593
5594impl HttpsRedirectMiddleware {
5595 #[must_use]
5597 pub fn new() -> Self {
5598 Self {
5599 config: HttpsRedirectConfig::default(),
5600 }
5601 }
5602
5603 #[must_use]
5605 pub fn redirect_enabled(mut self, enabled: bool) -> Self {
5606 self.config.redirect_enabled = enabled;
5607 self
5608 }
5609
5610 #[must_use]
5614 pub fn permanent_redirect(mut self, permanent: bool) -> Self {
5615 self.config.permanent_redirect = permanent;
5616 self
5617 }
5618
5619 #[must_use]
5624 pub fn hsts_max_age_secs(mut self, secs: u64) -> Self {
5625 self.config.hsts_max_age_secs = secs;
5626 self
5627 }
5628
5629 #[must_use]
5631 pub fn include_subdomains(mut self, include: bool) -> Self {
5632 self.config.hsts_include_subdomains = include;
5633 self
5634 }
5635
5636 #[must_use]
5641 pub fn preload(mut self, preload: bool) -> Self {
5642 self.config.hsts_preload = preload;
5643 self
5644 }
5645
5646 #[must_use]
5651 pub fn exclude_path(mut self, path: impl Into<String>) -> Self {
5652 self.config.exclude_paths.push(path.into());
5653 self
5654 }
5655
5656 #[must_use]
5658 pub fn exclude_paths(mut self, paths: Vec<String>) -> Self {
5659 self.config.exclude_paths = paths;
5660 self
5661 }
5662
5663 #[must_use]
5665 pub fn https_port(mut self, port: u16) -> Self {
5666 self.config.https_port = port;
5667 self
5668 }
5669
5670 fn is_secure(&self, req: &Request) -> bool {
5675 fn trim_ascii(mut bytes: &[u8]) -> &[u8] {
5676 while matches!(bytes.first(), Some(b' ' | b'\t')) {
5677 bytes = &bytes[1..];
5678 }
5679 while matches!(bytes.last(), Some(b' ' | b'\t')) {
5680 bytes = &bytes[..bytes.len() - 1];
5681 }
5682 bytes
5683 }
5684
5685 if let Some(info) = req.get_extension::<crate::request::ConnectionInfo>() {
5686 if info.is_tls {
5687 return true;
5688 }
5689 }
5690
5691 if let Some(forwarded) = req.headers().get("Forwarded") {
5693 if let Ok(s) = std::str::from_utf8(forwarded) {
5694 for entry in s.split(',') {
5695 for param in entry.split(';') {
5696 let param = param.trim();
5697 if let Some((k, v)) = param.split_once('=') {
5698 if k.trim().eq_ignore_ascii_case("proto") {
5699 let proto = v.trim().trim_matches('"');
5700 if proto.eq_ignore_ascii_case("https") {
5701 return true;
5702 }
5703 }
5704 }
5705 }
5706 }
5707 }
5708 }
5709
5710 if let Some(proto) = req.headers().get("X-Forwarded-Proto") {
5712 let first = proto.split(|&b| b == b',').next().unwrap_or(proto);
5713 return trim_ascii(first).eq_ignore_ascii_case(b"https");
5714 }
5715
5716 if let Some(ssl) = req.headers().get("X-Forwarded-Ssl") {
5718 return ssl.eq_ignore_ascii_case(b"on");
5719 }
5720
5721 if let Some(https) = req.headers().get("Front-End-Https") {
5723 return https.eq_ignore_ascii_case(b"on");
5724 }
5725
5726 false
5727 }
5728
5729 fn is_excluded(&self, path: &str) -> bool {
5731 self.config
5732 .exclude_paths
5733 .iter()
5734 .any(|p| path.starts_with(p))
5735 }
5736
5737 fn build_hsts_header(&self) -> Option<Vec<u8>> {
5739 if self.config.hsts_max_age_secs == 0 {
5740 return None;
5741 }
5742
5743 let mut value = format!("max-age={}", self.config.hsts_max_age_secs);
5744
5745 if self.config.hsts_include_subdomains {
5746 value.push_str("; includeSubDomains");
5747 }
5748
5749 if self.config.hsts_preload {
5750 value.push_str("; preload");
5751 }
5752
5753 Some(value.into_bytes())
5754 }
5755
5756 fn build_redirect_url(&self, req: &Request) -> String {
5758 let host = req
5759 .headers()
5760 .get("Host")
5761 .map(|h| String::from_utf8_lossy(h).to_string())
5762 .unwrap_or_else(|| "localhost".to_string());
5763
5764 let host_without_port = host.split(':').next().unwrap_or(&host);
5766
5767 let path = req.path();
5768 let query = req.query();
5769
5770 if self.config.https_port == 443 {
5771 match query {
5772 Some(q) => format!("https://{}{}?{}", host_without_port, path, q),
5773 None => format!("https://{}{}", host_without_port, path),
5774 }
5775 } else {
5776 match query {
5777 Some(q) => format!(
5778 "https://{}:{}{}?{}",
5779 host_without_port, self.config.https_port, path, q
5780 ),
5781 None => format!(
5782 "https://{}:{}{}",
5783 host_without_port, self.config.https_port, path
5784 ),
5785 }
5786 }
5787 }
5788}
5789
5790impl Middleware for HttpsRedirectMiddleware {
5791 fn before<'a>(
5792 &'a self,
5793 _ctx: &'a RequestContext,
5794 req: &'a mut Request,
5795 ) -> BoxFuture<'a, ControlFlow> {
5796 Box::pin(async move {
5797 if !self.config.redirect_enabled {
5799 return ControlFlow::Continue;
5800 }
5801
5802 if self.is_secure(req) {
5804 return ControlFlow::Continue;
5805 }
5806
5807 if self.is_excluded(req.path()) {
5809 return ControlFlow::Continue;
5810 }
5811
5812 let redirect_url = self.build_redirect_url(req);
5814
5815 let status = if self.config.permanent_redirect {
5817 crate::response::StatusCode::MOVED_PERMANENTLY
5818 } else {
5819 crate::response::StatusCode::TEMPORARY_REDIRECT
5820 };
5821
5822 let response = Response::with_status(status)
5824 .header("Location", redirect_url.into_bytes())
5825 .header("Content-Type", b"text/plain".to_vec())
5826 .body(crate::response::ResponseBody::Bytes(
5827 b"Redirecting to HTTPS...".to_vec(),
5828 ));
5829
5830 ControlFlow::Break(response)
5831 })
5832 }
5833
5834 fn after<'a>(
5835 &'a self,
5836 _ctx: &'a RequestContext,
5837 req: &'a Request,
5838 response: Response,
5839 ) -> BoxFuture<'a, Response> {
5840 Box::pin(async move {
5841 if !self.is_secure(req) {
5843 return response;
5844 }
5845
5846 if let Some(hsts_value) = self.build_hsts_header() {
5848 response.header("Strict-Transport-Security", hsts_value)
5849 } else {
5850 response
5851 }
5852 })
5853 }
5854
5855 fn name(&self) -> &'static str {
5856 "HttpsRedirect"
5857 }
5858}
5859
5860pub trait ResponseInterceptor: Send + Sync {
5899 fn intercept<'a>(
5910 &'a self,
5911 ctx: &'a ResponseInterceptorContext<'a>,
5912 response: Response,
5913 ) -> BoxFuture<'a, Response>;
5914
5915 fn name(&self) -> &'static str {
5917 std::any::type_name::<Self>()
5918 }
5919}
5920
5921#[derive(Debug)]
5926pub struct ResponseInterceptorContext<'a> {
5927 pub request: &'a Request,
5929 pub start_time: Instant,
5931 pub request_ctx: &'a RequestContext,
5933}
5934
5935impl<'a> ResponseInterceptorContext<'a> {
5936 pub fn new(request: &'a Request, request_ctx: &'a RequestContext, start_time: Instant) -> Self {
5938 Self {
5939 request,
5940 start_time,
5941 request_ctx,
5942 }
5943 }
5944
5945 pub fn elapsed(&self) -> std::time::Duration {
5947 self.start_time.elapsed()
5948 }
5949
5950 pub fn elapsed_ms(&self) -> u128 {
5952 self.start_time.elapsed().as_millis()
5953 }
5954}
5955
5956#[derive(Default)]
5971pub struct ResponseInterceptorStack {
5972 interceptors: Vec<Arc<dyn ResponseInterceptor>>,
5973}
5974
5975impl ResponseInterceptorStack {
5976 #[must_use]
5978 pub fn new() -> Self {
5979 Self {
5980 interceptors: Vec::new(),
5981 }
5982 }
5983
5984 #[must_use]
5986 pub fn with_capacity(capacity: usize) -> Self {
5987 Self {
5988 interceptors: Vec::with_capacity(capacity),
5989 }
5990 }
5991
5992 pub fn push<I: ResponseInterceptor + 'static>(&mut self, interceptor: I) {
5994 self.interceptors.push(Arc::new(interceptor));
5995 }
5996
5997 pub fn push_arc(&mut self, interceptor: Arc<dyn ResponseInterceptor>) {
5999 self.interceptors.push(interceptor);
6000 }
6001
6002 #[must_use]
6004 pub fn len(&self) -> usize {
6005 self.interceptors.len()
6006 }
6007
6008 #[must_use]
6010 pub fn is_empty(&self) -> bool {
6011 self.interceptors.is_empty()
6012 }
6013
6014 pub async fn process(
6016 &self,
6017 ctx: &ResponseInterceptorContext<'_>,
6018 mut response: Response,
6019 ) -> Response {
6020 for interceptor in &self.interceptors {
6021 let _ = ctx.request_ctx.checkpoint();
6022 response = interceptor.intercept(ctx, response).await;
6023 }
6024 response
6025 }
6026}
6027
6028#[derive(Debug, Clone)]
6045pub struct TimingInterceptor {
6046 header_name: String,
6048 include_server_timing: bool,
6050 server_timing_name: String,
6052}
6053
6054impl Default for TimingInterceptor {
6055 fn default() -> Self {
6056 Self::new()
6057 }
6058}
6059
6060impl TimingInterceptor {
6061 #[must_use]
6063 pub fn new() -> Self {
6064 Self {
6065 header_name: "X-Response-Time".to_string(),
6066 include_server_timing: false,
6067 server_timing_name: "total".to_string(),
6068 }
6069 }
6070
6071 #[must_use]
6073 pub fn with_server_timing(mut self, metric_name: impl Into<String>) -> Self {
6074 self.include_server_timing = true;
6075 self.server_timing_name = metric_name.into();
6076 self
6077 }
6078
6079 #[must_use]
6081 pub fn header_name(mut self, name: impl Into<String>) -> Self {
6082 self.header_name = name.into();
6083 self
6084 }
6085}
6086
6087impl ResponseInterceptor for TimingInterceptor {
6088 fn intercept<'a>(
6089 &'a self,
6090 ctx: &'a ResponseInterceptorContext<'a>,
6091 response: Response,
6092 ) -> BoxFuture<'a, Response> {
6093 Box::pin(async move {
6094 let elapsed_ms = ctx.elapsed_ms();
6095 let timing_value = format!("{}ms", elapsed_ms);
6096
6097 let response = response.header(&self.header_name, timing_value.clone().into_bytes());
6098
6099 if self.include_server_timing {
6100 let server_timing = format!("{};dur={}", self.server_timing_name, elapsed_ms);
6102 response.header("Server-Timing", server_timing.into_bytes())
6103 } else {
6104 response
6105 }
6106 })
6107 }
6108
6109 fn name(&self) -> &'static str {
6110 "TimingInterceptor"
6111 }
6112}
6113
6114#[derive(Debug, Clone)]
6138#[allow(clippy::struct_excessive_bools)]
6139pub struct DebugInfoInterceptor {
6140 include_path: bool,
6142 include_method: bool,
6144 include_request_id: bool,
6146 include_timing: bool,
6148 header_prefix: String,
6150}
6151
6152impl Default for DebugInfoInterceptor {
6153 fn default() -> Self {
6154 Self::new()
6155 }
6156}
6157
6158impl DebugInfoInterceptor {
6159 #[must_use]
6161 pub fn new() -> Self {
6162 Self {
6163 include_path: true,
6164 include_method: true,
6165 include_request_id: true,
6166 include_timing: true,
6167 header_prefix: "X-Debug-".to_string(),
6168 }
6169 }
6170
6171 #[must_use]
6173 pub fn include_path(mut self, include: bool) -> Self {
6174 self.include_path = include;
6175 self
6176 }
6177
6178 #[must_use]
6180 pub fn include_method(mut self, include: bool) -> Self {
6181 self.include_method = include;
6182 self
6183 }
6184
6185 #[must_use]
6187 pub fn include_request_id(mut self, include: bool) -> Self {
6188 self.include_request_id = include;
6189 self
6190 }
6191
6192 #[must_use]
6194 pub fn include_timing(mut self, include: bool) -> Self {
6195 self.include_timing = include;
6196 self
6197 }
6198
6199 #[must_use]
6201 pub fn header_prefix(mut self, prefix: impl Into<String>) -> Self {
6202 self.header_prefix = prefix.into();
6203 self
6204 }
6205}
6206
6207impl ResponseInterceptor for DebugInfoInterceptor {
6208 fn intercept<'a>(
6209 &'a self,
6210 ctx: &'a ResponseInterceptorContext<'a>,
6211 response: Response,
6212 ) -> BoxFuture<'a, Response> {
6213 Box::pin(async move {
6214 let mut resp = response;
6215
6216 if self.include_path {
6217 let header_name = format!("{}Path", self.header_prefix);
6218 resp = resp.header(header_name, ctx.request.path().as_bytes().to_vec());
6219 }
6220
6221 if self.include_method {
6222 let header_name = format!("{}Method", self.header_prefix);
6223 resp = resp.header(
6224 header_name,
6225 ctx.request.method().as_str().as_bytes().to_vec(),
6226 );
6227 }
6228
6229 if self.include_request_id {
6230 if let Some(request_id) = ctx.request.get_extension::<RequestId>() {
6231 let header_name = format!("{}Request-Id", self.header_prefix);
6232 resp = resp.header(header_name, request_id.0.as_bytes().to_vec());
6233 }
6234 }
6235
6236 if self.include_timing {
6237 let header_name = format!("{}Handler-Time", self.header_prefix);
6238 let timing = format!("{}ms", ctx.elapsed_ms());
6239 resp = resp.header(header_name, timing.into_bytes());
6240 }
6241
6242 resp
6243 })
6244 }
6245
6246 fn name(&self) -> &'static str {
6247 "DebugInfoInterceptor"
6248 }
6249}
6250
6251pub struct ResponseBodyTransform<F>
6272where
6273 F: Fn(Vec<u8>) -> Vec<u8> + Send + Sync,
6274{
6275 transform_fn: F,
6276 content_type_filter: Option<String>,
6278}
6279
6280impl<F> ResponseBodyTransform<F>
6281where
6282 F: Fn(Vec<u8>) -> Vec<u8> + Send + Sync,
6283{
6284 pub fn new(transform_fn: F) -> Self {
6286 Self {
6287 transform_fn,
6288 content_type_filter: None,
6289 }
6290 }
6291
6292 #[must_use]
6294 pub fn for_content_type(mut self, content_type: impl Into<String>) -> Self {
6295 self.content_type_filter = Some(content_type.into());
6296 self
6297 }
6298
6299 fn should_transform(&self, response: &Response) -> bool {
6300 match &self.content_type_filter {
6301 Some(filter) => response
6302 .headers()
6303 .iter()
6304 .find(|(name, _)| name.eq_ignore_ascii_case("content-type"))
6305 .and_then(|(_, ct)| std::str::from_utf8(ct).ok())
6306 .map(|ct| ct.starts_with(filter))
6307 .unwrap_or(false),
6308 None => true,
6309 }
6310 }
6311}
6312
6313impl<F> ResponseInterceptor for ResponseBodyTransform<F>
6314where
6315 F: Fn(Vec<u8>) -> Vec<u8> + Send + Sync,
6316{
6317 fn intercept<'a>(
6318 &'a self,
6319 _ctx: &'a ResponseInterceptorContext<'a>,
6320 response: Response,
6321 ) -> BoxFuture<'a, Response> {
6322 Box::pin(async move {
6323 if !self.should_transform(&response) {
6324 return response;
6325 }
6326
6327 let body_bytes = match response.body_ref() {
6329 crate::response::ResponseBody::Empty => Vec::new(),
6330 crate::response::ResponseBody::Bytes(b) => b.clone(),
6331 crate::response::ResponseBody::Stream(_) => {
6332 return response;
6334 }
6335 };
6336
6337 let transformed = (self.transform_fn)(body_bytes);
6339
6340 response.body(crate::response::ResponseBody::Bytes(transformed))
6342 })
6343 }
6344
6345 fn name(&self) -> &'static str {
6346 "ResponseBodyTransform"
6347 }
6348}
6349
6350#[derive(Debug, Clone, Default)]
6367pub struct HeaderTransformInterceptor {
6368 add_headers: Vec<(String, Vec<u8>)>,
6370 remove_headers: Vec<String>,
6372 rename_headers: Vec<(String, String)>,
6374}
6375
6376impl HeaderTransformInterceptor {
6377 #[must_use]
6379 pub fn new() -> Self {
6380 Self::default()
6381 }
6382
6383 #[must_use]
6385 pub fn add(mut self, name: impl Into<String>, value: impl Into<Vec<u8>>) -> Self {
6386 self.add_headers.push((name.into(), value.into()));
6387 self
6388 }
6389
6390 #[must_use]
6392 pub fn remove(mut self, name: impl Into<String>) -> Self {
6393 self.remove_headers.push(name.into());
6394 self
6395 }
6396
6397 #[must_use]
6399 pub fn rename(mut self, old_name: impl Into<String>, new_name: impl Into<String>) -> Self {
6400 self.rename_headers.push((old_name.into(), new_name.into()));
6401 self
6402 }
6403}
6404
6405impl ResponseInterceptor for HeaderTransformInterceptor {
6406 fn intercept<'a>(
6407 &'a self,
6408 _ctx: &'a ResponseInterceptorContext<'a>,
6409 response: Response,
6410 ) -> BoxFuture<'a, Response> {
6411 let add_headers = self.add_headers.clone();
6412 let remove_headers = self.remove_headers.clone();
6413 let rename_headers = self.rename_headers.clone();
6414
6415 Box::pin(async move {
6416 let mut resp = response;
6417
6418 for (old_name, new_name) in &rename_headers {
6420 let values: Vec<Vec<u8>> = resp
6421 .headers()
6422 .iter()
6423 .filter(|(name, _)| name.eq_ignore_ascii_case(old_name))
6424 .map(|(_, v)| v.clone())
6425 .collect();
6426
6427 if !values.is_empty() {
6428 resp = resp.remove_header(old_name);
6429 for v in values {
6430 resp = resp.header(new_name, v);
6431 }
6432 }
6433 }
6434
6435 for (name, value) in add_headers {
6437 resp = resp.header(name, value);
6438 }
6439
6440 for name in &remove_headers {
6442 resp = resp.remove_header(name);
6443 }
6444
6445 resp
6446 })
6447 }
6448
6449 fn name(&self) -> &'static str {
6450 "HeaderTransformInterceptor"
6451 }
6452}
6453
6454pub struct ConditionalInterceptor<I, F>
6470where
6471 I: ResponseInterceptor,
6472 F: Fn(&ResponseInterceptorContext, &Response) -> bool + Send + Sync,
6473{
6474 inner: I,
6475 condition: F,
6476}
6477
6478impl<I, F> ConditionalInterceptor<I, F>
6479where
6480 I: ResponseInterceptor,
6481 F: Fn(&ResponseInterceptorContext, &Response) -> bool + Send + Sync,
6482{
6483 pub fn new(inner: I, condition: F) -> Self {
6485 Self { inner, condition }
6486 }
6487}
6488
6489impl<I, F> ResponseInterceptor for ConditionalInterceptor<I, F>
6490where
6491 I: ResponseInterceptor,
6492 F: Fn(&ResponseInterceptorContext, &Response) -> bool + Send + Sync,
6493{
6494 fn intercept<'a>(
6495 &'a self,
6496 ctx: &'a ResponseInterceptorContext<'a>,
6497 response: Response,
6498 ) -> BoxFuture<'a, Response> {
6499 Box::pin(async move {
6500 if (self.condition)(ctx, &response) {
6501 self.inner.intercept(ctx, response).await
6502 } else {
6503 response
6504 }
6505 })
6506 }
6507
6508 fn name(&self) -> &'static str {
6509 "ConditionalInterceptor"
6510 }
6511}
6512
6513#[derive(Debug, Clone)]
6532pub struct ErrorResponseTransformer {
6533 status_codes: HashSet<u16>,
6535 replacement_body: Option<Vec<u8>>,
6537 add_error_id: bool,
6539}
6540
6541impl Default for ErrorResponseTransformer {
6542 fn default() -> Self {
6543 Self::new()
6544 }
6545}
6546
6547impl ErrorResponseTransformer {
6548 #[must_use]
6550 pub fn new() -> Self {
6551 Self {
6552 status_codes: HashSet::new(),
6553 replacement_body: None,
6554 add_error_id: false,
6555 }
6556 }
6557
6558 #[must_use]
6560 pub fn hide_details_for_status(mut self, status: crate::response::StatusCode) -> Self {
6561 self.status_codes.insert(status.as_u16());
6562 self
6563 }
6564
6565 #[must_use]
6567 pub fn with_replacement_body(mut self, body: impl Into<Vec<u8>>) -> Self {
6568 self.replacement_body = Some(body.into());
6569 self
6570 }
6571
6572 #[must_use]
6574 pub fn add_error_id(mut self, enable: bool) -> Self {
6575 self.add_error_id = enable;
6576 self
6577 }
6578}
6579
6580impl ResponseInterceptor for ErrorResponseTransformer {
6581 fn intercept<'a>(
6582 &'a self,
6583 ctx: &'a ResponseInterceptorContext<'a>,
6584 response: Response,
6585 ) -> BoxFuture<'a, Response> {
6586 Box::pin(async move {
6587 let status_code = response.status().as_u16();
6588
6589 if !self.status_codes.contains(&status_code) {
6590 return response;
6591 }
6592
6593 let mut resp = response;
6594
6595 if let Some(ref replacement) = self.replacement_body {
6597 resp = resp.body(crate::response::ResponseBody::Bytes(replacement.clone()));
6598 }
6599
6600 if self.add_error_id {
6602 let error_id = ctx
6604 .request
6605 .get_extension::<RequestId>()
6606 .map(|r| r.0.clone())
6607 .unwrap_or_else(|| format!("err-{}", ctx.elapsed_ms()));
6608 resp = resp.header("X-Error-Id", error_id.into_bytes());
6609 }
6610
6611 resp
6612 })
6613 }
6614
6615 fn name(&self) -> &'static str {
6616 "ErrorResponseTransformer"
6617 }
6618}
6619
6620pub struct ResponseInterceptorMiddleware<I>
6636where
6637 I: ResponseInterceptor,
6638{
6639 interceptor: I,
6640}
6641
6642impl<I> ResponseInterceptorMiddleware<I>
6643where
6644 I: ResponseInterceptor,
6645{
6646 pub fn new(interceptor: I) -> Self {
6648 Self { interceptor }
6649 }
6650}
6651
6652impl<I> Middleware for ResponseInterceptorMiddleware<I>
6653where
6654 I: ResponseInterceptor,
6655{
6656 fn before<'a>(
6657 &'a self,
6658 _ctx: &'a RequestContext,
6659 req: &'a mut Request,
6660 ) -> BoxFuture<'a, ControlFlow> {
6661 req.insert_extension(InterceptorStartTime(Instant::now()));
6663 Box::pin(async { ControlFlow::Continue })
6664 }
6665
6666 fn after<'a>(
6667 &'a self,
6668 ctx: &'a RequestContext,
6669 req: &'a Request,
6670 response: Response,
6671 ) -> BoxFuture<'a, Response> {
6672 Box::pin(async move {
6673 let start_time = req
6675 .get_extension::<InterceptorStartTime>()
6676 .map(|t| t.0)
6677 .unwrap_or_else(Instant::now);
6678
6679 let interceptor_ctx = ResponseInterceptorContext::new(req, ctx, start_time);
6680 self.interceptor.intercept(&interceptor_ctx, response).await
6681 })
6682 }
6683
6684 fn name(&self) -> &'static str {
6685 self.interceptor.name()
6686 }
6687}
6688
6689#[derive(Debug, Clone, Copy)]
6691struct InterceptorStartTime(Instant);
6692
6693#[derive(Debug, Clone)]
6725pub struct ServerTimingEntry {
6726 name: String,
6728 duration_ms: f64,
6730 description: Option<String>,
6732}
6733
6734impl ServerTimingEntry {
6735 #[must_use]
6737 pub fn new(name: impl Into<String>, duration_ms: f64) -> Self {
6738 Self {
6739 name: name.into(),
6740 duration_ms,
6741 description: None,
6742 }
6743 }
6744
6745 #[must_use]
6747 pub fn with_description(mut self, desc: impl Into<String>) -> Self {
6748 self.description = Some(desc.into());
6749 self
6750 }
6751
6752 #[must_use]
6754 pub fn to_header_value(&self) -> String {
6755 match &self.description {
6756 Some(desc) => format!(
6757 "{};dur={:.3};desc=\"{}\"",
6758 self.name, self.duration_ms, desc
6759 ),
6760 None => format!("{};dur={:.3}", self.name, self.duration_ms),
6761 }
6762 }
6763}
6764
6765#[derive(Debug, Clone, Default)]
6781pub struct ServerTimingBuilder {
6782 entries: Vec<ServerTimingEntry>,
6783}
6784
6785impl ServerTimingBuilder {
6786 #[must_use]
6788 pub fn new() -> Self {
6789 Self::default()
6790 }
6791
6792 #[must_use]
6794 pub fn add(mut self, name: impl Into<String>, duration_ms: f64) -> Self {
6795 self.entries.push(ServerTimingEntry::new(name, duration_ms));
6796 self
6797 }
6798
6799 #[must_use]
6801 pub fn add_with_desc(
6802 mut self,
6803 name: impl Into<String>,
6804 duration_ms: f64,
6805 description: impl Into<String>,
6806 ) -> Self {
6807 self.entries
6808 .push(ServerTimingEntry::new(name, duration_ms).with_description(description));
6809 self
6810 }
6811
6812 #[must_use]
6814 pub fn add_entry(mut self, entry: ServerTimingEntry) -> Self {
6815 self.entries.push(entry);
6816 self
6817 }
6818
6819 #[must_use]
6821 pub fn build(&self) -> String {
6822 self.entries
6823 .iter()
6824 .map(ServerTimingEntry::to_header_value)
6825 .collect::<Vec<_>>()
6826 .join(", ")
6827 }
6828
6829 #[must_use]
6831 pub fn is_empty(&self) -> bool {
6832 self.entries.is_empty()
6833 }
6834
6835 #[must_use]
6837 pub fn len(&self) -> usize {
6838 self.entries.len()
6839 }
6840}
6841
6842#[derive(Debug, Clone)]
6858pub struct TimingMetrics {
6859 pub start_time: Instant,
6861 pub first_byte_time: Option<Instant>,
6863 pub custom_metrics: Vec<(String, f64, Option<String>)>,
6865}
6866
6867impl TimingMetrics {
6868 #[must_use]
6870 pub fn new() -> Self {
6871 Self {
6872 start_time: Instant::now(),
6873 first_byte_time: None,
6874 custom_metrics: Vec::new(),
6875 }
6876 }
6877
6878 #[must_use]
6880 pub fn with_start_time(start_time: Instant) -> Self {
6881 Self {
6882 start_time,
6883 first_byte_time: None,
6884 custom_metrics: Vec::new(),
6885 }
6886 }
6887
6888 pub fn mark_first_byte(&mut self) {
6890 self.first_byte_time = Some(Instant::now());
6891 }
6892
6893 pub fn add_metric(&mut self, name: impl Into<String>, duration_ms: f64) {
6895 self.custom_metrics.push((name.into(), duration_ms, None));
6896 }
6897
6898 pub fn add_metric_with_desc(
6900 &mut self,
6901 name: impl Into<String>,
6902 duration_ms: f64,
6903 desc: impl Into<String>,
6904 ) {
6905 self.custom_metrics
6906 .push((name.into(), duration_ms, Some(desc.into())));
6907 }
6908
6909 #[must_use]
6911 pub fn total_ms(&self) -> f64 {
6912 self.start_time.elapsed().as_secs_f64() * 1000.0
6913 }
6914
6915 #[must_use]
6917 pub fn ttfb_ms(&self) -> Option<f64> {
6918 self.first_byte_time
6919 .map(|t| t.duration_since(self.start_time).as_secs_f64() * 1000.0)
6920 }
6921
6922 #[must_use]
6924 pub fn to_server_timing(&self) -> ServerTimingBuilder {
6925 let mut builder = ServerTimingBuilder::new().add_with_desc(
6926 "total",
6927 self.total_ms(),
6928 "Total request time",
6929 );
6930
6931 if let Some(ttfb) = self.ttfb_ms() {
6932 builder = builder.add_with_desc("ttfb", ttfb, "Time to first byte");
6933 }
6934
6935 for (name, duration, desc) in &self.custom_metrics {
6936 match desc {
6937 Some(d) => builder = builder.add_with_desc(name, *duration, d),
6938 None => builder = builder.add(name, *duration),
6939 }
6940 }
6941
6942 builder
6943 }
6944}
6945
6946impl Default for TimingMetrics {
6947 fn default() -> Self {
6948 Self::new()
6949 }
6950}
6951
6952#[derive(Debug, Clone)]
6954#[allow(clippy::struct_excessive_bools)]
6955pub struct TimingMetricsConfig {
6956 pub add_server_timing_header: bool,
6958 pub add_response_time_header: bool,
6960 pub response_time_header_name: String,
6962 pub include_custom_metrics: bool,
6964 pub include_ttfb: bool,
6966}
6967
6968impl Default for TimingMetricsConfig {
6969 fn default() -> Self {
6970 Self {
6971 add_server_timing_header: true,
6972 add_response_time_header: true,
6973 response_time_header_name: "X-Response-Time".to_string(),
6974 include_custom_metrics: true,
6975 include_ttfb: true,
6976 }
6977 }
6978}
6979
6980impl TimingMetricsConfig {
6981 #[must_use]
6983 pub fn new() -> Self {
6984 Self::default()
6985 }
6986
6987 #[must_use]
6989 pub fn server_timing(mut self, enabled: bool) -> Self {
6990 self.add_server_timing_header = enabled;
6991 self
6992 }
6993
6994 #[must_use]
6996 pub fn response_time(mut self, enabled: bool) -> Self {
6997 self.add_response_time_header = enabled;
6998 self
6999 }
7000
7001 #[must_use]
7003 pub fn response_time_header(mut self, name: impl Into<String>) -> Self {
7004 self.response_time_header_name = name.into();
7005 self
7006 }
7007
7008 #[must_use]
7010 pub fn custom_metrics(mut self, enabled: bool) -> Self {
7011 self.include_custom_metrics = enabled;
7012 self
7013 }
7014
7015 #[must_use]
7017 pub fn ttfb(mut self, enabled: bool) -> Self {
7018 self.include_ttfb = enabled;
7019 self
7020 }
7021
7022 #[must_use]
7024 pub fn production() -> Self {
7025 Self {
7026 add_server_timing_header: false,
7027 add_response_time_header: true,
7028 response_time_header_name: "X-Response-Time".to_string(),
7029 include_custom_metrics: false,
7030 include_ttfb: false,
7031 }
7032 }
7033
7034 #[must_use]
7036 pub fn development() -> Self {
7037 Self::default()
7038 }
7039}
7040
7041#[derive(Debug, Clone)]
7060pub struct TimingMetricsMiddleware {
7061 config: TimingMetricsConfig,
7062}
7063
7064impl TimingMetricsMiddleware {
7065 #[must_use]
7067 pub fn new() -> Self {
7068 Self {
7069 config: TimingMetricsConfig::default(),
7070 }
7071 }
7072
7073 #[must_use]
7075 pub fn with_config(config: TimingMetricsConfig) -> Self {
7076 Self { config }
7077 }
7078
7079 #[must_use]
7081 pub fn production() -> Self {
7082 Self {
7083 config: TimingMetricsConfig::production(),
7084 }
7085 }
7086
7087 #[must_use]
7089 pub fn development() -> Self {
7090 Self {
7091 config: TimingMetricsConfig::development(),
7092 }
7093 }
7094}
7095
7096impl Default for TimingMetricsMiddleware {
7097 fn default() -> Self {
7098 Self::new()
7099 }
7100}
7101
7102impl Middleware for TimingMetricsMiddleware {
7103 fn before<'a>(
7104 &'a self,
7105 _ctx: &'a RequestContext,
7106 req: &'a mut Request,
7107 ) -> BoxFuture<'a, ControlFlow> {
7108 req.insert_extension(TimingMetrics::new());
7110 Box::pin(async { ControlFlow::Continue })
7111 }
7112
7113 fn after<'a>(
7114 &'a self,
7115 _ctx: &'a RequestContext,
7116 req: &'a Request,
7117 response: Response,
7118 ) -> BoxFuture<'a, Response> {
7119 let config = self.config.clone();
7120
7121 Box::pin(async move {
7122 let mut resp = response;
7123
7124 let metrics = req.get_extension::<TimingMetrics>();
7126
7127 match metrics {
7128 Some(metrics) => {
7129 if config.add_response_time_header {
7131 let timing = format!("{:.3}ms", metrics.total_ms());
7132 resp = resp.header(&config.response_time_header_name, timing.into_bytes());
7133 }
7134
7135 if config.add_server_timing_header {
7137 let mut builder = ServerTimingBuilder::new().add_with_desc(
7138 "total",
7139 metrics.total_ms(),
7140 "Total request time",
7141 );
7142
7143 if config.include_ttfb {
7145 if let Some(ttfb) = metrics.ttfb_ms() {
7146 builder = builder.add_with_desc("ttfb", ttfb, "Time to first byte");
7147 }
7148 }
7149
7150 if config.include_custom_metrics {
7152 for (name, duration, desc) in &metrics.custom_metrics {
7153 match desc {
7154 Some(d) => builder = builder.add_with_desc(name, *duration, d),
7155 None => builder = builder.add(name, *duration),
7156 }
7157 }
7158 }
7159
7160 let header_value = builder.build();
7161 resp = resp.header("Server-Timing", header_value.into_bytes());
7162 }
7163 }
7164 None => {
7165 if config.add_response_time_header {
7168 resp = resp.header(&config.response_time_header_name, b"0.000ms".to_vec());
7169 }
7170 }
7171 }
7172
7173 resp
7174 })
7175 }
7176
7177 fn name(&self) -> &'static str {
7178 "TimingMetrics"
7179 }
7180}
7181
7182#[derive(Debug, Clone)]
7186pub struct TimingHistogramBucket {
7187 pub le: f64,
7189 pub count: u64,
7191}
7192
7193#[derive(Debug, Clone)]
7212pub struct TimingHistogram {
7213 bucket_bounds: Vec<f64>,
7215 bucket_counts: Vec<u64>,
7217 sum: f64,
7219 count: u64,
7221}
7222
7223impl TimingHistogram {
7224 #[must_use]
7228 pub fn with_buckets(bucket_bounds: Vec<f64>) -> Self {
7229 let bucket_counts = vec![0; bucket_bounds.len()];
7230 Self {
7231 bucket_bounds,
7232 bucket_counts,
7233 sum: 0.0,
7234 count: 0,
7235 }
7236 }
7237
7238 #[must_use]
7242 pub fn http_latency() -> Self {
7243 Self::with_buckets(vec![
7244 1.0, 5.0, 10.0, 25.0, 50.0, 100.0, 250.0, 500.0, 1000.0, 2500.0, 5000.0, 10000.0,
7245 ])
7246 }
7247
7248 pub fn observe(&mut self, value_ms: f64) {
7250 self.sum += value_ms;
7251 self.count += 1;
7252
7253 for (i, bound) in self.bucket_bounds.iter().enumerate() {
7255 if value_ms <= *bound {
7256 self.bucket_counts[i] += 1;
7257 }
7258 }
7259 }
7260
7261 #[must_use]
7263 pub fn count(&self) -> u64 {
7264 self.count
7265 }
7266
7267 #[must_use]
7269 pub fn sum(&self) -> f64 {
7270 self.sum
7271 }
7272
7273 #[must_use]
7275 pub fn mean(&self) -> f64 {
7276 if self.count == 0 {
7277 0.0
7278 } else {
7279 #[allow(clippy::cast_precision_loss)]
7280 {
7281 self.sum / self.count as f64
7282 }
7283 }
7284 }
7285
7286 #[must_use]
7288 pub fn buckets(&self) -> Vec<TimingHistogramBucket> {
7289 self.bucket_bounds
7290 .iter()
7291 .zip(&self.bucket_counts)
7292 .map(|(&le, &count)| TimingHistogramBucket { le, count })
7293 .collect()
7294 }
7295
7296 pub fn reset(&mut self) {
7298 self.sum = 0.0;
7299 self.count = 0;
7300 self.bucket_counts.fill(0);
7301 }
7302}
7303
7304impl Default for TimingHistogram {
7305 fn default() -> Self {
7306 Self::http_latency()
7307 }
7308}
7309
7310#[cfg(test)]
7315mod timing_metrics_tests {
7316 use super::*;
7317 use crate::request::Method;
7318 use crate::response::StatusCode;
7319
7320 fn test_context() -> RequestContext {
7321 RequestContext::new(asupersync::Cx::for_testing(), 1)
7322 }
7323
7324 fn test_request() -> Request {
7325 Request::new(Method::Get, "/test")
7326 }
7327
7328 fn run_middleware_before(mw: &impl Middleware, req: &mut Request) -> ControlFlow {
7329 let ctx = test_context();
7330 futures_executor::block_on(mw.before(&ctx, req))
7331 }
7332
7333 fn run_middleware_after(mw: &impl Middleware, req: &Request, resp: Response) -> Response {
7334 let ctx = test_context();
7335 futures_executor::block_on(mw.after(&ctx, req, resp))
7336 }
7337
7338 #[test]
7339 fn server_timing_entry_basic() {
7340 let entry = ServerTimingEntry::new("db", 42.5);
7341 assert_eq!(entry.to_header_value(), "db;dur=42.500");
7342 }
7343
7344 #[test]
7345 fn server_timing_entry_with_description() {
7346 let entry = ServerTimingEntry::new("db", 42.5).with_description("Database query");
7347 assert_eq!(
7348 entry.to_header_value(),
7349 "db;dur=42.500;desc=\"Database query\""
7350 );
7351 }
7352
7353 #[test]
7354 fn server_timing_builder_single_entry() {
7355 let timing = ServerTimingBuilder::new().add("total", 150.0).build();
7356 assert_eq!(timing, "total;dur=150.000");
7357 }
7358
7359 #[test]
7360 fn server_timing_builder_multiple_entries() {
7361 let timing = ServerTimingBuilder::new()
7362 .add("total", 150.0)
7363 .add_with_desc("db", 42.0, "Database")
7364 .add("cache", 5.0)
7365 .build();
7366
7367 assert!(timing.contains("total;dur=150.000"));
7368 assert!(timing.contains("db;dur=42.000;desc=\"Database\""));
7369 assert!(timing.contains("cache;dur=5.000"));
7370 assert!(timing.contains(", ")); }
7372
7373 #[test]
7374 fn server_timing_builder_empty() {
7375 let builder = ServerTimingBuilder::new();
7376 assert!(builder.is_empty());
7377 assert_eq!(builder.len(), 0);
7378 assert_eq!(builder.build(), "");
7379 }
7380
7381 #[test]
7382 fn timing_metrics_basic() {
7383 let metrics = TimingMetrics::new();
7384 std::thread::sleep(std::time::Duration::from_millis(5));
7385
7386 let total = metrics.total_ms();
7387 assert!(total >= 5.0, "Total should be at least 5ms");
7388 assert!(metrics.ttfb_ms().is_none(), "TTFB should not be set");
7389 }
7390
7391 #[test]
7392 fn timing_metrics_custom_metrics() {
7393 let mut metrics = TimingMetrics::new();
7394 metrics.add_metric("db", 42.5);
7395 metrics.add_metric_with_desc("cache", 5.0, "Cache lookup");
7396
7397 let timing = metrics.to_server_timing();
7398 assert_eq!(timing.len(), 3); let header = timing.build();
7401 assert!(header.contains("total"));
7402 assert!(header.contains("db;dur=42.500"));
7403 assert!(header.contains("cache;dur=5.000;desc=\"Cache lookup\""));
7404 }
7405
7406 #[test]
7407 fn timing_metrics_ttfb() {
7408 let mut metrics = TimingMetrics::new();
7409 std::thread::sleep(std::time::Duration::from_millis(5));
7410 metrics.mark_first_byte();
7411
7412 let ttfb = metrics.ttfb_ms().unwrap();
7413 assert!(ttfb >= 5.0, "TTFB should be at least 5ms");
7414 }
7415
7416 #[test]
7417 fn timing_metrics_config_default() {
7418 let config = TimingMetricsConfig::default();
7419 assert!(config.add_server_timing_header);
7420 assert!(config.add_response_time_header);
7421 assert!(config.include_custom_metrics);
7422 assert!(config.include_ttfb);
7423 }
7424
7425 #[test]
7426 fn timing_metrics_config_production() {
7427 let config = TimingMetricsConfig::production();
7428 assert!(!config.add_server_timing_header);
7429 assert!(config.add_response_time_header);
7430 assert!(!config.include_custom_metrics);
7431 }
7432
7433 #[test]
7434 fn timing_middleware_adds_metrics_to_request() {
7435 let mw = TimingMetricsMiddleware::new();
7436 let mut req = test_request();
7437
7438 let result = run_middleware_before(&mw, &mut req);
7440 assert!(result.is_continue());
7441
7442 let metrics = req.get_extension::<TimingMetrics>();
7443 assert!(metrics.is_some(), "TimingMetrics should be in extensions");
7444 }
7445
7446 #[test]
7447 fn timing_middleware_adds_response_time_header() {
7448 let mw = TimingMetricsMiddleware::new();
7449 let mut req = test_request();
7450
7451 run_middleware_before(&mw, &mut req);
7453
7454 let resp = Response::with_status(StatusCode::OK);
7455 let result = run_middleware_after(&mw, &req, resp);
7456
7457 let has_timing = result
7458 .headers()
7459 .iter()
7460 .any(|(name, _)| name == "X-Response-Time");
7461 assert!(has_timing, "Should have X-Response-Time header");
7462 }
7463
7464 #[test]
7465 fn timing_middleware_adds_server_timing_header() {
7466 let mw = TimingMetricsMiddleware::new();
7467 let mut req = test_request();
7468
7469 run_middleware_before(&mw, &mut req);
7470
7471 let resp = Response::with_status(StatusCode::OK);
7472 let result = run_middleware_after(&mw, &req, resp);
7473
7474 let server_timing = result
7475 .headers()
7476 .iter()
7477 .find(|(name, _)| name == "Server-Timing")
7478 .map(|(_, v)| String::from_utf8_lossy(v).to_string());
7479
7480 assert!(server_timing.is_some(), "Should have Server-Timing header");
7481 let header = server_timing.unwrap();
7482 assert!(header.contains("total"), "Should have total timing");
7483 }
7484
7485 #[test]
7486 fn timing_middleware_production_mode() {
7487 let mw = TimingMetricsMiddleware::production();
7488 let mut req = test_request();
7489
7490 run_middleware_before(&mw, &mut req);
7491
7492 let resp = Response::with_status(StatusCode::OK);
7493 let result = run_middleware_after(&mw, &req, resp);
7494
7495 let has_response_time = result
7497 .headers()
7498 .iter()
7499 .any(|(name, _)| name == "X-Response-Time");
7500 assert!(has_response_time);
7501
7502 let has_server_timing = result
7504 .headers()
7505 .iter()
7506 .any(|(name, _)| name == "Server-Timing");
7507 assert!(!has_server_timing);
7508 }
7509
7510 #[test]
7511 #[allow(clippy::float_cmp)]
7512 fn timing_histogram_basic() {
7513 let mut histogram = TimingHistogram::http_latency();
7514 assert_eq!(histogram.count(), 0);
7515 assert_eq!(histogram.sum(), 0.0);
7516
7517 histogram.observe(42.0);
7518 histogram.observe(150.0);
7519 histogram.observe(5.0);
7520
7521 assert_eq!(histogram.count(), 3);
7522 assert_eq!(histogram.sum(), 197.0);
7523 assert!((histogram.mean() - 65.666).abs() < 0.01);
7524 }
7525
7526 #[test]
7527 fn timing_histogram_buckets() {
7528 let mut histogram = TimingHistogram::with_buckets(vec![10.0, 50.0, 100.0]);
7529
7530 histogram.observe(5.0); histogram.observe(25.0); histogram.observe(75.0); histogram.observe(150.0); let buckets = histogram.buckets();
7536 assert_eq!(buckets.len(), 3);
7537
7538 assert_eq!(buckets[0].count, 1); assert_eq!(buckets[1].count, 2); assert_eq!(buckets[2].count, 3); }
7543
7544 #[test]
7545 #[allow(clippy::float_cmp)]
7546 fn timing_histogram_reset() {
7547 let mut histogram = TimingHistogram::http_latency();
7548 histogram.observe(100.0);
7549 histogram.observe(200.0);
7550
7551 assert_eq!(histogram.count(), 2);
7552
7553 histogram.reset();
7554
7555 assert_eq!(histogram.count(), 0);
7556 assert_eq!(histogram.sum(), 0.0);
7557 }
7558}
7559
7560#[cfg(test)]
7561mod response_interceptor_tests {
7562 use super::*;
7563 use crate::request::Method;
7564 use crate::response::StatusCode;
7565
7566 fn test_context() -> RequestContext {
7567 RequestContext::new(asupersync::Cx::for_testing(), 1)
7568 }
7569
7570 fn test_request() -> Request {
7571 Request::new(Method::Get, "/test")
7572 }
7573
7574 fn run_interceptor<I: ResponseInterceptor>(
7575 interceptor: &I,
7576 req: &Request,
7577 resp: Response,
7578 ) -> Response {
7579 let ctx = test_context();
7580 let start_time = Instant::now();
7581 let interceptor_ctx = ResponseInterceptorContext::new(req, &ctx, start_time);
7582 futures_executor::block_on(interceptor.intercept(&interceptor_ctx, resp))
7583 }
7584
7585 #[test]
7586 fn timing_interceptor_adds_header() {
7587 let interceptor = TimingInterceptor::new();
7588 let req = test_request();
7589 let resp = Response::with_status(StatusCode::OK);
7590
7591 let result = run_interceptor(&interceptor, &req, resp);
7592
7593 let has_timing = result
7594 .headers()
7595 .iter()
7596 .any(|(name, _)| name == "X-Response-Time");
7597 assert!(has_timing, "Should have X-Response-Time header");
7598 }
7599
7600 #[test]
7601 fn timing_interceptor_with_server_timing() {
7602 let interceptor = TimingInterceptor::new().with_server_timing("app");
7603 let req = test_request();
7604 let resp = Response::with_status(StatusCode::OK);
7605
7606 let result = run_interceptor(&interceptor, &req, resp);
7607
7608 let has_server_timing = result
7609 .headers()
7610 .iter()
7611 .any(|(name, _)| name == "Server-Timing");
7612 assert!(has_server_timing, "Should have Server-Timing header");
7613 }
7614
7615 #[test]
7616 fn timing_interceptor_custom_header_name() {
7617 let interceptor = TimingInterceptor::new().header_name("X-Custom-Time");
7618 let req = test_request();
7619 let resp = Response::with_status(StatusCode::OK);
7620
7621 let result = run_interceptor(&interceptor, &req, resp);
7622
7623 let has_custom = result
7624 .headers()
7625 .iter()
7626 .any(|(name, _)| name == "X-Custom-Time");
7627 assert!(has_custom, "Should have X-Custom-Time header");
7628 }
7629
7630 #[test]
7631 fn debug_info_interceptor_adds_headers() {
7632 let interceptor = DebugInfoInterceptor::new();
7633 let req = test_request();
7634 let resp = Response::with_status(StatusCode::OK);
7635
7636 let result = run_interceptor(&interceptor, &req, resp);
7637
7638 let has_path = result
7639 .headers()
7640 .iter()
7641 .any(|(name, _)| name == "X-Debug-Path");
7642 let has_method = result
7643 .headers()
7644 .iter()
7645 .any(|(name, _)| name == "X-Debug-Method");
7646 let has_timing = result
7647 .headers()
7648 .iter()
7649 .any(|(name, _)| name == "X-Debug-Handler-Time");
7650
7651 assert!(has_path, "Should have X-Debug-Path header");
7652 assert!(has_method, "Should have X-Debug-Method header");
7653 assert!(has_timing, "Should have X-Debug-Handler-Time header");
7654 }
7655
7656 #[test]
7657 fn debug_info_interceptor_custom_prefix() {
7658 let interceptor = DebugInfoInterceptor::new().header_prefix("X-Trace-");
7659 let req = test_request();
7660 let resp = Response::with_status(StatusCode::OK);
7661
7662 let result = run_interceptor(&interceptor, &req, resp);
7663
7664 let has_trace_path = result
7665 .headers()
7666 .iter()
7667 .any(|(name, _)| name == "X-Trace-Path");
7668 assert!(has_trace_path, "Should have X-Trace-Path header");
7669 }
7670
7671 #[test]
7672 fn debug_info_interceptor_selective_options() {
7673 let interceptor = DebugInfoInterceptor::new()
7674 .include_path(true)
7675 .include_method(false)
7676 .include_timing(false)
7677 .include_request_id(false);
7678 let req = test_request();
7679 let resp = Response::with_status(StatusCode::OK);
7680
7681 let result = run_interceptor(&interceptor, &req, resp);
7682
7683 let has_path = result
7684 .headers()
7685 .iter()
7686 .any(|(name, _)| name == "X-Debug-Path");
7687 let has_method = result
7688 .headers()
7689 .iter()
7690 .any(|(name, _)| name == "X-Debug-Method");
7691
7692 assert!(has_path, "Should have X-Debug-Path header");
7693 assert!(!has_method, "Should NOT have X-Debug-Method header");
7694 }
7695
7696 #[test]
7697 fn header_transform_adds_headers() {
7698 let interceptor = HeaderTransformInterceptor::new()
7699 .add("X-Powered-By", b"fastapi_rust".to_vec())
7700 .add("X-Version", b"1.0".to_vec());
7701 let req = test_request();
7702 let resp = Response::with_status(StatusCode::OK);
7703
7704 let result = run_interceptor(&interceptor, &req, resp);
7705
7706 let has_powered_by = result
7707 .headers()
7708 .iter()
7709 .any(|(name, _)| name == "X-Powered-By");
7710 let has_version = result.headers().iter().any(|(name, _)| name == "X-Version");
7711
7712 assert!(has_powered_by, "Should have X-Powered-By header");
7713 assert!(has_version, "Should have X-Version header");
7714 }
7715
7716 #[test]
7717 fn response_body_transform_modifies_body() {
7718 let transformer = ResponseBodyTransform::new(|body| {
7719 let mut result = b"[".to_vec();
7720 result.extend_from_slice(&body);
7721 result.extend_from_slice(b"]");
7722 result
7723 });
7724 let req = test_request();
7725 let resp = Response::with_status(StatusCode::OK)
7726 .body(crate::response::ResponseBody::Bytes(b"hello".to_vec()));
7727
7728 let result = run_interceptor(&transformer, &req, resp);
7729
7730 match result.body_ref() {
7731 crate::response::ResponseBody::Bytes(b) => {
7732 assert_eq!(b, b"[hello]");
7733 }
7734 _ => panic!("Expected bytes body"),
7735 }
7736 }
7737
7738 #[test]
7739 fn response_body_transform_with_content_type_filter() {
7740 let transformer =
7741 ResponseBodyTransform::new(|_| b"transformed".to_vec()).for_content_type("text/plain");
7742 let req = test_request();
7743
7744 let json_resp = Response::with_status(StatusCode::OK)
7746 .header("content-type", b"application/json".to_vec())
7747 .body(crate::response::ResponseBody::Bytes(b"original".to_vec()));
7748
7749 let result = run_interceptor(&transformer, &req, json_resp);
7750
7751 match result.body_ref() {
7752 crate::response::ResponseBody::Bytes(b) => {
7753 assert_eq!(b, b"original", "JSON should not be transformed");
7754 }
7755 _ => panic!("Expected bytes body"),
7756 }
7757
7758 let text_resp = Response::with_status(StatusCode::OK)
7760 .header("content-type", b"text/plain".to_vec())
7761 .body(crate::response::ResponseBody::Bytes(b"original".to_vec()));
7762
7763 let result = run_interceptor(&transformer, &req, text_resp);
7764
7765 match result.body_ref() {
7766 crate::response::ResponseBody::Bytes(b) => {
7767 assert_eq!(b, b"transformed", "Text should be transformed");
7768 }
7769 _ => panic!("Expected bytes body"),
7770 }
7771 }
7772
7773 #[test]
7774 fn error_response_transformer_hides_details() {
7775 let transformer = ErrorResponseTransformer::new()
7776 .hide_details_for_status(StatusCode::INTERNAL_SERVER_ERROR)
7777 .with_replacement_body(b"An error occurred");
7778
7779 let req = test_request();
7780
7781 let error_resp = Response::with_status(StatusCode::INTERNAL_SERVER_ERROR).body(
7783 crate::response::ResponseBody::Bytes(b"Sensitive error details".to_vec()),
7784 );
7785
7786 let result = run_interceptor(&transformer, &req, error_resp);
7787
7788 match result.body_ref() {
7789 crate::response::ResponseBody::Bytes(b) => {
7790 assert_eq!(b, b"An error occurred");
7791 }
7792 _ => panic!("Expected bytes body"),
7793 }
7794
7795 let ok_resp = Response::with_status(StatusCode::OK)
7797 .body(crate::response::ResponseBody::Bytes(b"Success".to_vec()));
7798
7799 let result = run_interceptor(&transformer, &req, ok_resp);
7800
7801 match result.body_ref() {
7802 crate::response::ResponseBody::Bytes(b) => {
7803 assert_eq!(b, b"Success");
7804 }
7805 _ => panic!("Expected bytes body"),
7806 }
7807 }
7808
7809 #[test]
7810 fn response_interceptor_stack_chains_interceptors() {
7811 let mut stack = ResponseInterceptorStack::new();
7812 stack.push(TimingInterceptor::new());
7813 stack.push(HeaderTransformInterceptor::new().add("X-Extra", b"value".to_vec()));
7814
7815 let req = test_request();
7816 let resp = Response::with_status(StatusCode::OK);
7817
7818 let ctx = test_context();
7819 let start_time = Instant::now();
7820 let interceptor_ctx = ResponseInterceptorContext::new(&req, &ctx, start_time);
7821 let result = futures_executor::block_on(stack.process(&interceptor_ctx, resp));
7822
7823 let has_timing = result
7824 .headers()
7825 .iter()
7826 .any(|(name, _)| name == "X-Response-Time");
7827 let has_extra = result.headers().iter().any(|(name, _)| name == "X-Extra");
7828
7829 assert!(
7830 has_timing,
7831 "Should have timing header from first interceptor"
7832 );
7833 assert!(
7834 has_extra,
7835 "Should have extra header from second interceptor"
7836 );
7837 }
7838
7839 #[test]
7840 fn response_interceptor_stack_empty_is_noop() {
7841 let stack = ResponseInterceptorStack::new();
7842 assert!(stack.is_empty());
7843 assert_eq!(stack.len(), 0);
7844
7845 let req = test_request();
7846 let resp = Response::with_status(StatusCode::OK)
7847 .body(crate::response::ResponseBody::Bytes(b"unchanged".to_vec()));
7848
7849 let ctx = test_context();
7850 let start_time = Instant::now();
7851 let interceptor_ctx = ResponseInterceptorContext::new(&req, &ctx, start_time);
7852 let result = futures_executor::block_on(stack.process(&interceptor_ctx, resp));
7853
7854 match result.body_ref() {
7855 crate::response::ResponseBody::Bytes(b) => {
7856 assert_eq!(b, b"unchanged");
7857 }
7858 _ => panic!("Expected bytes body"),
7859 }
7860 }
7861
7862 #[test]
7863 fn interceptor_context_provides_timing() {
7864 let ctx = test_context();
7865 let req = test_request();
7866 let start_time = Instant::now();
7867 std::thread::sleep(std::time::Duration::from_millis(5));
7868
7869 let interceptor_ctx = ResponseInterceptorContext::new(&req, &ctx, start_time);
7870
7871 assert!(
7872 interceptor_ctx.elapsed_ms() >= 5,
7873 "Elapsed time should be at least 5ms"
7874 );
7875 assert!(interceptor_ctx.elapsed().as_millis() >= 5);
7876 }
7877
7878 #[test]
7879 fn conditional_interceptor_applies_conditionally() {
7880 let inner = HeaderTransformInterceptor::new().add("X-Success", b"true".to_vec());
7882 let conditional =
7883 ConditionalInterceptor::new(inner, |_ctx, resp| resp.status().as_u16() == 200);
7884
7885 let req = test_request();
7886
7887 let ok_resp = Response::with_status(StatusCode::OK);
7889 let result = run_interceptor(&conditional, &req, ok_resp);
7890 let has_success = result.headers().iter().any(|(name, _)| name == "X-Success");
7891 assert!(has_success, "200 response should get X-Success header");
7892
7893 let not_found = Response::with_status(StatusCode::NOT_FOUND);
7895 let result = run_interceptor(&conditional, &req, not_found);
7896 let has_success = result.headers().iter().any(|(name, _)| name == "X-Success");
7897 assert!(!has_success, "404 response should NOT get X-Success header");
7898 }
7899}
7900
7901#[cfg(test)]
7902mod cache_control_tests {
7903 use super::*;
7904 use crate::request::Method;
7905 use crate::response::StatusCode;
7906
7907 fn test_context() -> RequestContext {
7908 RequestContext::new(asupersync::Cx::for_testing(), 1)
7909 }
7910
7911 fn run_after(mw: &CacheControlMiddleware, req: &Request, resp: Response) -> Response {
7912 let ctx = test_context();
7913 let fut = mw.after(&ctx, req, resp);
7914 futures_executor::block_on(fut)
7915 }
7916
7917 #[test]
7918 fn cache_directive_as_str_works() {
7919 assert_eq!(CacheDirective::Public.as_str(), "public");
7920 assert_eq!(CacheDirective::Private.as_str(), "private");
7921 assert_eq!(CacheDirective::NoStore.as_str(), "no-store");
7922 assert_eq!(CacheDirective::NoCache.as_str(), "no-cache");
7923 assert_eq!(CacheDirective::MustRevalidate.as_str(), "must-revalidate");
7924 assert_eq!(CacheDirective::Immutable.as_str(), "immutable");
7925 }
7926
7927 #[test]
7928 fn cache_control_builder_basic() {
7929 let cc = CacheControlBuilder::new()
7930 .public()
7931 .max_age_secs(3600)
7932 .build();
7933 assert!(cc.contains("public"));
7934 assert!(cc.contains("max-age=3600"));
7935 }
7936
7937 #[test]
7938 fn cache_control_builder_complex() {
7939 let cc = CacheControlBuilder::new()
7940 .public()
7941 .max_age_secs(60)
7942 .s_maxage_secs(3600)
7943 .stale_while_revalidate_secs(86400)
7944 .build();
7945 assert!(cc.contains("public"));
7946 assert!(cc.contains("max-age=60"));
7947 assert!(cc.contains("s-maxage=3600"));
7948 assert!(cc.contains("stale-while-revalidate=86400"));
7949 }
7950
7951 #[test]
7952 fn cache_control_builder_no_cache() {
7953 let cc = CacheControlBuilder::new()
7954 .no_store()
7955 .no_cache()
7956 .must_revalidate()
7957 .build();
7958 assert!(cc.contains("no-store"));
7959 assert!(cc.contains("no-cache"));
7960 assert!(cc.contains("must-revalidate"));
7961 }
7962
7963 #[test]
7964 fn cache_preset_no_cache() {
7965 let value = CachePreset::NoCache.to_header_value();
7966 assert!(value.contains("no-store"));
7967 assert!(value.contains("no-cache"));
7968 assert!(value.contains("must-revalidate"));
7969 }
7970
7971 #[test]
7972 fn cache_preset_immutable() {
7973 let value = CachePreset::Immutable.to_header_value();
7974 assert!(value.contains("public"));
7975 assert!(value.contains("max-age=31536000"));
7976 assert!(value.contains("immutable"));
7977 }
7978
7979 #[test]
7980 fn cache_preset_static_assets() {
7981 let value = CachePreset::StaticAssets.to_header_value();
7982 assert!(value.contains("public"));
7983 assert!(value.contains("max-age=86400"));
7984 }
7985
7986 #[test]
7987 fn middleware_adds_cache_control_header() {
7988 let mw = CacheControlMiddleware::with_preset(CachePreset::PublicOneHour);
7989 let req = Request::new(Method::Get, "/api/test");
7990 let resp = Response::with_status(StatusCode::OK);
7991
7992 let result = run_after(&mw, &req, resp);
7993 let headers = result.headers();
7994 let cc_header = headers
7995 .iter()
7996 .find(|(name, _)| name.eq_ignore_ascii_case("cache-control"));
7997 assert!(
7998 cc_header.is_some(),
7999 "Cache-Control header should be present"
8000 );
8001 let (_, value) = cc_header.unwrap();
8002 let value_str = String::from_utf8_lossy(value);
8003 assert!(value_str.contains("public"));
8004 assert!(value_str.contains("max-age=3600"));
8005 }
8006
8007 #[test]
8008 fn middleware_skips_post_requests() {
8009 let mw = CacheControlMiddleware::with_preset(CachePreset::PublicOneHour);
8010 let req = Request::new(Method::Post, "/api/test");
8011 let resp = Response::with_status(StatusCode::OK);
8012
8013 let result = run_after(&mw, &req, resp);
8014 let headers = result.headers();
8015 let cc_header = headers
8016 .iter()
8017 .find(|(name, _)| name.eq_ignore_ascii_case("cache-control"));
8018 assert!(
8019 cc_header.is_none(),
8020 "Cache-Control should not be added for POST"
8021 );
8022 }
8023
8024 #[test]
8025 fn middleware_skips_error_responses() {
8026 let mw = CacheControlMiddleware::with_preset(CachePreset::PublicOneHour);
8027 let req = Request::new(Method::Get, "/api/test");
8028 let resp = Response::with_status(StatusCode::INTERNAL_SERVER_ERROR);
8029
8030 let result = run_after(&mw, &req, resp);
8031 let headers = result.headers();
8032 let cc_header = headers
8033 .iter()
8034 .find(|(name, _)| name.eq_ignore_ascii_case("cache-control"));
8035 assert!(
8036 cc_header.is_none(),
8037 "Cache-Control should not be added for error responses"
8038 );
8039 }
8040
8041 #[test]
8042 fn middleware_with_vary_header() {
8043 let mw = CacheControlMiddleware::with_config(
8044 CacheControlConfig::from_preset(CachePreset::PublicOneHour)
8045 .vary("Accept-Encoding")
8046 .vary("Accept-Language"),
8047 );
8048 let req = Request::new(Method::Get, "/api/test");
8049 let resp = Response::with_status(StatusCode::OK);
8050
8051 let result = run_after(&mw, &req, resp);
8052 let headers = result.headers();
8053 let vary_header = headers
8054 .iter()
8055 .find(|(name, _)| name.eq_ignore_ascii_case("vary"));
8056 assert!(vary_header.is_some(), "Vary header should be present");
8057 let (_, value) = vary_header.unwrap();
8058 let value_str = String::from_utf8_lossy(value);
8059 assert!(value_str.contains("Accept-Encoding"));
8060 assert!(value_str.contains("Accept-Language"));
8061 }
8062
8063 #[test]
8064 fn middleware_preserves_existing_cache_control() {
8065 let mw = CacheControlMiddleware::with_config(
8066 CacheControlConfig::from_preset(CachePreset::PublicOneHour).preserve_existing(true),
8067 );
8068 let req = Request::new(Method::Get, "/api/test");
8069 let resp =
8070 Response::with_status(StatusCode::OK).header("Cache-Control", b"max-age=60".to_vec());
8071
8072 let result = run_after(&mw, &req, resp);
8073 let headers = result.headers();
8074 let cc_headers: Vec<_> = headers
8075 .iter()
8076 .filter(|(name, _)| name.eq_ignore_ascii_case("cache-control"))
8077 .collect();
8078 assert_eq!(cc_headers.len(), 1);
8080 let (_, value) = cc_headers[0];
8081 let value_str = String::from_utf8_lossy(value);
8082 assert_eq!(value_str, "max-age=60");
8083 }
8084
8085 #[test]
8086 fn path_pattern_matching_exact() {
8087 assert!(path_matches_pattern("/api/users", "/api/users"));
8088 assert!(!path_matches_pattern("/api/users", "/api/items"));
8089 }
8090
8091 #[test]
8092 fn path_pattern_matching_wildcard() {
8093 assert!(path_matches_pattern("/api/users/123", "/api/users/*"));
8094 assert!(path_matches_pattern("/static/css/style.css", "/static/*"));
8095 assert!(path_matches_pattern("/anything", "*"));
8096 }
8097
8098 #[test]
8099 fn date_formatting_works() {
8100 let now = std::time::SystemTime::now();
8102 let formatted = format_http_date(now);
8103 assert!(formatted.ends_with(" GMT"));
8105 let days = ["Sun", "Mon", "Tue", "Wed", "Thu", "Fri", "Sat"];
8107 assert!(days.iter().any(|d| formatted.starts_with(d)));
8108 }
8109
8110 #[test]
8111 fn leap_year_detection() {
8112 assert!(!is_leap_year(1900)); assert!(is_leap_year(2000)); assert!(is_leap_year(2024)); assert!(!is_leap_year(2023)); }
8117}
8118
8119#[cfg(test)]
8124mod trace_rejection_tests {
8125 use super::*;
8126 use crate::request::Method;
8127 use crate::response::StatusCode;
8128
8129 fn test_context() -> RequestContext {
8130 RequestContext::new(asupersync::Cx::for_testing(), 1)
8131 }
8132
8133 fn run_before(mw: &TraceRejectionMiddleware, req: &mut Request) -> ControlFlow {
8134 let ctx = test_context();
8135 let fut = mw.before(&ctx, req);
8136 futures_executor::block_on(fut)
8137 }
8138
8139 fn find_header<'a>(headers: &'a [(String, Vec<u8>)], name: &str) -> Option<&'a [u8]> {
8140 headers
8141 .iter()
8142 .find(|(n, _)| n.eq_ignore_ascii_case(name))
8143 .map(|(_, v)| v.as_slice())
8144 }
8145
8146 #[test]
8147 fn trace_request_rejected() {
8148 let mw = TraceRejectionMiddleware::new();
8149 let mut req = Request::new(Method::Trace, "/");
8150
8151 let result = run_before(&mw, &mut req);
8152
8153 match result {
8154 ControlFlow::Break(response) => {
8155 assert_eq!(response.status(), StatusCode::METHOD_NOT_ALLOWED);
8156 }
8157 ControlFlow::Continue => panic!("TRACE request should have been rejected"),
8158 }
8159 }
8160
8161 #[test]
8162 fn trace_request_with_path() {
8163 let mw = TraceRejectionMiddleware::new();
8164 let mut req = Request::new(Method::Trace, "/api/users/123");
8165
8166 let result = run_before(&mw, &mut req);
8167
8168 match result {
8169 ControlFlow::Break(response) => {
8170 assert_eq!(response.status(), StatusCode::METHOD_NOT_ALLOWED);
8171 }
8172 ControlFlow::Continue => panic!("TRACE request should have been rejected"),
8173 }
8174 }
8175
8176 #[test]
8177 fn get_request_allowed() {
8178 let mw = TraceRejectionMiddleware::new();
8179 let mut req = Request::new(Method::Get, "/");
8180
8181 let result = run_before(&mw, &mut req);
8182
8183 match result {
8184 ControlFlow::Continue => {} ControlFlow::Break(_) => panic!("GET request should be allowed"),
8186 }
8187 }
8188
8189 #[test]
8190 fn post_request_allowed() {
8191 let mw = TraceRejectionMiddleware::new();
8192 let mut req = Request::new(Method::Post, "/api/users");
8193
8194 let result = run_before(&mw, &mut req);
8195
8196 match result {
8197 ControlFlow::Continue => {} ControlFlow::Break(_) => panic!("POST request should be allowed"),
8199 }
8200 }
8201
8202 #[test]
8203 fn put_request_allowed() {
8204 let mw = TraceRejectionMiddleware::new();
8205 let mut req = Request::new(Method::Put, "/api/users/1");
8206
8207 let result = run_before(&mw, &mut req);
8208
8209 match result {
8210 ControlFlow::Continue => {} ControlFlow::Break(_) => panic!("PUT request should be allowed"),
8212 }
8213 }
8214
8215 #[test]
8216 fn delete_request_allowed() {
8217 let mw = TraceRejectionMiddleware::new();
8218 let mut req = Request::new(Method::Delete, "/api/users/1");
8219
8220 let result = run_before(&mw, &mut req);
8221
8222 match result {
8223 ControlFlow::Continue => {} ControlFlow::Break(_) => panic!("DELETE request should be allowed"),
8225 }
8226 }
8227
8228 #[test]
8229 fn patch_request_allowed() {
8230 let mw = TraceRejectionMiddleware::new();
8231 let mut req = Request::new(Method::Patch, "/api/users/1");
8232
8233 let result = run_before(&mw, &mut req);
8234
8235 match result {
8236 ControlFlow::Continue => {} ControlFlow::Break(_) => panic!("PATCH request should be allowed"),
8238 }
8239 }
8240
8241 #[test]
8242 fn options_request_allowed() {
8243 let mw = TraceRejectionMiddleware::new();
8244 let mut req = Request::new(Method::Options, "/api/users");
8245
8246 let result = run_before(&mw, &mut req);
8247
8248 match result {
8249 ControlFlow::Continue => {} ControlFlow::Break(_) => panic!("OPTIONS request should be allowed"),
8251 }
8252 }
8253
8254 #[test]
8255 fn head_request_allowed() {
8256 let mw = TraceRejectionMiddleware::new();
8257 let mut req = Request::new(Method::Head, "/");
8258
8259 let result = run_before(&mw, &mut req);
8260
8261 match result {
8262 ControlFlow::Continue => {} ControlFlow::Break(_) => panic!("HEAD request should be allowed"),
8264 }
8265 }
8266
8267 #[test]
8268 fn response_includes_allow_header() {
8269 let mw = TraceRejectionMiddleware::new();
8270 let mut req = Request::new(Method::Trace, "/");
8271
8272 let result = run_before(&mw, &mut req);
8273
8274 match result {
8275 ControlFlow::Break(response) => {
8276 let allow_header = find_header(response.headers(), "Allow");
8277 assert!(
8278 allow_header.is_some(),
8279 "Response should include Allow header"
8280 );
8281 }
8282 ControlFlow::Continue => panic!("TRACE request should have been rejected"),
8283 }
8284 }
8285
8286 #[test]
8287 fn response_has_json_content_type() {
8288 let mw = TraceRejectionMiddleware::new();
8289 let mut req = Request::new(Method::Trace, "/");
8290
8291 let result = run_before(&mw, &mut req);
8292
8293 match result {
8294 ControlFlow::Break(response) => {
8295 let ct_header = find_header(response.headers(), "Content-Type");
8296 assert_eq!(ct_header, Some(b"application/json".as_slice()));
8297 }
8298 ControlFlow::Continue => panic!("TRACE request should have been rejected"),
8299 }
8300 }
8301
8302 #[test]
8303 fn default_enables_logging() {
8304 let mw = TraceRejectionMiddleware::new();
8305 assert!(mw.log_attempts);
8306 }
8307
8308 #[test]
8309 fn log_attempts_can_be_disabled() {
8310 let mw = TraceRejectionMiddleware::new().log_attempts(false);
8311 assert!(!mw.log_attempts);
8312 }
8313
8314 #[test]
8315 fn middleware_name() {
8316 let mw = TraceRejectionMiddleware::new();
8317 assert_eq!(mw.name(), "TraceRejection");
8318 }
8319
8320 #[test]
8321 fn default_impl() {
8322 let mw = TraceRejectionMiddleware::default();
8323 assert!(mw.log_attempts);
8324 }
8325}
8326
8327#[cfg(test)]
8336mod https_redirect_tests {
8337 use super::*;
8338 use crate::request::Method;
8339 use crate::response::StatusCode;
8340
8341 fn test_context() -> RequestContext {
8342 RequestContext::new(asupersync::Cx::for_testing(), 1)
8343 }
8344
8345 fn run_before(mw: &HttpsRedirectMiddleware, req: &mut Request) -> ControlFlow {
8346 let ctx = test_context();
8347 let fut = mw.before(&ctx, req);
8348 futures_executor::block_on(fut)
8349 }
8350
8351 fn run_after(mw: &HttpsRedirectMiddleware, req: &Request, resp: Response) -> Response {
8352 let ctx = test_context();
8353 let fut = mw.after(&ctx, req, resp);
8354 futures_executor::block_on(fut)
8355 }
8356
8357 fn find_header<'a>(headers: &'a [(String, Vec<u8>)], name: &str) -> Option<&'a [u8]> {
8358 headers
8359 .iter()
8360 .find(|(n, _)| n.eq_ignore_ascii_case(name))
8361 .map(|(_, v)| v.as_slice())
8362 }
8363
8364 #[test]
8365 fn http_request_redirected() {
8366 let mw = HttpsRedirectMiddleware::new();
8367 let mut req = Request::new(Method::Get, "/");
8368 req.headers_mut().insert("Host", b"example.com".to_vec());
8369
8370 let result = run_before(&mw, &mut req);
8371
8372 match result {
8373 ControlFlow::Break(response) => {
8374 assert_eq!(response.status(), StatusCode::MOVED_PERMANENTLY);
8375 let location = find_header(response.headers(), "Location");
8376 assert_eq!(location, Some(b"https://example.com/".as_slice()));
8377 }
8378 ControlFlow::Continue => panic!("HTTP request should be redirected"),
8379 }
8380 }
8381
8382 #[test]
8383 fn http_request_with_path_and_query() {
8384 let mw = HttpsRedirectMiddleware::new();
8385 let mut req = Request::new(Method::Get, "/api/users?page=1");
8386 req.headers_mut().insert("Host", b"example.com".to_vec());
8387
8388 let result = run_before(&mw, &mut req);
8389
8390 match result {
8391 ControlFlow::Break(response) => {
8392 let location = find_header(response.headers(), "Location");
8393 assert_eq!(
8394 location,
8395 Some(b"https://example.com/api/users?page=1".as_slice())
8396 );
8397 }
8398 ControlFlow::Continue => panic!("HTTP request should be redirected"),
8399 }
8400 }
8401
8402 #[test]
8403 fn https_request_not_redirected() {
8404 let mw = HttpsRedirectMiddleware::new();
8405 let mut req = Request::new(Method::Get, "/");
8406 req.headers_mut().insert("Host", b"example.com".to_vec());
8407 req.headers_mut()
8408 .insert("X-Forwarded-Proto", b"https".to_vec());
8409
8410 let result = run_before(&mw, &mut req);
8411
8412 match result {
8413 ControlFlow::Continue => {} ControlFlow::Break(_) => panic!("HTTPS request should not be redirected"),
8415 }
8416 }
8417
8418 #[test]
8419 fn x_forwarded_ssl_recognized() {
8420 let mw = HttpsRedirectMiddleware::new();
8421 let mut req = Request::new(Method::Get, "/");
8422 req.headers_mut().insert("Host", b"example.com".to_vec());
8423 req.headers_mut().insert("X-Forwarded-Ssl", b"on".to_vec());
8424
8425 let result = run_before(&mw, &mut req);
8426
8427 match result {
8428 ControlFlow::Continue => {} ControlFlow::Break(_) => panic!("Request with X-Forwarded-Ssl=on should not redirect"),
8430 }
8431 }
8432
8433 #[test]
8434 fn excluded_path_not_redirected() {
8435 let mw = HttpsRedirectMiddleware::new().exclude_path("/health");
8436 let mut req = Request::new(Method::Get, "/health");
8437 req.headers_mut().insert("Host", b"example.com".to_vec());
8438
8439 let result = run_before(&mw, &mut req);
8440
8441 match result {
8442 ControlFlow::Continue => {} ControlFlow::Break(_) => panic!("Excluded path should not be redirected"),
8444 }
8445 }
8446
8447 #[test]
8448 fn excluded_path_prefix_matches() {
8449 let mw = HttpsRedirectMiddleware::new().exclude_path("/health");
8450 let mut req = Request::new(Method::Get, "/health/live");
8451 req.headers_mut().insert("Host", b"example.com".to_vec());
8452
8453 let result = run_before(&mw, &mut req);
8454
8455 match result {
8456 ControlFlow::Continue => {} ControlFlow::Break(_) => panic!("Path with excluded prefix should not be redirected"),
8458 }
8459 }
8460
8461 #[test]
8462 fn temporary_redirect_option() {
8463 let mw = HttpsRedirectMiddleware::new().permanent_redirect(false);
8464 let mut req = Request::new(Method::Get, "/");
8465 req.headers_mut().insert("Host", b"example.com".to_vec());
8466
8467 let result = run_before(&mw, &mut req);
8468
8469 match result {
8470 ControlFlow::Break(response) => {
8471 assert_eq!(response.status(), StatusCode::TEMPORARY_REDIRECT);
8472 }
8473 ControlFlow::Continue => panic!("HTTP request should be redirected"),
8474 }
8475 }
8476
8477 #[test]
8478 fn redirect_disabled() {
8479 let mw = HttpsRedirectMiddleware::new().redirect_enabled(false);
8480 let mut req = Request::new(Method::Get, "/");
8481 req.headers_mut().insert("Host", b"example.com".to_vec());
8482
8483 let result = run_before(&mw, &mut req);
8484
8485 match result {
8486 ControlFlow::Continue => {} ControlFlow::Break(_) => panic!("Redirects are disabled, should continue"),
8488 }
8489 }
8490
8491 #[test]
8492 fn hsts_header_on_https_response() {
8493 let mw = HttpsRedirectMiddleware::new();
8494 let mut req = Request::new(Method::Get, "/");
8495 req.headers_mut()
8496 .insert("X-Forwarded-Proto", b"https".to_vec());
8497
8498 let response = Response::with_status(StatusCode::OK);
8499 let result = run_after(&mw, &req, response);
8500
8501 let hsts = find_header(result.headers(), "Strict-Transport-Security");
8502 assert!(
8503 hsts.is_some(),
8504 "HSTS header should be present on HTTPS response"
8505 );
8506 let hsts_str = String::from_utf8_lossy(hsts.unwrap());
8507 assert!(hsts_str.contains("max-age=31536000"));
8508 }
8509
8510 #[test]
8511 fn hsts_header_not_on_http_response() {
8512 let mw = HttpsRedirectMiddleware::new().redirect_enabled(false);
8513 let req = Request::new(Method::Get, "/");
8514 let response = Response::with_status(StatusCode::OK);
8517 let result = run_after(&mw, &req, response);
8518
8519 let hsts = find_header(result.headers(), "Strict-Transport-Security");
8520 assert!(hsts.is_none(), "HSTS header should not be on HTTP response");
8521 }
8522
8523 #[test]
8524 fn hsts_with_include_subdomains() {
8525 let mw = HttpsRedirectMiddleware::new().include_subdomains(true);
8526 let mut req = Request::new(Method::Get, "/");
8527 req.headers_mut()
8528 .insert("X-Forwarded-Proto", b"https".to_vec());
8529
8530 let response = Response::with_status(StatusCode::OK);
8531 let result = run_after(&mw, &req, response);
8532
8533 let hsts = find_header(result.headers(), "Strict-Transport-Security");
8534 let hsts_str = String::from_utf8_lossy(hsts.unwrap());
8535 assert!(hsts_str.contains("includeSubDomains"));
8536 }
8537
8538 #[test]
8539 fn hsts_with_preload() {
8540 let mw = HttpsRedirectMiddleware::new().preload(true);
8541 let mut req = Request::new(Method::Get, "/");
8542 req.headers_mut()
8543 .insert("X-Forwarded-Proto", b"https".to_vec());
8544
8545 let response = Response::with_status(StatusCode::OK);
8546 let result = run_after(&mw, &req, response);
8547
8548 let hsts = find_header(result.headers(), "Strict-Transport-Security");
8549 let hsts_str = String::from_utf8_lossy(hsts.unwrap());
8550 assert!(hsts_str.contains("preload"));
8551 }
8552
8553 #[test]
8554 fn hsts_disabled_with_zero_max_age() {
8555 let mw = HttpsRedirectMiddleware::new().hsts_max_age_secs(0);
8556 let mut req = Request::new(Method::Get, "/");
8557 req.headers_mut()
8558 .insert("X-Forwarded-Proto", b"https".to_vec());
8559
8560 let response = Response::with_status(StatusCode::OK);
8561 let result = run_after(&mw, &req, response);
8562
8563 let hsts = find_header(result.headers(), "Strict-Transport-Security");
8564 assert!(hsts.is_none(), "HSTS should be disabled with max-age=0");
8565 }
8566
8567 #[test]
8568 fn custom_https_port() {
8569 let mw = HttpsRedirectMiddleware::new().https_port(8443);
8570 let mut req = Request::new(Method::Get, "/");
8571 req.headers_mut().insert("Host", b"example.com".to_vec());
8572
8573 let result = run_before(&mw, &mut req);
8574
8575 match result {
8576 ControlFlow::Break(response) => {
8577 let location = find_header(response.headers(), "Location");
8578 assert_eq!(location, Some(b"https://example.com:8443/".as_slice()));
8579 }
8580 ControlFlow::Continue => panic!("HTTP request should be redirected"),
8581 }
8582 }
8583
8584 #[test]
8585 fn host_with_port_stripped() {
8586 let mw = HttpsRedirectMiddleware::new();
8587 let mut req = Request::new(Method::Get, "/");
8588 req.headers_mut()
8589 .insert("Host", b"example.com:8080".to_vec());
8590
8591 let result = run_before(&mw, &mut req);
8592
8593 match result {
8594 ControlFlow::Break(response) => {
8595 let location = find_header(response.headers(), "Location");
8596 assert_eq!(location, Some(b"https://example.com/".as_slice()));
8598 }
8599 ControlFlow::Continue => panic!("HTTP request should be redirected"),
8600 }
8601 }
8602
8603 #[test]
8604 fn middleware_name() {
8605 let mw = HttpsRedirectMiddleware::new();
8606 assert_eq!(mw.name(), "HttpsRedirect");
8607 }
8608
8609 #[test]
8610 fn default_impl() {
8611 let mw = HttpsRedirectMiddleware::default();
8612 assert!(mw.config.redirect_enabled);
8613 assert!(mw.config.permanent_redirect);
8614 assert_eq!(mw.config.hsts_max_age_secs, 31_536_000);
8615 }
8616
8617 #[test]
8618 fn config_builder() {
8619 let mw = HttpsRedirectMiddleware::new()
8620 .redirect_enabled(false)
8621 .permanent_redirect(false)
8622 .hsts_max_age_secs(86400)
8623 .include_subdomains(true)
8624 .preload(true)
8625 .https_port(8443);
8626
8627 assert!(!mw.config.redirect_enabled);
8628 assert!(!mw.config.permanent_redirect);
8629 assert_eq!(mw.config.hsts_max_age_secs, 86400);
8630 assert!(mw.config.hsts_include_subdomains);
8631 assert!(mw.config.hsts_preload);
8632 assert_eq!(mw.config.https_port, 8443);
8633 }
8634
8635 #[test]
8636 fn exclude_paths_method() {
8637 let mw = HttpsRedirectMiddleware::new()
8638 .exclude_paths(vec!["/health".to_string(), "/ready".to_string()]);
8639
8640 assert_eq!(mw.config.exclude_paths.len(), 2);
8641 assert!(mw.config.exclude_paths.contains(&"/health".to_string()));
8642 assert!(mw.config.exclude_paths.contains(&"/ready".to_string()));
8643 }
8644}
8645
8646#[cfg(test)]
8655mod tests {
8656 use super::*;
8657 use crate::response::{ResponseBody, StatusCode};
8658
8659 #[allow(dead_code)]
8661 struct AddHeaderMiddleware {
8662 name: &'static str,
8663 value: &'static [u8],
8664 }
8665
8666 impl Middleware for AddHeaderMiddleware {
8667 fn after<'a>(
8668 &'a self,
8669 _ctx: &'a RequestContext,
8670 _req: &'a Request,
8671 response: Response,
8672 ) -> BoxFuture<'a, Response> {
8673 Box::pin(async move { response.header(self.name, self.value.to_vec()) })
8674 }
8675 }
8676
8677 #[allow(dead_code)]
8679 struct BlockingMiddleware;
8680
8681 impl Middleware for BlockingMiddleware {
8682 fn before<'a>(
8683 &'a self,
8684 _ctx: &'a RequestContext,
8685 _req: &'a mut Request,
8686 ) -> BoxFuture<'a, ControlFlow> {
8687 Box::pin(async {
8688 ControlFlow::Break(
8689 Response::with_status(StatusCode::FORBIDDEN)
8690 .body(ResponseBody::Bytes(b"blocked".to_vec())),
8691 )
8692 })
8693 }
8694 }
8695
8696 #[allow(dead_code)]
8698 struct TrackingMiddleware {
8699 before_count: std::sync::atomic::AtomicUsize,
8700 after_count: std::sync::atomic::AtomicUsize,
8701 }
8702
8703 #[allow(dead_code)]
8704 impl TrackingMiddleware {
8705 fn new() -> Self {
8706 Self {
8707 before_count: std::sync::atomic::AtomicUsize::new(0),
8708 after_count: std::sync::atomic::AtomicUsize::new(0),
8709 }
8710 }
8711
8712 fn before_count(&self) -> usize {
8713 self.before_count.load(std::sync::atomic::Ordering::SeqCst)
8714 }
8715
8716 fn after_count(&self) -> usize {
8717 self.after_count.load(std::sync::atomic::Ordering::SeqCst)
8718 }
8719 }
8720
8721 impl Middleware for TrackingMiddleware {
8722 fn before<'a>(
8723 &'a self,
8724 _ctx: &'a RequestContext,
8725 _req: &'a mut Request,
8726 ) -> BoxFuture<'a, ControlFlow> {
8727 self.before_count
8728 .fetch_add(1, std::sync::atomic::Ordering::SeqCst);
8729 Box::pin(async { ControlFlow::Continue })
8730 }
8731
8732 fn after<'a>(
8733 &'a self,
8734 _ctx: &'a RequestContext,
8735 _req: &'a Request,
8736 response: Response,
8737 ) -> BoxFuture<'a, Response> {
8738 self.after_count
8739 .fetch_add(1, std::sync::atomic::Ordering::SeqCst);
8740 Box::pin(async move { response })
8741 }
8742 }
8743
8744 #[test]
8745 fn control_flow_variants() {
8746 let cont = ControlFlow::Continue;
8747 assert!(cont.is_continue());
8748 assert!(!cont.is_break());
8749
8750 let brk = ControlFlow::Break(Response::ok());
8751 assert!(!brk.is_continue());
8752 assert!(brk.is_break());
8753 }
8754
8755 #[test]
8756 fn middleware_stack_empty() {
8757 let stack = MiddlewareStack::new();
8758 assert!(stack.is_empty());
8759 assert_eq!(stack.len(), 0);
8760 }
8761
8762 #[test]
8763 fn middleware_stack_push() {
8764 let mut stack = MiddlewareStack::new();
8765 stack.push(NoopMiddleware);
8766 stack.push(NoopMiddleware);
8767 assert_eq!(stack.len(), 2);
8768 assert!(!stack.is_empty());
8769 }
8770
8771 #[test]
8772 fn noop_middleware_name() {
8773 let mw = NoopMiddleware;
8774 assert_eq!(mw.name(), "Noop");
8775 }
8776
8777 #[test]
8778 fn logging_redacts_sensitive_headers() {
8779 let mut headers = crate::request::Headers::new();
8780 headers.insert("Authorization", b"secret".to_vec());
8781 headers.insert("X-Request-Id", b"abc123".to_vec());
8782
8783 let redacted = super::default_redacted_headers();
8784 let formatted = super::format_headers(headers.iter(), &redacted);
8785
8786 assert!(formatted.contains("authorization=<redacted>"));
8787 assert!(formatted.contains("x-request-id=abc123"));
8788 }
8789
8790 #[test]
8791 fn logging_body_truncation() {
8792 let body = b"abcdef";
8793 let preview = super::format_bytes(body, 4);
8794 assert_eq!(preview, "abcd...");
8795
8796 let preview_full = super::format_bytes(body, 10);
8797 assert_eq!(preview_full, "abcdef");
8798 }
8799
8800 fn test_context() -> RequestContext {
8801 let cx = asupersync::Cx::for_testing();
8802 RequestContext::new(cx, 1)
8803 }
8804
8805 fn header_value(response: &Response, name: &str) -> Option<String> {
8806 response
8807 .headers()
8808 .iter()
8809 .find(|(n, _)| n.eq_ignore_ascii_case(name))
8810 .and_then(|(_, v)| std::str::from_utf8(v).ok())
8811 .map(ToString::to_string)
8812 }
8813
8814 #[test]
8815 fn cors_exact_origin_allows() {
8816 let cors = Cors::new().allow_origin("https://example.com");
8817 let ctx = test_context();
8818 let mut req = Request::new(crate::request::Method::Get, "/");
8819 req.headers_mut()
8820 .insert("origin", b"https://example.com".to_vec());
8821
8822 let result = futures_executor::block_on(cors.before(&ctx, &mut req));
8823 assert!(matches!(result, ControlFlow::Continue));
8824
8825 let response = Response::ok().body(ResponseBody::Bytes(b"ok".to_vec()));
8826 let response = futures_executor::block_on(cors.after(&ctx, &req, response));
8827
8828 assert_eq!(
8829 header_value(&response, "access-control-allow-origin"),
8830 Some("https://example.com".to_string())
8831 );
8832 assert_eq!(header_value(&response, "vary"), Some("Origin".to_string()));
8833 }
8834
8835 #[test]
8836 fn cors_wildcard_origin_allows() {
8837 let cors = Cors::new().allow_origin_wildcard("https://*.example.com");
8838 let ctx = test_context();
8839 let mut req = Request::new(crate::request::Method::Get, "/");
8840 req.headers_mut()
8841 .insert("origin", b"https://api.example.com".to_vec());
8842
8843 let result = futures_executor::block_on(cors.before(&ctx, &mut req));
8844 assert!(matches!(result, ControlFlow::Continue));
8845 }
8846
8847 #[test]
8848 fn cors_regex_origin_allows() {
8849 let cors = Cors::new().allow_origin_regex(r"^https://.*\.example\.com$");
8850 let ctx = test_context();
8851 let mut req = Request::new(crate::request::Method::Get, "/");
8852 req.headers_mut()
8853 .insert("origin", b"https://svc.example.com".to_vec());
8854
8855 let result = futures_executor::block_on(cors.before(&ctx, &mut req));
8856 assert!(matches!(result, ControlFlow::Continue));
8857 }
8858
8859 #[test]
8860 fn cors_preflight_handled() {
8861 let cors = Cors::new()
8862 .allow_any_origin()
8863 .allow_headers(["x-test", "content-type"])
8864 .max_age(600);
8865 let ctx = test_context();
8866 let mut req = Request::new(crate::request::Method::Options, "/");
8867 req.headers_mut()
8868 .insert("origin", b"https://example.com".to_vec());
8869 req.headers_mut()
8870 .insert("access-control-request-method", b"POST".to_vec());
8871 req.headers_mut().insert(
8872 "access-control-request-headers",
8873 b"x-test, content-type".to_vec(),
8874 );
8875
8876 let result = futures_executor::block_on(cors.before(&ctx, &mut req));
8877 let ControlFlow::Break(response) = result else {
8878 panic!("expected preflight break");
8879 };
8880
8881 assert_eq!(response.status().as_u16(), 204);
8882 assert_eq!(
8883 header_value(&response, "access-control-allow-origin"),
8884 Some("*".to_string())
8885 );
8886 assert_eq!(
8887 header_value(&response, "access-control-allow-methods"),
8888 Some("GET, POST, PUT, PATCH, DELETE, OPTIONS, HEAD".to_string())
8889 );
8890 assert_eq!(
8891 header_value(&response, "access-control-allow-headers"),
8892 Some("x-test, content-type".to_string())
8893 );
8894 assert_eq!(
8895 header_value(&response, "access-control-max-age"),
8896 Some("600".to_string())
8897 );
8898 }
8899
8900 #[test]
8901 fn cors_credentials_echo_origin() {
8902 let cors = Cors::new().allow_any_origin().allow_credentials(true);
8903 let ctx = test_context();
8904 let mut req = Request::new(crate::request::Method::Get, "/");
8905 req.headers_mut()
8906 .insert("origin", b"https://example.com".to_vec());
8907
8908 let result = futures_executor::block_on(cors.before(&ctx, &mut req));
8909 assert!(matches!(result, ControlFlow::Continue));
8910
8911 let response = futures_executor::block_on(cors.after(&ctx, &req, Response::ok()));
8912 assert_eq!(
8913 header_value(&response, "access-control-allow-origin"),
8914 Some("https://example.com".to_string())
8915 );
8916 assert_eq!(
8917 header_value(&response, "access-control-allow-credentials"),
8918 Some("true".to_string())
8919 );
8920 }
8921
8922 #[test]
8927 fn cors_spec_compliance_credentials_never_wildcard_origin() {
8928 let cors = Cors::new().allow_any_origin().allow_credentials(true);
8931 let ctx = test_context();
8932
8933 for origin in &[
8935 "https://example.com",
8936 "https://api.example.com",
8937 "http://localhost:3000",
8938 ] {
8939 let mut req = Request::new(crate::request::Method::Get, "/");
8940 req.headers_mut()
8941 .insert("origin", origin.as_bytes().to_vec());
8942
8943 futures_executor::block_on(cors.before(&ctx, &mut req));
8944 let response = futures_executor::block_on(cors.after(&ctx, &req, Response::ok()));
8945
8946 let allow_origin = header_value(&response, "access-control-allow-origin");
8947 assert_eq!(
8948 allow_origin,
8949 Some((*origin).to_string()),
8950 "With credentials enabled, Access-Control-Allow-Origin must echo '{}', not '*'",
8951 origin
8952 );
8953 assert_ne!(
8954 allow_origin,
8955 Some("*".to_string()),
8956 "CORS spec violation: credentials + wildcard origin is forbidden"
8957 );
8958 }
8959 }
8960
8961 #[test]
8962 fn cors_spec_compliance_preflight_with_credentials() {
8963 let cors = Cors::new()
8965 .allow_any_origin()
8966 .allow_credentials(true)
8967 .allow_headers(["content-type", "x-custom-header"]);
8968 let ctx = test_context();
8969
8970 let mut req = Request::new(crate::request::Method::Options, "/");
8971 req.headers_mut()
8972 .insert("origin", b"https://example.com".to_vec());
8973 req.headers_mut()
8974 .insert("access-control-request-method", b"POST".to_vec());
8975 req.headers_mut()
8976 .insert("access-control-request-headers", b"content-type".to_vec());
8977
8978 let result = futures_executor::block_on(cors.before(&ctx, &mut req));
8979 let ControlFlow::Break(response) = result else {
8980 panic!("expected preflight break");
8981 };
8982
8983 let allow_origin = header_value(&response, "access-control-allow-origin");
8985 assert_eq!(allow_origin, Some("https://example.com".to_string()));
8986 assert_ne!(
8987 allow_origin,
8988 Some("*".to_string()),
8989 "CORS spec violation: preflight with credentials must not use wildcard origin"
8990 );
8991
8992 assert_eq!(
8994 header_value(&response, "access-control-allow-credentials"),
8995 Some("true".to_string())
8996 );
8997 }
8998
8999 #[test]
9000 fn cors_spec_without_credentials_allows_wildcard() {
9001 let cors = Cors::new().allow_any_origin();
9003 let ctx = test_context();
9004 let mut req = Request::new(crate::request::Method::Get, "/");
9005 req.headers_mut()
9006 .insert("origin", b"https://example.com".to_vec());
9007
9008 futures_executor::block_on(cors.before(&ctx, &mut req));
9009 let response = futures_executor::block_on(cors.after(&ctx, &req, Response::ok()));
9010
9011 assert_eq!(
9013 header_value(&response, "access-control-allow-origin"),
9014 Some("*".to_string())
9015 );
9016 assert!(header_value(&response, "access-control-allow-credentials").is_none());
9018 }
9019
9020 #[test]
9021 fn cors_disallowed_preflight_forbidden() {
9022 let cors = Cors::new().allow_origin("https://good.example");
9023 let ctx = test_context();
9024 let mut req = Request::new(crate::request::Method::Options, "/");
9025 req.headers_mut()
9026 .insert("origin", b"https://evil.example".to_vec());
9027 req.headers_mut()
9028 .insert("access-control-request-method", b"GET".to_vec());
9029
9030 let result = futures_executor::block_on(cors.before(&ctx, &mut req));
9031 let ControlFlow::Break(response) = result else {
9032 panic!("expected forbidden preflight");
9033 };
9034 assert_eq!(response.status().as_u16(), 403);
9035 }
9036
9037 #[test]
9038 fn cors_simple_request_disallowed_origin_no_headers() {
9039 let cors = Cors::new().allow_origin("https://good.example");
9041 let ctx = test_context();
9042 let mut req = Request::new(crate::request::Method::Get, "/");
9043 req.headers_mut()
9044 .insert("origin", b"https://evil.example".to_vec());
9045
9046 let result = futures_executor::block_on(cors.before(&ctx, &mut req));
9047 assert!(matches!(result, ControlFlow::Continue));
9049
9050 let response = futures_executor::block_on(cors.after(&ctx, &req, Response::ok()));
9051 assert!(header_value(&response, "access-control-allow-origin").is_none());
9053 }
9054
9055 #[test]
9056 fn cors_expose_headers_configuration() {
9057 let cors = Cors::new()
9058 .allow_any_origin()
9059 .expose_headers(["x-custom-header", "x-another-header"]);
9060 let ctx = test_context();
9061 let mut req = Request::new(crate::request::Method::Get, "/");
9062 req.headers_mut()
9063 .insert("origin", b"https://example.com".to_vec());
9064
9065 let result = futures_executor::block_on(cors.before(&ctx, &mut req));
9066 assert!(matches!(result, ControlFlow::Continue));
9067
9068 let response = futures_executor::block_on(cors.after(&ctx, &req, Response::ok()));
9069 assert_eq!(
9070 header_value(&response, "access-control-expose-headers"),
9071 Some("x-custom-header, x-another-header".to_string())
9072 );
9073 }
9074
9075 #[test]
9076 fn cors_any_origin_sets_wildcard() {
9077 let cors = Cors::new().allow_any_origin();
9078 let ctx = test_context();
9079 let mut req = Request::new(crate::request::Method::Get, "/");
9080 req.headers_mut()
9081 .insert("origin", b"https://any-site.com".to_vec());
9082
9083 let result = futures_executor::block_on(cors.before(&ctx, &mut req));
9084 assert!(matches!(result, ControlFlow::Continue));
9085
9086 let response = futures_executor::block_on(cors.after(&ctx, &req, Response::ok()));
9087 assert_eq!(
9088 header_value(&response, "access-control-allow-origin"),
9089 Some("*".to_string())
9090 );
9091 }
9092
9093 #[test]
9094 fn cors_config_allows_method_override() {
9095 let cors = Cors::new()
9097 .allow_any_origin()
9098 .allow_methods([crate::request::Method::Get, crate::request::Method::Post]);
9099 let ctx = test_context();
9100 let mut req = Request::new(crate::request::Method::Options, "/");
9101 req.headers_mut()
9102 .insert("origin", b"https://example.com".to_vec());
9103 req.headers_mut()
9104 .insert("access-control-request-method", b"POST".to_vec());
9105
9106 let result = futures_executor::block_on(cors.before(&ctx, &mut req));
9107 let ControlFlow::Break(response) = result else {
9108 panic!("expected preflight break");
9109 };
9110 assert_eq!(
9111 header_value(&response, "access-control-allow-methods"),
9112 Some("GET, POST".to_string())
9113 );
9114 }
9115
9116 #[test]
9117 fn cors_no_origin_header_skips_cors() {
9118 let cors = Cors::new().allow_any_origin();
9120 let ctx = test_context();
9121 let mut req = Request::new(crate::request::Method::Get, "/");
9122
9123 let result = futures_executor::block_on(cors.before(&ctx, &mut req));
9124 assert!(matches!(result, ControlFlow::Continue));
9125
9126 let response = futures_executor::block_on(cors.after(&ctx, &req, Response::ok()));
9127 assert!(header_value(&response, "access-control-allow-origin").is_none());
9128 }
9129
9130 #[test]
9131 fn cors_middleware_name() {
9132 let cors = Cors::new();
9133 assert_eq!(cors.name(), "Cors");
9134 }
9135
9136 #[test]
9137 fn cors_empty_allowed_headers_does_not_reflect_request_headers() {
9138 let cors = Cors::new().allow_any_origin(); let ctx = test_context();
9143 let mut req = Request::new(crate::request::Method::Options, "/api");
9144 req.headers_mut()
9145 .insert("origin", b"https://example.com".to_vec());
9146 req.headers_mut()
9147 .insert("access-control-request-method", b"GET".to_vec());
9148 req.headers_mut().insert(
9149 "access-control-request-headers",
9150 b"x-evil-custom, authorization".to_vec(),
9151 );
9152
9153 let result = futures_executor::block_on(cors.before(&ctx, &mut req));
9154 if let ControlFlow::Break(response) = result {
9155 assert_eq!(
9158 header_value(&response, "access-control-allow-headers"),
9159 None,
9160 "Empty allowed_headers must not reflect request headers"
9161 );
9162 } else {
9163 panic!("Preflight should have been handled (Break)");
9164 }
9165 }
9166
9167 #[test]
9168 fn cors_explicit_allowed_headers_returned_in_preflight() {
9169 let cors = Cors::new()
9170 .allow_any_origin()
9171 .allow_headers(["x-token", "content-type"]);
9172 let ctx = test_context();
9173 let mut req = Request::new(crate::request::Method::Options, "/api");
9174 req.headers_mut()
9175 .insert("origin", b"https://example.com".to_vec());
9176 req.headers_mut()
9177 .insert("access-control-request-method", b"POST".to_vec());
9178
9179 let result = futures_executor::block_on(cors.before(&ctx, &mut req));
9180 if let ControlFlow::Break(response) = result {
9181 let headers_val = header_value(&response, "access-control-allow-headers");
9182 assert!(headers_val.is_some());
9183 let val = headers_val.unwrap();
9184 assert!(val.contains("x-token"));
9185 assert!(val.contains("content-type"));
9186 } else {
9187 panic!("Preflight should have been handled (Break)");
9188 }
9189 }
9190
9191 #[test]
9196 fn request_id_generates_unique_ids() {
9197 let id1 = RequestId::generate();
9198 let id2 = RequestId::generate();
9199 let id3 = RequestId::generate();
9200
9201 assert_ne!(id1, id2);
9202 assert_ne!(id2, id3);
9203 assert_ne!(id1, id3);
9204
9205 assert!(!id1.as_str().is_empty());
9207 assert!(!id2.as_str().is_empty());
9208 assert!(!id3.as_str().is_empty());
9209 }
9210
9211 #[test]
9212 fn request_id_display() {
9213 let id = RequestId::new("test-request-123");
9214 assert_eq!(format!("{}", id), "test-request-123");
9215 }
9216
9217 #[test]
9218 fn request_id_from_string() {
9219 let id: RequestId = "my-id".into();
9220 assert_eq!(id.as_str(), "my-id");
9221
9222 let id2: RequestId = String::from("my-id-2").into();
9223 assert_eq!(id2.as_str(), "my-id-2");
9224 }
9225
9226 #[test]
9227 fn request_id_config_defaults() {
9228 let config = RequestIdConfig::default();
9229 assert_eq!(config.header_name, "x-request-id");
9230 assert!(config.accept_from_client);
9231 assert!(config.add_to_response);
9232 assert_eq!(config.max_client_id_length, 128);
9233 }
9234
9235 #[test]
9236 fn request_id_config_builder() {
9237 let config = RequestIdConfig::new()
9238 .header_name("X-Trace-ID")
9239 .accept_from_client(false)
9240 .add_to_response(false)
9241 .max_client_id_length(64);
9242
9243 assert_eq!(config.header_name, "X-Trace-ID");
9244 assert!(!config.accept_from_client);
9245 assert!(!config.add_to_response);
9246 assert_eq!(config.max_client_id_length, 64);
9247 }
9248
9249 #[test]
9250 fn request_id_middleware_generates_id() {
9251 let middleware = RequestIdMiddleware::new();
9252 let ctx = test_context();
9253 let mut req = Request::new(crate::request::Method::Get, "/");
9254
9255 let result = futures_executor::block_on(middleware.before(&ctx, &mut req));
9256 assert!(matches!(result, ControlFlow::Continue));
9257
9258 let stored_id = req.get_extension::<RequestId>();
9259 assert!(stored_id.is_some());
9260 assert!(!stored_id.unwrap().as_str().is_empty());
9261 }
9262
9263 #[test]
9264 fn request_id_middleware_accepts_client_id() {
9265 let middleware = RequestIdMiddleware::new();
9266 let ctx = test_context();
9267 let mut req = Request::new(crate::request::Method::Get, "/");
9268 req.headers_mut()
9269 .insert("x-request-id", b"client-provided-id-123".to_vec());
9270
9271 futures_executor::block_on(middleware.before(&ctx, &mut req));
9272
9273 let stored_id = req.get_extension::<RequestId>().unwrap();
9274 assert_eq!(stored_id.as_str(), "client-provided-id-123");
9275 }
9276
9277 #[test]
9278 fn request_id_middleware_rejects_invalid_client_id() {
9279 let middleware = RequestIdMiddleware::new();
9280 let ctx = test_context();
9281
9282 let mut req = Request::new(crate::request::Method::Get, "/");
9284 req.headers_mut()
9285 .insert("x-request-id", b"invalid<script>id".to_vec());
9286
9287 futures_executor::block_on(middleware.before(&ctx, &mut req));
9288
9289 let stored_id = req.get_extension::<RequestId>().unwrap();
9290 assert_ne!(stored_id.as_str(), "invalid<script>id");
9292 }
9293
9294 #[test]
9295 fn request_id_middleware_rejects_too_long_client_id() {
9296 let config = RequestIdConfig::new().max_client_id_length(10);
9297 let middleware = RequestIdMiddleware::with_config(config);
9298 let ctx = test_context();
9299
9300 let mut req = Request::new(crate::request::Method::Get, "/");
9301 req.headers_mut()
9302 .insert("x-request-id", b"this-id-is-way-too-long".to_vec());
9303
9304 futures_executor::block_on(middleware.before(&ctx, &mut req));
9305
9306 let stored_id = req.get_extension::<RequestId>().unwrap();
9307 assert_ne!(stored_id.as_str(), "this-id-is-way-too-long");
9309 }
9310
9311 #[test]
9312 fn request_id_middleware_adds_to_response() {
9313 let middleware = RequestIdMiddleware::new();
9314 let ctx = test_context();
9315 let mut req = Request::new(crate::request::Method::Get, "/");
9316
9317 futures_executor::block_on(middleware.before(&ctx, &mut req));
9318 let stored_id = req.get_extension::<RequestId>().unwrap().clone();
9319
9320 let response = Response::ok();
9321 let response = futures_executor::block_on(middleware.after(&ctx, &req, response));
9322
9323 let header = header_value(&response, "x-request-id");
9324 assert_eq!(header, Some(stored_id.0));
9325 }
9326
9327 #[test]
9328 fn request_id_middleware_respects_add_to_response_false() {
9329 let config = RequestIdConfig::new().add_to_response(false);
9330 let middleware = RequestIdMiddleware::with_config(config);
9331 let ctx = test_context();
9332 let mut req = Request::new(crate::request::Method::Get, "/");
9333
9334 futures_executor::block_on(middleware.before(&ctx, &mut req));
9335
9336 let response = Response::ok();
9337 let response = futures_executor::block_on(middleware.after(&ctx, &req, response));
9338
9339 let header = header_value(&response, "x-request-id");
9340 assert!(header.is_none());
9341 }
9342
9343 #[test]
9344 fn request_id_middleware_respects_accept_from_client_false() {
9345 let config = RequestIdConfig::new().accept_from_client(false);
9346 let middleware = RequestIdMiddleware::with_config(config);
9347 let ctx = test_context();
9348 let mut req = Request::new(crate::request::Method::Get, "/");
9349 req.headers_mut()
9350 .insert("x-request-id", b"client-id".to_vec());
9351
9352 futures_executor::block_on(middleware.before(&ctx, &mut req));
9353
9354 let stored_id = req.get_extension::<RequestId>().unwrap();
9355 assert_ne!(stored_id.as_str(), "client-id");
9357 }
9358
9359 #[test]
9360 fn request_id_middleware_custom_header_name() {
9361 let config = RequestIdConfig::new().header_name("X-Trace-ID");
9362 let middleware = RequestIdMiddleware::with_config(config);
9363 let ctx = test_context();
9364 let mut req = Request::new(crate::request::Method::Get, "/");
9365 req.headers_mut()
9366 .insert("X-Trace-ID", b"trace-123".to_vec());
9367
9368 futures_executor::block_on(middleware.before(&ctx, &mut req));
9369
9370 let stored_id = req.get_extension::<RequestId>().unwrap();
9371 assert_eq!(stored_id.as_str(), "trace-123");
9372
9373 let response = Response::ok();
9374 let response = futures_executor::block_on(middleware.after(&ctx, &req, response));
9375
9376 let header = header_value(&response, "X-Trace-ID");
9377 assert_eq!(header, Some("trace-123".to_string()));
9378 }
9379
9380 #[test]
9381 fn is_valid_request_id_accepts_valid() {
9382 assert!(super::is_valid_request_id("abc123"));
9383 assert!(super::is_valid_request_id("request-id-123"));
9384 assert!(super::is_valid_request_id("request_id_123"));
9385 assert!(super::is_valid_request_id("request.id.123"));
9386 assert!(super::is_valid_request_id("ABC123"));
9387 assert!(super::is_valid_request_id("a-b_c.D"));
9388 }
9389
9390 #[test]
9391 fn is_valid_request_id_rejects_invalid() {
9392 assert!(!super::is_valid_request_id(""));
9393 assert!(!super::is_valid_request_id("id with spaces"));
9394 assert!(!super::is_valid_request_id("id<script>"));
9395 assert!(!super::is_valid_request_id("id\nwith\nnewlines"));
9396 assert!(!super::is_valid_request_id("id;with;semicolons"));
9397 assert!(!super::is_valid_request_id("id/with/slashes"));
9398 }
9399
9400 #[test]
9401 fn request_id_middleware_name() {
9402 let middleware = RequestIdMiddleware::new();
9403 assert_eq!(middleware.name(), "RequestId");
9404 }
9405
9406 struct OrderTrackingMiddleware {
9412 id: &'static str,
9413 log: Arc<std::sync::Mutex<Vec<String>>>,
9414 }
9415
9416 impl OrderTrackingMiddleware {
9417 fn new(id: &'static str, log: Arc<std::sync::Mutex<Vec<String>>>) -> Self {
9418 Self { id, log }
9419 }
9420 }
9421
9422 impl Middleware for OrderTrackingMiddleware {
9423 fn before<'a>(
9424 &'a self,
9425 _ctx: &'a RequestContext,
9426 _req: &'a mut Request,
9427 ) -> BoxFuture<'a, ControlFlow> {
9428 self.log.lock().unwrap().push(format!("{}.before", self.id));
9429 Box::pin(async { ControlFlow::Continue })
9430 }
9431
9432 fn after<'a>(
9433 &'a self,
9434 _ctx: &'a RequestContext,
9435 _req: &'a Request,
9436 response: Response,
9437 ) -> BoxFuture<'a, Response> {
9438 self.log.lock().unwrap().push(format!("{}.after", self.id));
9439 Box::pin(async move { response })
9440 }
9441 }
9442
9443 struct ConditionalBreakMiddleware {
9445 id: &'static str,
9446 should_break: bool,
9447 log: Arc<std::sync::Mutex<Vec<String>>>,
9448 }
9449
9450 impl ConditionalBreakMiddleware {
9451 fn new(
9452 id: &'static str,
9453 should_break: bool,
9454 log: Arc<std::sync::Mutex<Vec<String>>>,
9455 ) -> Self {
9456 Self {
9457 id,
9458 should_break,
9459 log,
9460 }
9461 }
9462 }
9463
9464 impl Middleware for ConditionalBreakMiddleware {
9465 fn before<'a>(
9466 &'a self,
9467 _ctx: &'a RequestContext,
9468 _req: &'a mut Request,
9469 ) -> BoxFuture<'a, ControlFlow> {
9470 self.log.lock().unwrap().push(format!("{}.before", self.id));
9471 let should_break = self.should_break;
9472 Box::pin(async move {
9473 if should_break {
9474 ControlFlow::Break(
9475 Response::with_status(StatusCode::FORBIDDEN)
9476 .body(ResponseBody::Bytes(b"blocked".to_vec())),
9477 )
9478 } else {
9479 ControlFlow::Continue
9480 }
9481 })
9482 }
9483
9484 fn after<'a>(
9485 &'a self,
9486 _ctx: &'a RequestContext,
9487 _req: &'a Request,
9488 response: Response,
9489 ) -> BoxFuture<'a, Response> {
9490 self.log.lock().unwrap().push(format!("{}.after", self.id));
9491 Box::pin(async move { response })
9492 }
9493 }
9494
9495 struct OkHandler;
9497
9498 impl Handler for OkHandler {
9499 fn call<'a>(
9500 &'a self,
9501 _ctx: &'a RequestContext,
9502 _req: &'a mut Request,
9503 ) -> BoxFuture<'a, Response> {
9504 Box::pin(async move { Response::ok().body(ResponseBody::Bytes(b"handler".to_vec())) })
9505 }
9506 }
9507
9508 struct CheckHeaderHandler;
9510
9511 impl Handler for CheckHeaderHandler {
9512 fn call<'a>(
9513 &'a self,
9514 _ctx: &'a RequestContext,
9515 req: &'a mut Request,
9516 ) -> BoxFuture<'a, Response> {
9517 let has_header = req.headers().get("X-Modified-By").is_some();
9518 Box::pin(async move {
9519 if has_header {
9520 Response::ok().body(ResponseBody::Bytes(b"header-present".to_vec()))
9521 } else {
9522 Response::with_status(StatusCode::BAD_REQUEST)
9523 }
9524 })
9525 }
9526 }
9527
9528 struct ErrorHandler;
9530
9531 impl Handler for ErrorHandler {
9532 fn call<'a>(
9533 &'a self,
9534 _ctx: &'a RequestContext,
9535 _req: &'a mut Request,
9536 ) -> BoxFuture<'a, Response> {
9537 Box::pin(async move { Response::with_status(StatusCode::INTERNAL_SERVER_ERROR) })
9538 }
9539 }
9540
9541 #[test]
9542 fn middleware_stack_executes_in_correct_order() {
9543 let log = Arc::new(std::sync::Mutex::new(Vec::new()));
9546
9547 let mut stack = MiddlewareStack::new();
9548 stack.push(OrderTrackingMiddleware::new("mw1", log.clone()));
9549 stack.push(OrderTrackingMiddleware::new("mw2", log.clone()));
9550 stack.push(OrderTrackingMiddleware::new("mw3", log.clone()));
9551
9552 let ctx = test_context();
9553 let mut req = Request::new(crate::request::Method::Get, "/");
9554
9555 futures_executor::block_on(stack.execute(&OkHandler, &ctx, &mut req));
9556
9557 let calls = log.lock().unwrap().clone();
9558 assert_eq!(
9559 calls,
9560 vec![
9561 "mw1.before",
9562 "mw2.before",
9563 "mw3.before",
9564 "mw3.after",
9565 "mw2.after",
9566 "mw1.after",
9567 ]
9568 );
9569 }
9570
9571 #[test]
9572 fn middleware_stack_short_circuit_skips_later_middleware() {
9573 let log = Arc::new(std::sync::Mutex::new(Vec::new()));
9576
9577 let mut stack = MiddlewareStack::new();
9578 stack.push(OrderTrackingMiddleware::new("mw1", log.clone()));
9579 stack.push(ConditionalBreakMiddleware::new("mw2", true, log.clone()));
9580 stack.push(OrderTrackingMiddleware::new("mw3", log.clone()));
9581
9582 let ctx = test_context();
9583 let mut req = Request::new(crate::request::Method::Get, "/");
9584
9585 let response = futures_executor::block_on(stack.execute(&OkHandler, &ctx, &mut req));
9586
9587 assert_eq!(response.status().as_u16(), 403);
9589
9590 let calls = log.lock().unwrap().clone();
9591 assert_eq!(
9592 calls,
9593 vec![
9594 "mw1.before",
9595 "mw2.before",
9596 "mw1.after",
9599 ]
9600 );
9601 }
9602
9603 #[test]
9604 fn middleware_stack_first_middleware_breaks() {
9605 let log = Arc::new(std::sync::Mutex::new(Vec::new()));
9607
9608 let mut stack = MiddlewareStack::new();
9609 stack.push(ConditionalBreakMiddleware::new("mw1", true, log.clone()));
9610 stack.push(OrderTrackingMiddleware::new("mw2", log.clone()));
9611
9612 let ctx = test_context();
9613 let mut req = Request::new(crate::request::Method::Get, "/");
9614
9615 let response = futures_executor::block_on(stack.execute(&OkHandler, &ctx, &mut req));
9616
9617 assert_eq!(response.status().as_u16(), 403);
9618
9619 let calls = log.lock().unwrap().clone();
9620 assert_eq!(calls, vec!["mw1.before"]);
9621 }
9623
9624 #[test]
9625 fn middleware_stack_last_middleware_breaks() {
9626 let log = Arc::new(std::sync::Mutex::new(Vec::new()));
9628
9629 let mut stack = MiddlewareStack::new();
9630 stack.push(OrderTrackingMiddleware::new("mw1", log.clone()));
9631 stack.push(OrderTrackingMiddleware::new("mw2", log.clone()));
9632 stack.push(ConditionalBreakMiddleware::new("mw3", true, log.clone()));
9633
9634 let ctx = test_context();
9635 let mut req = Request::new(crate::request::Method::Get, "/");
9636
9637 let response = futures_executor::block_on(stack.execute(&OkHandler, &ctx, &mut req));
9638
9639 assert_eq!(response.status().as_u16(), 403);
9640
9641 let calls = log.lock().unwrap().clone();
9642 assert_eq!(
9643 calls,
9644 vec![
9645 "mw1.before",
9646 "mw2.before",
9647 "mw3.before",
9648 "mw2.after",
9650 "mw1.after",
9651 ]
9652 );
9653 }
9654
9655 #[test]
9656 fn middleware_stack_empty_executes_handler_directly() {
9657 let stack = MiddlewareStack::new();
9658 let ctx = test_context();
9659 let mut req = Request::new(crate::request::Method::Get, "/");
9660
9661 let response = futures_executor::block_on(stack.execute(&OkHandler, &ctx, &mut req));
9662
9663 assert_eq!(response.status().as_u16(), 200);
9664 }
9665
9666 #[test]
9667 fn middleware_stack_with_capacity() {
9668 let stack = MiddlewareStack::with_capacity(10);
9669 assert!(stack.is_empty());
9670 assert_eq!(stack.len(), 0);
9671 }
9672
9673 #[test]
9674 fn middleware_stack_push_arc() {
9675 let mut stack = MiddlewareStack::new();
9676 let mw: Arc<dyn Middleware> = Arc::new(NoopMiddleware);
9677 stack.push_arc(mw);
9678 assert_eq!(stack.len(), 1);
9679 }
9680
9681 #[test]
9686 fn add_response_header_adds_header() {
9687 let mw = AddResponseHeader::new("X-Custom", b"custom-value".to_vec());
9688 let ctx = test_context();
9689 let req = Request::new(crate::request::Method::Get, "/");
9690
9691 let response = Response::ok();
9692 let response = futures_executor::block_on(mw.after(&ctx, &req, response));
9693
9694 assert_eq!(
9695 header_value(&response, "X-Custom"),
9696 Some("custom-value".to_string())
9697 );
9698 }
9699
9700 #[test]
9701 fn add_response_header_preserves_existing_headers() {
9702 let mw = AddResponseHeader::new("X-New", b"new".to_vec());
9703 let ctx = test_context();
9704 let req = Request::new(crate::request::Method::Get, "/");
9705
9706 let response = Response::ok().header("X-Existing", b"existing".to_vec());
9707 let response = futures_executor::block_on(mw.after(&ctx, &req, response));
9708
9709 assert_eq!(
9710 header_value(&response, "X-Existing"),
9711 Some("existing".to_string())
9712 );
9713 assert_eq!(header_value(&response, "X-New"), Some("new".to_string()));
9714 }
9715
9716 #[test]
9717 fn add_response_header_name() {
9718 let mw = AddResponseHeader::new("X-Test", b"test".to_vec());
9719 assert_eq!(mw.name(), "AddResponseHeader");
9720 }
9721
9722 #[test]
9727 fn require_header_allows_with_header() {
9728 let mw = RequireHeader::new("X-Api-Key");
9729 let ctx = test_context();
9730 let mut req = Request::new(crate::request::Method::Get, "/");
9731 req.headers_mut()
9732 .insert("X-Api-Key", b"secret-key".to_vec());
9733
9734 let result = futures_executor::block_on(mw.before(&ctx, &mut req));
9735 assert!(matches!(result, ControlFlow::Continue));
9736 }
9737
9738 #[test]
9739 fn require_header_blocks_without_header() {
9740 let mw = RequireHeader::new("X-Api-Key");
9741 let ctx = test_context();
9742 let mut req = Request::new(crate::request::Method::Get, "/");
9743
9744 let result = futures_executor::block_on(mw.before(&ctx, &mut req));
9745
9746 match result {
9747 ControlFlow::Break(response) => {
9748 assert_eq!(response.status().as_u16(), 400);
9749 }
9750 ControlFlow::Continue => panic!("Expected Break, got Continue"),
9751 }
9752 }
9753
9754 #[test]
9755 fn require_header_name() {
9756 let mw = RequireHeader::new("X-Test");
9757 assert_eq!(mw.name(), "RequireHeader");
9758 }
9759
9760 #[test]
9765 fn path_prefix_filter_allows_matching_path() {
9766 let mw = PathPrefixFilter::new("/api");
9767 let ctx = test_context();
9768 let mut req = Request::new(crate::request::Method::Get, "/api/users");
9769
9770 let result = futures_executor::block_on(mw.before(&ctx, &mut req));
9771 assert!(matches!(result, ControlFlow::Continue));
9772 }
9773
9774 #[test]
9775 fn path_prefix_filter_allows_exact_prefix() {
9776 let mw = PathPrefixFilter::new("/api");
9777 let ctx = test_context();
9778 let mut req = Request::new(crate::request::Method::Get, "/api");
9779
9780 let result = futures_executor::block_on(mw.before(&ctx, &mut req));
9781 assert!(matches!(result, ControlFlow::Continue));
9782 }
9783
9784 #[test]
9785 fn path_prefix_filter_blocks_non_matching_path() {
9786 let mw = PathPrefixFilter::new("/api");
9787 let ctx = test_context();
9788 let mut req = Request::new(crate::request::Method::Get, "/admin/users");
9789
9790 let result = futures_executor::block_on(mw.before(&ctx, &mut req));
9791
9792 match result {
9793 ControlFlow::Break(response) => {
9794 assert_eq!(response.status().as_u16(), 404);
9795 }
9796 ControlFlow::Continue => panic!("Expected Break, got Continue"),
9797 }
9798 }
9799
9800 #[test]
9801 fn path_prefix_filter_name() {
9802 let mw = PathPrefixFilter::new("/api");
9803 assert_eq!(mw.name(), "PathPrefixFilter");
9804 }
9805
9806 #[test]
9811 fn conditional_status_applies_true_status() {
9812 let mw = ConditionalStatus::new(
9813 |req| req.path() == "/health",
9814 StatusCode::OK,
9815 StatusCode::NOT_FOUND,
9816 );
9817 let ctx = test_context();
9818 let req = Request::new(crate::request::Method::Get, "/health");
9819 let response = Response::with_status(StatusCode::INTERNAL_SERVER_ERROR);
9820
9821 let response = futures_executor::block_on(mw.after(&ctx, &req, response));
9822 assert_eq!(response.status().as_u16(), 200);
9823 }
9824
9825 #[test]
9826 fn conditional_status_applies_false_status() {
9827 let mw = ConditionalStatus::new(
9828 |req| req.path() == "/health",
9829 StatusCode::OK,
9830 StatusCode::NOT_FOUND,
9831 );
9832 let ctx = test_context();
9833 let req = Request::new(crate::request::Method::Get, "/other");
9834 let response = Response::with_status(StatusCode::INTERNAL_SERVER_ERROR);
9835
9836 let response = futures_executor::block_on(mw.after(&ctx, &req, response));
9837 assert_eq!(response.status().as_u16(), 404);
9838 }
9839
9840 #[test]
9841 fn conditional_status_name() {
9842 let mw = ConditionalStatus::new(|_| true, StatusCode::OK, StatusCode::NOT_FOUND);
9843 assert_eq!(mw.name(), "ConditionalStatus");
9844 }
9845
9846 #[derive(Clone)]
9851 struct LayerTestMiddleware {
9852 prefix: String,
9853 }
9854
9855 impl LayerTestMiddleware {
9856 fn new(prefix: impl Into<String>) -> Self {
9857 Self {
9858 prefix: prefix.into(),
9859 }
9860 }
9861 }
9862
9863 impl Middleware for LayerTestMiddleware {
9864 fn after<'a>(
9865 &'a self,
9866 _ctx: &'a RequestContext,
9867 _req: &'a Request,
9868 response: Response,
9869 ) -> BoxFuture<'a, Response> {
9870 let prefix = self.prefix.clone();
9871 Box::pin(async move { response.header("X-Layer", prefix.into_bytes()) })
9872 }
9873 }
9874
9875 #[test]
9876 fn layer_wraps_handler() {
9877 let layer = Layer::new(LayerTestMiddleware::new("wrapped"));
9878 let wrapped = layer.wrap(OkHandler);
9879
9880 let ctx = test_context();
9881 let mut req = Request::new(crate::request::Method::Get, "/");
9882
9883 let response = futures_executor::block_on(wrapped.call(&ctx, &mut req));
9884
9885 assert_eq!(response.status().as_u16(), 200);
9886 assert_eq!(
9887 header_value(&response, "X-Layer"),
9888 Some("wrapped".to_string())
9889 );
9890 }
9891
9892 #[test]
9893 fn layered_handles_break() {
9894 #[derive(Clone)]
9895 struct BreakingMiddleware;
9896
9897 impl Middleware for BreakingMiddleware {
9898 fn before<'a>(
9899 &'a self,
9900 _ctx: &'a RequestContext,
9901 _req: &'a mut Request,
9902 ) -> BoxFuture<'a, ControlFlow> {
9903 Box::pin(async {
9904 ControlFlow::Break(Response::with_status(StatusCode::UNAUTHORIZED))
9905 })
9906 }
9907
9908 fn after<'a>(
9909 &'a self,
9910 _ctx: &'a RequestContext,
9911 _req: &'a Request,
9912 response: Response,
9913 ) -> BoxFuture<'a, Response> {
9914 Box::pin(async move { response.header("X-After", b"ran".to_vec()) })
9915 }
9916 }
9917
9918 let layer = Layer::new(BreakingMiddleware);
9919 let wrapped = layer.wrap(OkHandler);
9920
9921 let ctx = test_context();
9922 let mut req = Request::new(crate::request::Method::Get, "/");
9923
9924 let response = futures_executor::block_on(wrapped.call(&ctx, &mut req));
9925
9926 assert_eq!(response.status().as_u16(), 401);
9928 assert_eq!(header_value(&response, "X-After"), Some("ran".to_string()));
9930 }
9931
9932 #[test]
9937 fn request_response_logger_default() {
9938 let logger = RequestResponseLogger::default();
9939 assert!(logger.log_request_headers);
9940 assert!(logger.log_response_headers);
9941 assert!(!logger.log_body);
9942 assert_eq!(logger.max_body_bytes, 1024);
9943 }
9944
9945 #[test]
9946 fn request_response_logger_builder() {
9947 let logger = RequestResponseLogger::new()
9948 .log_request_headers(false)
9949 .log_response_headers(false)
9950 .log_body(true)
9951 .max_body_bytes(2048)
9952 .redact_header("x-secret");
9953
9954 assert!(!logger.log_request_headers);
9955 assert!(!logger.log_response_headers);
9956 assert!(logger.log_body);
9957 assert_eq!(logger.max_body_bytes, 2048);
9958 assert!(logger.redact_headers.contains("x-secret"));
9959 }
9960
9961 #[test]
9962 fn request_response_logger_name() {
9963 let logger = RequestResponseLogger::new();
9964 assert_eq!(logger.name(), "RequestResponseLogger");
9965 }
9966
9967 #[test]
9972 fn middleware_stack_modifies_request_for_handler() {
9973 struct RequestModifier;
9975
9976 impl Middleware for RequestModifier {
9977 fn before<'a>(
9978 &'a self,
9979 _ctx: &'a RequestContext,
9980 req: &'a mut Request,
9981 ) -> BoxFuture<'a, ControlFlow> {
9982 req.headers_mut()
9983 .insert("X-Modified-By", b"middleware".to_vec());
9984 Box::pin(async { ControlFlow::Continue })
9985 }
9986 }
9987
9988 let mut stack = MiddlewareStack::new();
9989 stack.push(RequestModifier);
9990
9991 let ctx = test_context();
9992 let mut req = Request::new(crate::request::Method::Get, "/");
9993
9994 let response =
9995 futures_executor::block_on(stack.execute(&CheckHeaderHandler, &ctx, &mut req));
9996
9997 assert_eq!(response.status().as_u16(), 200);
9998 }
9999
10000 #[test]
10001 fn middleware_stack_multiple_response_modifications() {
10002 let mut stack = MiddlewareStack::new();
10003 stack.push(AddResponseHeader::new("X-First", b"1".to_vec()));
10004 stack.push(AddResponseHeader::new("X-Second", b"2".to_vec()));
10005 stack.push(AddResponseHeader::new("X-Third", b"3".to_vec()));
10006
10007 let ctx = test_context();
10008 let mut req = Request::new(crate::request::Method::Get, "/");
10009
10010 let response = futures_executor::block_on(stack.execute(&OkHandler, &ctx, &mut req));
10011
10012 assert_eq!(header_value(&response, "X-First"), Some("1".to_string()));
10014 assert_eq!(header_value(&response, "X-Second"), Some("2".to_string()));
10015 assert_eq!(header_value(&response, "X-Third"), Some("3".to_string()));
10016 }
10017
10018 #[test]
10019 fn middleware_stack_handler_receives_response_after_break() {
10020 let mut stack = MiddlewareStack::new();
10022 stack.push(ConditionalBreakMiddleware::new(
10023 "breaker",
10024 true,
10025 Arc::new(std::sync::Mutex::new(Vec::new())),
10026 ));
10027
10028 let ctx = test_context();
10029 let mut req = Request::new(crate::request::Method::Get, "/");
10030
10031 let response = futures_executor::block_on(stack.execute(&OkHandler, &ctx, &mut req));
10032
10033 assert_eq!(response.status().as_u16(), 403);
10034 match response.body_ref() {
10036 ResponseBody::Bytes(b) => assert_eq!(b, b"blocked"),
10037 _ => panic!("Expected Bytes body"),
10038 }
10039 }
10040
10041 #[test]
10046 fn middleware_after_can_change_status() {
10047 struct StatusChanger;
10048
10049 impl Middleware for StatusChanger {
10050 fn after<'a>(
10051 &'a self,
10052 _ctx: &'a RequestContext,
10053 _req: &'a Request,
10054 _response: Response,
10055 ) -> BoxFuture<'a, Response> {
10056 Box::pin(async { Response::with_status(StatusCode::SERVICE_UNAVAILABLE) })
10057 }
10058 }
10059
10060 let mut stack = MiddlewareStack::new();
10061 stack.push(StatusChanger);
10062
10063 let ctx = test_context();
10064 let mut req = Request::new(crate::request::Method::Get, "/");
10065
10066 let response = futures_executor::block_on(stack.execute(&OkHandler, &ctx, &mut req));
10067
10068 assert_eq!(response.status().as_u16(), 503);
10070 }
10071
10072 #[test]
10073 fn middleware_after_runs_even_on_error_status() {
10074 let log = Arc::new(std::sync::Mutex::new(Vec::new()));
10075 let mut stack = MiddlewareStack::new();
10076 stack.push(OrderTrackingMiddleware::new("mw1", log.clone()));
10077
10078 let ctx = test_context();
10079 let mut req = Request::new(crate::request::Method::Get, "/");
10080
10081 let response = futures_executor::block_on(stack.execute(&ErrorHandler, &ctx, &mut req));
10082
10083 assert_eq!(response.status().as_u16(), 500);
10084
10085 let calls = log.lock().unwrap().clone();
10086 assert_eq!(calls, vec!["mw1.before", "mw1.after"]);
10088 }
10089
10090 #[test]
10095 fn wildcard_match_simple() {
10096 assert!(super::wildcard_match("*.example.com", "api.example.com"));
10097 assert!(super::wildcard_match("*.example.com", "www.example.com"));
10098 assert!(!super::wildcard_match("*.example.com", "example.com"));
10099 }
10100
10101 #[test]
10102 fn wildcard_match_suffix_pattern() {
10103 assert!(super::wildcard_match("*.txt", "file.txt"));
10105 assert!(super::wildcard_match("*.txt", "document.txt"));
10106 assert!(!super::wildcard_match("*.txt", "file.doc"));
10107 assert!(super::wildcard_match("*-suffix", "any-suffix"));
10108 }
10109
10110 #[test]
10111 fn wildcard_match_no_wildcard() {
10112 assert!(super::wildcard_match("exact", "exact"));
10113 assert!(!super::wildcard_match("exact", "different"));
10114 }
10115
10116 #[test]
10117 fn regex_match_anchored() {
10118 assert!(super::regex_match("^hello$", "hello"));
10119 assert!(!super::regex_match("^hello$", "hello world"));
10120 assert!(!super::regex_match("^hello$", "say hello"));
10121 }
10122
10123 #[test]
10124 fn regex_match_dot_wildcard() {
10125 assert!(super::regex_match("h.llo", "hello"));
10126 assert!(super::regex_match("h.llo", "hallo"));
10127 }
10128
10129 #[test]
10130 fn regex_match_star() {
10131 assert!(super::regex_match("hel*o", "hello"));
10132 assert!(super::regex_match("hel*o", "helo"));
10133 assert!(super::regex_match("hel*o", "hellllllo"));
10134 }
10135
10136 #[test]
10141 fn middleware_default_before_continues() {
10142 struct DefaultBefore;
10143 impl Middleware for DefaultBefore {}
10144
10145 let mw = DefaultBefore;
10146 let ctx = test_context();
10147 let mut req = Request::new(crate::request::Method::Get, "/");
10148
10149 let result = futures_executor::block_on(mw.before(&ctx, &mut req));
10150 assert!(matches!(result, ControlFlow::Continue));
10151 }
10152
10153 #[test]
10154 fn middleware_default_after_passes_through() {
10155 struct DefaultAfter;
10156 impl Middleware for DefaultAfter {}
10157
10158 let mw = DefaultAfter;
10159 let ctx = test_context();
10160 let req = Request::new(crate::request::Method::Get, "/");
10161 let response = Response::with_status(StatusCode::CREATED);
10162
10163 let result = futures_executor::block_on(mw.after(&ctx, &req, response));
10164 assert_eq!(result.status().as_u16(), 201);
10165 }
10166
10167 #[test]
10168 fn middleware_default_name_is_type_name() {
10169 struct MyCustomMiddleware;
10170 impl Middleware for MyCustomMiddleware {}
10171
10172 let mw = MyCustomMiddleware;
10173 assert!(mw.name().contains("MyCustomMiddleware"));
10174 }
10175
10176 #[test]
10181 fn security_headers_default_config() {
10182 let config = SecurityHeadersConfig::default();
10183 assert_eq!(config.x_content_type_options, Some("nosniff"));
10184 assert_eq!(config.x_frame_options, Some(XFrameOptions::Deny));
10185 assert_eq!(config.x_xss_protection, Some("0"));
10186 assert!(config.content_security_policy.is_none());
10187 assert!(config.hsts.is_none());
10188 assert_eq!(
10189 config.referrer_policy,
10190 Some(ReferrerPolicy::StrictOriginWhenCrossOrigin)
10191 );
10192 assert!(config.permissions_policy.is_none());
10193 }
10194
10195 #[test]
10196 fn security_headers_none_config() {
10197 let config = SecurityHeadersConfig::none();
10198 assert!(config.x_content_type_options.is_none());
10199 assert!(config.x_frame_options.is_none());
10200 assert!(config.x_xss_protection.is_none());
10201 assert!(config.content_security_policy.is_none());
10202 assert!(config.hsts.is_none());
10203 assert!(config.referrer_policy.is_none());
10204 assert!(config.permissions_policy.is_none());
10205 }
10206
10207 #[test]
10208 fn security_headers_strict_config() {
10209 let config = SecurityHeadersConfig::strict();
10210 assert_eq!(config.x_content_type_options, Some("nosniff"));
10211 assert_eq!(config.x_frame_options, Some(XFrameOptions::Deny));
10212 assert_eq!(
10213 config.content_security_policy,
10214 Some("default-src 'self'".to_string())
10215 );
10216 assert_eq!(config.hsts, Some((31536000, true, false)));
10217 assert_eq!(config.referrer_policy, Some(ReferrerPolicy::NoReferrer));
10218 assert!(config.permissions_policy.is_some());
10219 }
10220
10221 #[test]
10222 fn security_headers_config_builder() {
10223 let config = SecurityHeadersConfig::new()
10224 .x_frame_options(Some(XFrameOptions::SameOrigin))
10225 .content_security_policy("default-src 'self'")
10226 .hsts(86400, false, false)
10227 .referrer_policy(Some(ReferrerPolicy::Origin));
10228
10229 assert_eq!(config.x_frame_options, Some(XFrameOptions::SameOrigin));
10230 assert_eq!(
10231 config.content_security_policy,
10232 Some("default-src 'self'".to_string())
10233 );
10234 assert_eq!(config.hsts, Some((86400, false, false)));
10235 assert_eq!(config.referrer_policy, Some(ReferrerPolicy::Origin));
10236 }
10237
10238 #[test]
10239 fn security_headers_hsts_value_format() {
10240 let config = SecurityHeadersConfig::none().hsts(3600, false, false);
10242 assert_eq!(config.build_hsts_value(), Some("max-age=3600".to_string()));
10243
10244 let config = SecurityHeadersConfig::none().hsts(3600, true, false);
10246 assert_eq!(
10247 config.build_hsts_value(),
10248 Some("max-age=3600; includeSubDomains".to_string())
10249 );
10250
10251 let config = SecurityHeadersConfig::none().hsts(3600, false, true);
10253 assert_eq!(
10254 config.build_hsts_value(),
10255 Some("max-age=3600; preload".to_string())
10256 );
10257
10258 let config = SecurityHeadersConfig::none().hsts(3600, true, true);
10260 assert_eq!(
10261 config.build_hsts_value(),
10262 Some("max-age=3600; includeSubDomains; preload".to_string())
10263 );
10264 }
10265
10266 #[test]
10267 fn security_headers_middleware_adds_default_headers() {
10268 let mw = SecurityHeaders::new();
10269 let ctx = test_context();
10270 let req = Request::new(crate::request::Method::Get, "/");
10271 let response = Response::ok();
10272
10273 let result = futures_executor::block_on(mw.after(&ctx, &req, response));
10274
10275 assert!(header_value(&result, "X-Content-Type-Options").is_some());
10277 assert!(header_value(&result, "X-Frame-Options").is_some());
10278 assert!(header_value(&result, "X-XSS-Protection").is_some());
10279 assert!(header_value(&result, "Referrer-Policy").is_some());
10280
10281 assert!(header_value(&result, "Content-Security-Policy").is_none());
10283 assert!(header_value(&result, "Strict-Transport-Security").is_none());
10284 assert!(header_value(&result, "Permissions-Policy").is_none());
10285 }
10286
10287 #[test]
10288 fn security_headers_middleware_with_csp() {
10289 let config = SecurityHeadersConfig::new()
10290 .content_security_policy("default-src 'self'; script-src 'self' 'unsafe-inline'");
10291 let mw = SecurityHeaders::with_config(config);
10292 let ctx = test_context();
10293 let req = Request::new(crate::request::Method::Get, "/");
10294 let response = Response::ok();
10295
10296 let result = futures_executor::block_on(mw.after(&ctx, &req, response));
10297
10298 let csp = header_value(&result, "Content-Security-Policy");
10299 assert!(csp.is_some());
10300 assert_eq!(
10301 csp.unwrap(),
10302 "default-src 'self'; script-src 'self' 'unsafe-inline'"
10303 );
10304 }
10305
10306 #[test]
10307 fn security_headers_middleware_with_hsts() {
10308 let config = SecurityHeadersConfig::new().hsts(31536000, true, false);
10309 let mw = SecurityHeaders::with_config(config);
10310 let ctx = test_context();
10311 let req = Request::new(crate::request::Method::Get, "/");
10312 let response = Response::ok();
10313
10314 let result = futures_executor::block_on(mw.after(&ctx, &req, response));
10315
10316 let hsts = header_value(&result, "Strict-Transport-Security");
10317 assert!(hsts.is_some());
10318 assert_eq!(hsts.unwrap(), "max-age=31536000; includeSubDomains");
10319 }
10320
10321 #[test]
10322 fn security_headers_middleware_name() {
10323 let mw = SecurityHeaders::new();
10324 assert_eq!(mw.name(), "SecurityHeaders");
10325 }
10326
10327 #[test]
10328 fn x_frame_options_values() {
10329 assert_eq!(XFrameOptions::Deny.as_bytes(), b"DENY");
10330 assert_eq!(XFrameOptions::SameOrigin.as_bytes(), b"SAMEORIGIN");
10331 }
10332
10333 #[test]
10334 fn referrer_policy_values() {
10335 assert_eq!(ReferrerPolicy::NoReferrer.as_bytes(), b"no-referrer");
10336 assert_eq!(
10337 ReferrerPolicy::NoReferrerWhenDowngrade.as_bytes(),
10338 b"no-referrer-when-downgrade"
10339 );
10340 assert_eq!(ReferrerPolicy::Origin.as_bytes(), b"origin");
10341 assert_eq!(
10342 ReferrerPolicy::OriginWhenCrossOrigin.as_bytes(),
10343 b"origin-when-cross-origin"
10344 );
10345 assert_eq!(ReferrerPolicy::SameOrigin.as_bytes(), b"same-origin");
10346 assert_eq!(ReferrerPolicy::StrictOrigin.as_bytes(), b"strict-origin");
10347 assert_eq!(
10348 ReferrerPolicy::StrictOriginWhenCrossOrigin.as_bytes(),
10349 b"strict-origin-when-cross-origin"
10350 );
10351 assert_eq!(ReferrerPolicy::UnsafeUrl.as_bytes(), b"unsafe-url");
10352 }
10353
10354 #[test]
10355 fn security_headers_strict_preset() {
10356 let mw = SecurityHeaders::strict();
10357 let ctx = test_context();
10358 let req = Request::new(crate::request::Method::Get, "/");
10359 let response = Response::ok();
10360
10361 let result = futures_executor::block_on(mw.after(&ctx, &req, response));
10362
10363 assert!(header_value(&result, "X-Content-Type-Options").is_some());
10365 assert!(header_value(&result, "X-Frame-Options").is_some());
10366 assert!(header_value(&result, "Content-Security-Policy").is_some());
10367 assert!(header_value(&result, "Strict-Transport-Security").is_some());
10368 assert!(header_value(&result, "Referrer-Policy").is_some());
10369 assert!(header_value(&result, "Permissions-Policy").is_some());
10370 }
10371
10372 #[test]
10373 fn security_headers_config_clearing_methods() {
10374 let config = SecurityHeadersConfig::strict()
10375 .no_content_security_policy()
10376 .no_hsts()
10377 .no_permissions_policy();
10378
10379 assert!(config.content_security_policy.is_none());
10380 assert!(config.hsts.is_none());
10381 assert!(config.permissions_policy.is_none());
10382 }
10383
10384 #[test]
10389 fn csrf_token_generate_produces_unique_tokens() {
10390 let token1 = CsrfToken::generate();
10391 let token2 = CsrfToken::generate();
10392 assert_ne!(token1, token2);
10393 assert!(!token1.as_str().is_empty());
10394 assert!(!token2.as_str().is_empty());
10395 }
10396
10397 #[test]
10398 fn csrf_token_display() {
10399 let token = CsrfToken::new("test-token-123");
10400 assert_eq!(format!("{}", token), "test-token-123");
10401 }
10402
10403 #[test]
10404 fn csrf_config_defaults() {
10405 let config = CsrfConfig::default();
10406 assert_eq!(config.cookie_name, "csrf_token");
10407 assert_eq!(config.header_name, "x-csrf-token");
10408 assert_eq!(config.mode, CsrfMode::DoubleSubmit);
10409 assert!(!config.rotate_token);
10410 assert!(config.production);
10411 assert!(config.error_message.is_none());
10412 }
10413
10414 #[test]
10415 fn csrf_config_builder() {
10416 let config = CsrfConfig::new()
10417 .cookie_name("XSRF-TOKEN")
10418 .header_name("X-XSRF-Token")
10419 .mode(CsrfMode::HeaderOnly)
10420 .rotate_token(true)
10421 .production(false)
10422 .error_message("Custom CSRF error");
10423
10424 assert_eq!(config.cookie_name, "XSRF-TOKEN");
10425 assert_eq!(config.header_name, "X-XSRF-Token");
10426 assert_eq!(config.mode, CsrfMode::HeaderOnly);
10427 assert!(config.rotate_token);
10428 assert!(!config.production);
10429 assert_eq!(config.error_message, Some("Custom CSRF error".to_string()));
10430 }
10431
10432 #[test]
10433 fn csrf_middleware_allows_get_without_token() {
10434 let csrf = CsrfMiddleware::new();
10435 let ctx = test_context();
10436 let mut req = Request::new(crate::request::Method::Get, "/");
10437
10438 let result = futures_executor::block_on(csrf.before(&ctx, &mut req));
10439 assert!(result.is_continue());
10440 assert!(req.get_extension::<CsrfToken>().is_some());
10442 }
10443
10444 #[test]
10445 fn csrf_middleware_allows_head_without_token() {
10446 let csrf = CsrfMiddleware::new();
10447 let ctx = test_context();
10448 let mut req = Request::new(crate::request::Method::Head, "/");
10449
10450 let result = futures_executor::block_on(csrf.before(&ctx, &mut req));
10451 assert!(result.is_continue());
10452 }
10453
10454 #[test]
10455 fn csrf_middleware_allows_options_without_token() {
10456 let csrf = CsrfMiddleware::new();
10457 let ctx = test_context();
10458 let mut req = Request::new(crate::request::Method::Options, "/");
10459
10460 let result = futures_executor::block_on(csrf.before(&ctx, &mut req));
10461 assert!(result.is_continue());
10462 }
10463
10464 #[test]
10465 fn csrf_middleware_blocks_post_without_token() {
10466 let csrf = CsrfMiddleware::new();
10467 let ctx = test_context();
10468 let mut req = Request::new(crate::request::Method::Post, "/");
10469
10470 let result = futures_executor::block_on(csrf.before(&ctx, &mut req));
10471 assert!(result.is_break());
10472
10473 if let ControlFlow::Break(response) = result {
10474 assert_eq!(response.status(), StatusCode::FORBIDDEN);
10475 }
10476 }
10477
10478 #[test]
10479 fn csrf_middleware_blocks_put_without_token() {
10480 let csrf = CsrfMiddleware::new();
10481 let ctx = test_context();
10482 let mut req = Request::new(crate::request::Method::Put, "/");
10483
10484 let result = futures_executor::block_on(csrf.before(&ctx, &mut req));
10485 assert!(result.is_break());
10486 }
10487
10488 #[test]
10489 fn csrf_middleware_blocks_delete_without_token() {
10490 let csrf = CsrfMiddleware::new();
10491 let ctx = test_context();
10492 let mut req = Request::new(crate::request::Method::Delete, "/");
10493
10494 let result = futures_executor::block_on(csrf.before(&ctx, &mut req));
10495 assert!(result.is_break());
10496 }
10497
10498 #[test]
10499 fn csrf_middleware_blocks_patch_without_token() {
10500 let csrf = CsrfMiddleware::new();
10501 let ctx = test_context();
10502 let mut req = Request::new(crate::request::Method::Patch, "/");
10503
10504 let result = futures_executor::block_on(csrf.before(&ctx, &mut req));
10505 assert!(result.is_break());
10506 }
10507
10508 #[test]
10509 fn csrf_middleware_allows_post_with_matching_tokens() {
10510 let csrf = CsrfMiddleware::new();
10511 let ctx = test_context();
10512 let mut req = Request::new(crate::request::Method::Post, "/");
10513
10514 let token = "valid-csrf-token-12345";
10516 req.headers_mut()
10517 .insert("cookie", format!("csrf_token={}", token).into_bytes());
10518 req.headers_mut()
10519 .insert("x-csrf-token", token.as_bytes().to_vec());
10520
10521 let result = futures_executor::block_on(csrf.before(&ctx, &mut req));
10522 assert!(result.is_continue());
10523
10524 let stored_token = req.get_extension::<CsrfToken>().unwrap();
10526 assert_eq!(stored_token.as_str(), token);
10527 }
10528
10529 #[test]
10530 fn csrf_middleware_blocks_post_with_mismatched_tokens() {
10531 let csrf = CsrfMiddleware::new();
10532 let ctx = test_context();
10533 let mut req = Request::new(crate::request::Method::Post, "/");
10534
10535 req.headers_mut()
10537 .insert("cookie", b"csrf_token=token-in-cookie".to_vec());
10538 req.headers_mut()
10539 .insert("x-csrf-token", b"different-token".to_vec());
10540
10541 let result = futures_executor::block_on(csrf.before(&ctx, &mut req));
10542 assert!(result.is_break());
10543
10544 if let ControlFlow::Break(response) = result {
10545 assert_eq!(response.status(), StatusCode::FORBIDDEN);
10546 }
10547 }
10548
10549 #[test]
10550 fn csrf_middleware_blocks_post_with_header_only_in_double_submit_mode() {
10551 let csrf = CsrfMiddleware::new();
10552 let ctx = test_context();
10553 let mut req = Request::new(crate::request::Method::Post, "/");
10554
10555 req.headers_mut()
10557 .insert("x-csrf-token", b"some-token".to_vec());
10558
10559 let result = futures_executor::block_on(csrf.before(&ctx, &mut req));
10560 assert!(result.is_break());
10561 }
10562
10563 #[test]
10564 fn csrf_middleware_blocks_post_with_cookie_only_in_double_submit_mode() {
10565 let csrf = CsrfMiddleware::new();
10566 let ctx = test_context();
10567 let mut req = Request::new(crate::request::Method::Post, "/");
10568
10569 req.headers_mut()
10571 .insert("cookie", b"csrf_token=some-token".to_vec());
10572
10573 let result = futures_executor::block_on(csrf.before(&ctx, &mut req));
10574 assert!(result.is_break());
10575 }
10576
10577 #[test]
10578 fn csrf_middleware_header_only_mode_accepts_header_token() {
10579 let csrf = CsrfMiddleware::with_config(CsrfConfig::new().mode(CsrfMode::HeaderOnly));
10580 let ctx = test_context();
10581 let mut req = Request::new(crate::request::Method::Post, "/");
10582
10583 req.headers_mut()
10584 .insert("x-csrf-token", b"valid-token".to_vec());
10585
10586 let result = futures_executor::block_on(csrf.before(&ctx, &mut req));
10587 assert!(result.is_continue());
10588 }
10589
10590 #[test]
10591 fn csrf_middleware_header_only_mode_rejects_empty_header() {
10592 let csrf = CsrfMiddleware::with_config(CsrfConfig::new().mode(CsrfMode::HeaderOnly));
10593 let ctx = test_context();
10594 let mut req = Request::new(crate::request::Method::Post, "/");
10595
10596 req.headers_mut().insert("x-csrf-token", b"".to_vec());
10597
10598 let result = futures_executor::block_on(csrf.before(&ctx, &mut req));
10599 assert!(result.is_break());
10600 }
10601
10602 #[test]
10603 fn csrf_middleware_sets_cookie_on_get() {
10604 let csrf = CsrfMiddleware::new();
10605 let ctx = test_context();
10606 let mut req = Request::new(crate::request::Method::Get, "/");
10607
10608 let _ = futures_executor::block_on(csrf.before(&ctx, &mut req));
10610
10611 let response = Response::ok();
10613 let result = futures_executor::block_on(csrf.after(&ctx, &req, response));
10614
10615 let cookie_value = header_value(&result, "set-cookie");
10617 assert!(cookie_value.is_some());
10618
10619 let cookie_value = cookie_value.unwrap();
10620 assert!(cookie_value.starts_with("csrf_token="));
10621 assert!(cookie_value.contains("SameSite=Strict"));
10622 assert!(cookie_value.contains("Secure")); }
10624
10625 #[test]
10626 fn csrf_middleware_no_secure_in_dev_mode() {
10627 let csrf = CsrfMiddleware::with_config(CsrfConfig::new().production(false));
10628 let ctx = test_context();
10629 let mut req = Request::new(crate::request::Method::Get, "/");
10630
10631 let _ = futures_executor::block_on(csrf.before(&ctx, &mut req));
10632
10633 let response = Response::ok();
10634 let result = futures_executor::block_on(csrf.after(&ctx, &req, response));
10635
10636 let cookie_value = header_value(&result, "set-cookie").unwrap();
10637 assert!(!cookie_value.contains("Secure")); }
10639
10640 #[test]
10641 fn csrf_middleware_does_not_set_cookie_if_already_present() {
10642 let csrf = CsrfMiddleware::new();
10643 let ctx = test_context();
10644 let mut req = Request::new(crate::request::Method::Get, "/");
10645
10646 req.headers_mut()
10648 .insert("cookie", b"csrf_token=existing-token".to_vec());
10649
10650 let _ = futures_executor::block_on(csrf.before(&ctx, &mut req));
10651
10652 let response = Response::ok();
10653 let result = futures_executor::block_on(csrf.after(&ctx, &req, response));
10654
10655 assert!(header_value(&result, "set-cookie").is_none());
10657 }
10658
10659 #[test]
10660 fn csrf_middleware_rotates_token_when_configured() {
10661 let csrf = CsrfMiddleware::with_config(CsrfConfig::new().rotate_token(true));
10662 let ctx = test_context();
10663 let mut req = Request::new(crate::request::Method::Get, "/");
10664
10665 req.headers_mut()
10667 .insert("cookie", b"csrf_token=old-token".to_vec());
10668
10669 let _ = futures_executor::block_on(csrf.before(&ctx, &mut req));
10670
10671 let response = Response::ok();
10672 let result = futures_executor::block_on(csrf.after(&ctx, &req, response));
10673
10674 assert!(header_value(&result, "set-cookie").is_some());
10676 }
10677
10678 #[test]
10679 fn csrf_middleware_custom_header_name() {
10680 let csrf = CsrfMiddleware::with_config(
10681 CsrfConfig::new()
10682 .header_name("X-XSRF-Token")
10683 .cookie_name("XSRF-TOKEN"),
10684 );
10685 let ctx = test_context();
10686 let mut req = Request::new(crate::request::Method::Post, "/");
10687
10688 let token = "custom-token-value";
10689 req.headers_mut()
10690 .insert("cookie", format!("XSRF-TOKEN={}", token).into_bytes());
10691 req.headers_mut()
10692 .insert("x-xsrf-token", token.as_bytes().to_vec());
10693
10694 let result = futures_executor::block_on(csrf.before(&ctx, &mut req));
10695 assert!(result.is_continue());
10696 }
10697
10698 #[test]
10699 fn csrf_middleware_error_response_is_json() {
10700 let csrf = CsrfMiddleware::new();
10701 let ctx = test_context();
10702 let mut req = Request::new(crate::request::Method::Post, "/");
10703
10704 let result = futures_executor::block_on(csrf.before(&ctx, &mut req));
10705
10706 if let ControlFlow::Break(response) = result {
10707 let content_type = header_value(&response, "content-type");
10708 assert_eq!(content_type, Some("application/json".to_string()));
10709
10710 if let ResponseBody::Bytes(body) = response.body_ref() {
10712 let body_str = std::str::from_utf8(body).unwrap();
10713 assert!(body_str.contains("csrf_error"));
10714 assert!(body_str.contains("x-csrf-token"));
10715 } else {
10716 panic!("Expected Bytes body");
10717 }
10718 } else {
10719 panic!("Expected Break");
10720 }
10721 }
10722
10723 #[test]
10724 fn csrf_middleware_custom_error_message() {
10725 let csrf = CsrfMiddleware::with_config(
10726 CsrfConfig::new().error_message("Access denied: invalid security token"),
10727 );
10728 let ctx = test_context();
10729 let mut req = Request::new(crate::request::Method::Post, "/");
10730
10731 let result = futures_executor::block_on(csrf.before(&ctx, &mut req));
10732
10733 if let ControlFlow::Break(response) = result {
10734 if let ResponseBody::Bytes(body) = response.body_ref() {
10735 let body_str = std::str::from_utf8(body).unwrap();
10736 assert!(body_str.contains("Access denied: invalid security token"));
10737 }
10738 }
10739 }
10740
10741 #[test]
10742 fn csrf_middleware_name() {
10743 let csrf = CsrfMiddleware::new();
10744 assert_eq!(csrf.name(), "CSRF");
10745 }
10746
10747 #[test]
10748 fn csrf_middleware_parses_cookie_with_multiple_cookies() {
10749 let csrf = CsrfMiddleware::new();
10750 let ctx = test_context();
10751 let mut req = Request::new(crate::request::Method::Post, "/");
10752
10753 let token = "the-csrf-token";
10755 req.headers_mut().insert(
10756 "cookie",
10757 format!("session=abc123; csrf_token={}; user=test", token).into_bytes(),
10758 );
10759 req.headers_mut()
10760 .insert("x-csrf-token", token.as_bytes().to_vec());
10761
10762 let result = futures_executor::block_on(csrf.before(&ctx, &mut req));
10763 assert!(result.is_continue());
10764 }
10765
10766 #[test]
10767 fn csrf_middleware_handles_empty_token_value() {
10768 let csrf = CsrfMiddleware::new();
10769 let ctx = test_context();
10770 let mut req = Request::new(crate::request::Method::Post, "/");
10771
10772 req.headers_mut().insert("cookie", b"csrf_token=".to_vec());
10774 req.headers_mut().insert("x-csrf-token", b"".to_vec());
10775
10776 let result = futures_executor::block_on(csrf.before(&ctx, &mut req));
10777 assert!(result.is_break()); }
10779
10780 #[test]
10783 fn csrf_token_generate_many_unique() {
10784 let mut tokens = std::collections::HashSet::new();
10786 for _ in 0..100 {
10787 let token = CsrfToken::generate();
10788 assert!(
10789 tokens.insert(token.0.clone()),
10790 "Duplicate token generated: {}",
10791 token.0
10792 );
10793 }
10794 assert_eq!(tokens.len(), 100);
10795 }
10796
10797 #[test]
10798 fn csrf_token_generate_format_is_hex() {
10799 let token = CsrfToken::generate();
10800 let s = token.as_str();
10801 assert!(
10803 s.len() >= 64,
10804 "Expected at least 64 hex characters, got {} in '{s}'",
10805 s.len()
10806 );
10807 assert!(
10808 s.chars().all(|c| c.is_ascii_hexdigit()),
10809 "Non-hex character in token: {s}"
10810 );
10811 }
10812
10813 #[test]
10814 fn csrf_token_generate_minimum_length() {
10815 let token = CsrfToken::generate();
10816 assert!(
10818 token.as_str().len() >= 64,
10819 "Token too short: {} (len={})",
10820 token.as_str(),
10821 token.as_str().len()
10822 );
10823 }
10824
10825 #[test]
10826 fn csrf_token_from_str() {
10827 let token: CsrfToken = "my-token".into();
10828 assert_eq!(token.as_str(), "my-token");
10829 assert_eq!(token.0, "my-token");
10830 }
10831
10832 #[test]
10833 fn csrf_token_clone_eq() {
10834 let t1 = CsrfToken::new("abc");
10835 let t2 = t1.clone();
10836 assert_eq!(t1, t2);
10837 assert_eq!(t1.as_str(), t2.as_str());
10838 }
10839
10840 #[test]
10841 fn csrf_middleware_allows_trace_without_token() {
10842 let csrf = CsrfMiddleware::new();
10843 let ctx = test_context();
10844 let mut req = Request::new(crate::request::Method::Trace, "/");
10845
10846 let result = futures_executor::block_on(csrf.before(&ctx, &mut req));
10847 assert!(result.is_continue());
10848 assert!(req.get_extension::<CsrfToken>().is_some());
10850 }
10851
10852 #[test]
10853 fn csrf_safe_method_generates_token_into_extension() {
10854 let csrf = CsrfMiddleware::new();
10855 let ctx = test_context();
10856
10857 for method in [
10858 crate::request::Method::Get,
10859 crate::request::Method::Head,
10860 crate::request::Method::Options,
10861 crate::request::Method::Trace,
10862 ] {
10863 let mut req = Request::new(method, "/test");
10864 let result = futures_executor::block_on(csrf.before(&ctx, &mut req));
10865 assert!(result.is_continue());
10866 let token = req.get_extension::<CsrfToken>().expect("token missing");
10867 assert!(!token.as_str().is_empty());
10868 }
10869 }
10870
10871 #[test]
10872 fn csrf_safe_method_preserves_existing_cookie_token() {
10873 let csrf = CsrfMiddleware::new();
10874 let ctx = test_context();
10875 let mut req = Request::new(crate::request::Method::Get, "/");
10876 req.headers_mut()
10877 .insert("cookie", b"csrf_token=my-existing-token".to_vec());
10878
10879 let _ = futures_executor::block_on(csrf.before(&ctx, &mut req));
10880
10881 let token = req.get_extension::<CsrfToken>().unwrap();
10883 assert_eq!(token.as_str(), "my-existing-token");
10884 }
10885
10886 #[test]
10887 fn csrf_valid_post_stores_token_in_extension() {
10888 let csrf = CsrfMiddleware::new();
10889 let ctx = test_context();
10890 let mut req = Request::new(crate::request::Method::Post, "/submit");
10891
10892 let tk = "valid-token-xyz";
10893 req.headers_mut()
10894 .insert("cookie", format!("csrf_token={}", tk).into_bytes());
10895 req.headers_mut()
10896 .insert("x-csrf-token", tk.as_bytes().to_vec());
10897
10898 let result = futures_executor::block_on(csrf.before(&ctx, &mut req));
10899 assert!(result.is_continue());
10900 let stored = req.get_extension::<CsrfToken>().unwrap();
10901 assert_eq!(stored.as_str(), tk);
10902 }
10903
10904 #[test]
10905 fn csrf_double_submit_both_empty_strings_rejected() {
10906 let csrf = CsrfMiddleware::new();
10907 let ctx = test_context();
10908 let mut req = Request::new(crate::request::Method::Post, "/");
10909
10910 req.headers_mut().insert("cookie", b"csrf_token=".to_vec());
10912 req.headers_mut().insert("x-csrf-token", b"".to_vec());
10913
10914 let result = futures_executor::block_on(csrf.before(&ctx, &mut req));
10915 assert!(result.is_break());
10916 }
10917
10918 #[test]
10919 fn csrf_double_submit_matching_empty_rejected() {
10920 let csrf = CsrfMiddleware::new();
10922 let ctx = test_context();
10923 let mut req = Request::new(crate::request::Method::Post, "/");
10924
10925 req.headers_mut().insert("cookie", b"csrf_token=".to_vec());
10926 req.headers_mut().insert("x-csrf-token", b"".to_vec());
10927
10928 let result = futures_executor::block_on(csrf.before(&ctx, &mut req));
10929 assert!(
10930 result.is_break(),
10931 "Empty matching tokens should be rejected"
10932 );
10933 }
10934
10935 #[test]
10936 fn csrf_header_only_mode_does_not_need_cookie() {
10937 let csrf = CsrfMiddleware::with_config(CsrfConfig::new().mode(CsrfMode::HeaderOnly));
10938 let ctx = test_context();
10939 let mut req = Request::new(crate::request::Method::Post, "/");
10940
10941 req.headers_mut()
10943 .insert("x-csrf-token", b"header-only-token".to_vec());
10944
10945 let result = futures_executor::block_on(csrf.before(&ctx, &mut req));
10946 assert!(result.is_continue());
10947 let token = req.get_extension::<CsrfToken>().unwrap();
10948 assert_eq!(token.as_str(), "header-only-token");
10949 }
10950
10951 #[test]
10952 fn csrf_header_only_mode_ignores_mismatched_cookie() {
10953 let csrf = CsrfMiddleware::with_config(CsrfConfig::new().mode(CsrfMode::HeaderOnly));
10955 let ctx = test_context();
10956 let mut req = Request::new(crate::request::Method::Post, "/");
10957
10958 req.headers_mut()
10959 .insert("cookie", b"csrf_token=different-value".to_vec());
10960 req.headers_mut()
10961 .insert("x-csrf-token", b"header-value".to_vec());
10962
10963 let result = futures_executor::block_on(csrf.before(&ctx, &mut req));
10964 assert!(result.is_continue(), "HeaderOnly should ignore cookie");
10965 }
10966
10967 #[test]
10968 fn csrf_header_only_mode_rejects_no_header() {
10969 let csrf = CsrfMiddleware::with_config(CsrfConfig::new().mode(CsrfMode::HeaderOnly));
10970 let ctx = test_context();
10971 let mut req = Request::new(crate::request::Method::Post, "/");
10972 let result = futures_executor::block_on(csrf.before(&ctx, &mut req));
10974 assert!(result.is_break());
10975 }
10976
10977 #[test]
10978 fn csrf_header_only_error_message_mentions_header() {
10979 let csrf = CsrfMiddleware::with_config(CsrfConfig::new().mode(CsrfMode::HeaderOnly));
10980 let ctx = test_context();
10981 let mut req = Request::new(crate::request::Method::Post, "/");
10982
10983 let result = futures_executor::block_on(csrf.before(&ctx, &mut req));
10984 if let ControlFlow::Break(response) = result {
10985 if let ResponseBody::Bytes(body) = response.body_ref() {
10986 let body_str = std::str::from_utf8(body).unwrap();
10987 assert!(
10988 body_str.contains("missing in header"),
10989 "Expected 'missing in header' in: {}",
10990 body_str
10991 );
10992 }
10993 } else {
10994 panic!("Expected Break");
10995 }
10996 }
10997
10998 #[test]
10999 fn csrf_mismatch_error_differs_from_missing_error() {
11000 let csrf = CsrfMiddleware::new();
11001 let ctx = test_context();
11002
11003 let mut req_missing = Request::new(crate::request::Method::Post, "/");
11005 let missing_result = futures_executor::block_on(csrf.before(&ctx, &mut req_missing));
11006 let missing_body = match missing_result {
11007 ControlFlow::Break(r) => match r.body_ref() {
11008 ResponseBody::Bytes(b) => std::str::from_utf8(b).unwrap().to_string(),
11009 ResponseBody::Empty | ResponseBody::Stream(_) => panic!("Expected Bytes"),
11010 },
11011 ControlFlow::Continue => panic!("Expected Break"),
11012 };
11013
11014 let mut req_mismatch = Request::new(crate::request::Method::Post, "/");
11016 req_mismatch
11017 .headers_mut()
11018 .insert("cookie", b"csrf_token=aaa".to_vec());
11019 req_mismatch
11020 .headers_mut()
11021 .insert("x-csrf-token", b"bbb".to_vec());
11022 let mismatch_result = futures_executor::block_on(csrf.before(&ctx, &mut req_mismatch));
11023 let mismatch_body = match mismatch_result {
11024 ControlFlow::Break(r) => match r.body_ref() {
11025 ResponseBody::Bytes(b) => std::str::from_utf8(b).unwrap().to_string(),
11026 ResponseBody::Empty | ResponseBody::Stream(_) => panic!("Expected Bytes"),
11027 },
11028 ControlFlow::Continue => panic!("Expected Break"),
11029 };
11030
11031 assert_ne!(
11033 missing_body, mismatch_body,
11034 "Missing vs mismatch should have different error messages"
11035 );
11036 assert!(missing_body.contains("missing"));
11037 assert!(mismatch_body.contains("mismatch"));
11038 }
11039
11040 #[test]
11041 fn csrf_cookie_not_httponly() {
11042 let csrf = CsrfMiddleware::new();
11044 let ctx = test_context();
11045 let mut req = Request::new(crate::request::Method::Get, "/");
11046
11047 let _ = futures_executor::block_on(csrf.before(&ctx, &mut req));
11048 let response = Response::ok();
11049 let result = futures_executor::block_on(csrf.after(&ctx, &req, response));
11050
11051 let cookie_value = header_value(&result, "set-cookie").unwrap();
11052 assert!(
11053 !cookie_value.to_lowercase().contains("httponly"),
11054 "CSRF cookie must NOT be HttpOnly (needs JS access), got: {}",
11055 cookie_value
11056 );
11057 }
11058
11059 #[test]
11060 fn csrf_cookie_has_path_slash() {
11061 let csrf = CsrfMiddleware::new();
11062 let ctx = test_context();
11063 let mut req = Request::new(crate::request::Method::Get, "/");
11064
11065 let _ = futures_executor::block_on(csrf.before(&ctx, &mut req));
11066 let response = Response::ok();
11067 let result = futures_executor::block_on(csrf.after(&ctx, &req, response));
11068
11069 let cookie_value = header_value(&result, "set-cookie").unwrap();
11070 assert!(
11071 cookie_value.contains("Path=/"),
11072 "Cookie should have Path=/, got: {}",
11073 cookie_value
11074 );
11075 }
11076
11077 #[test]
11078 fn csrf_cookie_has_samesite_strict() {
11079 let csrf = CsrfMiddleware::new();
11080 let ctx = test_context();
11081 let mut req = Request::new(crate::request::Method::Get, "/");
11082
11083 let _ = futures_executor::block_on(csrf.before(&ctx, &mut req));
11084 let response = Response::ok();
11085 let result = futures_executor::block_on(csrf.after(&ctx, &req, response));
11086
11087 let cookie_value = header_value(&result, "set-cookie").unwrap();
11088 assert!(
11089 cookie_value.contains("SameSite=Strict"),
11090 "Cookie should have SameSite=Strict, got: {}",
11091 cookie_value
11092 );
11093 }
11094
11095 #[test]
11096 fn csrf_production_mode_sets_secure_flag() {
11097 let csrf = CsrfMiddleware::with_config(CsrfConfig::new().production(true));
11098 let ctx = test_context();
11099 let mut req = Request::new(crate::request::Method::Get, "/");
11100
11101 let _ = futures_executor::block_on(csrf.before(&ctx, &mut req));
11102 let response = Response::ok();
11103 let result = futures_executor::block_on(csrf.after(&ctx, &req, response));
11104
11105 let cookie_value = header_value(&result, "set-cookie").unwrap();
11106 assert!(
11107 cookie_value.contains("Secure"),
11108 "Production cookie must have Secure flag, got: {}",
11109 cookie_value
11110 );
11111 }
11112
11113 #[test]
11114 fn csrf_no_set_cookie_on_post_response() {
11115 let csrf = CsrfMiddleware::new();
11117 let ctx = test_context();
11118 let mut req = Request::new(crate::request::Method::Post, "/");
11119
11120 let token = "valid-token";
11121 req.headers_mut()
11122 .insert("cookie", format!("csrf_token={}", token).into_bytes());
11123 req.headers_mut()
11124 .insert("x-csrf-token", token.as_bytes().to_vec());
11125
11126 let _ = futures_executor::block_on(csrf.before(&ctx, &mut req));
11127 let response = Response::ok();
11128 let result = futures_executor::block_on(csrf.after(&ctx, &req, response));
11129
11130 assert!(
11131 header_value(&result, "set-cookie").is_none(),
11132 "POST response should not set CSRF cookie"
11133 );
11134 }
11135
11136 #[test]
11137 fn csrf_head_method_sets_cookie() {
11138 let csrf = CsrfMiddleware::new();
11139 let ctx = test_context();
11140 let mut req = Request::new(crate::request::Method::Head, "/");
11141
11142 let _ = futures_executor::block_on(csrf.before(&ctx, &mut req));
11143 let response = Response::ok();
11144 let result = futures_executor::block_on(csrf.after(&ctx, &req, response));
11145
11146 assert!(
11147 header_value(&result, "set-cookie").is_some(),
11148 "HEAD response should set CSRF cookie"
11149 );
11150 }
11151
11152 #[test]
11153 fn csrf_options_method_sets_cookie() {
11154 let csrf = CsrfMiddleware::new();
11155 let ctx = test_context();
11156 let mut req = Request::new(crate::request::Method::Options, "/");
11157
11158 let _ = futures_executor::block_on(csrf.before(&ctx, &mut req));
11159 let response = Response::ok();
11160 let result = futures_executor::block_on(csrf.after(&ctx, &req, response));
11161
11162 assert!(
11163 header_value(&result, "set-cookie").is_some(),
11164 "OPTIONS response should set CSRF cookie"
11165 );
11166 }
11167
11168 #[test]
11169 fn csrf_rotation_produces_different_token_in_cookie() {
11170 let csrf = CsrfMiddleware::with_config(CsrfConfig::new().rotate_token(true));
11171 let ctx = test_context();
11172 let mut req = Request::new(crate::request::Method::Get, "/");
11173
11174 let old_token = "old-token-value";
11175 req.headers_mut()
11176 .insert("cookie", format!("csrf_token={}", old_token).into_bytes());
11177
11178 let _ = futures_executor::block_on(csrf.before(&ctx, &mut req));
11179 let response = Response::ok();
11180 let result = futures_executor::block_on(csrf.after(&ctx, &req, response));
11181
11182 let cookie_value = header_value(&result, "set-cookie").unwrap();
11183 assert!(cookie_value.starts_with("csrf_token="));
11188 }
11189
11190 #[test]
11191 fn csrf_no_rotation_skips_set_cookie_when_present() {
11192 let csrf = CsrfMiddleware::with_config(CsrfConfig::new().rotate_token(false));
11193 let ctx = test_context();
11194 let mut req = Request::new(crate::request::Method::Get, "/");
11195
11196 req.headers_mut()
11197 .insert("cookie", b"csrf_token=existing".to_vec());
11198
11199 let _ = futures_executor::block_on(csrf.before(&ctx, &mut req));
11200 let response = Response::ok();
11201 let result = futures_executor::block_on(csrf.after(&ctx, &req, response));
11202
11203 assert!(
11204 header_value(&result, "set-cookie").is_none(),
11205 "Without rotation, should not re-set existing cookie"
11206 );
11207 }
11208
11209 #[test]
11210 fn csrf_custom_cookie_name_in_set_cookie_response() {
11211 let csrf = CsrfMiddleware::with_config(CsrfConfig::new().cookie_name("XSRF-TOKEN"));
11212 let ctx = test_context();
11213 let mut req = Request::new(crate::request::Method::Get, "/");
11214
11215 let _ = futures_executor::block_on(csrf.before(&ctx, &mut req));
11216 let response = Response::ok();
11217 let result = futures_executor::block_on(csrf.after(&ctx, &req, response));
11218
11219 let cookie_value = header_value(&result, "set-cookie").unwrap();
11220 assert!(
11221 cookie_value.starts_with("XSRF-TOKEN="),
11222 "Custom cookie name should appear in Set-Cookie, got: {}",
11223 cookie_value
11224 );
11225 }
11226
11227 #[test]
11228 fn csrf_custom_header_name_validated() {
11229 let csrf = CsrfMiddleware::with_config(
11230 CsrfConfig::new()
11231 .header_name("X-Custom-CSRF")
11232 .cookie_name("my_csrf"),
11233 );
11234 let ctx = test_context();
11235 let mut req = Request::new(crate::request::Method::Post, "/");
11236
11237 let token = "custom-tok";
11238 req.headers_mut()
11239 .insert("cookie", format!("my_csrf={}", token).into_bytes());
11240 req.headers_mut()
11241 .insert("x-custom-csrf", token.as_bytes().to_vec());
11242
11243 let result = futures_executor::block_on(csrf.before(&ctx, &mut req));
11244 assert!(result.is_continue());
11245 }
11246
11247 #[test]
11248 fn csrf_custom_header_name_wrong_header_rejected() {
11249 let csrf = CsrfMiddleware::with_config(CsrfConfig::new().header_name("X-Custom-CSRF"));
11250 let ctx = test_context();
11251 let mut req = Request::new(crate::request::Method::Post, "/");
11252
11253 let token = "some-token";
11254 req.headers_mut()
11255 .insert("cookie", format!("csrf_token={}", token).into_bytes());
11256 req.headers_mut()
11258 .insert("x-csrf-token", token.as_bytes().to_vec());
11259
11260 let result = futures_executor::block_on(csrf.before(&ctx, &mut req));
11261 assert!(result.is_break(), "Wrong header name should be rejected");
11262 }
11263
11264 #[test]
11265 fn csrf_cookie_parsing_multiple_cookies_picks_correct() {
11266 let csrf = CsrfMiddleware::new();
11267 let ctx = test_context();
11268 let mut req = Request::new(crate::request::Method::Post, "/");
11269
11270 let token = "correct-csrf";
11271 req.headers_mut().insert(
11272 "cookie",
11273 format!("session=abc; other=xyz; csrf_token={}; tracking=123", token).into_bytes(),
11274 );
11275 req.headers_mut()
11276 .insert("x-csrf-token", token.as_bytes().to_vec());
11277
11278 let result = futures_executor::block_on(csrf.before(&ctx, &mut req));
11279 assert!(result.is_continue());
11280 }
11281
11282 #[test]
11283 fn csrf_cookie_parsing_spaces_around_semicolons() {
11284 let csrf = CsrfMiddleware::new();
11285 let ctx = test_context();
11286 let mut req = Request::new(crate::request::Method::Post, "/");
11287
11288 let token = "spaced-token";
11289 req.headers_mut().insert(
11290 "cookie",
11291 format!("session=abc ; csrf_token={} ; other=xyz", token).into_bytes(),
11292 );
11293 req.headers_mut()
11294 .insert("x-csrf-token", token.as_bytes().to_vec());
11295
11296 let result = futures_executor::block_on(csrf.before(&ctx, &mut req));
11297 assert!(result.is_continue());
11298 }
11299
11300 #[test]
11301 fn csrf_error_response_status_is_403() {
11302 let csrf = CsrfMiddleware::new();
11303 let ctx = test_context();
11304
11305 for method in [
11307 crate::request::Method::Post,
11308 crate::request::Method::Put,
11309 crate::request::Method::Delete,
11310 crate::request::Method::Patch,
11311 ] {
11312 let mut req = Request::new(method, "/");
11313 let result = futures_executor::block_on(csrf.before(&ctx, &mut req));
11314 match result {
11315 ControlFlow::Break(response) => {
11316 assert_eq!(
11317 response.status(),
11318 StatusCode::FORBIDDEN,
11319 "Expected 403 for {:?}",
11320 method
11321 );
11322 }
11323 ControlFlow::Continue => panic!("Expected Break for {:?}", method),
11324 }
11325 }
11326 }
11327
11328 #[test]
11329 fn csrf_error_body_json_structure() {
11330 let csrf = CsrfMiddleware::new();
11331 let ctx = test_context();
11332 let mut req = Request::new(crate::request::Method::Post, "/");
11333
11334 let result = futures_executor::block_on(csrf.before(&ctx, &mut req));
11335 if let ControlFlow::Break(response) = result {
11336 if let ResponseBody::Bytes(body) = response.body_ref() {
11337 let body_str = std::str::from_utf8(body).unwrap();
11338 let parsed: serde_json::Value = serde_json::from_str(body_str)
11340 .unwrap_or_else(|e| panic!("Invalid JSON: {}: {}", body_str, e));
11341 assert!(parsed["detail"].is_array());
11342 let detail = &parsed["detail"][0];
11343 assert_eq!(detail["type"], "csrf_error");
11344 assert!(detail["loc"].is_array());
11345 assert_eq!(detail["loc"][0], "header");
11346 assert_eq!(detail["loc"][1], "x-csrf-token");
11347 assert!(detail["msg"].is_string());
11348 } else {
11349 panic!("Expected Bytes body");
11350 }
11351 } else {
11352 panic!("Expected Break");
11353 }
11354 }
11355
11356 #[test]
11357 fn csrf_default_trait() {
11358 let csrf = CsrfMiddleware::default();
11359 assert_eq!(csrf.name(), "CSRF");
11360 let ctx = test_context();
11362 let mut req = Request::new(crate::request::Method::Get, "/");
11363 let result = futures_executor::block_on(csrf.before(&ctx, &mut req));
11364 assert!(result.is_continue());
11365 }
11366
11367 #[test]
11368 fn csrf_mode_default_is_double_submit() {
11369 assert_eq!(CsrfMode::default(), CsrfMode::DoubleSubmit);
11370 }
11371
11372 #[test]
11373 fn csrf_double_submit_both_present_same_non_empty_passes() {
11374 let csrf = CsrfMiddleware::new();
11376 let ctx = test_context();
11377
11378 let token = "a1b2c3d4e5f6";
11379 let mut req = Request::new(crate::request::Method::Delete, "/resource/1");
11380 req.headers_mut()
11381 .insert("cookie", format!("csrf_token={}", token).into_bytes());
11382 req.headers_mut()
11383 .insert("x-csrf-token", token.as_bytes().to_vec());
11384
11385 let result = futures_executor::block_on(csrf.before(&ctx, &mut req));
11386 assert!(result.is_continue());
11387 }
11388
11389 #[test]
11390 fn csrf_double_submit_case_sensitive() {
11391 let csrf = CsrfMiddleware::new();
11393 let ctx = test_context();
11394 let mut req = Request::new(crate::request::Method::Post, "/");
11395
11396 req.headers_mut()
11397 .insert("cookie", b"csrf_token=AbCdEf".to_vec());
11398 req.headers_mut().insert("x-csrf-token", b"abcdef".to_vec());
11399
11400 let result = futures_executor::block_on(csrf.before(&ctx, &mut req));
11401 assert!(
11402 result.is_break(),
11403 "Token comparison should be case-sensitive"
11404 );
11405 }
11406
11407 #[test]
11408 fn csrf_token_cookie_extractor_reads_csrf_cookie() {
11409 use crate::extract::{CookieName, CsrfTokenCookie};
11411 assert_eq!(CsrfTokenCookie::NAME, "csrf_token");
11412 }
11413
11414 #[test]
11415 fn csrf_make_set_cookie_header_value_production() {
11416 let value = CsrfMiddleware::make_set_cookie_header_value("csrf_token", "tok123", true);
11417 let s = std::str::from_utf8(&value).unwrap();
11418 assert!(s.contains("csrf_token=tok123"));
11419 assert!(s.contains("Path=/"));
11420 assert!(s.contains("SameSite=Strict"));
11421 assert!(s.contains("Secure"));
11422 assert!(!s.to_lowercase().contains("httponly"));
11423 }
11424
11425 #[test]
11426 fn csrf_make_set_cookie_header_value_development() {
11427 let value = CsrfMiddleware::make_set_cookie_header_value("csrf_token", "tok123", false);
11428 let s = std::str::from_utf8(&value).unwrap();
11429 assert!(s.contains("csrf_token=tok123"));
11430 assert!(s.contains("Path=/"));
11431 assert!(s.contains("SameSite=Strict"));
11432 assert!(!s.contains("Secure"));
11433 }
11434
11435 #[test]
11436 fn csrf_before_after_full_cycle_get_then_post() {
11437 let csrf = CsrfMiddleware::new();
11439 let ctx = test_context();
11440
11441 let mut get_req = Request::new(crate::request::Method::Get, "/form");
11443 let _ = futures_executor::block_on(csrf.before(&ctx, &mut get_req));
11444 let get_response = Response::ok();
11445 let get_result = futures_executor::block_on(csrf.after(&ctx, &get_req, get_response));
11446
11447 let set_cookie = header_value(&get_result, "set-cookie").expect("GET should set cookie");
11448 let token_value = set_cookie
11450 .strip_prefix("csrf_token=")
11451 .unwrap()
11452 .split(';')
11453 .next()
11454 .unwrap();
11455 assert!(!token_value.is_empty());
11456
11457 let mut post_req = Request::new(crate::request::Method::Post, "/form");
11459 post_req
11460 .headers_mut()
11461 .insert("cookie", format!("csrf_token={}", token_value).into_bytes());
11462 post_req
11463 .headers_mut()
11464 .insert("x-csrf-token", token_value.as_bytes().to_vec());
11465
11466 let result = futures_executor::block_on(csrf.before(&ctx, &mut post_req));
11467 assert!(result.is_continue(), "POST with valid token should pass");
11468 }
11469
11470 #[test]
11471 fn csrf_all_state_changing_methods_require_token() {
11472 let csrf = CsrfMiddleware::new();
11473 let ctx = test_context();
11474
11475 for method in [
11476 crate::request::Method::Post,
11477 crate::request::Method::Put,
11478 crate::request::Method::Delete,
11479 crate::request::Method::Patch,
11480 ] {
11481 let mut req = Request::new(method, "/resource");
11482 let result = futures_executor::block_on(csrf.before(&ctx, &mut req));
11483 assert!(
11484 result.is_break(),
11485 "{:?} without token should be rejected",
11486 method
11487 );
11488 }
11489 }
11490
11491 #[test]
11492 fn csrf_all_safe_methods_pass_without_token() {
11493 let csrf = CsrfMiddleware::new();
11494 let ctx = test_context();
11495
11496 for method in [
11497 crate::request::Method::Get,
11498 crate::request::Method::Head,
11499 crate::request::Method::Options,
11500 crate::request::Method::Trace,
11501 ] {
11502 let mut req = Request::new(method, "/resource");
11503 let result = futures_executor::block_on(csrf.before(&ctx, &mut req));
11504 assert!(
11505 result.is_continue(),
11506 "{:?} should be allowed without token",
11507 method
11508 );
11509 }
11510 }
11511
11512 #[derive(Clone)]
11519 struct OrderRecordingMiddleware {
11520 id: &'static str,
11521 log: std::sync::Arc<std::sync::Mutex<Vec<String>>>,
11522 }
11523
11524 impl OrderRecordingMiddleware {
11525 fn new(id: &'static str, log: std::sync::Arc<std::sync::Mutex<Vec<String>>>) -> Self {
11526 Self { id, log }
11527 }
11528 }
11529
11530 impl Middleware for OrderRecordingMiddleware {
11531 fn before<'a>(
11532 &'a self,
11533 _ctx: &'a RequestContext,
11534 _req: &'a mut Request,
11535 ) -> BoxFuture<'a, ControlFlow> {
11536 let id = self.id;
11537 let log = self.log.clone();
11538 Box::pin(async move {
11539 log.lock().unwrap().push(format!("{id}:before"));
11540 ControlFlow::Continue
11541 })
11542 }
11543
11544 fn after<'a>(
11545 &'a self,
11546 _ctx: &'a RequestContext,
11547 _req: &'a Request,
11548 response: Response,
11549 ) -> BoxFuture<'a, Response> {
11550 let id = self.id;
11551 let log = self.log.clone();
11552 Box::pin(async move {
11553 log.lock().unwrap().push(format!("{id}:after"));
11554 response
11555 })
11556 }
11557
11558 fn name(&self) -> &'static str {
11559 "OrderRecording"
11560 }
11561 }
11562
11563 struct ShortCircuitMiddleware {
11565 id: &'static str,
11566 log: std::sync::Arc<std::sync::Mutex<Vec<String>>>,
11567 }
11568
11569 impl ShortCircuitMiddleware {
11570 fn new(id: &'static str, log: std::sync::Arc<std::sync::Mutex<Vec<String>>>) -> Self {
11571 Self { id, log }
11572 }
11573 }
11574
11575 impl Middleware for ShortCircuitMiddleware {
11576 fn before<'a>(
11577 &'a self,
11578 _ctx: &'a RequestContext,
11579 _req: &'a mut Request,
11580 ) -> BoxFuture<'a, ControlFlow> {
11581 let id = self.id;
11582 let log = self.log.clone();
11583 Box::pin(async move {
11584 log.lock().unwrap().push(format!("{id}:before:break"));
11585 ControlFlow::Break(
11586 Response::with_status(StatusCode::FORBIDDEN)
11587 .body(ResponseBody::Bytes(b"short-circuited".to_vec())),
11588 )
11589 })
11590 }
11591
11592 fn after<'a>(
11593 &'a self,
11594 _ctx: &'a RequestContext,
11595 _req: &'a Request,
11596 response: Response,
11597 ) -> BoxFuture<'a, Response> {
11598 let id = self.id;
11599 let log = self.log.clone();
11600 Box::pin(async move {
11601 log.lock().unwrap().push(format!("{id}:after"));
11602 response
11603 })
11604 }
11605
11606 fn name(&self) -> &'static str {
11607 "ShortCircuit"
11608 }
11609 }
11610
11611 struct RecordingHandler {
11613 log: std::sync::Arc<std::sync::Mutex<Vec<String>>>,
11614 }
11615
11616 impl RecordingHandler {
11617 fn new(log: std::sync::Arc<std::sync::Mutex<Vec<String>>>) -> Self {
11618 Self { log }
11619 }
11620 }
11621
11622 impl Handler for RecordingHandler {
11623 fn call<'a>(
11624 &'a self,
11625 _ctx: &'a RequestContext,
11626 _req: &'a mut Request,
11627 ) -> BoxFuture<'a, Response> {
11628 let log = self.log.clone();
11629 Box::pin(async move {
11630 log.lock().unwrap().push("handler".to_string());
11631 Response::ok().body(ResponseBody::Bytes(b"ok".to_vec()))
11632 })
11633 }
11634 }
11635
11636 #[test]
11637 fn middleware_stack_three_middleware_onion_order() {
11638 let log = std::sync::Arc::new(std::sync::Mutex::new(Vec::new()));
11642
11643 let mut stack = MiddlewareStack::new();
11644 stack.push(OrderRecordingMiddleware::new("mw1", log.clone()));
11645 stack.push(OrderRecordingMiddleware::new("mw2", log.clone()));
11646 stack.push(OrderRecordingMiddleware::new("mw3", log.clone()));
11647
11648 let handler = RecordingHandler::new(log.clone());
11649 let ctx = test_context();
11650 let mut req = Request::new(crate::request::Method::Get, "/");
11651
11652 let _response = futures_executor::block_on(stack.execute(&handler, &ctx, &mut req));
11653
11654 let execution_log = log.lock().unwrap().clone();
11655 assert_eq!(
11656 execution_log,
11657 vec![
11658 "mw1:before",
11659 "mw2:before",
11660 "mw3:before",
11661 "handler",
11662 "mw3:after",
11663 "mw2:after",
11664 "mw1:after",
11665 ]
11666 );
11667 }
11668
11669 #[test]
11670 fn middleware_stack_short_circuit_runs_prior_after_hooks() {
11671 let log = std::sync::Arc::new(std::sync::Mutex::new(Vec::new()));
11679
11680 let mut stack = MiddlewareStack::new();
11681 stack.push(OrderRecordingMiddleware::new("mw1", log.clone()));
11682 stack.push(ShortCircuitMiddleware::new("mw2", log.clone()));
11683 stack.push(OrderRecordingMiddleware::new("mw3", log.clone()));
11684
11685 let handler = RecordingHandler::new(log.clone());
11686 let ctx = test_context();
11687 let mut req = Request::new(crate::request::Method::Get, "/");
11688
11689 let response = futures_executor::block_on(stack.execute(&handler, &ctx, &mut req));
11690
11691 assert_eq!(response.status().as_u16(), 403);
11693
11694 let execution_log = log.lock().unwrap().clone();
11695 assert_eq!(
11698 execution_log,
11699 vec!["mw1:before", "mw2:before:break", "mw1:after",]
11700 );
11701 }
11702
11703 #[test]
11704 fn middleware_stack_first_middleware_short_circuits() {
11705 let log = std::sync::Arc::new(std::sync::Mutex::new(Vec::new()));
11709
11710 let mut stack = MiddlewareStack::new();
11711 stack.push(ShortCircuitMiddleware::new("mw1", log.clone()));
11712 stack.push(OrderRecordingMiddleware::new("mw2", log.clone()));
11713
11714 let handler = RecordingHandler::new(log.clone());
11715 let ctx = test_context();
11716 let mut req = Request::new(crate::request::Method::Get, "/");
11717
11718 let response = futures_executor::block_on(stack.execute(&handler, &ctx, &mut req));
11719 assert_eq!(response.status().as_u16(), 403);
11720
11721 let execution_log = log.lock().unwrap().clone();
11722 assert_eq!(execution_log, vec!["mw1:before:break",]);
11724 }
11725
11726 #[test]
11727 fn middleware_stack_empty_runs_handler_only() {
11728 let log = std::sync::Arc::new(std::sync::Mutex::new(Vec::new()));
11730
11731 let stack = MiddlewareStack::new();
11732 let handler = RecordingHandler::new(log.clone());
11733 let ctx = test_context();
11734 let mut req = Request::new(crate::request::Method::Get, "/");
11735
11736 let response = futures_executor::block_on(stack.execute(&handler, &ctx, &mut req));
11737 assert_eq!(response.status().as_u16(), 200);
11738
11739 let execution_log = log.lock().unwrap().clone();
11740 assert_eq!(execution_log, vec!["handler"]);
11741 }
11742
11743 #[test]
11744 fn middleware_stack_single_middleware_ordering() {
11745 let log = std::sync::Arc::new(std::sync::Mutex::new(Vec::new()));
11747
11748 let mut stack = MiddlewareStack::new();
11749 stack.push(OrderRecordingMiddleware::new("mw1", log.clone()));
11750
11751 let handler = RecordingHandler::new(log.clone());
11752 let ctx = test_context();
11753 let mut req = Request::new(crate::request::Method::Get, "/");
11754
11755 let _response = futures_executor::block_on(stack.execute(&handler, &ctx, &mut req));
11756
11757 let execution_log = log.lock().unwrap().clone();
11758 assert_eq!(execution_log, vec!["mw1:before", "handler", "mw1:after",]);
11759 }
11760
11761 #[test]
11762 fn middleware_stack_five_middleware_onion_order() {
11763 let log = std::sync::Arc::new(std::sync::Mutex::new(Vec::new()));
11765
11766 let mut stack = MiddlewareStack::new();
11767 stack.push(OrderRecordingMiddleware::new("a", log.clone()));
11768 stack.push(OrderRecordingMiddleware::new("b", log.clone()));
11769 stack.push(OrderRecordingMiddleware::new("c", log.clone()));
11770 stack.push(OrderRecordingMiddleware::new("d", log.clone()));
11771 stack.push(OrderRecordingMiddleware::new("e", log.clone()));
11772
11773 let handler = RecordingHandler::new(log.clone());
11774 let ctx = test_context();
11775 let mut req = Request::new(crate::request::Method::Get, "/");
11776
11777 let _response = futures_executor::block_on(stack.execute(&handler, &ctx, &mut req));
11778
11779 let execution_log = log.lock().unwrap().clone();
11780 assert_eq!(
11781 execution_log,
11782 vec![
11783 "a:before", "b:before", "c:before", "d:before", "e:before", "handler", "e:after",
11784 "d:after", "c:after", "b:after", "a:after",
11785 ]
11786 );
11787 }
11788
11789 #[test]
11790 fn middleware_stack_short_circuit_at_end_runs_prior_afters() {
11791 let log = std::sync::Arc::new(std::sync::Mutex::new(Vec::new()));
11798
11799 let mut stack = MiddlewareStack::new();
11800 stack.push(OrderRecordingMiddleware::new("mw1", log.clone()));
11801 stack.push(OrderRecordingMiddleware::new("mw2", log.clone()));
11802 stack.push(ShortCircuitMiddleware::new("mw3", log.clone()));
11803
11804 let handler = RecordingHandler::new(log.clone());
11805 let ctx = test_context();
11806 let mut req = Request::new(crate::request::Method::Get, "/");
11807
11808 let response = futures_executor::block_on(stack.execute(&handler, &ctx, &mut req));
11809 assert_eq!(response.status().as_u16(), 403);
11810
11811 let execution_log = log.lock().unwrap().clone();
11812 assert_eq!(
11814 execution_log,
11815 vec![
11816 "mw1:before",
11817 "mw2:before",
11818 "mw3:before:break",
11819 "mw2:after",
11820 "mw1:after",
11821 ]
11822 );
11823 }
11824
11825 struct ModifyingMiddleware {
11827 id: &'static str,
11828 log: std::sync::Arc<std::sync::Mutex<Vec<String>>>,
11829 }
11830
11831 impl ModifyingMiddleware {
11832 fn new(id: &'static str, log: std::sync::Arc<std::sync::Mutex<Vec<String>>>) -> Self {
11833 Self { id, log }
11834 }
11835 }
11836
11837 impl Middleware for ModifyingMiddleware {
11838 fn before<'a>(
11839 &'a self,
11840 _ctx: &'a RequestContext,
11841 req: &'a mut Request,
11842 ) -> BoxFuture<'a, ControlFlow> {
11843 let id = self.id;
11844 let log = self.log.clone();
11845 Box::pin(async move {
11846 req.headers_mut()
11848 .insert(format!("x-{id}-before"), b"true".to_vec());
11849 log.lock().unwrap().push(format!("{id}:before"));
11850 ControlFlow::Continue
11851 })
11852 }
11853
11854 fn after<'a>(
11855 &'a self,
11856 _ctx: &'a RequestContext,
11857 _req: &'a Request,
11858 response: Response,
11859 ) -> BoxFuture<'a, Response> {
11860 let id = self.id;
11861 let log = self.log.clone();
11862 Box::pin(async move {
11863 log.lock().unwrap().push(format!("{id}:after"));
11864 response.header(format!("x-{id}-after"), b"true".to_vec())
11866 })
11867 }
11868
11869 fn name(&self) -> &'static str {
11870 "Modifying"
11871 }
11872 }
11873
11874 #[test]
11875 fn middleware_stack_modifications_accumulate_correctly() {
11876 let log = std::sync::Arc::new(std::sync::Mutex::new(Vec::new()));
11879
11880 let mut stack = MiddlewareStack::new();
11881 stack.push(ModifyingMiddleware::new("mw1", log.clone()));
11882 stack.push(ModifyingMiddleware::new("mw2", log.clone()));
11883 stack.push(ModifyingMiddleware::new("mw3", log.clone()));
11884
11885 let handler = RecordingHandler::new(log.clone());
11886 let ctx = test_context();
11887 let mut req = Request::new(crate::request::Method::Get, "/");
11888
11889 let response = futures_executor::block_on(stack.execute(&handler, &ctx, &mut req));
11890
11891 assert!(header_value(&response, "x-mw1-after").is_some());
11893 assert!(header_value(&response, "x-mw2-after").is_some());
11894 assert!(header_value(&response, "x-mw3-after").is_some());
11895
11896 assert!(req.headers().contains("x-mw1-before"));
11898 assert!(req.headers().contains("x-mw2-before"));
11899 assert!(req.headers().contains("x-mw3-before"));
11900 }
11901
11902 #[test]
11903 fn layer_wrap_maintains_middleware_order() {
11904 let log = std::sync::Arc::new(std::sync::Mutex::new(Vec::new()));
11906
11907 let layer = Layer::new(OrderRecordingMiddleware::new("layer", log.clone()));
11909
11910 let handler = RecordingHandler::new(log.clone());
11912 let layered_handler = layer.wrap(handler);
11913
11914 let ctx = test_context();
11915 let mut req = Request::new(crate::request::Method::Get, "/");
11916
11917 let _response = futures_executor::block_on(layered_handler.call(&ctx, &mut req));
11919
11920 let execution_log = log.lock().unwrap().clone();
11921 assert_eq!(
11922 execution_log,
11923 vec!["layer:before", "handler", "layer:after",]
11924 );
11925 }
11926}
11927
11928#[cfg(all(test, feature = "compression"))]
11933mod compression_tests {
11934 use super::*;
11935 use crate::request::Method;
11936 use crate::response::ResponseBody;
11937
11938 fn test_context() -> RequestContext {
11939 RequestContext::new(asupersync::Cx::for_testing(), 1)
11940 }
11941
11942 #[test]
11943 fn compression_config_defaults() {
11944 let config = CompressionConfig::default();
11945 assert_eq!(config.min_size, 1024);
11946 assert_eq!(config.level, 6);
11947 assert!(!config.skip_content_types.is_empty());
11948 }
11949
11950 #[test]
11951 fn compression_config_builder() {
11952 let config = CompressionConfig::new().min_size(512).level(9);
11953 assert_eq!(config.min_size, 512);
11954 assert_eq!(config.level, 9);
11955 }
11956
11957 #[test]
11958 fn compression_level_clamped() {
11959 let config = CompressionConfig::new().level(100);
11960 assert_eq!(config.level, 9);
11961
11962 let config = CompressionConfig::new().level(0);
11963 assert_eq!(config.level, 1);
11964 }
11965
11966 #[test]
11967 fn skip_content_type_exact_match() {
11968 let config = CompressionConfig::default();
11969 assert!(config.should_skip_content_type("image/jpeg"));
11970 assert!(config.should_skip_content_type("image/jpeg; charset=utf-8"));
11971 assert!(!config.should_skip_content_type("text/html"));
11972 }
11973
11974 #[test]
11975 fn skip_content_type_prefix_match() {
11976 let config = CompressionConfig::default();
11977 assert!(config.should_skip_content_type("video/mp4"));
11979 assert!(config.should_skip_content_type("video/webm"));
11980 assert!(config.should_skip_content_type("audio/mpeg"));
11981 }
11982
11983 #[test]
11984 fn compression_skips_small_responses() {
11985 let middleware = CompressionMiddleware::new();
11986 let ctx = test_context();
11987
11988 let mut req = Request::new(Method::Get, "/");
11990 req.headers_mut()
11991 .insert("accept-encoding", b"gzip".to_vec());
11992
11993 let response = Response::ok()
11995 .header("content-type", b"text/plain".to_vec())
11996 .body(ResponseBody::Bytes(b"Hello, World!".to_vec()));
11997
11998 let result = futures_executor::block_on(middleware.after(&ctx, &req, response));
12000
12001 let has_encoding = result
12003 .headers()
12004 .iter()
12005 .any(|(name, _)| name.eq_ignore_ascii_case("content-encoding"));
12006 assert!(!has_encoding, "Small response should not be compressed");
12007 }
12008
12009 #[test]
12010 fn compression_works_for_large_responses() {
12011 let config = CompressionConfig::new().min_size(10); let middleware = CompressionMiddleware::with_config(config);
12013 let ctx = test_context();
12014
12015 let mut req = Request::new(Method::Get, "/");
12017 req.headers_mut()
12018 .insert("accept-encoding", b"gzip".to_vec());
12019
12020 let body = "Hello, World! ".repeat(100);
12022 let original_size = body.len();
12023
12024 let response = Response::ok()
12025 .header("content-type", b"text/plain".to_vec())
12026 .body(ResponseBody::Bytes(body.into_bytes()));
12027
12028 let result = futures_executor::block_on(middleware.after(&ctx, &req, response));
12030
12031 let encoding = result
12033 .headers()
12034 .iter()
12035 .find(|(name, _)| name.eq_ignore_ascii_case("content-encoding"));
12036 assert!(encoding.is_some(), "Large response should be compressed");
12037
12038 let (_, value) = encoding.unwrap();
12039 assert_eq!(value, b"gzip");
12040
12041 let vary = result
12043 .headers()
12044 .iter()
12045 .find(|(name, _)| name.eq_ignore_ascii_case("vary"));
12046 assert!(vary.is_some(), "Should have Vary header");
12047
12048 if let ResponseBody::Bytes(compressed) = result.body_ref() {
12050 assert!(
12051 compressed.len() < original_size,
12052 "Compressed size should be smaller"
12053 );
12054 } else {
12055 panic!("Expected Bytes body");
12056 }
12057 }
12058
12059 #[test]
12060 fn compression_skips_without_accept_encoding() {
12061 let config = CompressionConfig::new().min_size(10);
12062 let middleware = CompressionMiddleware::with_config(config);
12063 let ctx = test_context();
12064
12065 let req = Request::new(Method::Get, "/");
12067
12068 let body = "Hello, World! ".repeat(100);
12069 let response = Response::ok()
12070 .header("content-type", b"text/plain".to_vec())
12071 .body(ResponseBody::Bytes(body.into_bytes()));
12072
12073 let result = futures_executor::block_on(middleware.after(&ctx, &req, response));
12074
12075 let has_encoding = result
12077 .headers()
12078 .iter()
12079 .any(|(name, _)| name.eq_ignore_ascii_case("content-encoding"));
12080 assert!(!has_encoding, "Should not compress without Accept-Encoding");
12081 }
12082
12083 #[test]
12084 fn compression_skips_already_compressed_content() {
12085 let config = CompressionConfig::new().min_size(10);
12086 let middleware = CompressionMiddleware::with_config(config);
12087 let ctx = test_context();
12088
12089 let mut req = Request::new(Method::Get, "/");
12091 req.headers_mut()
12092 .insert("accept-encoding", b"gzip".to_vec());
12093
12094 let body = "Some image data".repeat(100);
12096 let response = Response::ok()
12097 .header("content-type", b"image/jpeg".to_vec())
12098 .body(ResponseBody::Bytes(body.into_bytes()));
12099
12100 let result = futures_executor::block_on(middleware.after(&ctx, &req, response));
12101
12102 let has_encoding = result
12104 .headers()
12105 .iter()
12106 .any(|(name, _)| name.eq_ignore_ascii_case("content-encoding"));
12107 assert!(
12108 !has_encoding,
12109 "Should not compress already-compressed content types"
12110 );
12111 }
12112
12113 #[test]
12114 fn compression_skips_if_already_has_content_encoding() {
12115 let config = CompressionConfig::new().min_size(10);
12116 let middleware = CompressionMiddleware::with_config(config);
12117 let ctx = test_context();
12118
12119 let mut req = Request::new(Method::Get, "/");
12121 req.headers_mut()
12122 .insert("accept-encoding", b"gzip".to_vec());
12123
12124 let body = "Hello, World! ".repeat(100);
12126 let response = Response::ok()
12127 .header("content-type", b"text/plain".to_vec())
12128 .header("content-encoding", b"br".to_vec())
12129 .body(ResponseBody::Bytes(body.into_bytes()));
12130
12131 let result = futures_executor::block_on(middleware.after(&ctx, &req, response));
12132
12133 let encodings: Vec<_> = result
12135 .headers()
12136 .iter()
12137 .filter(|(name, _)| name.eq_ignore_ascii_case("content-encoding"))
12138 .collect();
12139
12140 assert_eq!(encodings.len(), 1);
12142 assert_eq!(encodings[0].1, b"br");
12143 }
12144
12145 #[test]
12146 fn accepts_gzip_parses_header_correctly() {
12147 let mut req = Request::new(Method::Get, "/");
12151 req.headers_mut()
12152 .insert("accept-encoding", b"gzip".to_vec());
12153 assert!(CompressionMiddleware::accepts_gzip(&req));
12154
12155 let mut req = Request::new(Method::Get, "/");
12157 req.headers_mut()
12158 .insert("accept-encoding", b"deflate, gzip, br".to_vec());
12159 assert!(CompressionMiddleware::accepts_gzip(&req));
12160
12161 let mut req = Request::new(Method::Get, "/");
12163 req.headers_mut()
12164 .insert("accept-encoding", b"gzip;q=1.0, identity;q=0.5".to_vec());
12165 assert!(CompressionMiddleware::accepts_gzip(&req));
12166
12167 let mut req = Request::new(Method::Get, "/");
12169 req.headers_mut().insert("accept-encoding", b"*".to_vec());
12170 assert!(CompressionMiddleware::accepts_gzip(&req));
12171
12172 let mut req = Request::new(Method::Get, "/");
12174 req.headers_mut()
12175 .insert("accept-encoding", b"deflate, br".to_vec());
12176 assert!(!CompressionMiddleware::accepts_gzip(&req));
12177
12178 let req_no_header = Request::new(Method::Get, "/");
12180 assert!(!CompressionMiddleware::accepts_gzip(&req_no_header));
12181 }
12182
12183 #[test]
12184 fn compression_middleware_name() {
12185 let middleware = CompressionMiddleware::new();
12186 assert_eq!(middleware.name(), "Compression");
12187 }
12188}
12189
12190#[cfg(test)]
12195mod request_inspection_tests {
12196 use super::*;
12197 use crate::request::Method;
12198 use crate::response::ResponseBody;
12199
12200 fn test_context() -> RequestContext {
12201 RequestContext::new(asupersync::Cx::for_testing(), 1)
12202 }
12203
12204 #[test]
12205 fn inspection_middleware_default_creates_normal_verbosity() {
12206 let mw = RequestInspectionMiddleware::new();
12207 assert_eq!(mw.verbosity, InspectionVerbosity::Normal);
12208 assert_eq!(mw.slow_threshold_ms, 1000);
12209 assert_eq!(mw.max_body_preview, 2048);
12210 assert_eq!(mw.name(), "RequestInspection");
12211 }
12212
12213 #[test]
12214 fn inspection_middleware_builder_methods() {
12215 let mw = RequestInspectionMiddleware::new()
12216 .verbosity(InspectionVerbosity::Verbose)
12217 .slow_threshold_ms(500)
12218 .max_body_preview(4096)
12219 .log_config(LogConfig::development())
12220 .redact_header("x-api-key");
12221
12222 assert_eq!(mw.verbosity, InspectionVerbosity::Verbose);
12223 assert_eq!(mw.slow_threshold_ms, 500);
12224 assert_eq!(mw.max_body_preview, 4096);
12225 assert!(mw.redact_headers.contains("x-api-key"));
12226 assert!(mw.redact_headers.contains("authorization"));
12228 assert!(mw.redact_headers.contains("cookie"));
12229 }
12230
12231 #[test]
12232 fn inspection_before_continues_processing() {
12233 let mw = RequestInspectionMiddleware::new().verbosity(InspectionVerbosity::Minimal);
12234 let ctx = test_context();
12235 let mut req = Request::new(Method::Post, "/api/users");
12236
12237 let result = futures_executor::block_on(mw.before(&ctx, &mut req));
12238 assert!(result.is_continue());
12239 }
12240
12241 #[test]
12242 fn inspection_after_returns_response_unchanged() {
12243 let mw = RequestInspectionMiddleware::new().verbosity(InspectionVerbosity::Minimal);
12244 let ctx = test_context();
12245 let mut req = Request::new(Method::Get, "/health");
12246
12247 let _ = futures_executor::block_on(mw.before(&ctx, &mut req));
12249
12250 let response = Response::ok().body(ResponseBody::Bytes(b"OK".to_vec()));
12251
12252 let result = futures_executor::block_on(mw.after(&ctx, &req, response));
12253 assert_eq!(result.status().as_u16(), 200);
12254 assert_eq!(result.body_ref().len(), 2);
12255 }
12256
12257 #[test]
12258 fn inspection_stores_start_extension() {
12259 let mw = RequestInspectionMiddleware::new();
12260 let ctx = test_context();
12261 let mut req = Request::new(Method::Get, "/");
12262
12263 let _ = futures_executor::block_on(mw.before(&ctx, &mut req));
12264
12265 assert!(req.get_extension::<InspectionStart>().is_some());
12267 }
12268
12269 #[test]
12270 fn inspection_all_verbosity_levels_continue() {
12271 for verbosity in [
12272 InspectionVerbosity::Minimal,
12273 InspectionVerbosity::Normal,
12274 InspectionVerbosity::Verbose,
12275 ] {
12276 let mw = RequestInspectionMiddleware::new().verbosity(verbosity);
12277 let ctx = test_context();
12278 let mut req = Request::new(Method::Get, "/test");
12279 req.headers_mut()
12280 .insert("content-type", b"text/plain".to_vec());
12281
12282 let result = futures_executor::block_on(mw.before(&ctx, &mut req));
12283 assert!(
12284 result.is_continue(),
12285 "Verbosity {verbosity:?} should continue"
12286 );
12287 }
12288 }
12289
12290 #[test]
12291 fn inspection_verbose_with_json_body() {
12292 let mw = RequestInspectionMiddleware::new().verbosity(InspectionVerbosity::Verbose);
12293 let ctx = test_context();
12294 let body = br#"{"name":"Alice","age":30}"#;
12295 let mut req = Request::new(Method::Post, "/api/users");
12296 req.headers_mut()
12297 .insert("content-type", b"application/json".to_vec());
12298 req.set_body(Body::Bytes(body.to_vec()));
12299
12300 let result = futures_executor::block_on(mw.before(&ctx, &mut req));
12301 assert!(result.is_continue());
12302 }
12303
12304 #[test]
12305 fn inspection_verbose_after_with_json_response() {
12306 let mw = RequestInspectionMiddleware::new().verbosity(InspectionVerbosity::Verbose);
12307 let ctx = test_context();
12308 let mut req = Request::new(Method::Get, "/api/users/1");
12309
12310 let _ = futures_executor::block_on(mw.before(&ctx, &mut req));
12311
12312 let response = Response::ok()
12313 .header("content-type", b"application/json".to_vec())
12314 .body(ResponseBody::Bytes(br#"{"id":1,"name":"Alice"}"#.to_vec()));
12315
12316 let result = futures_executor::block_on(mw.after(&ctx, &req, response));
12317 assert_eq!(result.status().as_u16(), 200);
12318 }
12319
12320 #[test]
12321 fn inspection_redacts_sensitive_headers() {
12322 let mw = RequestInspectionMiddleware::new();
12323
12324 assert!(mw.redact_headers.contains("authorization"));
12326 assert!(mw.redact_headers.contains("proxy-authorization"));
12327 assert!(mw.redact_headers.contains("cookie"));
12328 assert!(mw.redact_headers.contains("set-cookie"));
12329 }
12330
12331 #[test]
12332 fn inspection_format_headers_redacts() {
12333 let mw = RequestInspectionMiddleware::new().redact_header("x-secret");
12334
12335 let headers = vec![
12336 ("content-type", b"text/plain".as_slice()),
12337 ("x-secret", b"my-secret-value".as_slice()),
12338 ("x-normal", b"visible".as_slice()),
12339 ];
12340
12341 let output = mw.format_inspection_headers(headers.into_iter());
12342 assert!(output.contains("content-type: text/plain"));
12343 assert!(output.contains("x-secret: [REDACTED]"));
12344 assert!(output.contains("x-normal: visible"));
12345 assert!(!output.contains("my-secret-value"));
12346 }
12347
12348 #[test]
12349 fn inspection_format_body_preview_truncates() {
12350 let mw = RequestInspectionMiddleware::new().max_body_preview(10);
12351
12352 let body = b"Hello, World! This is a long body.";
12353 let result = mw.format_body_preview(body, None);
12354 assert!(result.is_some());
12355 let text = result.unwrap();
12356 assert!(text.ends_with("..."));
12357 assert!(text.len() <= 15); }
12359
12360 #[test]
12361 fn inspection_format_body_preview_empty() {
12362 let mw = RequestInspectionMiddleware::new();
12363 assert!(mw.format_body_preview(b"", None).is_none());
12364 }
12365
12366 #[test]
12367 fn inspection_format_body_preview_zero_max() {
12368 let mw = RequestInspectionMiddleware::new().max_body_preview(0);
12369 assert!(mw.format_body_preview(b"hello", None).is_none());
12370 }
12371
12372 #[test]
12373 fn inspection_format_body_preview_json_pretty() {
12374 let mw = RequestInspectionMiddleware::new();
12375 let body = br#"{"key":"value","num":42}"#;
12376 let ct = b"application/json".as_slice();
12377 let result = mw.format_body_preview(body, Some(ct));
12378 assert!(result.is_some());
12379 let text = result.unwrap();
12380 assert!(text.contains('\n'));
12382 assert!(text.contains("\"key\": \"value\""));
12383 }
12384
12385 #[test]
12386 fn inspection_format_body_preview_non_json() {
12387 let mw = RequestInspectionMiddleware::new();
12388 let body = b"Hello, World!";
12389 let ct = b"text/plain".as_slice();
12390 let result = mw.format_body_preview(body, Some(ct));
12391 assert_eq!(result.unwrap(), "Hello, World!");
12392 }
12393
12394 #[test]
12395 fn inspection_format_body_preview_binary() {
12396 let mw = RequestInspectionMiddleware::new();
12397 let body: &[u8] = &[0xFF, 0xFE, 0xFD, 0x00];
12398 let result = mw.format_body_preview(body, None);
12399 assert!(result.is_some());
12400 assert!(result.unwrap().contains("binary"));
12401 }
12402
12403 #[test]
12404 fn try_pretty_json_valid_object() {
12405 let result = try_pretty_json(r#"{"a":"b","c":1}"#);
12406 assert!(result.is_some());
12407 let pretty = result.unwrap();
12408 assert!(pretty.contains('\n'));
12409 assert!(pretty.contains(" \"a\": \"b\""));
12410 }
12411
12412 #[test]
12413 fn try_pretty_json_valid_array() {
12414 let result = try_pretty_json(r"[1,2,3]");
12415 assert!(result.is_some());
12416 let pretty = result.unwrap();
12417 assert!(pretty.contains('\n'));
12418 }
12419
12420 #[test]
12421 fn try_pretty_json_empty_object() {
12422 let result = try_pretty_json("{}");
12423 assert!(result.is_some());
12424 assert_eq!(result.unwrap(), "{}");
12425 }
12426
12427 #[test]
12428 fn try_pretty_json_empty_array() {
12429 let result = try_pretty_json("[]");
12430 assert!(result.is_some());
12431 assert_eq!(result.unwrap(), "[]");
12432 }
12433
12434 #[test]
12435 fn try_pretty_json_not_json() {
12436 assert!(try_pretty_json("hello world").is_none());
12437 assert!(try_pretty_json("12345").is_none());
12438 }
12439
12440 #[test]
12441 fn try_pretty_json_nested() {
12442 let input = r#"{"user":{"name":"Alice","roles":["admin","user"]}}"#;
12443 let result = try_pretty_json(input);
12444 assert!(result.is_some());
12445 let pretty = result.unwrap();
12446 assert!(pretty.contains("\"user\":"));
12447 assert!(pretty.contains("\"name\": \"Alice\""));
12448 assert!(pretty.contains("\"roles\":"));
12449 }
12450
12451 #[test]
12452 fn try_pretty_json_with_escapes() {
12453 let input = r#"{"msg":"hello \"world\""}"#;
12454 let result = try_pretty_json(input);
12455 assert!(result.is_some());
12456 let pretty = result.unwrap();
12457 assert!(pretty.contains(r#"\"world\""#));
12458 }
12459
12460 #[test]
12461 fn inspection_name() {
12462 let mw = RequestInspectionMiddleware::new();
12463 assert_eq!(mw.name(), "RequestInspection");
12464 }
12465
12466 #[test]
12467 fn inspection_default_via_default_trait() {
12468 let mw = RequestInspectionMiddleware::default();
12469 assert_eq!(mw.verbosity, InspectionVerbosity::Normal);
12470 assert_eq!(mw.slow_threshold_ms, 1000);
12471 }
12472
12473 #[test]
12474 fn inspection_with_query_string() {
12475 let mw = RequestInspectionMiddleware::new().verbosity(InspectionVerbosity::Minimal);
12476 let ctx = test_context();
12477 let mut req = Request::new(Method::Get, "/search");
12478 req.set_query(Some("q=rust&page=1".to_string()));
12479
12480 let result = futures_executor::block_on(mw.before(&ctx, &mut req));
12481 assert!(result.is_continue());
12482 }
12483
12484 #[test]
12485 fn inspection_response_body_stream() {
12486 let mw = RequestInspectionMiddleware::new();
12487 let result = mw.format_response_preview(&ResponseBody::Empty, None);
12488 assert!(result.is_none());
12489 }
12490}
12491
12492#[cfg(test)]
12497mod rate_limit_tests {
12498 use super::*;
12499 use crate::request::Method;
12500 use crate::response::{ResponseBody, StatusCode};
12501 use std::time::{Duration, Instant};
12502
12503 fn test_context() -> RequestContext {
12504 RequestContext::new(asupersync::Cx::for_testing(), 1)
12505 }
12506
12507 fn run_rate_limit_before(mw: &RateLimitMiddleware, req: &mut Request) -> ControlFlow {
12508 let ctx = test_context();
12509 let fut = mw.before(&ctx, req);
12510 futures_executor::block_on(fut)
12511 }
12512
12513 fn run_rate_limit_after(mw: &RateLimitMiddleware, req: &Request, resp: Response) -> Response {
12514 let ctx = test_context();
12515 let fut = mw.after(&ctx, req, resp);
12516 futures_executor::block_on(fut)
12517 }
12518
12519 fn request_with_ip(key: &str) -> Request {
12520 let mut req = Request::new(Method::Get, "/");
12521 req.headers_mut()
12522 .insert("x-forwarded-for", key.as_bytes().to_vec());
12523 req
12524 }
12525
12526 fn rate_limit_entry_keys(
12527 store: &InMemoryRateLimitStore,
12528 algorithm: RateLimitAlgorithm,
12529 ) -> Vec<String> {
12530 match algorithm {
12531 RateLimitAlgorithm::TokenBucket => {
12532 store.token_buckets.lock().entries.keys().cloned().collect()
12533 }
12534 RateLimitAlgorithm::FixedWindow => {
12535 store.fixed_windows.lock().entries.keys().cloned().collect()
12536 }
12537 RateLimitAlgorithm::SlidingWindow => store
12538 .sliding_windows
12539 .lock()
12540 .entries
12541 .keys()
12542 .cloned()
12543 .collect(),
12544 }
12545 }
12546
12547 #[test]
12548 fn rate_limit_default_allows_requests() {
12549 let mw = RateLimitMiddleware::new();
12550 let mut req = Request::new(Method::Get, "/api/test");
12551 req.headers_mut()
12552 .insert("x-forwarded-for", b"192.168.1.1".to_vec());
12553
12554 let result = run_rate_limit_before(&mw, &mut req);
12555 assert!(result.is_continue(), "first request should be allowed");
12556 }
12557
12558 #[test]
12559 fn rate_limit_fixed_window_blocks_after_limit() {
12560 let mw = RateLimitMiddleware::builder()
12561 .requests(3)
12562 .per(Duration::from_secs(60))
12563 .algorithm(RateLimitAlgorithm::FixedWindow)
12564 .key_extractor(IpKeyExtractor)
12565 .build();
12566
12567 for i in 0..3 {
12568 let mut req = Request::new(Method::Get, "/api/test");
12569 req.headers_mut()
12570 .insert("x-forwarded-for", b"10.0.0.1".to_vec());
12571 let result = run_rate_limit_before(&mw, &mut req);
12572 assert!(
12573 result.is_continue(),
12574 "request {i} should be allowed within limit"
12575 );
12576 }
12577
12578 let mut req = Request::new(Method::Get, "/api/test");
12580 req.headers_mut()
12581 .insert("x-forwarded-for", b"10.0.0.1".to_vec());
12582 let result = run_rate_limit_before(&mw, &mut req);
12583 assert!(result.is_break(), "fourth request should be blocked");
12584
12585 if let ControlFlow::Break(resp) = result {
12587 assert_eq!(resp.status(), StatusCode::TOO_MANY_REQUESTS);
12588 }
12589 }
12590
12591 #[test]
12592 fn rate_limit_different_keys_independent() {
12593 let mw = RateLimitMiddleware::builder()
12594 .requests(2)
12595 .per(Duration::from_secs(60))
12596 .algorithm(RateLimitAlgorithm::FixedWindow)
12597 .key_extractor(IpKeyExtractor)
12598 .build();
12599
12600 for _ in 0..2 {
12602 let mut req = Request::new(Method::Get, "/");
12603 req.headers_mut()
12604 .insert("x-forwarded-for", b"1.1.1.1".to_vec());
12605 assert!(run_rate_limit_before(&mw, &mut req).is_continue());
12606 }
12607
12608 let mut req = Request::new(Method::Get, "/");
12610 req.headers_mut()
12611 .insert("x-forwarded-for", b"1.1.1.1".to_vec());
12612 assert!(run_rate_limit_before(&mw, &mut req).is_break());
12613
12614 let mut req = Request::new(Method::Get, "/");
12616 req.headers_mut()
12617 .insert("x-forwarded-for", b"2.2.2.2".to_vec());
12618 assert!(run_rate_limit_before(&mw, &mut req).is_continue());
12619 }
12620
12621 #[test]
12622 fn rate_limit_token_bucket_allows_burst() {
12623 let mw = RateLimitMiddleware::builder()
12624 .requests(5)
12625 .per(Duration::from_secs(60))
12626 .algorithm(RateLimitAlgorithm::TokenBucket)
12627 .key_extractor(IpKeyExtractor)
12628 .build();
12629
12630 for i in 0..5 {
12632 let mut req = Request::new(Method::Get, "/");
12633 req.headers_mut()
12634 .insert("x-forwarded-for", b"10.0.0.1".to_vec());
12635 let result = run_rate_limit_before(&mw, &mut req);
12636 assert!(result.is_continue(), "burst request {i} should be allowed");
12637 }
12638
12639 let mut req = Request::new(Method::Get, "/");
12641 req.headers_mut()
12642 .insert("x-forwarded-for", b"10.0.0.1".to_vec());
12643 assert!(run_rate_limit_before(&mw, &mut req).is_break());
12644 }
12645
12646 #[test]
12647 fn rate_limit_sliding_window_basic() {
12648 let mw = RateLimitMiddleware::builder()
12649 .requests(3)
12650 .per(Duration::from_secs(60))
12651 .algorithm(RateLimitAlgorithm::SlidingWindow)
12652 .key_extractor(IpKeyExtractor)
12653 .build();
12654
12655 for i in 0..3 {
12656 let mut req = Request::new(Method::Get, "/");
12657 req.headers_mut()
12658 .insert("x-forwarded-for", b"10.0.0.1".to_vec());
12659 assert!(
12660 run_rate_limit_before(&mw, &mut req).is_continue(),
12661 "sliding window request {i} should be allowed"
12662 );
12663 }
12664
12665 let mut req = Request::new(Method::Get, "/");
12667 req.headers_mut()
12668 .insert("x-forwarded-for", b"10.0.0.1".to_vec());
12669 assert!(run_rate_limit_before(&mw, &mut req).is_break());
12670 }
12671
12672 #[test]
12673 fn rate_limit_header_key_extractor() {
12674 let mw = RateLimitMiddleware::builder()
12675 .requests(2)
12676 .per(Duration::from_secs(60))
12677 .algorithm(RateLimitAlgorithm::FixedWindow)
12678 .key_extractor(HeaderKeyExtractor::new("x-api-key"))
12679 .build();
12680
12681 for _ in 0..2 {
12683 let mut req = Request::new(Method::Get, "/");
12684 req.headers_mut().insert("x-api-key", b"key-abc".to_vec());
12685 assert!(run_rate_limit_before(&mw, &mut req).is_continue());
12686 }
12687
12688 let mut req = Request::new(Method::Get, "/");
12690 req.headers_mut().insert("x-api-key", b"key-abc".to_vec());
12691 assert!(run_rate_limit_before(&mw, &mut req).is_break());
12692
12693 let mut req = Request::new(Method::Get, "/");
12695 req.headers_mut().insert("x-api-key", b"key-xyz".to_vec());
12696 assert!(run_rate_limit_before(&mw, &mut req).is_continue());
12697 }
12698
12699 #[test]
12700 fn rate_limit_path_key_extractor() {
12701 let mw = RateLimitMiddleware::builder()
12702 .requests(1)
12703 .per(Duration::from_secs(60))
12704 .algorithm(RateLimitAlgorithm::FixedWindow)
12705 .key_extractor(PathKeyExtractor)
12706 .build();
12707
12708 let mut req = Request::new(Method::Get, "/api/a");
12709 assert!(run_rate_limit_before(&mw, &mut req).is_continue());
12710
12711 let mut req = Request::new(Method::Get, "/api/a");
12713 assert!(run_rate_limit_before(&mw, &mut req).is_break());
12714
12715 let mut req = Request::new(Method::Get, "/api/b");
12717 assert!(run_rate_limit_before(&mw, &mut req).is_continue());
12718 }
12719
12720 #[test]
12721 fn rate_limit_no_key_skips_limiting() {
12722 let mw = RateLimitMiddleware::builder()
12723 .requests(1)
12724 .per(Duration::from_secs(60))
12725 .algorithm(RateLimitAlgorithm::FixedWindow)
12726 .key_extractor(HeaderKeyExtractor::new("x-api-key"))
12727 .build();
12728
12729 let mut req = Request::new(Method::Get, "/");
12731 assert!(run_rate_limit_before(&mw, &mut req).is_continue());
12732
12733 for _ in 0..10 {
12735 let mut req = Request::new(Method::Get, "/");
12736 assert!(run_rate_limit_before(&mw, &mut req).is_continue());
12737 }
12738 }
12739
12740 #[test]
12741 fn rate_limit_response_headers_on_success() {
12742 let mw = RateLimitMiddleware::builder()
12743 .requests(10)
12744 .per(Duration::from_secs(60))
12745 .algorithm(RateLimitAlgorithm::FixedWindow)
12746 .key_extractor(IpKeyExtractor)
12747 .build();
12748
12749 let mut req = Request::new(Method::Get, "/");
12750 req.headers_mut()
12751 .insert("x-forwarded-for", b"10.0.0.1".to_vec());
12752 let cf = run_rate_limit_before(&mw, &mut req);
12753 assert!(cf.is_continue());
12754
12755 let resp = Response::with_status(StatusCode::OK);
12756 let resp = run_rate_limit_after(&mw, &req, resp);
12757
12758 let headers = resp.headers();
12760 let has_limit = headers
12761 .iter()
12762 .any(|(n, _)| n.eq_ignore_ascii_case("x-ratelimit-limit"));
12763 let has_remaining = headers
12764 .iter()
12765 .any(|(n, _)| n.eq_ignore_ascii_case("x-ratelimit-remaining"));
12766 let has_reset = headers
12767 .iter()
12768 .any(|(n, _)| n.eq_ignore_ascii_case("x-ratelimit-reset"));
12769
12770 assert!(has_limit, "should have X-RateLimit-Limit header");
12771 assert!(has_remaining, "should have X-RateLimit-Remaining header");
12772 assert!(has_reset, "should have X-RateLimit-Reset header");
12773
12774 let limit_val = headers
12776 .iter()
12777 .find(|(n, _)| n.eq_ignore_ascii_case("x-ratelimit-limit"))
12778 .map(|(_, v)| std::str::from_utf8(v).unwrap().to_string())
12779 .unwrap();
12780 assert_eq!(limit_val, "10");
12781 }
12782
12783 #[test]
12784 fn rate_limit_429_response_has_retry_after() {
12785 let mw = RateLimitMiddleware::builder()
12786 .requests(1)
12787 .per(Duration::from_secs(60))
12788 .algorithm(RateLimitAlgorithm::FixedWindow)
12789 .key_extractor(IpKeyExtractor)
12790 .build();
12791
12792 let mut req = Request::new(Method::Get, "/");
12794 req.headers_mut()
12795 .insert("x-forwarded-for", b"10.0.0.1".to_vec());
12796 assert!(run_rate_limit_before(&mw, &mut req).is_continue());
12797
12798 let mut req = Request::new(Method::Get, "/");
12800 req.headers_mut()
12801 .insert("x-forwarded-for", b"10.0.0.1".to_vec());
12802 let result = run_rate_limit_before(&mw, &mut req);
12803
12804 if let ControlFlow::Break(resp) = result {
12805 assert_eq!(resp.status(), StatusCode::TOO_MANY_REQUESTS);
12806
12807 let has_retry = resp
12809 .headers()
12810 .iter()
12811 .any(|(n, _)| n.eq_ignore_ascii_case("retry-after"));
12812 assert!(has_retry, "429 response should have Retry-After header");
12813
12814 let has_ct = resp
12816 .headers()
12817 .iter()
12818 .any(|(n, v)| n.eq_ignore_ascii_case("content-type") && v == b"application/json");
12819 assert!(has_ct, "429 response should have JSON content type");
12820 } else {
12821 panic!("expected Break(429)");
12822 }
12823 }
12824
12825 #[test]
12826 fn rate_limit_no_headers_when_disabled() {
12827 let mw = RateLimitMiddleware::builder()
12828 .requests(10)
12829 .per(Duration::from_secs(60))
12830 .algorithm(RateLimitAlgorithm::FixedWindow)
12831 .key_extractor(IpKeyExtractor)
12832 .include_headers(false)
12833 .build();
12834
12835 let mut req = Request::new(Method::Get, "/");
12836 req.headers_mut()
12837 .insert("x-forwarded-for", b"10.0.0.1".to_vec());
12838 assert!(run_rate_limit_before(&mw, &mut req).is_continue());
12839
12840 let resp = Response::with_status(StatusCode::OK);
12841 let resp = run_rate_limit_after(&mw, &req, resp);
12842
12843 let has_limit = resp
12844 .headers()
12845 .iter()
12846 .any(|(n, _)| n.eq_ignore_ascii_case("x-ratelimit-limit"));
12847 assert!(
12848 !has_limit,
12849 "should NOT have rate limit headers when disabled"
12850 );
12851 }
12852
12853 #[test]
12854 fn rate_limit_custom_retry_message() {
12855 let mw = RateLimitMiddleware::builder()
12856 .requests(1)
12857 .per(Duration::from_secs(60))
12858 .algorithm(RateLimitAlgorithm::FixedWindow)
12859 .key_extractor(IpKeyExtractor)
12860 .retry_message("Slow down, partner!")
12861 .build();
12862
12863 let mut req = Request::new(Method::Get, "/");
12865 req.headers_mut()
12866 .insert("x-forwarded-for", b"10.0.0.1".to_vec());
12867 run_rate_limit_before(&mw, &mut req);
12868
12869 let mut req = Request::new(Method::Get, "/");
12871 req.headers_mut()
12872 .insert("x-forwarded-for", b"10.0.0.1".to_vec());
12873 if let ControlFlow::Break(resp) = run_rate_limit_before(&mw, &mut req) {
12874 if let ResponseBody::Bytes(body) = resp.body_ref() {
12875 let body_str = std::str::from_utf8(body).unwrap();
12876 assert!(
12877 body_str.contains("Slow down, partner!"),
12878 "expected custom message in body, got: {body_str}"
12879 );
12880 } else {
12881 panic!("expected Bytes body");
12882 }
12883 } else {
12884 panic!("expected Break(429)");
12885 }
12886 }
12887
12888 #[test]
12889 fn rate_limit_ip_extractor_x_forwarded_for() {
12890 let extractor = IpKeyExtractor;
12891 let mut req = Request::new(Method::Get, "/");
12892 req.headers_mut()
12893 .insert("x-forwarded-for", b"1.2.3.4, 5.6.7.8".to_vec());
12894 assert_eq!(extractor.extract_key(&req), Some("1.2.3.4".to_string()));
12895 }
12896
12897 #[test]
12898 fn rate_limit_ip_extractor_x_real_ip() {
12899 let extractor = IpKeyExtractor;
12900 let mut req = Request::new(Method::Get, "/");
12901 req.headers_mut().insert("x-real-ip", b"9.8.7.6".to_vec());
12902 assert_eq!(extractor.extract_key(&req), Some("9.8.7.6".to_string()));
12903 }
12904
12905 #[test]
12906 fn rate_limit_ip_extractor_fallback() {
12907 let extractor = IpKeyExtractor;
12908 let req = Request::new(Method::Get, "/");
12909 assert_eq!(extractor.extract_key(&req), Some("unknown".to_string()));
12910 }
12911
12912 #[test]
12914 fn connected_ip_extractor_with_remote_addr() {
12915 use std::net::{IpAddr, Ipv4Addr};
12916
12917 let extractor = ConnectedIpKeyExtractor;
12918 let mut req = Request::new(Method::Get, "/");
12919 req.insert_extension(RemoteAddr(IpAddr::V4(Ipv4Addr::new(192, 168, 1, 100))));
12920
12921 assert_eq!(
12922 extractor.extract_key(&req),
12923 Some("192.168.1.100".to_string())
12924 );
12925 }
12926
12927 #[test]
12928 fn connected_ip_extractor_without_remote_addr() {
12929 let extractor = ConnectedIpKeyExtractor;
12930 let req = Request::new(Method::Get, "/");
12931
12932 assert_eq!(extractor.extract_key(&req), None);
12934 }
12935
12936 #[test]
12937 fn connected_ip_extractor_ignores_headers() {
12938 use std::net::{IpAddr, Ipv4Addr};
12939
12940 let extractor = ConnectedIpKeyExtractor;
12941 let mut req = Request::new(Method::Get, "/");
12942 req.insert_extension(RemoteAddr(IpAddr::V4(Ipv4Addr::new(10, 0, 0, 1))));
12943 req.headers_mut()
12945 .insert("x-forwarded-for", b"1.2.3.4".to_vec());
12946
12947 assert_eq!(extractor.extract_key(&req), Some("10.0.0.1".to_string()));
12949 }
12950
12951 #[test]
12953 fn trusted_proxy_extractor_from_trusted_proxy() {
12954 use std::net::{IpAddr, Ipv4Addr};
12955
12956 let extractor = TrustedProxyIpKeyExtractor::new().trust_cidr("10.0.0.0/8");
12957
12958 let mut req = Request::new(Method::Get, "/");
12959 req.insert_extension(RemoteAddr(IpAddr::V4(Ipv4Addr::new(10, 0, 0, 1))));
12961 req.headers_mut()
12963 .insert("x-forwarded-for", b"203.0.113.50".to_vec());
12964
12965 assert_eq!(
12967 extractor.extract_key(&req),
12968 Some("203.0.113.50".to_string())
12969 );
12970 }
12971
12972 #[test]
12973 fn trusted_proxy_extractor_from_untrusted_direct() {
12974 use std::net::{IpAddr, Ipv4Addr};
12975
12976 let extractor = TrustedProxyIpKeyExtractor::new().trust_cidr("10.0.0.0/8");
12977
12978 let mut req = Request::new(Method::Get, "/");
12979 req.insert_extension(RemoteAddr(IpAddr::V4(Ipv4Addr::new(203, 0, 113, 50))));
12981 req.headers_mut()
12983 .insert("x-forwarded-for", b"1.2.3.4".to_vec());
12984
12985 assert_eq!(
12987 extractor.extract_key(&req),
12988 Some("203.0.113.50".to_string())
12989 );
12990 }
12991
12992 #[test]
12993 fn trusted_proxy_extractor_no_remote_addr() {
12994 let extractor = TrustedProxyIpKeyExtractor::new().trust_loopback();
12995
12996 let mut req = Request::new(Method::Get, "/");
12997 req.headers_mut()
12999 .insert("x-forwarded-for", b"1.2.3.4".to_vec());
13000
13001 assert_eq!(extractor.extract_key(&req), None);
13002 }
13003
13004 #[test]
13005 fn trusted_proxy_extractor_loopback_ipv4() {
13006 use std::net::{IpAddr, Ipv4Addr};
13007
13008 let extractor = TrustedProxyIpKeyExtractor::new().trust_loopback();
13009
13010 let mut req = Request::new(Method::Get, "/");
13011 req.insert_extension(RemoteAddr(IpAddr::V4(Ipv4Addr::LOCALHOST)));
13012 req.headers_mut()
13013 .insert("x-forwarded-for", b"8.8.8.8".to_vec());
13014
13015 assert_eq!(extractor.extract_key(&req), Some("8.8.8.8".to_string()));
13016 }
13017
13018 #[test]
13019 fn trusted_proxy_extractor_loopback_ipv6() {
13020 use std::net::{IpAddr, Ipv6Addr};
13021
13022 let extractor = TrustedProxyIpKeyExtractor::new().trust_loopback();
13023
13024 let mut req = Request::new(Method::Get, "/");
13025 req.insert_extension(RemoteAddr(IpAddr::V6(Ipv6Addr::LOCALHOST)));
13026 req.headers_mut()
13027 .insert("x-forwarded-for", b"8.8.8.8".to_vec());
13028
13029 assert_eq!(extractor.extract_key(&req), Some("8.8.8.8".to_string()));
13030 }
13031
13032 #[test]
13033 fn cidr_parsing() {
13034 assert!(parse_cidr("10.0.0.0/8").is_some());
13036 assert!(parse_cidr("192.168.1.0/24").is_some());
13037 assert!(parse_cidr("0.0.0.0/0").is_some());
13038 assert!(parse_cidr("::1/128").is_some());
13039 assert!(parse_cidr("::/0").is_some());
13040
13041 assert!(parse_cidr("10.0.0.0/33").is_none()); assert!(parse_cidr("invalid").is_none());
13044 assert!(parse_cidr("10.0.0.0").is_none()); }
13046
13047 #[test]
13048 fn ip_in_cidr_matching() {
13049 use std::net::{IpAddr, Ipv4Addr};
13050
13051 let cidr_10 = IpAddr::V4(Ipv4Addr::new(10, 0, 0, 0));
13052
13053 assert!(ip_in_cidr(
13055 IpAddr::V4(Ipv4Addr::new(10, 0, 0, 1)),
13056 cidr_10,
13057 8
13058 ));
13059 assert!(ip_in_cidr(
13060 IpAddr::V4(Ipv4Addr::new(10, 255, 255, 255)),
13061 cidr_10,
13062 8
13063 ));
13064
13065 assert!(!ip_in_cidr(
13067 IpAddr::V4(Ipv4Addr::new(11, 0, 0, 1)),
13068 cidr_10,
13069 8
13070 ));
13071 assert!(!ip_in_cidr(
13072 IpAddr::V4(Ipv4Addr::new(192, 168, 1, 1)),
13073 cidr_10,
13074 8
13075 ));
13076 }
13077
13078 #[test]
13079 fn rate_limit_composite_key_extractor() {
13080 let extractor =
13081 CompositeKeyExtractor::new(vec![Box::new(IpKeyExtractor), Box::new(PathKeyExtractor)]);
13082
13083 let mut req = Request::new(Method::Get, "/api/users");
13084 req.headers_mut()
13085 .insert("x-forwarded-for", b"10.0.0.1".to_vec());
13086
13087 let key = extractor.extract_key(&req);
13088 assert_eq!(key, Some("10.0.0.1:/api/users".to_string()));
13089 }
13090
13091 #[test]
13092 fn rate_limit_builder_defaults() {
13093 let mw = RateLimitMiddleware::builder().build();
13094 assert_eq!(DEFAULT_RATE_LIMIT_MAX_KEYS, 65_536);
13095 assert_eq!(mw.config.max_requests, 100);
13096 assert_eq!(mw.config.window, Duration::from_secs(60));
13097 assert_eq!(mw.config.algorithm, RateLimitAlgorithm::TokenBucket);
13098 assert_eq!(mw.config.max_keys, DEFAULT_RATE_LIMIT_MAX_KEYS);
13099 assert_eq!(mw.store.max_keys, DEFAULT_RATE_LIMIT_MAX_KEYS);
13100 assert!(mw.config.include_headers);
13101 }
13102
13103 #[test]
13104 fn rate_limit_builder_applies_custom_key_bound() {
13105 let mw = RateLimitMiddleware::builder().max_keys(7).build();
13106 assert_eq!(mw.config.max_keys, 7);
13107 assert_eq!(mw.store.max_keys, 7);
13108 }
13109
13110 #[test]
13111 fn rate_limit_store_bounds_every_algorithm_and_fails_closed() {
13112 let window = Duration::from_secs(60);
13113
13114 for algorithm in [
13115 RateLimitAlgorithm::TokenBucket,
13116 RateLimitAlgorithm::FixedWindow,
13117 RateLimitAlgorithm::SlidingWindow,
13118 ] {
13119 let store = InMemoryRateLimitStore::with_max_keys(2);
13120 assert!(store.check("resident-a", algorithm, 10, window).allowed);
13121 assert!(store.check("resident-b", algorithm, 10, window).allowed);
13122
13123 let saturated = store.check("unseen", algorithm, 10, window);
13124 assert!(!saturated.allowed, "{algorithm:?} must fail closed");
13125 assert_eq!(saturated.limit, 10);
13126 assert_eq!(saturated.remaining, 0);
13127 assert_eq!(saturated.reset_after_secs, 60);
13128
13129 let keys = rate_limit_entry_keys(&store, algorithm);
13130 assert_eq!(keys.len(), 2, "{algorithm:?} exceeded its key bound");
13131 assert!(keys.contains(&"resident-a".to_string()));
13132 assert!(keys.contains(&"resident-b".to_string()));
13133 assert!(!keys.contains(&"unseen".to_string()));
13134 }
13135 }
13136
13137 #[test]
13138 fn saturated_retry_after_rounds_fractional_windows_up() {
13139 let store = InMemoryRateLimitStore::with_max_keys(0);
13140 let result = store.check(
13141 "unseen",
13142 RateLimitAlgorithm::FixedWindow,
13143 10,
13144 Duration::from_millis(1_500),
13145 );
13146 assert!(!result.allowed);
13147 assert_eq!(result.reset_after_secs, 2);
13148 }
13149
13150 #[test]
13151 fn rate_limit_saturation_returns_429_without_resetting_resident_counter() {
13152 for algorithm in [
13153 RateLimitAlgorithm::TokenBucket,
13154 RateLimitAlgorithm::FixedWindow,
13155 RateLimitAlgorithm::SlidingWindow,
13156 ] {
13157 let mw = RateLimitMiddleware::builder()
13158 .requests(2)
13159 .per(Duration::from_secs(60))
13160 .algorithm(algorithm)
13161 .key_extractor(IpKeyExtractor)
13162 .max_keys(1)
13163 .build();
13164
13165 for _ in 0..2 {
13166 let mut resident = request_with_ip("resident");
13167 assert!(run_rate_limit_before(&mw, &mut resident).is_continue());
13168 }
13169
13170 let mut unseen = request_with_ip("unseen");
13171 let ControlFlow::Break(response) = run_rate_limit_before(&mw, &mut unseen) else {
13172 panic!("{algorithm:?} unseen key must fail closed while saturated");
13173 };
13174 assert_eq!(response.status(), StatusCode::TOO_MANY_REQUESTS);
13175 assert!(response.headers().iter().any(|(name, value)| {
13176 name.eq_ignore_ascii_case("retry-after") && value.as_slice() == b"60"
13177 }));
13178
13179 let mut resident = request_with_ip("resident");
13180 assert!(
13181 run_rate_limit_before(&mw, &mut resident).is_break(),
13182 "{algorithm:?} resident counter must remain exhausted"
13183 );
13184
13185 let keys = rate_limit_entry_keys(&mw.store, algorithm);
13186 assert_eq!(keys, vec!["resident".to_string()]);
13187 }
13188 }
13189
13190 #[test]
13191 fn rate_limit_store_reclaims_algorithm_specific_stale_entries() {
13192 let window = Duration::from_secs(3_600);
13193
13194 for algorithm in [
13195 RateLimitAlgorithm::TokenBucket,
13196 RateLimitAlgorithm::FixedWindow,
13197 RateLimitAlgorithm::SlidingWindow,
13198 ] {
13199 let store = InMemoryRateLimitStore::with_max_keys(1);
13200 let start = Instant::now();
13201 assert!(
13202 store
13203 .check_at("stale", algorithm, 10, window, start)
13204 .allowed
13205 );
13206
13207 let replacement_at = if algorithm == RateLimitAlgorithm::SlidingWindow {
13208 start + Duration::from_secs(7_201)
13209 } else {
13210 start + Duration::from_secs(3_601)
13211 };
13212
13213 assert!(
13214 store
13215 .check_at("replacement", algorithm, 10, window, replacement_at)
13216 .allowed,
13217 "{algorithm:?} should reclaim its stale entry"
13218 );
13219 assert_eq!(
13220 rate_limit_entry_keys(&store, algorithm),
13221 vec!["replacement".to_string()]
13222 );
13223 }
13224 }
13225
13226 #[test]
13227 fn sliding_window_retains_entries_until_two_windows_are_inactive() {
13228 let algorithm = RateLimitAlgorithm::SlidingWindow;
13229 let window = Duration::from_secs(3_600);
13230 let store = InMemoryRateLimitStore::with_max_keys(1);
13231 let start = Instant::now();
13232 assert!(
13233 store
13234 .check_at("resident", algorithm, 10, window, start)
13235 .allowed
13236 );
13237
13238 assert!(
13239 !store
13240 .check_at(
13241 "too-early",
13242 algorithm,
13243 10,
13244 window,
13245 start + Duration::from_secs(3_601),
13246 )
13247 .allowed
13248 );
13249 assert_eq!(
13250 rate_limit_entry_keys(&store, algorithm),
13251 vec!["resident".to_string()]
13252 );
13253
13254 assert!(
13255 store
13256 .check_at(
13257 "replacement",
13258 algorithm,
13259 10,
13260 window,
13261 start + Duration::from_secs(7_202),
13262 )
13263 .allowed
13264 );
13265 assert_eq!(
13266 rate_limit_entry_keys(&store, algorithm),
13267 vec!["replacement".to_string()]
13268 );
13269 }
13270
13271 #[test]
13272 fn sliding_window_preserves_fractional_progress_when_rotating() {
13273 let algorithm = RateLimitAlgorithm::SlidingWindow;
13274 let window = Duration::from_secs(100);
13275 let store = InMemoryRateLimitStore::with_max_keys(1);
13276 let start = Instant::now();
13277
13278 for _ in 0..10 {
13279 assert!(
13280 store
13281 .check_at("resident", algorithm, 10, window, start)
13282 .allowed
13283 );
13284 }
13285 assert!(
13286 !store
13287 .check_at("resident", algorithm, 10, window, start)
13288 .allowed
13289 );
13290
13291 let half_into_next_window = store.check_at(
13292 "resident",
13293 algorithm,
13294 10,
13295 window,
13296 start + Duration::from_secs(150),
13297 );
13298 assert!(half_into_next_window.allowed);
13299 assert_eq!(half_into_next_window.remaining, 4);
13300 assert_eq!(half_into_next_window.reset_after_secs, 50);
13301 }
13302
13303 #[test]
13304 fn cleanup_uses_each_resident_entry_retention_policy() {
13305 let short_window = Duration::from_secs(1);
13306 let long_window = Duration::from_secs(3_600);
13307
13308 for algorithm in [
13309 RateLimitAlgorithm::TokenBucket,
13310 RateLimitAlgorithm::FixedWindow,
13311 RateLimitAlgorithm::SlidingWindow,
13312 ] {
13313 let store = InMemoryRateLimitStore::with_max_keys(1);
13314 let start = Instant::now();
13315 assert!(
13316 store
13317 .check_at("long-lived", algorithm, 10, long_window, start)
13318 .allowed
13319 );
13320
13321 let result = store.check_at(
13322 "short-window-unseen",
13323 algorithm,
13324 10,
13325 short_window,
13326 start + Duration::from_secs(2),
13327 );
13328 assert!(
13329 !result.allowed,
13330 "{algorithm:?} must not evict a resident using the incoming short window"
13331 );
13332 assert_eq!(
13333 rate_limit_entry_keys(&store, algorithm),
13334 vec!["long-lived".to_string()]
13335 );
13336 }
13337 }
13338
13339 #[test]
13340 fn cleanup_reclaims_short_resident_despite_long_incoming_window() {
13341 let short_window = Duration::from_secs(1);
13342 let long_window = Duration::from_secs(3_600);
13343
13344 for algorithm in [
13345 RateLimitAlgorithm::TokenBucket,
13346 RateLimitAlgorithm::FixedWindow,
13347 RateLimitAlgorithm::SlidingWindow,
13348 ] {
13349 let store = InMemoryRateLimitStore::with_max_keys(1);
13350 let start = Instant::now();
13351 assert!(
13352 store
13353 .check_at("short-lived", algorithm, 10, short_window, start)
13354 .allowed
13355 );
13356
13357 let result = store.check_at(
13358 "long-window-replacement",
13359 algorithm,
13360 10,
13361 long_window,
13362 start + Duration::from_secs(3),
13363 );
13364 assert!(
13365 result.allowed,
13366 "{algorithm:?} must reclaim a stale resident independently of the incoming window"
13367 );
13368 assert_eq!(
13369 rate_limit_entry_keys(&store, algorithm),
13370 vec!["long-window-replacement".to_string()]
13371 );
13372 }
13373 }
13374
13375 #[test]
13376 fn rate_limit_cleanup_is_throttled_to_one_sweep_per_second() {
13377 let algorithm = RateLimitAlgorithm::SlidingWindow;
13378 let window = Duration::from_secs(100);
13379 let store = InMemoryRateLimitStore::with_max_keys(1);
13380 let start = Instant::now();
13381 assert!(
13382 store
13383 .check_at("stale", algorithm, 10, window, start)
13384 .allowed
13385 );
13386
13387 assert!(
13388 !store
13389 .check_at(
13390 "initial-sweep",
13391 algorithm,
13392 10,
13393 window,
13394 start + Duration::from_millis(199_500),
13395 )
13396 .allowed
13397 );
13398 assert!(
13399 !store
13400 .check_at(
13401 "first-unseen",
13402 algorithm,
13403 10,
13404 window,
13405 start + Duration::from_millis(200_100),
13406 )
13407 .allowed
13408 );
13409 assert!(
13410 !store
13411 .check_at(
13412 "second-unseen",
13413 algorithm,
13414 10,
13415 window,
13416 start + Duration::from_millis(200_400),
13417 )
13418 .allowed
13419 );
13420 assert_eq!(
13421 rate_limit_entry_keys(&store, algorithm),
13422 vec!["stale".to_string()]
13423 );
13424
13425 assert!(
13426 store
13427 .check_at(
13428 "replacement",
13429 algorithm,
13430 10,
13431 window,
13432 start + Duration::from_millis(200_500),
13433 )
13434 .allowed
13435 );
13436 assert_eq!(
13437 rate_limit_entry_keys(&store, algorithm),
13438 vec!["replacement".to_string()]
13439 );
13440 }
13441
13442 #[test]
13443 fn zero_window_fails_closed_without_retaining_keys() {
13444 let store = InMemoryRateLimitStore::with_max_keys(1);
13445
13446 for algorithm in [
13447 RateLimitAlgorithm::TokenBucket,
13448 RateLimitAlgorithm::FixedWindow,
13449 RateLimitAlgorithm::SlidingWindow,
13450 ] {
13451 let result = store.check("never-retained", algorithm, 10, Duration::ZERO);
13452 assert!(
13453 !result.allowed,
13454 "{algorithm:?} zero window must fail closed"
13455 );
13456 assert_eq!(result.reset_after_secs, 1);
13457 assert!(rate_limit_entry_keys(&store, algorithm).is_empty());
13458 }
13459 }
13460
13461 #[test]
13462 fn rate_limit_builder_per_minute() {
13463 let mw = RateLimitMiddleware::builder()
13464 .requests(50)
13465 .per_minute(2)
13466 .algorithm(RateLimitAlgorithm::SlidingWindow)
13467 .build();
13468 assert_eq!(mw.config.max_requests, 50);
13469 assert_eq!(mw.config.window, Duration::from_secs(120));
13470 assert_eq!(mw.config.algorithm, RateLimitAlgorithm::SlidingWindow);
13471 }
13472
13473 #[test]
13474 fn rate_limit_builder_per_hour() {
13475 let mw = RateLimitMiddleware::builder()
13476 .requests(1000)
13477 .per_hour(1)
13478 .build();
13479 assert_eq!(mw.config.window, Duration::from_secs(3600));
13480 }
13481
13482 #[test]
13483 fn rate_limit_middleware_name() {
13484 let mw = RateLimitMiddleware::new();
13485 assert_eq!(mw.name(), "RateLimit");
13486 }
13487
13488 #[test]
13489 fn rate_limit_default_via_default_trait() {
13490 let mw = RateLimitMiddleware::default();
13491 assert_eq!(mw.config.max_requests, 100);
13492 }
13493
13494 #[test]
13499 fn etag_middleware_generates_etag_for_get() {
13500 let mw = ETagMiddleware::new();
13501 let ctx = test_context();
13502 let req = Request::new(crate::request::Method::Get, "/resource");
13503
13504 let response = Response::ok()
13506 .header("content-type", b"application/json".to_vec())
13507 .body(ResponseBody::Bytes(br#"{"status":"ok"}"#.to_vec()));
13508
13509 let response = futures_executor::block_on(mw.after(&ctx, &req, response));
13510
13511 let etag = response
13513 .headers()
13514 .iter()
13515 .find(|(name, _)| name.eq_ignore_ascii_case("etag"));
13516 assert!(etag.is_some(), "Response should have ETag header");
13517
13518 let etag_value = std::str::from_utf8(&etag.unwrap().1).unwrap();
13520 assert!(etag_value.starts_with('"'), "ETag should start with quote");
13521 assert!(etag_value.ends_with('"'), "ETag should end with quote");
13522 }
13523
13524 #[test]
13525 fn etag_middleware_returns_304_on_match() {
13526 let mw = ETagMiddleware::new();
13527 let ctx = test_context();
13528
13529 let req1 = Request::new(crate::request::Method::Get, "/resource");
13531 let body = br#"{"status":"ok"}"#.to_vec();
13532 let response1 = Response::ok().body(ResponseBody::Bytes(body.clone()));
13533 let response1 = futures_executor::block_on(mw.after(&ctx, &req1, response1));
13534
13535 let etag = response1
13536 .headers()
13537 .iter()
13538 .find(|(name, _)| name.eq_ignore_ascii_case("etag"))
13539 .map(|(_, v)| std::str::from_utf8(v).unwrap().to_string())
13540 .unwrap();
13541
13542 let mut req2 = Request::new(crate::request::Method::Get, "/resource");
13544 req2.headers_mut()
13545 .insert("if-none-match", etag.as_bytes().to_vec());
13546
13547 let response2 = Response::ok().body(ResponseBody::Bytes(body));
13548 let response2 = futures_executor::block_on(mw.after(&ctx, &req2, response2));
13549
13550 assert_eq!(response2.status().as_u16(), 304);
13552 assert!(response2.body_ref().is_empty());
13553 }
13554
13555 #[test]
13556 fn etag_middleware_returns_full_response_on_mismatch() {
13557 let mw = ETagMiddleware::new();
13558 let ctx = test_context();
13559
13560 let mut req = Request::new(crate::request::Method::Get, "/resource");
13561 req.headers_mut()
13562 .insert("if-none-match", b"\"old-etag\"".to_vec());
13563
13564 let body = br#"{"status":"updated"}"#.to_vec();
13565 let response = Response::ok().body(ResponseBody::Bytes(body.clone()));
13566 let response = futures_executor::block_on(mw.after(&ctx, &req, response));
13567
13568 assert_eq!(response.status().as_u16(), 200);
13570 assert!(!response.body_ref().is_empty());
13571 }
13572
13573 #[test]
13574 fn etag_middleware_weak_etag_generation() {
13575 let config = ETagConfig::new().weak(true);
13576 let mw = ETagMiddleware::with_config(config);
13577 let ctx = test_context();
13578 let req = Request::new(crate::request::Method::Get, "/resource");
13579
13580 let response = Response::ok().body(ResponseBody::Bytes(b"data".to_vec()));
13581 let response = futures_executor::block_on(mw.after(&ctx, &req, response));
13582
13583 let etag = response
13584 .headers()
13585 .iter()
13586 .find(|(name, _)| name.eq_ignore_ascii_case("etag"))
13587 .map(|(_, v)| std::str::from_utf8(v).unwrap().to_string())
13588 .unwrap();
13589
13590 assert!(etag.starts_with("W/"), "Weak ETag should start with W/");
13591 }
13592
13593 #[test]
13594 fn etag_middleware_skips_post_requests() {
13595 let mw = ETagMiddleware::new();
13596 let ctx = test_context();
13597 let req = Request::new(crate::request::Method::Post, "/resource");
13598
13599 let response = Response::ok().body(ResponseBody::Bytes(b"created".to_vec()));
13600 let response = futures_executor::block_on(mw.after(&ctx, &req, response));
13601
13602 let etag = response
13604 .headers()
13605 .iter()
13606 .find(|(name, _)| name.eq_ignore_ascii_case("etag"));
13607 assert!(etag.is_none(), "POST should not have ETag");
13608 }
13609
13610 #[test]
13611 fn etag_middleware_handles_head_requests() {
13612 let mw = ETagMiddleware::new();
13613 let ctx = test_context();
13614 let req = Request::new(crate::request::Method::Head, "/resource");
13615
13616 let response = Response::ok().body(ResponseBody::Bytes(b"data".to_vec()));
13617 let response = futures_executor::block_on(mw.after(&ctx, &req, response));
13618
13619 let etag = response
13621 .headers()
13622 .iter()
13623 .find(|(name, _)| name.eq_ignore_ascii_case("etag"));
13624 assert!(etag.is_some(), "HEAD should have ETag");
13625 }
13626
13627 #[test]
13628 fn etag_middleware_disabled_mode() {
13629 let config = ETagConfig::new().mode(ETagMode::Disabled);
13630 let mw = ETagMiddleware::with_config(config);
13631 let ctx = test_context();
13632 let req = Request::new(crate::request::Method::Get, "/resource");
13633
13634 let response = Response::ok().body(ResponseBody::Bytes(b"data".to_vec()));
13635 let response = futures_executor::block_on(mw.after(&ctx, &req, response));
13636
13637 let etag = response
13639 .headers()
13640 .iter()
13641 .find(|(name, _)| name.eq_ignore_ascii_case("etag"));
13642 assert!(etag.is_none(), "Disabled mode should not add ETag");
13643 }
13644
13645 #[test]
13646 fn etag_middleware_min_size_filter() {
13647 let config = ETagConfig::new().min_size(1000);
13648 let mw = ETagMiddleware::with_config(config);
13649 let ctx = test_context();
13650 let req = Request::new(crate::request::Method::Get, "/resource");
13651
13652 let response = Response::ok().body(ResponseBody::Bytes(b"small".to_vec()));
13654 let response = futures_executor::block_on(mw.after(&ctx, &req, response));
13655
13656 let etag = response
13658 .headers()
13659 .iter()
13660 .find(|(name, _)| name.eq_ignore_ascii_case("etag"));
13661 assert!(etag.is_none(), "Small body should not get ETag");
13662 }
13663
13664 #[test]
13665 fn etag_middleware_preserves_existing_etag() {
13666 let config = ETagConfig::new().mode(ETagMode::Manual);
13667 let mw = ETagMiddleware::with_config(config);
13668 let ctx = test_context();
13669
13670 let mut req = Request::new(crate::request::Method::Get, "/resource");
13672 req.headers_mut()
13673 .insert("if-none-match", b"\"custom-etag\"".to_vec());
13674
13675 let response = Response::ok()
13677 .header("etag", b"\"custom-etag\"".to_vec())
13678 .body(ResponseBody::Bytes(b"data".to_vec()));
13679 let response = futures_executor::block_on(mw.after(&ctx, &req, response));
13680
13681 assert_eq!(response.status().as_u16(), 304);
13683 }
13684
13685 #[test]
13686 fn etag_middleware_wildcard_if_none_match() {
13687 let mw = ETagMiddleware::new();
13688 let ctx = test_context();
13689 let mut req = Request::new(crate::request::Method::Get, "/resource");
13690 req.headers_mut().insert("if-none-match", b"*".to_vec());
13691
13692 let response = Response::ok().body(ResponseBody::Bytes(b"data".to_vec()));
13693 let response = futures_executor::block_on(mw.after(&ctx, &req, response));
13694
13695 assert_eq!(response.status().as_u16(), 304);
13697 }
13698
13699 #[test]
13700 fn etag_middleware_weak_comparison_matches() {
13701 let mw = ETagMiddleware::new();
13702 let ctx = test_context();
13703
13704 let req1 = Request::new(crate::request::Method::Get, "/resource");
13706 let body = b"test data".to_vec();
13707 let response1 = Response::ok().body(ResponseBody::Bytes(body.clone()));
13708 let response1 = futures_executor::block_on(mw.after(&ctx, &req1, response1));
13709
13710 let etag = response1
13711 .headers()
13712 .iter()
13713 .find(|(name, _)| name.eq_ignore_ascii_case("etag"))
13714 .map(|(_, v)| std::str::from_utf8(v).unwrap().to_string())
13715 .unwrap();
13716
13717 let mut req2 = Request::new(crate::request::Method::Get, "/resource");
13719 let weak_etag = format!("W/{}", etag);
13720 req2.headers_mut()
13721 .insert("if-none-match", weak_etag.as_bytes().to_vec());
13722
13723 let response2 = Response::ok().body(ResponseBody::Bytes(body));
13724 let response2 = futures_executor::block_on(mw.after(&ctx, &req2, response2));
13725
13726 assert_eq!(response2.status().as_u16(), 304);
13728 }
13729
13730 #[test]
13731 fn etag_middleware_name() {
13732 let mw = ETagMiddleware::new();
13733 assert_eq!(mw.name(), "ETagMiddleware");
13734 }
13735
13736 #[test]
13737 fn etag_config_builder() {
13738 let config = ETagConfig::new()
13739 .mode(ETagMode::Auto)
13740 .weak(true)
13741 .min_size(512);
13742
13743 assert_eq!(config.mode, ETagMode::Auto);
13744 assert!(config.weak);
13745 assert_eq!(config.min_size, 512);
13746 }
13747
13748 #[test]
13749 fn etag_generates_consistent_hash() {
13750 let etag1 = ETagMiddleware::generate_etag(b"hello world", false);
13752 let etag2 = ETagMiddleware::generate_etag(b"hello world", false);
13753 assert_eq!(etag1, etag2);
13754
13755 let etag3 = ETagMiddleware::generate_etag(b"hello world!", false);
13757 assert_ne!(etag1, etag3);
13758 }
13759}