1use async_trait::async_trait;
7use bytes::Bytes;
8use pingora::http::ResponseHeader;
9use pingora::prelude::*;
10use pingora::protocols::Digest;
11use pingora::proxy::{ProxyHttp, Session};
12use pingora::upstreams::peer::Peer;
13use pingora_cache::{
14 CacheKey, CacheMeta, ForcedFreshness, HitHandler, NoCacheReason, RespCacheable,
15};
16use pingora_timeout::sleep;
17use std::os::unix::io::RawFd;
18use std::time::Duration;
19use tracing::{debug, error, info, trace, warn};
20
21use crate::cache::{get_cache_eviction, get_cache_lock, get_cache_storage};
22use crate::disk_cache::DiskHitHandler;
23use crate::hybrid_cache::HybridHitHandler;
24use crate::inference::{
25 extract_inference_content, is_sse_response, PromptInjectionResult, StreamingTokenCounter,
26};
27use crate::logging::{AccessLogEntry, AuditEventType, AuditLogEntry};
28use crate::rate_limit::HeaderAccessor;
29use crate::routing::RequestInfo;
30
31use super::context::{FallbackReason, RequestContext};
32use super::fallback::FallbackEvaluator;
33use super::fallback_metrics::get_fallback_metrics;
34use super::listener_addr::listener_for_addr;
35use super::model_routing;
36use super::model_routing_metrics::get_model_routing_metrics;
37use super::ZentinelProxy;
38
39fn to_socket_addr(
44 addr: &pingora::protocols::l4::socket::SocketAddr,
45) -> Option<std::net::SocketAddr> {
46 addr.as_inet().copied()
47}
48
49struct NoHeaderAccessor;
51impl HeaderAccessor for NoHeaderAccessor {
52 fn get_header(&self, _name: &str) -> Option<String> {
53 None
54 }
55}
56
57impl ZentinelProxy {
58 fn listener_matcher_for(
64 &self,
65 session: &Session,
66 ) -> Option<std::sync::Arc<crate::routing::RouteMatcher>> {
67 let matchers = self.listener_matchers.read();
68 if matchers.is_empty() {
69 return None;
70 }
71 let addr = to_socket_addr(session.downstream_session.server_addr()?)?;
75 matchers.get(addr).cloned()
76 }
77}
78
79fn cache_status_member(cache_name: &str, status: &super::context::CacheStatus) -> String {
86 use super::context::CacheStatus;
87 match status {
88 CacheStatus::HitMemory => format!("{cache_name}; hit; detail=memory"),
89 CacheStatus::HitDisk => format!("{cache_name}; hit; detail=disk"),
90 CacheStatus::Hit => format!("{cache_name}; hit"),
91 CacheStatus::HitStale => format!("{cache_name}; fwd=stale"),
92 CacheStatus::Miss => format!("{cache_name}; fwd=miss"),
93 CacheStatus::Bypass(reason) => match *reason {
94 "method" => format!("{cache_name}; fwd=bypass; detail=method"),
95 "disabled" => format!("{cache_name}; fwd=bypass; detail=disabled"),
96 "no-route" => format!("{cache_name}; fwd=bypass; detail=no-route"),
97 _ => format!("{cache_name}; fwd=bypass"),
98 },
99 }
100}
101
102fn apply_cache_status(
109 response: &mut pingora::http::ResponseHeader,
110 cache_name: &str,
111 status: &super::context::CacheStatus,
112) {
113 let member = cache_status_member(cache_name, status);
114 response.append_header("Cache-Status", &member).ok();
115}
116
117#[async_trait]
118impl ProxyHttp for ZentinelProxy {
119 type CTX = RequestContext;
120
121 fn new_ctx(&self) -> Self::CTX {
122 RequestContext::new()
123 }
124
125 fn should_retry_response(
135 &self,
136 session: &Session,
137 resp: &ResponseHeader,
138 ctx: &mut Self::CTX,
139 ) -> bool {
140 let Some(route) = ctx.route_config() else {
141 return false;
142 };
143 let Some(policy) = route.retry_policy.as_ref() else {
144 return false;
145 };
146
147 if !policy.is_retryable_status(resp.status.as_u16()) {
148 return false;
149 }
150
151 let method = session.req_header().method.as_str();
154 if !policy.may_retry_method(method) {
155 debug!(
156 correlation_id = %ctx.trace_id,
157 method = method,
158 status = resp.status.as_u16(),
159 "Not retrying a non-idempotent request; set retry-non-idempotent to allow it"
160 );
161 return false;
162 }
163
164 if ctx.request_attempts >= policy.max_attempts {
167 debug!(
168 correlation_id = %ctx.trace_id,
169 status = resp.status.as_u16(),
170 attempts = ctx.request_attempts,
171 "Retry budget exhausted; forwarding the upstream response"
172 );
173 return false;
174 }
175
176 info!(
177 correlation_id = %ctx.trace_id,
178 status = resp.status.as_u16(),
179 attempt = ctx.request_attempts,
180 max_attempts = policy.max_attempts,
181 "Retrying request after a retryable upstream status"
182 );
183 true
184 }
185
186 fn fail_to_connect(
187 &self,
188 _session: &mut Session,
189 peer: &HttpPeer,
190 ctx: &mut Self::CTX,
191 e: Box<Error>,
192 ) -> Box<Error> {
193 error!(
194 correlation_id = %ctx.trace_id,
195 route_id = ctx.route_id.as_deref().unwrap_or("unknown"),
196 upstream = ctx.upstream.as_deref().unwrap_or("unknown"),
197 peer_address = %peer.address(),
198 error = %e,
199 "Failed to connect to upstream peer"
200 );
201 self.log_manager.log_request_error(
202 "error",
203 "Failed to connect to upstream peer",
204 &ctx.trace_id,
205 ctx.route_id.as_deref(),
206 ctx.upstream.as_deref(),
207 Some(format!("peer={} error={}", peer.address(), e)),
208 );
209 e
211 }
212
213 async fn early_request_filter(
216 &self,
217 session: &mut Session,
218 ctx: &mut Self::CTX,
219 ) -> Result<(), Box<Error>> {
220 self.reload_coordinator.inc_requests();
223
224 let req_header = session.req_header();
226 let method = req_header.method.as_str();
227 let path = req_header.uri.path();
228 let host = crate::http_helpers::extract_request_host(req_header);
229
230 if let Some(ref challenge_manager) = self.acme_challenges {
232 if let Some(token) = crate::acme::ChallengeManager::extract_token(path) {
233 if let Some(key_authorization) = challenge_manager.get_response(token) {
234 debug!(
235 token = %token,
236 "Serving ACME HTTP-01 challenge response"
237 );
238
239 let mut resp = ResponseHeader::build(200, None)?;
241 resp.insert_header("Content-Type", "text/plain")?;
242 resp.insert_header("Content-Length", key_authorization.len().to_string())?;
243
244 session.write_response_header(Box::new(resp), false).await?;
246 session
247 .write_response_body(Some(Bytes::from(key_authorization)), true)
248 .await?;
249
250 return Err(Error::explain(
252 ErrorType::InternalError,
253 "ACME challenge served",
254 ));
255 } else {
256 warn!(
258 token = %token,
259 "ACME challenge token not found"
260 );
261 }
262 }
263 }
264
265 ctx.method = method.to_string();
266 ctx.path = path.to_string();
267 ctx.host = Some(host.to_string());
268
269 let listener_matcher = self.listener_matcher_for(session);
273
274 let route_match = {
276 let mut request_info = RequestInfo::new(method, path, host);
277 let matched = if let Some(ref matcher) = listener_matcher {
278 if matcher.needs_headers() {
280 request_info = request_info
281 .with_headers(RequestInfo::build_headers(req_header.headers.iter()));
282 }
283 matcher.match_request(&request_info)
284 } else {
285 let route_matcher = self.route_matcher.read();
286 if route_matcher.needs_headers() {
287 request_info = request_info
288 .with_headers(RequestInfo::build_headers(req_header.headers.iter()));
289 }
290 route_matcher.match_request(&request_info)
291 };
292
293 match matched {
294 Some(m) => m,
295 None => return Ok(()), }
297 };
298
299 ctx.trace_id = self.get_trace_id(session);
300 ctx.route_id = Some(route_match.route_id.to_string());
301 ctx.route_config = Some(route_match.config.clone());
302
303 if let Some(traceparent) = req_header.headers.get(crate::otel::TRACEPARENT_HEADER) {
305 if let Ok(s) = traceparent.to_str() {
306 ctx.trace_context = crate::otel::TraceContext::parse_traceparent(s);
307 }
308 }
309
310 if let Some(tracer) = crate::otel::get_tracer() {
312 ctx.otel_span = Some(tracer.start_span(method, path, ctx.trace_context.as_ref()));
313 }
314
315 if route_match.config.service_type == zentinel_config::ServiceType::Builtin {
317 trace!(
318 correlation_id = %ctx.trace_id,
319 route_id = %route_match.route_id,
320 builtin_handler = ?route_match.config.builtin_handler,
321 "Handling builtin route in early_request_filter"
322 );
323
324 let handled = self
326 .handle_builtin_route(session, ctx, &route_match)
327 .await?;
328
329 if handled {
330 return Err(Error::explain(
332 ErrorType::InternalError,
333 "Builtin handler complete",
334 ));
335 }
336 }
337
338 Ok(())
339 }
340
341 async fn upstream_peer(
342 &self,
343 session: &mut Session,
344 ctx: &mut Self::CTX,
345 ) -> Result<Box<HttpPeer>, Box<Error>> {
346 if ctx.config.is_none() {
348 ctx.config = Some(self.config_manager.current());
349 }
350
351 if ctx.client_ip.is_empty() {
353 ctx.client_ip = session
354 .client_addr()
355 .map(|a| a.to_string())
356 .unwrap_or_else(|| "unknown".to_string());
357 }
358
359 let req_header = session.req_header();
360
361 if ctx.method.is_empty() {
363 ctx.method = req_header.method.to_string();
364 ctx.path = req_header.uri.path().to_string();
365 ctx.query = req_header.uri.query().map(|q| q.to_string());
366 ctx.host = Some(crate::http_helpers::extract_request_host(req_header).to_string());
367 }
368 ctx.user_agent = req_header
369 .headers
370 .get("user-agent")
371 .and_then(|v| v.to_str().ok())
372 .map(|s| s.to_string());
373 ctx.referer = req_header
374 .headers
375 .get("referer")
376 .and_then(|v| v.to_str().ok())
377 .map(|s| s.to_string());
378
379 trace!(
380 correlation_id = %ctx.trace_id,
381 client_ip = %ctx.client_ip,
382 "Request received, initializing context"
383 );
384
385 let route_match = if let Some(ref route_config) = ctx.route_config {
387 let route_id = ctx.route_id.as_deref().unwrap_or("");
388 crate::routing::RouteMatch {
389 route_id: zentinel_common::RouteId::new(route_id),
390 config: route_config.clone(),
391 }
392 } else {
393 let listener_matcher = self.listener_matcher_for(session);
397 let (match_result, route_duration) = {
398 let host = ctx.host.as_deref().unwrap_or("");
399
400 let mut request_info = RequestInfo::new(&ctx.method, &ctx.path, host);
402
403 let route_start = std::time::Instant::now();
404 let matched = if let Some(ref matcher) = listener_matcher {
405 if matcher.needs_headers() {
406 request_info = request_info
407 .with_headers(RequestInfo::build_headers(req_header.headers.iter()));
408 }
409 if matcher.needs_query_params() {
410 request_info = request_info
411 .with_query_params(RequestInfo::parse_query_params(&ctx.path));
412 }
413 matcher.match_request(&request_info)
414 } else {
415 let route_matcher = self.route_matcher.read();
416 if route_matcher.needs_headers() {
418 request_info = request_info
419 .with_headers(RequestInfo::build_headers(req_header.headers.iter()));
420 }
421 if route_matcher.needs_query_params() {
423 request_info = request_info
424 .with_query_params(RequestInfo::parse_query_params(&ctx.path));
425 }
426 route_matcher.match_request(&request_info)
427 };
428
429 let route_match = matched.ok_or_else(|| {
430 warn!(
431 correlation_id = %ctx.trace_id,
432 method = %request_info.method,
433 path = %request_info.path,
434 host = %request_info.host,
435 "No matching route found for request"
436 );
437 self.log_manager.log_request_error(
438 "warn",
439 "No matching route found for request",
440 &ctx.trace_id,
441 None,
442 None,
443 Some(format!(
444 "method={} path={} host={}",
445 request_info.method, request_info.path, request_info.host
446 )),
447 );
448 Error::explain(ErrorType::HTTPStatus(404), "No matching route found")
449 })?;
450 let route_duration = route_start.elapsed();
451 (route_match, route_duration)
453 };
454
455 ctx.route_id = Some(match_result.route_id.to_string());
456 ctx.route_config = Some(match_result.config.clone());
457
458 if ctx.trace_id.is_empty() {
460 ctx.trace_id = self.get_trace_id(session);
461
462 if let Some(traceparent) = req_header.headers.get(crate::otel::TRACEPARENT_HEADER) {
464 if let Ok(s) = traceparent.to_str() {
465 ctx.trace_context = crate::otel::TraceContext::parse_traceparent(s);
466 }
467 }
468
469 if let Some(tracer) = crate::otel::get_tracer() {
471 ctx.otel_span =
472 Some(tracer.start_span(&ctx.method, &ctx.path, ctx.trace_context.as_ref()));
473 }
474 }
475
476 trace!(
477 correlation_id = %ctx.trace_id,
478 route_id = %match_result.route_id,
479 route_duration_us = route_duration.as_micros(),
480 service_type = ?match_result.config.service_type,
481 "Route matched"
482 );
483 match_result
484 };
485
486 if route_match.config.service_type == zentinel_config::ServiceType::Builtin {
488 trace!(
489 correlation_id = %ctx.trace_id,
490 route_id = %route_match.route_id,
491 builtin_handler = ?route_match.config.builtin_handler,
492 "Route type is builtin, skipping upstream"
493 );
494 ctx.upstream = Some(format!("_builtin_{}", route_match.route_id));
496 return Err(Error::explain(
498 ErrorType::InternalError,
499 "Builtin handler handled in request_filter",
500 ));
501 }
502
503 if route_match.config.service_type == zentinel_config::ServiceType::Static {
505 trace!(
506 correlation_id = %ctx.trace_id,
507 route_id = %route_match.route_id,
508 "Route type is static, checking for static server"
509 );
510 if self
512 .static_servers
513 .get(route_match.route_id.as_str())
514 .await
515 .is_some()
516 {
517 ctx.upstream = Some(format!("_static_{}", route_match.route_id));
519 info!(
520 correlation_id = %ctx.trace_id,
521 route_id = %route_match.route_id,
522 path = %ctx.path,
523 "Serving static file"
524 );
525 return Err(Error::explain(
527 ErrorType::InternalError,
528 "Static file serving handled in request_filter",
529 ));
530 }
531 }
532
533 let mut model_routing_applied = false;
536 if let Some(ref inference) = route_match.config.inference {
537 if let Some(ref model_routing) = inference.model_routing {
538 let model = model_routing::extract_model_from_headers(&req_header.headers);
540
541 if let Some(ref model_name) = model {
542 if let Some(routing_result) =
544 model_routing::find_upstream_for_model(model_routing, model_name)
545 {
546 debug!(
547 correlation_id = %ctx.trace_id,
548 route_id = %route_match.route_id,
549 model = %model_name,
550 upstream = %routing_result.upstream,
551 is_default = routing_result.is_default,
552 provider_override = ?routing_result.provider,
553 "Model-based routing selected upstream"
554 );
555
556 ctx.record_model_routing(
557 &routing_result.upstream,
558 Some(model_name.clone()),
559 routing_result.provider,
560 );
561 model_routing_applied = true;
562
563 if let Some(metrics) = get_model_routing_metrics() {
565 metrics.record_model_routed(
566 route_match.route_id.as_str(),
567 model_name,
568 &routing_result.upstream,
569 );
570 if routing_result.is_default {
571 metrics.record_default_upstream(route_match.route_id.as_str());
572 }
573 if let Some(provider) = routing_result.provider {
574 metrics.record_provider_override(
575 route_match.route_id.as_str(),
576 &routing_result.upstream,
577 provider.as_str(),
578 );
579 }
580 }
581 }
582 } else if let Some(ref default_upstream) = model_routing.default_upstream {
583 debug!(
585 correlation_id = %ctx.trace_id,
586 route_id = %route_match.route_id,
587 upstream = %default_upstream,
588 "Model-based routing using default upstream (no model header)"
589 );
590 ctx.record_model_routing(default_upstream, None, None);
591 model_routing_applied = true;
592
593 if let Some(metrics) = get_model_routing_metrics() {
595 metrics.record_no_model_header(route_match.route_id.as_str());
596 }
597 }
598 }
599 }
600
601 if !model_routing_applied {
603 if let Some(ref upstream) = route_match.config.upstream {
604 ctx.upstream = Some(upstream.clone());
605 trace!(
606 correlation_id = %ctx.trace_id,
607 route_id = %route_match.route_id,
608 upstream = %upstream,
609 "Upstream configured for route"
610 );
611 } else {
612 warn!(
616 correlation_id = %ctx.trace_id,
617 route_id = %route_match.route_id,
618 "Route has no upstream configured, returning 500"
619 );
620 crate::http_helpers::write_error(
621 session,
622 500,
623 "Internal Server Error",
624 "text/plain",
625 )
626 .await?;
627 return Err(Error::explain(
628 ErrorType::HTTPStatus(500),
629 "Route has no valid upstream",
630 ));
631 }
632 }
633
634 if let Some(ref fallback_config) = route_match.config.fallback {
637 let upstream_name = ctx.upstream.as_ref().unwrap();
638
639 let is_healthy = if let Some(pool) = self.upstream_pools.get(upstream_name).await {
641 pool.has_healthy_targets().await
642 } else {
643 false };
645
646 let is_budget_exhausted = ctx.inference_budget_exhausted;
648
649 let current_model = ctx.inference_model.as_deref();
651
652 let evaluator = FallbackEvaluator::new(
654 fallback_config,
655 ctx.tried_upstreams(),
656 ctx.fallback_attempt,
657 );
658
659 if let Some(decision) = evaluator.should_fallback_before_request(
661 upstream_name,
662 is_healthy,
663 is_budget_exhausted,
664 current_model,
665 ) {
666 info!(
667 correlation_id = %ctx.trace_id,
668 route_id = %route_match.route_id,
669 from_upstream = %upstream_name,
670 to_upstream = %decision.next_upstream,
671 reason = %decision.reason,
672 fallback_attempt = ctx.fallback_attempt + 1,
673 "Triggering fallback routing"
674 );
675
676 if let Some(metrics) = get_fallback_metrics() {
678 metrics.record_fallback_attempt(
679 route_match.route_id.as_str(),
680 upstream_name,
681 &decision.next_upstream,
682 &decision.reason,
683 );
684 }
685
686 ctx.record_fallback(decision.reason, &decision.next_upstream);
688
689 if let Some((original, mapped)) = decision.model_mapping {
691 if let Some(metrics) = get_fallback_metrics() {
693 metrics.record_model_mapping(
694 route_match.route_id.as_str(),
695 &original,
696 &mapped,
697 );
698 }
699
700 ctx.record_model_mapping(original, mapped);
701 trace!(
702 correlation_id = %ctx.trace_id,
703 original_model = ?ctx.model_mapping_applied().map(|m| &m.0),
704 mapped_model = ?ctx.model_mapping_applied().map(|m| &m.1),
705 "Applied model mapping for fallback"
706 );
707 }
708 }
709 }
710
711 debug!(
712 correlation_id = %ctx.trace_id,
713 route_id = %route_match.route_id,
714 upstream = ?ctx.upstream,
715 method = %req_header.method,
716 path = %req_header.uri.path(),
717 host = ctx.host.as_deref().unwrap_or("-"),
718 client_ip = %ctx.client_ip,
719 "Processing request"
720 );
721
722 if ctx
724 .upstream
725 .as_ref()
726 .is_some_and(|u| u.starts_with("_static_"))
727 {
728 return Err(Error::explain(
730 ErrorType::InternalError,
731 "Static route should be handled in request_filter",
732 ));
733 }
734
735 let upstream_name = ctx
736 .upstream
737 .as_ref()
738 .ok_or_else(|| Error::explain(ErrorType::InternalError, "No upstream configured"))?;
739
740 trace!(
741 correlation_id = %ctx.trace_id,
742 upstream = %upstream_name,
743 "Looking up upstream pool"
744 );
745
746 let pool = self
747 .upstream_pools
748 .get(upstream_name)
749 .await
750 .ok_or_else(|| {
751 error!(
752 correlation_id = %ctx.trace_id,
753 upstream = %upstream_name,
754 "Upstream pool not found"
755 );
756 self.log_manager.log_request_error(
757 "error",
758 "Upstream pool not found",
759 &ctx.trace_id,
760 ctx.route_id.as_deref(),
761 Some(upstream_name),
762 None,
763 );
764 Error::explain(
765 ErrorType::InternalError,
766 format!("Upstream pool '{}' not found", upstream_name),
767 )
768 })?;
769
770 const PEER_SELECTION_ATTEMPTS: u32 = 2;
777 let max_retries = PEER_SELECTION_ATTEMPTS;
778
779 ctx.request_attempts += 1;
785 if ctx.request_attempts > 1 {
786 if let Some(policy) = route_match.config.retry_policy.as_ref() {
787 let backoff = policy.backoff_for(ctx.request_attempts);
788 if !backoff.is_zero() {
789 trace!(
790 correlation_id = %ctx.trace_id,
791 attempt = ctx.request_attempts,
792 backoff_ms = backoff.as_millis(),
793 "Backing off before retrying the request"
794 );
795 tokio::time::sleep(backoff).await;
796 }
797 }
798 }
799
800 trace!(
801 correlation_id = %ctx.trace_id,
802 upstream = %upstream_name,
803 max_retries = max_retries,
804 "Starting upstream peer selection"
805 );
806
807 let mut last_error = None;
808 let selection_start = std::time::Instant::now();
809
810 for attempt in 1..=max_retries {
811 ctx.upstream_attempts = attempt;
812
813 trace!(
814 correlation_id = %ctx.trace_id,
815 upstream = %upstream_name,
816 attempt = attempt,
817 max_retries = max_retries,
818 "Attempting to select upstream peer"
819 );
820
821 match pool.select_peer_with_metadata(None).await {
822 Ok((mut peer, metadata)) => {
823 let selection_duration = selection_start.elapsed();
824 pool.increment_active();
826 let peer_addr = peer.address().to_string();
828 ctx.selected_upstream_address = Some(peer_addr.clone());
829
830 if metadata.contains_key("sticky_session_new") {
832 ctx.sticky_session_new_assignment = true;
833 ctx.sticky_session_set_cookie =
834 metadata.get("sticky_set_cookie_header").cloned();
835 ctx.sticky_target_index = metadata
836 .get("sticky_target_index")
837 .and_then(|s| s.parse().ok());
838
839 trace!(
840 correlation_id = %ctx.trace_id,
841 sticky_target_index = ?ctx.sticky_target_index,
842 "New sticky session assignment, will set cookie"
843 );
844 }
845
846 debug!(
847 correlation_id = %ctx.trace_id,
848 upstream = %upstream_name,
849 peer_address = %peer_addr,
850 attempt = attempt,
851 selection_duration_us = selection_duration.as_micros(),
852 sticky_session_hit = metadata.contains_key("sticky_session_hit"),
853 sticky_session_new = ctx.sticky_session_new_assignment,
854 "Selected upstream peer"
855 );
856 if let Some(ref rc) = ctx.route_config {
858 if let Some(timeout_secs) = rc.policies.timeout_secs {
859 peer.options.read_timeout = Some(Duration::from_secs(timeout_secs));
860 }
861 }
862
863 if let Some(connect_secs) = ctx.filter_connect_timeout_secs {
865 peer.options.connection_timeout = Some(Duration::from_secs(connect_secs));
866 }
867 if let Some(upstream_secs) = ctx.filter_upstream_timeout_secs {
868 peer.options.read_timeout = Some(Duration::from_secs(upstream_secs));
869 }
870
871 if let Some(per_attempt) = route_match
877 .config
878 .retry_policy
879 .as_ref()
880 .and_then(|p| p.per_attempt_timeout)
881 {
882 peer.options.connection_timeout = Some(per_attempt);
883 }
884
885 return Ok(Box::new(peer));
886 }
887 Err(e) => {
888 warn!(
889 correlation_id = %ctx.trace_id,
890 upstream = %upstream_name,
891 attempt = attempt,
892 max_retries = max_retries,
893 error = %e,
894 "Failed to select upstream peer"
895 );
896 last_error = Some(e);
897
898 if attempt < max_retries {
899 let backoff = Duration::from_millis(100 * 2_u64.pow(attempt - 1));
901 trace!(
902 correlation_id = %ctx.trace_id,
903 backoff_ms = backoff.as_millis(),
904 "Backing off before retry"
905 );
906 sleep(backoff).await;
907 }
908 }
909 }
910 }
911
912 let selection_duration = selection_start.elapsed();
913 error!(
914 correlation_id = %ctx.trace_id,
915 upstream = %upstream_name,
916 attempts = max_retries,
917 selection_duration_ms = selection_duration.as_millis(),
918 last_error = ?last_error,
919 "All upstream selection attempts failed"
920 );
921 self.log_manager.log_request_error(
922 "error",
923 "All upstream selection attempts failed",
924 &ctx.trace_id,
925 ctx.route_id.as_deref(),
926 Some(upstream_name),
927 Some(format!("attempts={} error={:?}", max_retries, last_error)),
928 );
929
930 if ctx.used_fallback() {
932 if let Some(metrics) = get_fallback_metrics() {
933 metrics.record_fallback_exhausted(ctx.route_id.as_deref().unwrap_or("unknown"));
934 }
935 }
936
937 Err(Error::explain(
938 ErrorType::InternalError,
939 format!("All upstream attempts failed: {:?}", last_error),
940 ))
941 }
942
943 async fn request_filter(
944 &self,
945 session: &mut Session,
946 ctx: &mut Self::CTX,
947 ) -> Result<bool, Box<Error>> {
948 trace!(
949 correlation_id = %ctx.trace_id,
950 route_id = ctx.route_id.as_deref().unwrap_or("unknown"),
951 "Starting request filter phase"
952 );
953
954 if let Some(local_addr) = session
958 .downstream_session
959 .server_addr()
960 .and_then(to_socket_addr)
961 {
962 let config = ctx
963 .config
964 .get_or_insert_with(|| self.config_manager.current());
965 if let Some(listener) = listener_for_addr(&config.listeners, local_addr) {
966 let request_timeout_secs = listener.request_timeout_secs;
967 let keepalive_timeout_secs = listener.keepalive_timeout_secs;
968 session
970 .downstream_session
971 .set_read_timeout(Some(std::time::Duration::from_secs(request_timeout_secs)));
972 ctx.listener_keepalive_timeout_secs = Some(keepalive_timeout_secs);
974 }
975 }
976
977 if let Some(route_id) = ctx.route_id.as_deref() {
980 if self.rate_limit_manager.has_route_limiter(route_id) {
981 let rate_result = self.rate_limit_manager.check(
982 route_id,
983 &ctx.client_ip,
984 &ctx.path,
985 Option::<&NoHeaderAccessor>::None,
986 );
987
988 if rate_result.limit > 0 {
990 ctx.rate_limit_info = Some(super::context::RateLimitHeaderInfo {
991 limit: rate_result.limit,
992 remaining: rate_result.remaining,
993 reset_at: rate_result.reset_at,
994 });
995 }
996
997 if !rate_result.allowed {
998 use zentinel_config::RateLimitAction;
999
1000 match rate_result.action {
1001 RateLimitAction::Reject => {
1002 warn!(
1003 correlation_id = %ctx.trace_id,
1004 route_id = route_id,
1005 client_ip = %ctx.client_ip,
1006 limiter = %rate_result.limiter,
1007 limit = rate_result.limit,
1008 remaining = rate_result.remaining,
1009 "Request rate limited"
1010 );
1011 self.metrics.record_blocked_request("rate_limited");
1012
1013 let audit_entry = AuditLogEntry::rate_limited(
1015 &ctx.trace_id,
1016 &ctx.method,
1017 &ctx.path,
1018 &ctx.client_ip,
1019 &rate_result.limiter,
1020 )
1021 .with_route_id(route_id)
1022 .with_status_code(rate_result.status_code);
1023 self.log_manager.log_audit(&audit_entry);
1024
1025 let body = rate_result
1027 .message
1028 .unwrap_or_else(|| "Rate limit exceeded".to_string());
1029
1030 let retry_after = rate_result.reset_at.saturating_sub(
1032 std::time::SystemTime::now()
1033 .duration_since(std::time::UNIX_EPOCH)
1034 .unwrap_or_default()
1035 .as_secs(),
1036 );
1037 crate::http_helpers::write_rate_limit_error(
1038 session,
1039 rate_result.status_code,
1040 &body,
1041 rate_result.limit,
1042 rate_result.remaining,
1043 rate_result.reset_at,
1044 retry_after,
1045 )
1046 .await?;
1047 return Ok(true); }
1049 RateLimitAction::LogOnly => {
1050 debug!(
1051 correlation_id = %ctx.trace_id,
1052 route_id = route_id,
1053 "Rate limit exceeded (log only mode)"
1054 );
1055 }
1057 RateLimitAction::Delay => {
1058 if let Some(delay_ms) = rate_result.suggested_delay_ms {
1060 let actual_delay = delay_ms.min(rate_result.max_delay_ms);
1062
1063 if actual_delay > 0 {
1064 debug!(
1065 correlation_id = %ctx.trace_id,
1066 route_id = route_id,
1067 suggested_delay_ms = delay_ms,
1068 max_delay_ms = rate_result.max_delay_ms,
1069 actual_delay_ms = actual_delay,
1070 "Applying rate limit delay"
1071 );
1072
1073 tokio::time::sleep(std::time::Duration::from_millis(
1074 actual_delay,
1075 ))
1076 .await;
1077 }
1078 }
1079 }
1081 }
1082 }
1083 }
1084 }
1085
1086 if let Some(route_id) = ctx.route_id.as_deref() {
1089 if let Some(ref route_config) = ctx.route_config {
1090 if route_config.service_type == zentinel_config::ServiceType::Inference
1091 && self.inference_rate_limit_manager.has_route(route_id)
1092 {
1093 let headers = &session.req_header().headers;
1096
1097 let body = ctx.body_buffer.as_slice();
1099
1100 let rate_limit_key = &ctx.client_ip;
1102
1103 if let Some(check_result) = self.inference_rate_limit_manager.check(
1104 route_id,
1105 rate_limit_key,
1106 headers,
1107 body,
1108 ) {
1109 ctx.inference_rate_limit_enabled = true;
1111 ctx.inference_estimated_tokens = check_result.estimated_tokens;
1112 ctx.inference_rate_limit_key = Some(rate_limit_key.to_string());
1113 ctx.inference_model = check_result.model.clone();
1114
1115 if !check_result.is_allowed() {
1116 let retry_after_ms = check_result.retry_after_ms();
1117 let retry_after_secs = retry_after_ms.div_ceil(1000);
1118
1119 warn!(
1120 correlation_id = %ctx.trace_id,
1121 route_id = route_id,
1122 client_ip = %ctx.client_ip,
1123 estimated_tokens = check_result.estimated_tokens,
1124 model = ?check_result.model,
1125 retry_after_ms = retry_after_ms,
1126 "Inference rate limit exceeded (tokens)"
1127 );
1128 self.metrics
1129 .record_blocked_request("inference_rate_limited");
1130
1131 let audit_entry = AuditLogEntry::new(
1133 &ctx.trace_id,
1134 AuditEventType::RateLimitExceeded,
1135 &ctx.method,
1136 &ctx.path,
1137 &ctx.client_ip,
1138 )
1139 .with_route_id(route_id)
1140 .with_status_code(429)
1141 .with_reason(format!(
1142 "Token rate limit exceeded: estimated {} tokens, model={:?}",
1143 check_result.estimated_tokens, check_result.model
1144 ));
1145 self.log_manager.log_audit(&audit_entry);
1146
1147 let body = "Token rate limit exceeded";
1149 let reset_at = std::time::SystemTime::now()
1150 .duration_since(std::time::UNIX_EPOCH)
1151 .unwrap_or_default()
1152 .as_secs()
1153 + retry_after_secs;
1154
1155 crate::http_helpers::write_rate_limit_error(
1157 session,
1158 429,
1159 body,
1160 0, 0, reset_at,
1163 retry_after_secs,
1164 )
1165 .await?;
1166 return Ok(true); }
1168
1169 trace!(
1170 correlation_id = %ctx.trace_id,
1171 route_id = route_id,
1172 estimated_tokens = check_result.estimated_tokens,
1173 model = ?check_result.model,
1174 "Inference rate limit check passed"
1175 );
1176
1177 if self.inference_rate_limit_manager.has_budget(route_id) {
1179 ctx.inference_budget_enabled = true;
1180
1181 if let Some(budget_result) =
1182 self.inference_rate_limit_manager.check_budget(
1183 route_id,
1184 rate_limit_key,
1185 check_result.estimated_tokens,
1186 )
1187 {
1188 if !budget_result.is_allowed() {
1189 let retry_after_secs = budget_result.retry_after_secs();
1190
1191 warn!(
1192 correlation_id = %ctx.trace_id,
1193 route_id = route_id,
1194 client_ip = %ctx.client_ip,
1195 estimated_tokens = check_result.estimated_tokens,
1196 retry_after_secs = retry_after_secs,
1197 "Token budget exhausted"
1198 );
1199
1200 ctx.inference_budget_exhausted = true;
1201 self.metrics.record_blocked_request("budget_exhausted");
1202
1203 let audit_entry = AuditLogEntry::new(
1205 &ctx.trace_id,
1206 AuditEventType::RateLimitExceeded,
1207 &ctx.method,
1208 &ctx.path,
1209 &ctx.client_ip,
1210 )
1211 .with_route_id(route_id)
1212 .with_status_code(429)
1213 .with_reason("Token budget exhausted".to_string());
1214 self.log_manager.log_audit(&audit_entry);
1215
1216 let body = "Token budget exhausted";
1218 let reset_at = std::time::SystemTime::now()
1219 .duration_since(std::time::UNIX_EPOCH)
1220 .unwrap_or_default()
1221 .as_secs()
1222 + retry_after_secs;
1223
1224 crate::http_helpers::write_rate_limit_error(
1225 session,
1226 429,
1227 body,
1228 0,
1229 0,
1230 reset_at,
1231 retry_after_secs,
1232 )
1233 .await?;
1234 return Ok(true);
1235 }
1236
1237 let remaining = match &budget_result {
1239 zentinel_common::budget::BudgetCheckResult::Allowed {
1240 remaining,
1241 } => *remaining as i64,
1242 zentinel_common::budget::BudgetCheckResult::Soft {
1243 remaining,
1244 ..
1245 } => *remaining,
1246 _ => 0,
1247 };
1248 ctx.inference_budget_remaining = Some(remaining);
1249
1250 if let Some(status) = self
1252 .inference_rate_limit_manager
1253 .budget_status(route_id, rate_limit_key)
1254 {
1255 ctx.inference_budget_period_reset = Some(status.period_end);
1256 }
1257
1258 trace!(
1259 correlation_id = %ctx.trace_id,
1260 route_id = route_id,
1261 budget_remaining = remaining,
1262 "Token budget check passed"
1263 );
1264 }
1265 }
1266
1267 if self
1269 .inference_rate_limit_manager
1270 .has_cost_attribution(route_id)
1271 {
1272 ctx.inference_cost_enabled = true;
1273 }
1274 }
1275 }
1276 }
1277 }
1278
1279 if let Some(ref route_config) = ctx.route_config {
1281 if let Some(ref inference) = route_config.inference {
1282 if let Some(ref guardrails) = inference.guardrails {
1283 if let Some(ref pi_config) = guardrails.prompt_injection {
1284 if pi_config.enabled && !ctx.body_buffer.is_empty() {
1285 ctx.guardrails_enabled = true;
1286
1287 if let Some(content) = extract_inference_content(&ctx.body_buffer) {
1289 let result = self
1290 .guardrail_processor
1291 .check_prompt_injection(
1292 pi_config,
1293 &content,
1294 ctx.inference_model.as_deref(),
1295 ctx.route_id.as_deref(),
1296 &ctx.trace_id,
1297 )
1298 .await;
1299
1300 match result {
1301 PromptInjectionResult::Blocked {
1302 status,
1303 message,
1304 detections,
1305 } => {
1306 warn!(
1307 correlation_id = %ctx.trace_id,
1308 route_id = ctx.route_id.as_deref().unwrap_or("unknown"),
1309 detection_count = detections.len(),
1310 "Prompt injection detected, blocking request"
1311 );
1312
1313 self.metrics.record_blocked_request("prompt_injection");
1314
1315 ctx.guardrail_detection_categories =
1317 detections.iter().map(|d| d.category.clone()).collect();
1318
1319 let audit_entry = AuditLogEntry::new(
1321 &ctx.trace_id,
1322 AuditEventType::Blocked,
1323 &ctx.method,
1324 &ctx.path,
1325 &ctx.client_ip,
1326 )
1327 .with_route_id(ctx.route_id.as_deref().unwrap_or("unknown"))
1328 .with_status_code(status)
1329 .with_reason("Prompt injection detected".to_string());
1330 self.log_manager.log_audit(&audit_entry);
1331
1332 crate::http_helpers::write_json_error(
1334 session,
1335 status,
1336 "prompt_injection_blocked",
1337 Some(&message),
1338 )
1339 .await?;
1340 return Ok(true);
1341 }
1342 PromptInjectionResult::Detected { detections } => {
1343 warn!(
1345 correlation_id = %ctx.trace_id,
1346 route_id = ctx.route_id.as_deref().unwrap_or("unknown"),
1347 detection_count = detections.len(),
1348 "Prompt injection detected (logged only)"
1349 );
1350 ctx.guardrail_detection_categories =
1351 detections.iter().map(|d| d.category.clone()).collect();
1352 }
1353 PromptInjectionResult::Warning { detections } => {
1354 ctx.guardrail_warning = true;
1356 ctx.guardrail_detection_categories =
1357 detections.iter().map(|d| d.category.clone()).collect();
1358 debug!(
1359 correlation_id = %ctx.trace_id,
1360 "Prompt injection warning set"
1361 );
1362 }
1363 PromptInjectionResult::Clean => {
1364 trace!(
1365 correlation_id = %ctx.trace_id,
1366 "No prompt injection detected"
1367 );
1368 }
1369 PromptInjectionResult::Error { message } => {
1370 trace!(
1372 correlation_id = %ctx.trace_id,
1373 error = %message,
1374 "Prompt injection check error (failure mode applied)"
1375 );
1376 }
1377 }
1378 }
1379 }
1380 }
1381 }
1382 }
1383 }
1384
1385 if let Some(route_id) = ctx.route_id.as_deref() {
1387 if let Some(ref route_config) = ctx.route_config {
1388 for filter_id in &route_config.filters {
1389 if let Some(result) = self.geo_filter_manager.check(filter_id, &ctx.client_ip) {
1390 ctx.geo_country_code = result.country_code.clone();
1392 ctx.geo_lookup_performed = true;
1393
1394 if !result.allowed {
1395 warn!(
1396 correlation_id = %ctx.trace_id,
1397 route_id = route_id,
1398 client_ip = %ctx.client_ip,
1399 country = ?result.country_code,
1400 filter_id = %filter_id,
1401 "Request blocked by geo filter"
1402 );
1403 self.metrics.record_blocked_request("geo_blocked");
1404
1405 let audit_entry = AuditLogEntry::new(
1407 &ctx.trace_id,
1408 AuditEventType::Blocked,
1409 &ctx.method,
1410 &ctx.path,
1411 &ctx.client_ip,
1412 )
1413 .with_route_id(route_id)
1414 .with_status_code(result.status_code)
1415 .with_reason(format!(
1416 "Geo blocked: country={}, filter={}",
1417 result.country_code.as_deref().unwrap_or("unknown"),
1418 filter_id
1419 ));
1420 self.log_manager.log_audit(&audit_entry);
1421
1422 let body = result
1424 .block_message
1425 .unwrap_or_else(|| "Access denied".to_string());
1426
1427 crate::http_helpers::write_error(
1428 session,
1429 result.status_code,
1430 &body,
1431 "text/plain",
1432 )
1433 .await?;
1434 return Ok(true); }
1436
1437 break;
1439 }
1440 }
1441 }
1442 }
1443
1444 let config_for_filters = std::sync::Arc::clone(
1447 ctx.config
1448 .get_or_insert_with(|| self.config_manager.current()),
1449 );
1450 if super::filters::apply_request_filters(session, ctx, &config_for_filters).await? {
1451 return Ok(true); }
1453
1454 let is_websocket_upgrade = session
1456 .req_header()
1457 .headers
1458 .get(http::header::UPGRADE)
1459 .map(|v| v.as_bytes().eq_ignore_ascii_case(b"websocket"))
1460 .unwrap_or(false);
1461
1462 if is_websocket_upgrade {
1463 ctx.is_websocket_upgrade = true;
1464
1465 if let Some(ref route_config) = ctx.route_config {
1467 if !route_config.websocket {
1468 warn!(
1469 correlation_id = %ctx.trace_id,
1470 route_id = ctx.route_id.as_deref().unwrap_or("unknown"),
1471 client_ip = %ctx.client_ip,
1472 "WebSocket upgrade rejected: not enabled for route"
1473 );
1474
1475 self.metrics.record_blocked_request("websocket_not_enabled");
1476
1477 let audit_entry = AuditLogEntry::new(
1479 &ctx.trace_id,
1480 AuditEventType::Blocked,
1481 &ctx.method,
1482 &ctx.path,
1483 &ctx.client_ip,
1484 )
1485 .with_route_id(ctx.route_id.as_deref().unwrap_or("unknown"))
1486 .with_action("websocket_rejected")
1487 .with_reason("WebSocket not enabled for route");
1488 self.log_manager.log_audit(&audit_entry);
1489
1490 crate::http_helpers::write_error(
1492 session,
1493 403,
1494 "WebSocket not enabled for this route",
1495 "text/plain",
1496 )
1497 .await?;
1498 return Ok(true); }
1500
1501 debug!(
1502 correlation_id = %ctx.trace_id,
1503 route_id = ctx.route_id.as_deref().unwrap_or("unknown"),
1504 "WebSocket upgrade request allowed"
1505 );
1506
1507 if route_config.websocket_inspection {
1509 let has_compression = session
1511 .req_header()
1512 .headers
1513 .get("Sec-WebSocket-Extensions")
1514 .and_then(|v| v.to_str().ok())
1515 .map(|s| s.contains("permessage-deflate"))
1516 .unwrap_or(false);
1517
1518 if has_compression {
1519 debug!(
1520 correlation_id = %ctx.trace_id,
1521 route_id = ctx.route_id.as_deref().unwrap_or("unknown"),
1522 "WebSocket inspection skipped: permessage-deflate negotiated"
1523 );
1524 ctx.websocket_skip_inspection = true;
1525 } else {
1526 ctx.websocket_inspection_enabled = true;
1527
1528 ctx.websocket_inspection_agents = self.agent_manager.get_agents_for_event(
1530 zentinel_agent_protocol::EventType::WebSocketFrame,
1531 );
1532
1533 debug!(
1534 correlation_id = %ctx.trace_id,
1535 route_id = ctx.route_id.as_deref().unwrap_or("unknown"),
1536 agent_count = ctx.websocket_inspection_agents.len(),
1537 "WebSocket frame inspection enabled"
1538 );
1539 }
1540 }
1541 }
1542 }
1543
1544 if let Some(route_config) = ctx.route_config.clone() {
1547 if route_config.service_type == zentinel_config::ServiceType::Static {
1548 trace!(
1549 correlation_id = %ctx.trace_id,
1550 route_id = ctx.route_id.as_deref().unwrap_or("unknown"),
1551 "Handling static file route"
1552 );
1553 let route_match = crate::routing::RouteMatch {
1555 route_id: zentinel_common::RouteId::new(ctx.route_id.as_deref().unwrap_or("")),
1556 config: route_config.clone(),
1557 };
1558 return self.handle_static_route(session, ctx, &route_match).await;
1559 } else if route_config.service_type == zentinel_config::ServiceType::Builtin {
1560 trace!(
1561 correlation_id = %ctx.trace_id,
1562 route_id = ctx.route_id.as_deref().unwrap_or("unknown"),
1563 builtin_handler = ?route_config.builtin_handler,
1564 "Handling builtin route"
1565 );
1566 let route_match = crate::routing::RouteMatch {
1568 route_id: zentinel_common::RouteId::new(ctx.route_id.as_deref().unwrap_or("")),
1569 config: route_config.clone(),
1570 };
1571 return self.handle_builtin_route(session, ctx, &route_match).await;
1572 }
1573 }
1574
1575 if let Some(route_id) = ctx.route_id.clone() {
1577 if let Some(validator) = self.validators.get(&route_id).await {
1578 trace!(
1579 correlation_id = %ctx.trace_id,
1580 route_id = %route_id,
1581 "Running API schema validation"
1582 );
1583 if let Some(result) = self
1584 .validate_api_request(session, ctx, &route_id, &validator)
1585 .await?
1586 {
1587 debug!(
1588 correlation_id = %ctx.trace_id,
1589 route_id = %route_id,
1590 validation_passed = result,
1591 "API validation complete"
1592 );
1593 return Ok(result);
1594 }
1595 }
1596 }
1597
1598 let client_addr = session
1600 .client_addr()
1601 .map(|a| format!("{}", a))
1602 .unwrap_or_else(|| "unknown".to_string());
1603 let client_port = session.client_addr().map(|_| 0).unwrap_or(0);
1604
1605 let req_header = session.req_header_mut();
1606
1607 req_header
1609 .insert_header("X-Correlation-Id", &ctx.trace_id)
1610 .ok();
1611 req_header.insert_header("X-Forwarded-By", "Zentinel").ok();
1612
1613 let config = ctx
1615 .config
1616 .get_or_insert_with(|| self.config_manager.current());
1617
1618 const HEADER_LIMIT_THRESHOLD: usize = 1024 * 1024; let header_count = req_header.headers.len();
1623 if config.limits.max_header_count < HEADER_LIMIT_THRESHOLD
1624 && header_count > config.limits.max_header_count
1625 {
1626 warn!(
1627 correlation_id = %ctx.trace_id,
1628 header_count = header_count,
1629 limit = config.limits.max_header_count,
1630 "Request blocked: exceeds header count limit"
1631 );
1632
1633 self.metrics.record_blocked_request("header_count_exceeded");
1634 return Err(Error::explain(ErrorType::InternalError, "Too many headers"));
1635 }
1636
1637 if config.limits.max_header_size_bytes < HEADER_LIMIT_THRESHOLD {
1639 let total_header_size: usize = req_header
1640 .headers
1641 .iter()
1642 .map(|(k, v)| k.as_str().len() + v.len())
1643 .sum();
1644
1645 if total_header_size > config.limits.max_header_size_bytes {
1646 warn!(
1647 correlation_id = %ctx.trace_id,
1648 header_size = total_header_size,
1649 limit = config.limits.max_header_size_bytes,
1650 "Request blocked: exceeds header size limit"
1651 );
1652
1653 self.metrics.record_blocked_request("header_size_exceeded");
1654 return Err(Error::explain(
1655 ErrorType::InternalError,
1656 "Headers too large",
1657 ));
1658 }
1659 }
1660
1661 trace!(
1663 correlation_id = %ctx.trace_id,
1664 "Processing request through agents"
1665 );
1666 if let Err(e) = self
1667 .process_agents(session, ctx, &client_addr, client_port)
1668 .await
1669 {
1670 if let ErrorType::HTTPStatus(status) = e.etype() {
1673 let error_msg = e.to_string();
1675 let body = error_msg
1676 .split("context:")
1677 .nth(1)
1678 .map(|s| s.trim())
1679 .unwrap_or("Request blocked");
1680 debug!(
1681 correlation_id = %ctx.trace_id,
1682 status = status,
1683 body = %body,
1684 "Sending HTTP error response for agent block"
1685 );
1686 crate::http_helpers::write_error(session, *status, body, "text/plain").await?;
1687 return Ok(true); }
1689 return Err(e);
1691 }
1692
1693 trace!(
1694 correlation_id = %ctx.trace_id,
1695 "Request filter phase complete, forwarding to upstream"
1696 );
1697
1698 Ok(false) }
1700
1701 async fn request_body_filter(
1708 &self,
1709 session: &mut Session,
1710 body: &mut Option<Bytes>,
1711 end_of_stream: bool,
1712 ctx: &mut Self::CTX,
1713 ) -> Result<(), Box<Error>> {
1714 use zentinel_config::BodyStreamingMode;
1715
1716 self.evaluate_agentic_policy(session, body.as_ref(), end_of_stream, ctx)?;
1720
1721 if ctx.is_websocket_upgrade {
1723 if let Some(ref handler) = ctx.websocket_handler {
1724 let result = handler.process_client_data(body.take()).await;
1725 match result {
1726 crate::websocket::ProcessResult::Forward(data) => {
1727 *body = data;
1728 }
1729 crate::websocket::ProcessResult::Close(reason) => {
1730 warn!(
1731 correlation_id = %ctx.trace_id,
1732 code = reason.code,
1733 reason = %reason.reason,
1734 "WebSocket connection closed by agent (client->server)"
1735 );
1736 return Err(Error::explain(
1738 ErrorType::InternalError,
1739 format!("WebSocket closed: {} {}", reason.code, reason.reason),
1740 ));
1741 }
1742 }
1743 }
1744 return Ok(());
1746 }
1747
1748 let chunk_len = body.as_ref().map(|b| b.len()).unwrap_or(0);
1750 if chunk_len > 0 {
1751 ctx.request_body_bytes += chunk_len as u64;
1752
1753 trace!(
1754 correlation_id = %ctx.trace_id,
1755 chunk_size = chunk_len,
1756 total_body_bytes = ctx.request_body_bytes,
1757 end_of_stream = end_of_stream,
1758 streaming_mode = ?ctx.request_body_streaming_mode,
1759 "Processing request body chunk"
1760 );
1761
1762 let config = ctx
1764 .config
1765 .get_or_insert_with(|| self.config_manager.current());
1766 if ctx.request_body_bytes > config.limits.max_body_size_bytes as u64 {
1767 warn!(
1768 correlation_id = %ctx.trace_id,
1769 body_bytes = ctx.request_body_bytes,
1770 limit = config.limits.max_body_size_bytes,
1771 "Request body size limit exceeded"
1772 );
1773 self.metrics.record_blocked_request("body_size_exceeded");
1774 return Err(Error::explain(
1775 ErrorType::InternalError,
1776 "Request body too large",
1777 ));
1778 }
1779 }
1780
1781 if ctx.body_inspection_enabled && !ctx.body_inspection_agents.is_empty() {
1783 let config = ctx
1784 .config
1785 .get_or_insert_with(|| self.config_manager.current());
1786 let max_inspection_bytes = config
1787 .waf
1788 .as_ref()
1789 .map(|w| w.body_inspection.max_inspection_bytes as u64)
1790 .unwrap_or(1024 * 1024);
1791
1792 match ctx.request_body_streaming_mode {
1793 BodyStreamingMode::Stream => {
1794 if body.is_some() {
1796 self.process_body_chunk_streaming(body, end_of_stream, ctx)
1797 .await?;
1798 } else if end_of_stream && ctx.agent_needs_more {
1799 self.process_body_chunk_streaming(body, end_of_stream, ctx)
1801 .await?;
1802 }
1803 }
1804 BodyStreamingMode::Hybrid { buffer_threshold } => {
1805 if ctx.body_bytes_inspected < buffer_threshold as u64 {
1807 if let Some(ref chunk) = body {
1809 let bytes_to_buffer = std::cmp::min(
1810 chunk.len(),
1811 (buffer_threshold as u64 - ctx.body_bytes_inspected) as usize,
1812 );
1813 ctx.body_buffer.extend_from_slice(&chunk[..bytes_to_buffer]);
1814 ctx.body_bytes_inspected += bytes_to_buffer as u64;
1815
1816 if ctx.body_bytes_inspected >= buffer_threshold as u64 || end_of_stream
1818 {
1819 self.send_buffered_body_to_agents(
1821 end_of_stream && chunk.len() == bytes_to_buffer,
1822 ctx,
1823 )
1824 .await?;
1825 ctx.body_buffer.clear();
1826
1827 if bytes_to_buffer < chunk.len() {
1829 let remaining = chunk.slice(bytes_to_buffer..);
1830 let mut remaining_body = Some(remaining);
1831 self.process_body_chunk_streaming(
1832 &mut remaining_body,
1833 end_of_stream,
1834 ctx,
1835 )
1836 .await?;
1837 }
1838 }
1839 }
1840 } else {
1841 self.process_body_chunk_streaming(body, end_of_stream, ctx)
1843 .await?;
1844 }
1845 }
1846 BodyStreamingMode::Buffer => {
1847 if let Some(ref chunk) = body {
1849 if ctx.body_bytes_inspected < max_inspection_bytes {
1850 let bytes_to_inspect = std::cmp::min(
1851 chunk.len() as u64,
1852 max_inspection_bytes - ctx.body_bytes_inspected,
1853 ) as usize;
1854
1855 ctx.body_buffer
1856 .extend_from_slice(&chunk[..bytes_to_inspect]);
1857 ctx.body_bytes_inspected += bytes_to_inspect as u64;
1858
1859 trace!(
1860 correlation_id = %ctx.trace_id,
1861 bytes_inspected = ctx.body_bytes_inspected,
1862 max_inspection_bytes = max_inspection_bytes,
1863 buffer_size = ctx.body_buffer.len(),
1864 "Buffering body for agent inspection"
1865 );
1866 }
1867 }
1868
1869 let should_send =
1871 end_of_stream || ctx.body_bytes_inspected >= max_inspection_bytes;
1872 if should_send && !ctx.body_buffer.is_empty() {
1873 self.send_buffered_body_to_agents(end_of_stream, ctx)
1874 .await?;
1875 ctx.body_buffer.clear();
1876 }
1877 }
1878 }
1879 }
1880
1881 if end_of_stream {
1882 trace!(
1883 correlation_id = %ctx.trace_id,
1884 total_body_bytes = ctx.request_body_bytes,
1885 bytes_inspected = ctx.body_bytes_inspected,
1886 "Request body complete"
1887 );
1888 }
1889
1890 Ok(())
1891 }
1892
1893 async fn response_filter(
1894 &self,
1895 session: &mut Session,
1896 upstream_response: &mut ResponseHeader,
1897 ctx: &mut Self::CTX,
1898 ) -> Result<(), Box<Error>> {
1899 let status = upstream_response.status.as_u16();
1900 let duration = ctx.elapsed();
1901
1902 trace!(
1903 correlation_id = %ctx.trace_id,
1904 status = status,
1905 "Starting response filter phase"
1906 );
1907
1908 if status == 101 && ctx.is_websocket_upgrade {
1910 if ctx.websocket_inspection_enabled && !ctx.websocket_skip_inspection {
1911 let inspector = crate::websocket::WebSocketInspector::with_metrics(
1913 self.agent_manager.clone(),
1914 ctx.route_id
1915 .clone()
1916 .unwrap_or_else(|| "unknown".to_string()),
1917 ctx.trace_id.clone(),
1918 ctx.client_ip.clone(),
1919 100, Some(self.metrics.clone()),
1921 );
1922
1923 let handler = crate::websocket::WebSocketHandler::new(
1924 std::sync::Arc::new(inspector),
1925 1024 * 1024, );
1927
1928 ctx.websocket_handler = Some(std::sync::Arc::new(handler));
1929
1930 info!(
1931 correlation_id = %ctx.trace_id,
1932 route_id = ctx.route_id.as_deref().unwrap_or("unknown"),
1933 agent_count = ctx.websocket_inspection_agents.len(),
1934 "WebSocket upgrade successful, frame inspection enabled"
1935 );
1936 } else if ctx.websocket_skip_inspection {
1937 debug!(
1938 correlation_id = %ctx.trace_id,
1939 route_id = ctx.route_id.as_deref().unwrap_or("unknown"),
1940 "WebSocket upgrade successful, inspection skipped (compression negotiated)"
1941 );
1942 } else {
1943 debug!(
1944 correlation_id = %ctx.trace_id,
1945 route_id = ctx.route_id.as_deref().unwrap_or("unknown"),
1946 "WebSocket upgrade successful"
1947 );
1948 }
1949 }
1950
1951 upstream_response.insert_header("X-Correlation-Id", &ctx.trace_id)?;
1953
1954 if let Some(ref rate_info) = ctx.rate_limit_info {
1956 upstream_response.insert_header("X-RateLimit-Limit", rate_info.limit.to_string())?;
1957 upstream_response
1958 .insert_header("X-RateLimit-Remaining", rate_info.remaining.to_string())?;
1959 upstream_response.insert_header("X-RateLimit-Reset", rate_info.reset_at.to_string())?;
1960 }
1961
1962 if ctx.inference_budget_enabled {
1964 if let Some(remaining) = ctx.inference_budget_remaining {
1965 upstream_response.insert_header("X-Budget-Remaining", remaining.to_string())?;
1966 }
1967 if let Some(period_reset) = ctx.inference_budget_period_reset {
1968 let reset_datetime = chrono::DateTime::from_timestamp(period_reset as i64, 0)
1970 .map(|dt| dt.format("%Y-%m-%dT%H:%M:%SZ").to_string())
1971 .unwrap_or_else(|| period_reset.to_string());
1972 upstream_response.insert_header("X-Budget-Period-Reset", reset_datetime)?;
1973 }
1974 }
1975
1976 if let Some(ref country_code) = ctx.geo_country_code {
1978 upstream_response.insert_header("X-GeoIP-Country", country_code)?;
1979 }
1980
1981 if let Some(ref route_config) = ctx.route_config {
1983 let mods = &route_config.policies.response_headers;
1984 for (old_name, new_name) in &mods.rename {
1986 if let Some(value) = upstream_response
1987 .headers
1988 .get(old_name)
1989 .and_then(|v| v.to_str().ok())
1990 {
1991 let owned = value.to_string();
1992 upstream_response
1993 .insert_header(new_name.clone(), &owned)
1994 .ok();
1995 upstream_response.remove_header(old_name);
1996 }
1997 }
1998 for (name, value) in &mods.set {
1999 upstream_response
2000 .insert_header(name.clone(), value.as_str())
2001 .ok();
2002 }
2003 for (name, value) in &mods.add {
2004 upstream_response
2005 .append_header(name.clone(), value.as_str())
2006 .ok();
2007 }
2008 for name in &mods.remove {
2009 upstream_response.remove_header(name);
2010 }
2011 }
2012
2013 if let Some(ref cache_status) = ctx.cache_status {
2015 let status_header_enabled = ctx
2016 .config
2017 .as_ref()
2018 .and_then(|c| c.cache.as_ref())
2019 .map(|c| c.status_header)
2020 .unwrap_or(false);
2021
2022 if status_header_enabled {
2023 let cache_name = ctx
2024 .config
2025 .as_ref()
2026 .and_then(|c| c.cache.as_ref())
2027 .map(|c| c.status_header_name.as_str())
2028 .unwrap_or("zentinel");
2029
2030 apply_cache_status(upstream_response, cache_name, cache_status);
2031 }
2032 }
2033
2034 if let Some(config) = ctx.config.as_ref().map(std::sync::Arc::clone) {
2036 super::filters::apply_response_filters(upstream_response, ctx, &config);
2037 }
2038
2039 if ctx.compress_enabled {
2041 session.upstream_compression.adjust_level(6);
2042 }
2043
2044 if let Some(keepalive_secs) = ctx.listener_keepalive_timeout_secs {
2046 session
2047 .downstream_session
2048 .set_keepalive(Some(keepalive_secs));
2049 }
2050
2051 if ctx.sticky_session_new_assignment {
2053 if let Some(ref set_cookie_header) = ctx.sticky_session_set_cookie {
2054 upstream_response.insert_header("Set-Cookie", set_cookie_header)?;
2055 trace!(
2056 correlation_id = %ctx.trace_id,
2057 sticky_target_index = ?ctx.sticky_target_index,
2058 "Added sticky session Set-Cookie header"
2059 );
2060 }
2061 }
2062
2063 if ctx.guardrail_warning {
2065 upstream_response.insert_header("X-Guardrail-Warning", "prompt-injection-detected")?;
2066 }
2067
2068 if ctx.used_fallback() {
2070 upstream_response.insert_header("X-Fallback-Used", "true")?;
2071
2072 if let Some(ref upstream) = ctx.upstream {
2073 upstream_response.insert_header("X-Fallback-Upstream", upstream)?;
2074 }
2075
2076 if let Some(ref reason) = ctx.fallback_reason {
2077 upstream_response.insert_header("X-Fallback-Reason", reason.to_string())?;
2078 }
2079
2080 if let Some(ref original) = ctx.original_upstream {
2081 upstream_response.insert_header("X-Original-Upstream", original)?;
2082 }
2083
2084 if let Some(ref mapping) = ctx.model_mapping_applied {
2085 upstream_response
2086 .insert_header("X-Model-Mapping", format!("{} -> {}", mapping.0, mapping.1))?;
2087 }
2088
2089 trace!(
2090 correlation_id = %ctx.trace_id,
2091 fallback_attempt = ctx.fallback_attempt,
2092 fallback_upstream = ctx.upstream.as_deref().unwrap_or("unknown"),
2093 original_upstream = ctx.original_upstream.as_deref().unwrap_or("unknown"),
2094 "Added fallback response headers"
2095 );
2096
2097 if status < 400 {
2099 if let Some(metrics) = get_fallback_metrics() {
2100 metrics.record_fallback_success(
2101 ctx.route_id.as_deref().unwrap_or("unknown"),
2102 ctx.upstream.as_deref().unwrap_or("unknown"),
2103 );
2104 }
2105 }
2106 }
2107
2108 if ctx.inference_rate_limit_enabled {
2110 let content_type = upstream_response
2112 .headers
2113 .get("content-type")
2114 .and_then(|ct| ct.to_str().ok());
2115
2116 if is_sse_response(content_type) {
2117 let provider = ctx
2119 .route_config
2120 .as_ref()
2121 .and_then(|r| r.inference.as_ref())
2122 .map(|i| i.provider)
2123 .unwrap_or_default();
2124
2125 ctx.inference_streaming_response = true;
2126 ctx.inference_streaming_counter = Some(StreamingTokenCounter::new(
2127 provider,
2128 ctx.inference_model.clone(),
2129 ));
2130
2131 trace!(
2132 correlation_id = %ctx.trace_id,
2133 content_type = ?content_type,
2134 model = ?ctx.inference_model,
2135 "Initialized streaming token counter for SSE response"
2136 );
2137 }
2138 }
2139
2140 if !ctx.route_agent_ids.is_empty() {
2142 let agent_ids = ctx.route_agent_ids.clone();
2143 let mut resp_headers_map: std::collections::HashMap<String, Vec<String>> =
2144 std::collections::HashMap::with_capacity(upstream_response.headers.len());
2145 for (name, value) in upstream_response.headers.iter() {
2146 resp_headers_map
2147 .entry(name.as_str().to_string())
2148 .or_default()
2149 .push(value.to_str().unwrap_or("").to_string());
2150 }
2151
2152 let agent_ctx = crate::agents::AgentCallContext {
2153 correlation_id: zentinel_common::CorrelationId::from_string(&ctx.trace_id),
2154 metadata: zentinel_agent_protocol::RequestMetadata {
2155 correlation_id: ctx.trace_id.clone(),
2156 request_id: uuid::Uuid::new_v4().to_string(),
2157 client_ip: ctx.client_ip.clone(),
2158 client_port: 0,
2159 server_name: ctx.host.clone(),
2160 protocol: "HTTP/1.1".to_string(),
2161 tls_version: None,
2162 tls_cipher: None,
2163 route_id: ctx.route_id.clone(),
2164 upstream_id: ctx.upstream.clone(),
2165 timestamp: chrono::Utc::now().to_rfc3339(),
2166 traceparent: ctx.traceparent(),
2167 },
2168 route_id: ctx.route_id.clone(),
2169 upstream_id: ctx.upstream.clone(),
2170 request_body: None,
2171 response_body: None,
2172 };
2173
2174 match self
2175 .agent_manager
2176 .process_response_headers(&agent_ctx, status, &resp_headers_map, &agent_ids)
2177 .await
2178 {
2179 Ok(decision) => {
2180 for op in &decision.response_headers {
2182 match op {
2183 zentinel_agent_protocol::HeaderOp::Set { name, value } => {
2184 upstream_response
2185 .insert_header(name.clone(), value.as_str())
2186 .ok();
2187 }
2188 zentinel_agent_protocol::HeaderOp::Add { name, value } => {
2189 upstream_response
2190 .append_header(name.clone(), value.as_str())
2191 .ok();
2192 }
2193 zentinel_agent_protocol::HeaderOp::Remove { name } => {
2194 upstream_response.remove_header(name);
2195 }
2196 }
2197 }
2198
2199 let has_body_agents = self
2201 .agent_manager
2202 .any_agent_handles_event(
2203 &agent_ids,
2204 zentinel_agent_protocol::EventType::ResponseBodyChunk,
2205 )
2206 .await;
2207 if has_body_agents {
2208 ctx.response_agent_processing_enabled = true;
2209 upstream_response.insert_header("Connection", "close").ok();
2212 session.downstream_session.set_keepalive(None);
2213 debug!(
2214 correlation_id = %ctx.trace_id,
2215 "Enabling response body agent processing (agent subscribes to ResponseBody)"
2216 );
2217 }
2218
2219 debug!(
2220 correlation_id = %ctx.trace_id,
2221 response_headers_modified = !decision.response_headers.is_empty(),
2222 needs_body = ctx.response_agent_processing_enabled,
2223 "Response headers processed through agents"
2224 );
2225 }
2226 Err(e) => {
2227 warn!(
2228 correlation_id = %ctx.trace_id,
2229 error = %e,
2230 "Agent response header processing failed, continuing without agent"
2231 );
2232 }
2233 }
2234 }
2235
2236 if status >= 400 {
2238 trace!(
2239 correlation_id = %ctx.trace_id,
2240 status = status,
2241 "Handling error response"
2242 );
2243 self.handle_error_response(upstream_response, ctx).await?;
2244 }
2245
2246 self.metrics.record_request(
2248 ctx.route_id.as_deref().unwrap_or("unknown"),
2249 &ctx.method,
2250 status,
2251 duration,
2252 );
2253
2254 if let Some(ref mut span) = ctx.otel_span {
2256 span.set_status(status);
2257 if let Some(ref upstream) = ctx.upstream {
2258 span.set_upstream(upstream, "");
2259 }
2260 if status >= 500 {
2261 span.record_error(&format!("HTTP {}", status));
2262 }
2263 }
2264
2265 if let Some(ref upstream) = ctx.upstream {
2267 let success = status < 500;
2268
2269 trace!(
2270 correlation_id = %ctx.trace_id,
2271 upstream = %upstream,
2272 success = success,
2273 status = status,
2274 "Recording passive health check result"
2275 );
2276
2277 let error_msg = if !success {
2278 Some(format!("HTTP {}", status))
2279 } else {
2280 None
2281 };
2282 self.passive_health
2283 .record_outcome(upstream, success, error_msg.as_deref())
2284 .await;
2285
2286 if let Some(pool) = self.upstream_pools.get(upstream).await {
2288 pool.report_result(upstream, success).await;
2289 }
2290
2291 if !success {
2292 warn!(
2293 correlation_id = %ctx.trace_id,
2294 upstream = %upstream,
2295 status = status,
2296 "Upstream returned error status"
2297 );
2298 }
2299 }
2300
2301 if status >= 500 {
2303 error!(
2304 correlation_id = %ctx.trace_id,
2305 route_id = ctx.route_id.as_deref().unwrap_or("unknown"),
2306 upstream = ctx.upstream.as_deref().unwrap_or("none"),
2307 method = %ctx.method,
2308 path = %ctx.path,
2309 status = status,
2310 duration_ms = duration.as_millis(),
2311 attempts = ctx.upstream_attempts,
2312 "Request completed with server error"
2313 );
2314 self.log_manager.log_request_error(
2315 "error",
2316 "Request completed with server error",
2317 &ctx.trace_id,
2318 ctx.route_id.as_deref(),
2319 ctx.upstream.as_deref(),
2320 Some(format!(
2321 "status={} method={} path={} duration_ms={}",
2322 status,
2323 ctx.method,
2324 ctx.path,
2325 duration.as_millis()
2326 )),
2327 );
2328 } else if status >= 400 {
2329 warn!(
2330 correlation_id = %ctx.trace_id,
2331 route_id = ctx.route_id.as_deref().unwrap_or("unknown"),
2332 upstream = ctx.upstream.as_deref().unwrap_or("none"),
2333 method = %ctx.method,
2334 path = %ctx.path,
2335 status = status,
2336 duration_ms = duration.as_millis(),
2337 "Request completed with client error"
2338 );
2339 self.log_manager.log_request_error(
2340 "warn",
2341 "Request completed with client error",
2342 &ctx.trace_id,
2343 ctx.route_id.as_deref(),
2344 ctx.upstream.as_deref(),
2345 Some(format!(
2346 "status={} method={} path={}",
2347 status, ctx.method, ctx.path
2348 )),
2349 );
2350 } else {
2351 debug!(
2352 correlation_id = %ctx.trace_id,
2353 route_id = ctx.route_id.as_deref().unwrap_or("unknown"),
2354 upstream = ctx.upstream.as_deref().unwrap_or("none"),
2355 method = %ctx.method,
2356 path = %ctx.path,
2357 status = status,
2358 duration_ms = duration.as_millis(),
2359 attempts = ctx.upstream_attempts,
2360 "Request completed"
2361 );
2362 }
2363
2364 Ok(())
2365 }
2366
2367 async fn upstream_request_filter(
2370 &self,
2371 _session: &mut Session,
2372 upstream_request: &mut pingora::http::RequestHeader,
2373 ctx: &mut Self::CTX,
2374 ) -> Result<()>
2375 where
2376 Self::CTX: Send + Sync,
2377 {
2378 trace!(
2379 correlation_id = %ctx.trace_id,
2380 route_id = ctx.route_id.as_deref().unwrap_or("unknown"),
2381 "Applying upstream request modifications"
2382 );
2383
2384 upstream_request
2386 .insert_header("X-Trace-Id", &ctx.trace_id)
2387 .ok();
2388
2389 if let Some(ref span) = ctx.otel_span {
2391 let sampled = ctx
2392 .trace_context
2393 .as_ref()
2394 .map(|c| c.sampled)
2395 .unwrap_or(true);
2396 let traceparent =
2397 crate::otel::create_traceparent(&span.trace_id, &span.span_id, sampled);
2398 upstream_request
2399 .insert_header(crate::otel::TRACEPARENT_HEADER, &traceparent)
2400 .ok();
2401 }
2402
2403 upstream_request
2405 .insert_header("X-Forwarded-By", "Zentinel")
2406 .ok();
2407
2408 if let Some(ref route_config) = ctx.route_config {
2412 let mods = &route_config.policies.request_headers;
2413
2414 for (old_name, new_name) in &mods.rename {
2416 if let Some(value) = upstream_request
2417 .headers
2418 .get(old_name)
2419 .and_then(|v| v.to_str().ok())
2420 {
2421 let owned = value.to_string();
2422 upstream_request
2423 .insert_header(new_name.clone(), &owned)
2424 .ok();
2425 upstream_request.remove_header(old_name);
2426 }
2427 }
2428
2429 for (name, value) in &mods.set {
2431 upstream_request
2432 .insert_header(name.clone(), value.as_str())
2433 .ok();
2434 }
2435
2436 for (name, value) in &mods.add {
2438 upstream_request
2439 .append_header(name.clone(), value.as_str())
2440 .ok();
2441 }
2442
2443 for name in &mods.remove {
2445 upstream_request.remove_header(name);
2446 }
2447
2448 trace!(
2449 correlation_id = %ctx.trace_id,
2450 "Applied request header modifications"
2451 );
2452 }
2453
2454 if let Some(ref config) = ctx.config {
2456 super::filters::apply_request_headers_filters(upstream_request, ctx, config);
2457 }
2458
2459 upstream_request.remove_header("X-Internal-Token");
2461 upstream_request.remove_header("Authorization-Internal");
2462
2463 if let Some(ref route_config) = ctx.route_config {
2466 if let Some(ref shadow_config) = route_config.shadow {
2467 let pools_snapshot = self.upstream_pools.snapshot().await;
2469 let upstream_pools = std::sync::Arc::new(pools_snapshot);
2470
2471 let route_id = ctx
2473 .route_id
2474 .clone()
2475 .unwrap_or_else(|| "unknown".to_string());
2476
2477 let shadow_manager = crate::shadow::ShadowManager::new(
2479 upstream_pools,
2480 shadow_config.clone(),
2481 Some(std::sync::Arc::clone(&self.metrics)),
2482 route_id,
2483 );
2484
2485 if shadow_manager.should_shadow(upstream_request) {
2487 trace!(
2488 correlation_id = %ctx.trace_id,
2489 shadow_upstream = %shadow_config.upstream,
2490 percentage = shadow_config.percentage,
2491 "Shadowing request"
2492 );
2493
2494 let shadow_headers = upstream_request.clone();
2496
2497 let shadow_ctx = crate::upstream::RequestContext {
2499 client_ip: ctx.client_ip.parse().ok(),
2500 headers: std::collections::HashMap::new(), path: ctx.path.clone(),
2502 method: ctx.method.clone(),
2503 };
2504
2505 let buffer_body = shadow_config.buffer_body
2507 && crate::shadow::should_buffer_method(&ctx.method);
2508
2509 if buffer_body {
2510 trace!(
2514 correlation_id = %ctx.trace_id,
2515 "Deferring shadow request until body is buffered"
2516 );
2517 ctx.shadow_pending = Some(crate::proxy::context::ShadowPendingRequest {
2518 headers: shadow_headers,
2519 manager: std::sync::Arc::new(shadow_manager),
2520 request_ctx: shadow_ctx,
2521 include_body: true,
2522 });
2523 if !ctx.body_inspection_enabled {
2526 ctx.body_inspection_enabled = true;
2527 }
2530 } else {
2531 shadow_manager.shadow_request(shadow_headers, None, shadow_ctx);
2533 ctx.shadow_sent = true;
2534 }
2535 }
2536 }
2537 }
2538
2539 Ok(())
2540 }
2541
2542 fn response_body_filter(
2548 &self,
2549 _session: &mut Session,
2550 body: &mut Option<Bytes>,
2551 end_of_stream: bool,
2552 ctx: &mut Self::CTX,
2553 ) -> Result<Option<Duration>, Box<Error>> {
2554 if ctx.is_websocket_upgrade {
2557 if let Some(ref handler) = ctx.websocket_handler {
2558 let handler = handler.clone();
2559 let data = body.take();
2560
2561 let result = tokio::task::block_in_place(|| {
2564 tokio::runtime::Handle::current()
2565 .block_on(async { handler.process_server_data(data).await })
2566 });
2567
2568 match result {
2569 crate::websocket::ProcessResult::Forward(data) => {
2570 *body = data;
2571 }
2572 crate::websocket::ProcessResult::Close(reason) => {
2573 warn!(
2574 correlation_id = %ctx.trace_id,
2575 code = reason.code,
2576 reason = %reason.reason,
2577 "WebSocket connection closed by agent (server->client)"
2578 );
2579 let close_frame =
2582 crate::websocket::WebSocketFrame::close(reason.code, &reason.reason);
2583 let codec = crate::websocket::WebSocketCodec::new(1024 * 1024);
2584 if let Ok(encoded) = codec.encode_frame(&close_frame, false) {
2585 *body = Some(Bytes::from(encoded));
2586 }
2587 }
2588 }
2589 }
2590 return Ok(None);
2592 }
2593
2594 if ctx.response_agent_processing_enabled && !ctx.route_agent_ids.is_empty() {
2596 if let Some(ref chunk) = body {
2597 ctx.response_agent_body_buffer.extend_from_slice(chunk);
2598 }
2599
2600 if end_of_stream {
2601 let agent_ids = ctx.route_agent_ids.clone();
2602 let buffer = std::mem::take(&mut ctx.response_agent_body_buffer);
2603 let chunk_index = 0u32;
2604 let total_size = Some(buffer.len());
2605 let trace_id = ctx.trace_id.clone();
2606 let client_ip = ctx.client_ip.clone();
2607 let host = ctx.host.clone();
2608 let route_id = ctx.route_id.clone();
2609 let upstream_id = ctx.upstream.clone();
2610 let traceparent = ctx.traceparent();
2611 let agent_mgr = self.agent_manager.clone();
2612
2613 let result = tokio::task::block_in_place(|| {
2616 tokio::runtime::Handle::current().block_on(async {
2617 let agent_ctx = crate::agents::AgentCallContext {
2618 correlation_id: zentinel_common::CorrelationId::from_string(&trace_id),
2619 metadata: zentinel_agent_protocol::RequestMetadata {
2620 correlation_id: trace_id.clone(),
2621 request_id: uuid::Uuid::new_v4().to_string(),
2622 client_ip,
2623 client_port: 0,
2624 server_name: host,
2625 protocol: "HTTP/1.1".to_string(),
2626 tls_version: None,
2627 tls_cipher: None,
2628 route_id: route_id.clone(),
2629 upstream_id: upstream_id.clone(),
2630 timestamp: chrono::Utc::now().to_rfc3339(),
2631 traceparent,
2632 },
2633 route_id,
2634 upstream_id,
2635 request_body: None,
2636 response_body: None,
2637 };
2638
2639 agent_mgr
2640 .process_response_body_streaming(
2641 &agent_ctx,
2642 &buffer,
2643 true, chunk_index,
2645 buffer.len(),
2646 total_size,
2647 &agent_ids,
2648 )
2649 .await
2650 })
2651 });
2652
2653 match result {
2654 Ok(decision) => {
2655 if let Some(mutation) = decision.response_body_mutation {
2657 if let Some(ref data) = mutation.data {
2658 if !data.is_empty() {
2659 if let Ok(decoded) = base64::Engine::decode(
2661 &base64::engine::general_purpose::STANDARD,
2662 data,
2663 ) {
2664 debug!(
2665 correlation_id = %ctx.trace_id,
2666 original_size = buffer.len(),
2667 new_size = decoded.len(),
2668 "Agent replaced response body"
2669 );
2670 *body = Some(Bytes::from(decoded));
2671 ctx.response_agent_body_complete = true;
2672 } else {
2673 warn!(
2674 correlation_id = %ctx.trace_id,
2675 "Failed to decode agent response body mutation (invalid base64)"
2676 );
2677 }
2678 }
2679 }
2681 }
2683
2684 }
2688 Err(e) => {
2689 warn!(
2690 correlation_id = %ctx.trace_id,
2691 error = %e,
2692 "Agent response body processing failed, passing through original"
2693 );
2694 }
2695 }
2696 } else if !end_of_stream {
2697 *body = None;
2699 return Ok(None);
2700 }
2701 }
2702
2703 if let Some(ref chunk) = body {
2705 ctx.response_bytes += chunk.len() as u64;
2706
2707 trace!(
2708 correlation_id = %ctx.trace_id,
2709 chunk_size = chunk.len(),
2710 total_response_bytes = ctx.response_bytes,
2711 end_of_stream = end_of_stream,
2712 "Processing response body chunk"
2713 );
2714
2715 if let Some(ref mut counter) = ctx.inference_streaming_counter {
2717 let result = counter.process_chunk(chunk);
2718
2719 if result.content.is_some() || result.is_done {
2720 trace!(
2721 correlation_id = %ctx.trace_id,
2722 has_content = result.content.is_some(),
2723 is_done = result.is_done,
2724 chunks_processed = counter.chunks_processed(),
2725 accumulated_content_len = counter.content().len(),
2726 "Processed SSE chunk for token counting"
2727 );
2728 }
2729 }
2730
2731 if ctx.response_body_inspection_enabled
2735 && !ctx.response_body_inspection_agents.is_empty()
2736 {
2737 let config = ctx
2738 .config
2739 .get_or_insert_with(|| self.config_manager.current());
2740 let max_inspection_bytes = config
2741 .waf
2742 .as_ref()
2743 .map(|w| w.body_inspection.max_inspection_bytes as u64)
2744 .unwrap_or(1024 * 1024);
2745
2746 if ctx.response_body_bytes_inspected < max_inspection_bytes {
2747 let bytes_to_inspect = std::cmp::min(
2748 chunk.len() as u64,
2749 max_inspection_bytes - ctx.response_body_bytes_inspected,
2750 ) as usize;
2751
2752 ctx.response_body_bytes_inspected += bytes_to_inspect as u64;
2756 ctx.response_body_chunk_index += 1;
2757
2758 trace!(
2759 correlation_id = %ctx.trace_id,
2760 bytes_inspected = ctx.response_body_bytes_inspected,
2761 max_inspection_bytes = max_inspection_bytes,
2762 chunk_index = ctx.response_body_chunk_index,
2763 "Tracking response body for inspection"
2764 );
2765 }
2766 }
2767 }
2768
2769 if end_of_stream {
2770 trace!(
2771 correlation_id = %ctx.trace_id,
2772 total_response_bytes = ctx.response_bytes,
2773 response_bytes_inspected = ctx.response_body_bytes_inspected,
2774 "Response body complete"
2775 );
2776 }
2777
2778 Ok(None)
2780 }
2781
2782 async fn connected_to_upstream(
2785 &self,
2786 _session: &mut Session,
2787 reused: bool,
2788 peer: &HttpPeer,
2789 #[cfg(unix)] _fd: RawFd,
2790 #[cfg(windows)] _sock: std::os::windows::io::RawSocket,
2791 digest: Option<&Digest>,
2792 ctx: &mut Self::CTX,
2793 ) -> Result<(), Box<Error>> {
2794 ctx.connection_reused = reused;
2796
2797 if reused {
2799 trace!(
2800 correlation_id = %ctx.trace_id,
2801 upstream = ctx.upstream.as_deref().unwrap_or("unknown"),
2802 peer_address = %peer.address(),
2803 "Reusing existing upstream connection"
2804 );
2805 } else {
2806 debug!(
2807 correlation_id = %ctx.trace_id,
2808 upstream = ctx.upstream.as_deref().unwrap_or("unknown"),
2809 peer_address = %peer.address(),
2810 ssl = digest.as_ref().map(|d| d.ssl_digest.is_some()).unwrap_or(false),
2811 "Established new upstream connection"
2812 );
2813 }
2814
2815 Ok(())
2816 }
2817
2818 fn request_cache_filter(&self, session: &mut Session, ctx: &mut Self::CTX) -> Result<()> {
2828 let route_id = match ctx.route_id.as_deref() {
2830 Some(id) => id,
2831 None => {
2832 trace!(
2833 correlation_id = %ctx.trace_id,
2834 "Cache filter: no route ID, skipping cache"
2835 );
2836 return Ok(());
2837 }
2838 };
2839
2840 if !self.cache_manager.is_enabled(route_id) {
2842 ctx.cache_status = Some(super::context::CacheStatus::Bypass("disabled"));
2843 trace!(
2844 correlation_id = %ctx.trace_id,
2845 route_id = %route_id,
2846 "Cache disabled for route"
2847 );
2848 return Ok(());
2849 }
2850
2851 if !self
2853 .cache_manager
2854 .is_method_cacheable(route_id, &ctx.method)
2855 {
2856 ctx.cache_status = Some(super::context::CacheStatus::Bypass("method"));
2857 trace!(
2858 correlation_id = %ctx.trace_id,
2859 route_id = %route_id,
2860 method = %ctx.method,
2861 "Method not cacheable"
2862 );
2863 return Ok(());
2864 }
2865
2866 if !self.cache_manager.is_path_cacheable(route_id, &ctx.path) {
2868 ctx.cache_status = Some(super::context::CacheStatus::Bypass("excluded"));
2869 trace!(
2870 correlation_id = %ctx.trace_id,
2871 route_id = %route_id,
2872 path = %ctx.path,
2873 "Path excluded from caching"
2874 );
2875 return Ok(());
2876 }
2877
2878 debug!(
2880 correlation_id = %ctx.trace_id,
2881 route_id = %route_id,
2882 method = %ctx.method,
2883 path = %ctx.path,
2884 "Enabling HTTP caching for request"
2885 );
2886
2887 let storage = get_cache_storage();
2889 let eviction = get_cache_eviction();
2890 let cache_lock = get_cache_lock();
2891
2892 session.cache.enable(
2894 storage,
2895 Some(eviction),
2896 None, Some(cache_lock),
2898 None, );
2900
2901 ctx.cache_eligible = true;
2903
2904 trace!(
2905 correlation_id = %ctx.trace_id,
2906 route_id = %route_id,
2907 cache_enabled = session.cache.enabled(),
2908 "Cache enabled for request"
2909 );
2910
2911 Ok(())
2912 }
2913
2914 fn cache_key_callback(&self, session: &Session, ctx: &mut Self::CTX) -> Result<CacheKey> {
2919 let req_header = session.req_header();
2920 let method = req_header.method.as_str();
2921 let path = req_header.uri.path();
2922 let host = ctx.host.as_deref().unwrap_or("unknown");
2923 let query = req_header.uri.query();
2924
2925 let key_string = crate::cache::CacheManager::generate_cache_key(method, host, path, query);
2927
2928 trace!(
2929 correlation_id = %ctx.trace_id,
2930 cache_key = %key_string,
2931 "Generated cache key"
2932 );
2933
2934 Ok(CacheKey::new("", format!("{}", req_header.uri), ""))
2937 }
2938
2939 fn cache_miss(&self, session: &mut Session, ctx: &mut Self::CTX) {
2944 session.cache.cache_miss();
2946
2947 ctx.cache_status = Some(super::context::CacheStatus::Miss);
2948
2949 if let Some(route_id) = ctx.route_id.as_deref() {
2951 self.cache_manager.stats().record_miss();
2952
2953 trace!(
2954 correlation_id = %ctx.trace_id,
2955 route_id = %route_id,
2956 path = %ctx.path,
2957 "Cache miss"
2958 );
2959 }
2960 }
2961
2962 async fn cache_hit_filter(
2968 &self,
2969 session: &mut Session,
2970 meta: &CacheMeta,
2971 hit_handler: &mut HitHandler,
2972 is_fresh: bool,
2973 ctx: &mut Self::CTX,
2974 ) -> Result<Option<ForcedFreshness>>
2975 where
2976 Self::CTX: Send + Sync,
2977 {
2978 let req_header = session.req_header();
2980 let method = req_header.method.as_str();
2981 let path = req_header.uri.path();
2982 let host = req_header.uri.host().unwrap_or("localhost");
2983 let query = req_header.uri.query();
2984
2985 let cache_key = crate::cache::CacheManager::generate_cache_key(method, host, path, query);
2987
2988 if self.cache_manager.should_invalidate(&cache_key) {
2990 info!(
2991 correlation_id = %ctx.trace_id,
2992 route_id = ctx.route_id.as_deref().unwrap_or("unknown"),
2993 cache_key = %cache_key,
2994 "Cache entry invalidated by purge request"
2995 );
2996 return Ok(Some(ForcedFreshness::ForceExpired));
2998 }
2999
3000 if is_fresh {
3002 let is_disk_hit = hit_handler
3004 .as_any()
3005 .downcast_ref::<HybridHitHandler>()
3006 .is_some()
3007 || hit_handler
3008 .as_any()
3009 .downcast_ref::<DiskHitHandler>()
3010 .is_some();
3011
3012 let stats = self.cache_manager.stats();
3013 if is_disk_hit {
3014 ctx.cache_status = Some(super::context::CacheStatus::HitDisk);
3015 stats.record_disk_hit();
3016 } else {
3017 ctx.cache_status = Some(super::context::CacheStatus::HitMemory);
3018 stats.record_memory_hit();
3019 }
3020 stats.record_hit();
3021
3022 debug!(
3023 correlation_id = %ctx.trace_id,
3024 route_id = ctx.route_id.as_deref().unwrap_or("unknown"),
3025 is_fresh = is_fresh,
3026 tier = if is_disk_hit { "disk" } else { "memory" },
3027 "Cache hit (fresh)"
3028 );
3029 } else {
3030 ctx.cache_status = Some(super::context::CacheStatus::HitStale);
3031
3032 trace!(
3033 correlation_id = %ctx.trace_id,
3034 route_id = ctx.route_id.as_deref().unwrap_or("unknown"),
3035 is_fresh = is_fresh,
3036 "Cache hit (stale)"
3037 );
3038 }
3039
3040 Ok(None)
3042 }
3043
3044 fn response_cache_filter(
3049 &self,
3050 _session: &Session,
3051 resp: &ResponseHeader,
3052 ctx: &mut Self::CTX,
3053 ) -> Result<RespCacheable> {
3054 let route_id = match ctx.route_id.as_deref() {
3055 Some(id) => id,
3056 None => {
3057 return Ok(RespCacheable::Uncacheable(NoCacheReason::Custom(
3058 "no_route",
3059 )));
3060 }
3061 };
3062
3063 if !self.cache_manager.is_enabled(route_id) {
3065 return Ok(RespCacheable::Uncacheable(NoCacheReason::Custom(
3066 "disabled",
3067 )));
3068 }
3069
3070 let status = resp.status.as_u16();
3071
3072 if !self.cache_manager.is_status_cacheable(route_id, status) {
3074 trace!(
3075 correlation_id = %ctx.trace_id,
3076 route_id = %route_id,
3077 status = status,
3078 "Status code not cacheable"
3079 );
3080 return Ok(RespCacheable::Uncacheable(NoCacheReason::OriginNotCache));
3081 }
3082
3083 if let Some(cache_control) = resp.headers.get("cache-control") {
3085 if let Ok(cc_str) = cache_control.to_str() {
3086 if crate::cache::CacheManager::is_no_cache(cc_str) {
3087 trace!(
3088 correlation_id = %ctx.trace_id,
3089 route_id = %route_id,
3090 cache_control = %cc_str,
3091 "Response has no-cache directive"
3092 );
3093 return Ok(RespCacheable::Uncacheable(NoCacheReason::OriginNotCache));
3094 }
3095 }
3096 }
3097
3098 let cache_control = resp
3100 .headers
3101 .get("cache-control")
3102 .and_then(|v| v.to_str().ok());
3103 let ttl = self.cache_manager.calculate_ttl(route_id, cache_control);
3104
3105 if ttl.is_zero() {
3106 trace!(
3107 correlation_id = %ctx.trace_id,
3108 route_id = %route_id,
3109 "TTL is zero, not caching"
3110 );
3111 return Ok(RespCacheable::Uncacheable(NoCacheReason::OriginNotCache));
3112 }
3113
3114 let config = self
3116 .cache_manager
3117 .get_route_config(route_id)
3118 .unwrap_or_default();
3119
3120 let now = std::time::SystemTime::now();
3122 let fresh_until = now + ttl;
3123
3124 let header = resp.clone();
3126
3127 let cache_meta = CacheMeta::new(
3129 fresh_until,
3130 now,
3131 config.stale_while_revalidate_secs as u32,
3132 config.stale_if_error_secs as u32,
3133 header,
3134 );
3135
3136 self.cache_manager.stats().record_store();
3138
3139 debug!(
3140 correlation_id = %ctx.trace_id,
3141 route_id = %route_id,
3142 status = status,
3143 ttl_secs = ttl.as_secs(),
3144 stale_while_revalidate_secs = config.stale_while_revalidate_secs,
3145 stale_if_error_secs = config.stale_if_error_secs,
3146 "Caching response"
3147 );
3148
3149 Ok(RespCacheable::Cacheable(cache_meta))
3150 }
3151
3152 fn should_serve_stale(
3156 &self,
3157 _session: &mut Session,
3158 ctx: &mut Self::CTX,
3159 error: Option<&Error>,
3160 ) -> bool {
3161 let route_id = match ctx.route_id.as_deref() {
3162 Some(id) => id,
3163 None => return false,
3164 };
3165
3166 let config = match self.cache_manager.get_route_config(route_id) {
3168 Some(c) => c,
3169 None => return false,
3170 };
3171
3172 if let Some(e) = error {
3174 if e.esource() == &pingora::ErrorSource::Upstream {
3176 debug!(
3177 correlation_id = %ctx.trace_id,
3178 route_id = %route_id,
3179 error = %e,
3180 stale_if_error_secs = config.stale_if_error_secs,
3181 "Considering stale-if-error"
3182 );
3183 return config.stale_if_error_secs > 0;
3184 }
3185 }
3186
3187 if error.is_none() && config.stale_while_revalidate_secs > 0 {
3189 trace!(
3190 correlation_id = %ctx.trace_id,
3191 route_id = %route_id,
3192 stale_while_revalidate_secs = config.stale_while_revalidate_secs,
3193 "Allowing stale-while-revalidate"
3194 );
3195 return true;
3196 }
3197
3198 false
3199 }
3200
3201 fn range_header_filter(
3211 &self,
3212 session: &mut Session,
3213 response: &mut ResponseHeader,
3214 ctx: &mut Self::CTX,
3215 ) -> pingora_proxy::RangeType
3216 where
3217 Self::CTX: Send + Sync,
3218 {
3219 let supports_range = ctx.route_config.as_ref().is_none_or(|config| {
3221 matches!(
3223 config.service_type,
3224 zentinel_config::ServiceType::Static | zentinel_config::ServiceType::Web
3225 )
3226 });
3227
3228 if !supports_range {
3229 trace!(
3230 correlation_id = %ctx.trace_id,
3231 route_id = ctx.route_id.as_deref().unwrap_or("unknown"),
3232 "Range request not supported for this route type"
3233 );
3234 return pingora_proxy::RangeType::None;
3235 }
3236
3237 let range_type = pingora_proxy::range_header_filter(session.req_header(), response, None);
3239
3240 match &range_type {
3241 pingora_proxy::RangeType::None => {
3242 trace!(
3243 correlation_id = %ctx.trace_id,
3244 route_id = ctx.route_id.as_deref().unwrap_or("unknown"),
3245 "No range request or not applicable"
3246 );
3247 }
3248 pingora_proxy::RangeType::Single(range) => {
3249 trace!(
3250 correlation_id = %ctx.trace_id,
3251 route_id = ctx.route_id.as_deref().unwrap_or("unknown"),
3252 range_start = range.start,
3253 range_end = range.end,
3254 "Processing single-range request"
3255 );
3256 }
3257 pingora_proxy::RangeType::Multi(multi) => {
3258 trace!(
3259 correlation_id = %ctx.trace_id,
3260 route_id = ctx.route_id.as_deref().unwrap_or("unknown"),
3261 range_count = multi.ranges.len(),
3262 "Processing multi-range request"
3263 );
3264 }
3265 pingora_proxy::RangeType::Invalid => {
3266 debug!(
3267 correlation_id = %ctx.trace_id,
3268 route_id = ctx.route_id.as_deref().unwrap_or("unknown"),
3269 "Invalid range header"
3270 );
3271 }
3272 }
3273
3274 range_type
3275 }
3276
3277 async fn fail_to_proxy(
3280 &self,
3281 session: &mut Session,
3282 e: &Error,
3283 ctx: &mut Self::CTX,
3284 ) -> pingora_proxy::FailToProxy
3285 where
3286 Self::CTX: Send + Sync,
3287 {
3288 let error_code = match e.etype() {
3289 ErrorType::ConnectRefused => 503,
3291 ErrorType::ConnectTimedout => 504,
3292 ErrorType::ConnectNoRoute => 502,
3293
3294 ErrorType::ReadTimedout => 504,
3296 ErrorType::WriteTimedout => 504,
3297
3298 ErrorType::TLSHandshakeFailure => 502,
3300 ErrorType::InvalidCert => 502,
3301
3302 ErrorType::InvalidHTTPHeader => 400,
3304 ErrorType::H2Error => 502,
3305
3306 ErrorType::ConnectProxyFailure => 502,
3308 ErrorType::ConnectionClosed => 502,
3309
3310 ErrorType::HTTPStatus(status) => *status,
3312
3313 ErrorType::InternalError => 500,
3316
3317 _ => 502,
3319 };
3320
3321 error!(
3322 correlation_id = %ctx.trace_id,
3323 route_id = ctx.route_id.as_deref().unwrap_or("unknown"),
3324 upstream = ctx.upstream.as_deref().unwrap_or("unknown"),
3325 error_type = ?e.etype(),
3326 error = %e,
3327 error_code = error_code,
3328 "Proxy error occurred"
3329 );
3330
3331 self.metrics
3333 .record_blocked_request(&format!("proxy_error_{}", error_code));
3334
3335 let error_message = match error_code {
3339 400 => "Bad Request",
3340 502 => "Bad Gateway",
3341 503 => "Service Unavailable",
3342 504 => "Gateway Timeout",
3343 _ => "Internal Server Error",
3344 };
3345
3346 let body = format!(
3348 r#"{{"error":"{} {}","trace_id":"{}"}}"#,
3349 error_code, error_message, ctx.trace_id
3350 );
3351
3352 let mut header = pingora::http::ResponseHeader::build(error_code, None).unwrap();
3354 header
3355 .insert_header("Content-Type", "application/json")
3356 .ok();
3357 header
3358 .insert_header("Content-Length", body.len().to_string())
3359 .ok();
3360 header
3361 .insert_header("X-Correlation-Id", ctx.trace_id.as_str())
3362 .ok();
3363 header.insert_header("Connection", "close").ok();
3364
3365 if let Err(write_err) = session.write_response_header(Box::new(header), false).await {
3367 warn!(
3368 correlation_id = %ctx.trace_id,
3369 error = %write_err,
3370 "Failed to write error response header"
3371 );
3372 } else {
3373 if let Err(write_err) = session
3375 .write_response_body(Some(bytes::Bytes::from(body)), true)
3376 .await
3377 {
3378 warn!(
3379 correlation_id = %ctx.trace_id,
3380 error = %write_err,
3381 "Failed to write error response body"
3382 );
3383 }
3384 }
3385
3386 pingora_proxy::FailToProxy {
3389 error_code,
3390 can_reuse_downstream: false,
3391 }
3392 }
3393
3394 fn error_while_proxy(
3400 &self,
3401 peer: &HttpPeer,
3402 session: &mut Session,
3403 e: Box<Error>,
3404 ctx: &mut Self::CTX,
3405 client_reused: bool,
3406 ) -> Box<Error> {
3407 let error_type = e.etype().clone();
3408 let upstream_id = ctx.upstream.as_deref().unwrap_or("unknown");
3409
3410 let is_retryable = matches!(
3412 error_type,
3413 ErrorType::ConnectTimedout
3414 | ErrorType::ReadTimedout
3415 | ErrorType::WriteTimedout
3416 | ErrorType::ConnectionClosed
3417 | ErrorType::ConnectRefused
3418 );
3419
3420 warn!(
3422 correlation_id = %ctx.trace_id,
3423 route_id = ctx.route_id.as_deref().unwrap_or("unknown"),
3424 upstream = %upstream_id,
3425 peer_address = %peer.address(),
3426 error_type = ?error_type,
3427 error = %e,
3428 client_reused = client_reused,
3429 is_retryable = is_retryable,
3430 "Error during proxy operation"
3431 );
3432
3433 let peer_address = peer.address().to_string();
3436 let upstream_pools = self.upstream_pools.clone();
3437 let upstream_id_owned = upstream_id.to_string();
3438 tokio::spawn(async move {
3439 if let Some(pool) = upstream_pools.get(&upstream_id_owned).await {
3440 pool.report_result(&peer_address, false).await;
3441 }
3442 });
3443
3444 self.metrics
3446 .record_blocked_request(&format!("proxy_error_{:?}", error_type));
3447
3448 let mut enhanced_error = e.more_context(format!(
3450 "Upstream: {}, Peer: {}, Attempts: {}",
3451 upstream_id,
3452 peer.address(),
3453 ctx.upstream_attempts
3454 ));
3455
3456 if is_retryable {
3461 let can_retry = if client_reused {
3462 !session.as_ref().retry_buffer_truncated()
3464 } else {
3465 true
3467 };
3468
3469 enhanced_error.retry.decide_reuse(can_retry);
3470
3471 if can_retry {
3472 debug!(
3473 correlation_id = %ctx.trace_id,
3474 upstream = %upstream_id,
3475 error_type = ?error_type,
3476 "Error is retryable, will attempt retry"
3477 );
3478 }
3479 } else {
3480 enhanced_error.retry.decide_reuse(false);
3482 }
3483
3484 enhanced_error
3485 }
3486
3487 async fn logging(&self, session: &mut Session, _error: Option<&Error>, ctx: &mut Self::CTX) {
3488 self.reload_coordinator.dec_requests();
3490
3491 if !ctx.route_agent_ids.is_empty()
3494 || !ctx.body_inspection_agents.is_empty()
3495 || !ctx.websocket_inspection_agents.is_empty()
3496 {
3497 self.agent_manager.end_request(&ctx.trace_id).await;
3498 }
3499
3500 if !ctx.shadow_sent {
3502 if let Some(shadow_pending) = ctx.shadow_pending.take() {
3503 let body = if shadow_pending.include_body && !ctx.body_buffer.is_empty() {
3504 Some(ctx.body_buffer.clone())
3506 } else {
3507 None
3508 };
3509
3510 trace!(
3511 correlation_id = %ctx.trace_id,
3512 body_size = body.as_ref().map(|b| b.len()).unwrap_or(0),
3513 "Firing deferred shadow request with buffered body"
3514 );
3515
3516 shadow_pending.manager.shadow_request(
3517 shadow_pending.headers,
3518 body,
3519 shadow_pending.request_ctx,
3520 );
3521 ctx.shadow_sent = true;
3522 }
3523 }
3524
3525 let duration = ctx.elapsed();
3526
3527 let status = session
3529 .response_written()
3530 .map(|r| r.status.as_u16())
3531 .unwrap_or(0);
3532
3533 if let (Some(ref peer_addr), Some(ref upstream_id)) =
3536 (&ctx.selected_upstream_address, &ctx.upstream)
3537 {
3538 let success = status > 0 && status < 500;
3540
3541 if let Some(pool) = self.upstream_pools.get(upstream_id).await {
3542 pool.report_result_with_latency(peer_addr, success, Some(duration))
3543 .await;
3544 pool.decrement_active();
3545 trace!(
3546 correlation_id = %ctx.trace_id,
3547 upstream = %upstream_id,
3548 peer_address = %peer_addr,
3549 success = success,
3550 duration_ms = duration.as_millis(),
3551 status = status,
3552 "Reported result to adaptive load balancer"
3553 );
3554 }
3555
3556 if ctx.inference_rate_limit_enabled && success {
3558 let cold_detected = self.warmth_tracker.record_request(peer_addr, duration);
3559 if cold_detected {
3560 debug!(
3561 correlation_id = %ctx.trace_id,
3562 upstream = %upstream_id,
3563 peer_address = %peer_addr,
3564 duration_ms = duration.as_millis(),
3565 "Cold model detected on inference upstream"
3566 );
3567 }
3568 }
3569 }
3570
3571 if ctx.inference_rate_limit_enabled {
3574 if let (Some(route_id), Some(ref rate_limit_key)) =
3575 (ctx.route_id.as_deref(), &ctx.inference_rate_limit_key)
3576 {
3577 let response_headers = session
3579 .response_written()
3580 .map(|r| &r.headers)
3581 .cloned()
3582 .unwrap_or_default();
3583
3584 let streaming_result = if ctx.inference_streaming_response {
3586 ctx.inference_streaming_counter
3587 .as_ref()
3588 .map(|counter| counter.finalize())
3589 } else {
3590 None
3591 };
3592
3593 if let Some(ref result) = streaming_result {
3595 debug!(
3596 correlation_id = %ctx.trace_id,
3597 output_tokens = result.output_tokens,
3598 input_tokens = ?result.input_tokens,
3599 source = ?result.source,
3600 content_length = result.content_length,
3601 "Finalized streaming token count"
3602 );
3603 }
3604
3605 if ctx.inference_streaming_response {
3607 if let Some(ref route_config) = ctx.route_config {
3608 if let Some(ref inference) = route_config.inference {
3609 if let Some(ref guardrails) = inference.guardrails {
3610 if let Some(ref pii_config) = guardrails.pii_detection {
3611 if pii_config.enabled {
3612 if let Some(ref counter) = ctx.inference_streaming_counter {
3614 let response_content = counter.content();
3615 if !response_content.is_empty() {
3616 let pii_result = self
3617 .guardrail_processor
3618 .check_pii(
3619 pii_config,
3620 response_content,
3621 ctx.route_id.as_deref(),
3622 &ctx.trace_id,
3623 )
3624 .await;
3625
3626 match pii_result {
3627 crate::inference::PiiCheckResult::Detected {
3628 detections,
3629 redacted_content: _,
3630 } => {
3631 warn!(
3632 correlation_id = %ctx.trace_id,
3633 route_id = ctx.route_id.as_deref().unwrap_or("unknown"),
3634 detection_count = detections.len(),
3635 "PII detected in inference response"
3636 );
3637
3638 ctx.pii_detection_categories = detections
3640 .iter()
3641 .map(|d| d.category.clone())
3642 .collect();
3643
3644 for detection in &detections {
3646 self.metrics.record_pii_detected(
3647 ctx.route_id.as_deref().unwrap_or("unknown"),
3648 &detection.category,
3649 );
3650 }
3651 }
3652 crate::inference::PiiCheckResult::Clean => {
3653 trace!(
3654 correlation_id = %ctx.trace_id,
3655 "No PII detected in response"
3656 );
3657 }
3658 crate::inference::PiiCheckResult::Error { message } => {
3659 debug!(
3660 correlation_id = %ctx.trace_id,
3661 error = %message,
3662 "PII detection check failed"
3663 );
3664 }
3665 }
3666 }
3667 }
3668 }
3669 }
3670 }
3671 }
3672 }
3673 }
3674
3675 let empty_body: &[u8] = &[];
3679
3680 if let Some(actual_estimate) = self.inference_rate_limit_manager.record_actual(
3681 route_id,
3682 rate_limit_key,
3683 &response_headers,
3684 empty_body,
3685 ctx.inference_estimated_tokens,
3686 ) {
3687 let (actual_tokens, source_info) = if let Some(ref streaming) = streaming_result
3689 {
3690 if let Some(total_tokens) = streaming.total_tokens {
3692 (total_tokens, "streaming_api")
3693 } else if actual_estimate.source == crate::inference::TokenSource::Estimated
3694 {
3695 let total = ctx.inference_input_tokens + streaming.output_tokens;
3698 (total, "streaming_tiktoken")
3699 } else {
3700 (actual_estimate.tokens, "headers")
3701 }
3702 } else {
3703 (actual_estimate.tokens, "headers")
3704 };
3705
3706 ctx.inference_actual_tokens = Some(actual_tokens);
3707
3708 debug!(
3709 correlation_id = %ctx.trace_id,
3710 route_id = route_id,
3711 estimated_tokens = ctx.inference_estimated_tokens,
3712 actual_tokens = actual_tokens,
3713 source = source_info,
3714 streaming_response = ctx.inference_streaming_response,
3715 model = ?ctx.inference_model,
3716 "Recorded actual inference tokens"
3717 );
3718
3719 if ctx.inference_budget_enabled {
3721 let alerts = self.inference_rate_limit_manager.record_budget(
3722 route_id,
3723 rate_limit_key,
3724 actual_tokens,
3725 );
3726
3727 for alert in alerts.iter() {
3729 warn!(
3730 correlation_id = %ctx.trace_id,
3731 route_id = route_id,
3732 tenant = %alert.tenant,
3733 threshold_pct = alert.threshold * 100.0,
3734 tokens_used = alert.tokens_used,
3735 tokens_limit = alert.tokens_limit,
3736 "Token budget alert threshold crossed"
3737 );
3738 }
3739
3740 if let Some(status) = self
3742 .inference_rate_limit_manager
3743 .budget_status(route_id, rate_limit_key)
3744 {
3745 ctx.inference_budget_remaining = Some(status.tokens_remaining as i64);
3746 }
3747 }
3748
3749 if ctx.inference_cost_enabled {
3751 if let Some(model) = ctx.inference_model.as_deref() {
3752 let (input_tokens, output_tokens) = if let Some(ref streaming) =
3754 streaming_result
3755 {
3756 let input =
3758 streaming.input_tokens.unwrap_or(ctx.inference_input_tokens);
3759 let output = streaming.output_tokens;
3760 (input, output)
3761 } else {
3762 let input = ctx.inference_input_tokens;
3764 let output = actual_tokens.saturating_sub(input);
3765 (input, output)
3766 };
3767 ctx.inference_output_tokens = output_tokens;
3768
3769 if let Some(cost_result) = self
3770 .inference_rate_limit_manager
3771 .calculate_cost(route_id, model, input_tokens, output_tokens)
3772 {
3773 ctx.inference_request_cost = Some(cost_result.total_cost);
3774
3775 trace!(
3776 correlation_id = %ctx.trace_id,
3777 route_id = route_id,
3778 model = model,
3779 input_tokens = input_tokens,
3780 output_tokens = output_tokens,
3781 total_cost = cost_result.total_cost,
3782 currency = %cost_result.currency,
3783 "Calculated inference request cost"
3784 );
3785 }
3786 }
3787 }
3788 }
3789 }
3790 }
3791
3792 if self.log_manager.should_log_access(status) {
3794 let access_entry = AccessLogEntry {
3795 timestamp: chrono::Utc::now().to_rfc3339(),
3796 trace_id: ctx.trace_id.clone(),
3797 method: ctx.method.clone(),
3798 path: ctx.path.clone(),
3799 query: ctx.query.clone(),
3800 protocol: "HTTP/1.1".to_string(),
3801 status,
3802 body_bytes: ctx.response_bytes,
3803 duration_ms: duration.as_millis() as u64,
3804 client_ip: ctx.client_ip.clone(),
3805 user_agent: ctx.user_agent.clone(),
3806 referer: ctx.referer.clone(),
3807 host: ctx.host.clone(),
3808 route_id: ctx.route_id.clone(),
3809 upstream: ctx.upstream.clone(),
3810 upstream_attempts: ctx.upstream_attempts,
3811 instance_id: self.app_state.instance_id.clone(),
3812 namespace: ctx.namespace.clone(),
3813 service: ctx.service.clone(),
3814 body_bytes_sent: ctx.response_bytes,
3816 upstream_addr: ctx.selected_upstream_address.clone(),
3817 connection_reused: ctx.connection_reused,
3818 rate_limit_hit: status == 429,
3819 geo_country: ctx.geo_country_code.clone(),
3820 };
3821 self.log_manager.log_access(&access_entry);
3822 }
3823
3824 if tracing::enabled!(tracing::Level::DEBUG) {
3826 let write_pending_ms = session.upstream_write_pending_time().as_millis() as u64;
3828 debug!(
3829 trace_id = %ctx.trace_id,
3830 method = %ctx.method,
3831 path = %ctx.path,
3832 route_id = ?ctx.route_id,
3833 upstream = ?ctx.upstream,
3834 status = status,
3835 duration_ms = duration.as_millis() as u64,
3836 upstream_write_pending_ms = write_pending_ms,
3837 upstream_attempts = ctx.upstream_attempts,
3838 error = ?_error.map(|e| e.to_string()),
3839 "Request completed"
3840 );
3841 }
3842
3843 if ctx.is_websocket_upgrade && status == 101 {
3845 info!(
3846 trace_id = %ctx.trace_id,
3847 route_id = ?ctx.route_id,
3848 upstream = ?ctx.upstream,
3849 client_ip = %ctx.client_ip,
3850 "WebSocket connection established"
3851 );
3852 }
3853
3854 if let Some(span) = ctx.otel_span.take() {
3856 span.end();
3857 }
3858 }
3859}
3860
3861impl ZentinelProxy {
3866 fn evaluate_agentic_policy(
3877 &self,
3878 session: &Session,
3879 chunk: Option<&Bytes>,
3880 end_of_stream: bool,
3881 ctx: &mut RequestContext,
3882 ) -> Result<(), Box<Error>> {
3883 use crate::agentic::{self, jsonrpc};
3884
3885 let Some(route) = ctx.route_config.clone() else {
3886 return Ok(());
3887 };
3888 if route.mcp.is_none() && route.a2a.is_none() {
3889 return Ok(());
3890 }
3891
3892 if let Some(chunk) = chunk {
3896 let remaining = jsonrpc::MAX_ENVELOPE_BYTES.saturating_sub(ctx.agentic_body.len());
3897 if chunk.len() > remaining {
3898 ctx.agentic_body_oversize = true;
3899 ctx.agentic_body.extend_from_slice(&chunk[..remaining]);
3900 } else {
3901 ctx.agentic_body.extend_from_slice(chunk);
3902 }
3903 }
3904
3905 if !end_of_stream {
3906 return Ok(());
3907 }
3908
3909 let headers: Vec<(String, String)> = session
3913 .req_header()
3914 .headers
3915 .iter()
3916 .filter_map(|(k, v)| {
3917 v.to_str()
3918 .ok()
3919 .map(|v| (k.as_str().to_ascii_lowercase(), v.to_string()))
3920 })
3921 .collect();
3922
3923 match agentic::decide(&route, &headers, &ctx.agentic_body) {
3924 None => {}
3925 Some(agentic::Outcome::Allow {
3926 mcp_method,
3927 mcp_target,
3928 a2a_method,
3929 }) => {
3930 trace!(
3931 correlation_id = %ctx.trace_id,
3932 route_id = ?ctx.route_id,
3933 mcp_method = ?mcp_method,
3934 mcp_target = ?mcp_target,
3935 a2a_method = ?a2a_method,
3936 "Agentic request permitted"
3937 );
3938 ctx.mcp_method = mcp_method;
3939 ctx.mcp_target = mcp_target;
3940 ctx.a2a_method = a2a_method;
3941 }
3942 Some(agentic::Outcome::Deny { reason, kind }) => {
3943 warn!(
3944 correlation_id = %ctx.trace_id,
3945 route_id = ?ctx.route_id,
3946 policy = kind,
3947 reason = %reason,
3948 "Agentic request denied"
3949 );
3950 self.metrics.record_blocked_request(kind);
3951 return Err(Error::explain(ErrorType::HTTPStatus(403), reason));
3952 }
3953 }
3954
3955 Ok(())
3956 }
3957
3958 async fn process_body_chunk_streaming(
3960 &self,
3961 body: &mut Option<Bytes>,
3962 end_of_stream: bool,
3963 ctx: &mut RequestContext,
3964 ) -> Result<(), Box<Error>> {
3965 let chunk_data: Vec<u8> = body.as_ref().map(|b| b.to_vec()).unwrap_or_default();
3967 let chunk_index = ctx.request_body_chunk_index;
3968 ctx.request_body_chunk_index += 1;
3969 ctx.body_bytes_inspected += chunk_data.len() as u64;
3970
3971 debug!(
3972 correlation_id = %ctx.trace_id,
3973 chunk_index = chunk_index,
3974 chunk_size = chunk_data.len(),
3975 end_of_stream = end_of_stream,
3976 "Streaming body chunk to agents"
3977 );
3978
3979 let agent_ctx = crate::agents::AgentCallContext {
3981 correlation_id: zentinel_common::CorrelationId::from_string(&ctx.trace_id),
3982 metadata: zentinel_agent_protocol::RequestMetadata {
3983 correlation_id: ctx.trace_id.clone(),
3984 request_id: ctx.trace_id.clone(),
3985 client_ip: ctx.client_ip.clone(),
3986 client_port: 0,
3987 server_name: ctx.host.clone(),
3988 protocol: "HTTP/1.1".to_string(),
3989 tls_version: None,
3990 tls_cipher: None,
3991 route_id: ctx.route_id.clone(),
3992 upstream_id: ctx.upstream.clone(),
3993 timestamp: chrono::Utc::now().to_rfc3339(),
3994 traceparent: ctx.traceparent(),
3995 },
3996 route_id: ctx.route_id.clone(),
3997 upstream_id: ctx.upstream.clone(),
3998 request_body: None, response_body: None,
4000 };
4001
4002 let agent_ids = ctx.body_inspection_agents.clone();
4003 let total_size = None; match self
4006 .agent_manager
4007 .process_request_body_streaming(
4008 &agent_ctx,
4009 &chunk_data,
4010 end_of_stream,
4011 chunk_index,
4012 ctx.body_bytes_inspected as usize,
4013 total_size,
4014 &agent_ids,
4015 )
4016 .await
4017 {
4018 Ok(decision) => {
4019 ctx.agent_needs_more = decision.needs_more;
4021
4022 if let Some(ref mutation) = decision.request_body_mutation {
4024 if !mutation.is_pass_through() {
4025 if mutation.is_drop() {
4026 *body = None;
4028 trace!(
4029 correlation_id = %ctx.trace_id,
4030 chunk_index = chunk_index,
4031 "Agent dropped body chunk"
4032 );
4033 } else if let Some(ref new_data) = mutation.data {
4034 *body = Some(Bytes::from(new_data.clone()));
4036 trace!(
4037 correlation_id = %ctx.trace_id,
4038 chunk_index = chunk_index,
4039 original_size = chunk_data.len(),
4040 new_size = new_data.len(),
4041 "Agent mutated body chunk"
4042 );
4043 }
4044 }
4045 }
4046
4047 if !decision.needs_more && !decision.is_allow() {
4049 warn!(
4050 correlation_id = %ctx.trace_id,
4051 agent_id = decision.decided_by.as_deref().unwrap_or("unknown"),
4052 action = ?decision.action,
4053 "Agent blocked request body"
4054 );
4055 self.metrics.record_blocked_request("agent_body_inspection");
4056
4057 let (status, message) = match &decision.action {
4058 crate::agents::AgentAction::Block { status, body, .. } => (
4059 *status,
4060 body.clone().unwrap_or_else(|| "Blocked".to_string()),
4061 ),
4062 _ => (403, "Forbidden".to_string()),
4063 };
4064
4065 return Err(Error::explain(ErrorType::HTTPStatus(status), message));
4066 }
4067
4068 trace!(
4069 correlation_id = %ctx.trace_id,
4070 needs_more = decision.needs_more,
4071 "Agent processed body chunk"
4072 );
4073 }
4074 Err(e) => {
4075 let fail_closed = ctx
4076 .route_config
4077 .as_ref()
4078 .map(|r| r.policies.failure_mode == zentinel_config::FailureMode::Closed)
4079 .unwrap_or(false);
4080
4081 if fail_closed {
4082 error!(
4083 correlation_id = %ctx.trace_id,
4084 error = %e,
4085 "Agent streaming body inspection failed, blocking (fail-closed)"
4086 );
4087 self.log_manager.log_request_error(
4088 "error",
4089 "Agent streaming body inspection failed, blocking (fail-closed)",
4090 &ctx.trace_id,
4091 ctx.route_id.as_deref(),
4092 ctx.upstream.as_deref(),
4093 Some(format!("error={}", e)),
4094 );
4095 return Err(Error::explain(
4096 ErrorType::HTTPStatus(503),
4097 "Service unavailable",
4098 ));
4099 } else {
4100 warn!(
4101 correlation_id = %ctx.trace_id,
4102 error = %e,
4103 "Agent streaming body inspection failed, allowing (fail-open)"
4104 );
4105 self.log_manager.log_request_error(
4106 "warn",
4107 "Agent streaming body inspection failed, allowing (fail-open)",
4108 &ctx.trace_id,
4109 ctx.route_id.as_deref(),
4110 ctx.upstream.as_deref(),
4111 Some(format!("error={}", e)),
4112 );
4113 }
4114 }
4115 }
4116
4117 Ok(())
4118 }
4119
4120 async fn send_buffered_body_to_agents(
4122 &self,
4123 end_of_stream: bool,
4124 ctx: &mut RequestContext,
4125 ) -> Result<(), Box<Error>> {
4126 debug!(
4127 correlation_id = %ctx.trace_id,
4128 buffer_size = ctx.body_buffer.len(),
4129 end_of_stream = end_of_stream,
4130 agent_count = ctx.body_inspection_agents.len(),
4131 decompression_enabled = ctx.decompression_enabled,
4132 "Sending buffered body to agents for inspection"
4133 );
4134
4135 let body_for_inspection = if ctx.decompression_enabled {
4137 if let Some(ref encoding) = ctx.body_content_encoding {
4138 let config = crate::decompression::DecompressionConfig {
4139 max_ratio: ctx.max_decompression_ratio,
4140 max_output_bytes: ctx.max_decompression_bytes,
4141 };
4142
4143 match crate::decompression::decompress_body(&ctx.body_buffer, encoding, &config) {
4144 Ok(result) => {
4145 ctx.body_was_decompressed = true;
4146 self.metrics
4147 .record_decompression_success(encoding, result.ratio);
4148 debug!(
4149 correlation_id = %ctx.trace_id,
4150 encoding = %encoding,
4151 compressed_size = result.compressed_size,
4152 decompressed_size = result.decompressed_size,
4153 ratio = result.ratio,
4154 "Body decompressed for agent inspection"
4155 );
4156 result.data
4157 }
4158 Err(e) => {
4159 let failure_reason = match &e {
4161 crate::decompression::DecompressionError::RatioExceeded { .. } => {
4162 "ratio_exceeded"
4163 }
4164 crate::decompression::DecompressionError::SizeExceeded { .. } => {
4165 "size_exceeded"
4166 }
4167 crate::decompression::DecompressionError::InvalidData { .. } => {
4168 "invalid_data"
4169 }
4170 crate::decompression::DecompressionError::UnsupportedEncoding {
4171 ..
4172 } => "unsupported",
4173 crate::decompression::DecompressionError::IoError(_) => "io_error",
4174 };
4175 self.metrics
4176 .record_decompression_failure(encoding, failure_reason);
4177
4178 let fail_closed = ctx
4180 .route_config
4181 .as_ref()
4182 .map(|r| {
4183 r.policies.failure_mode == zentinel_config::FailureMode::Closed
4184 })
4185 .unwrap_or(false);
4186
4187 if fail_closed {
4188 error!(
4189 correlation_id = %ctx.trace_id,
4190 error = %e,
4191 encoding = %encoding,
4192 "Decompression failed, blocking (fail-closed)"
4193 );
4194 return Err(Error::explain(
4195 ErrorType::HTTPStatus(400),
4196 "Invalid compressed body",
4197 ));
4198 } else {
4199 warn!(
4200 correlation_id = %ctx.trace_id,
4201 error = %e,
4202 encoding = %encoding,
4203 "Decompression failed, sending compressed body (fail-open)"
4204 );
4205 ctx.body_buffer.clone()
4206 }
4207 }
4208 }
4209 } else {
4210 ctx.body_buffer.clone()
4211 }
4212 } else {
4213 ctx.body_buffer.clone()
4214 };
4215
4216 let agent_ctx = crate::agents::AgentCallContext {
4217 correlation_id: zentinel_common::CorrelationId::from_string(&ctx.trace_id),
4218 metadata: zentinel_agent_protocol::RequestMetadata {
4219 correlation_id: ctx.trace_id.clone(),
4220 request_id: ctx.trace_id.clone(),
4221 client_ip: ctx.client_ip.clone(),
4222 client_port: 0,
4223 server_name: ctx.host.clone(),
4224 protocol: "HTTP/1.1".to_string(),
4225 tls_version: None,
4226 tls_cipher: None,
4227 route_id: ctx.route_id.clone(),
4228 upstream_id: ctx.upstream.clone(),
4229 timestamp: chrono::Utc::now().to_rfc3339(),
4230 traceparent: ctx.traceparent(),
4231 },
4232 route_id: ctx.route_id.clone(),
4233 upstream_id: ctx.upstream.clone(),
4234 request_body: Some(body_for_inspection.clone()),
4235 response_body: None,
4236 };
4237
4238 let agent_ids = ctx.body_inspection_agents.clone();
4239 match self
4240 .agent_manager
4241 .process_request_body(&agent_ctx, &body_for_inspection, end_of_stream, &agent_ids)
4242 .await
4243 {
4244 Ok(decision) => {
4245 if !decision.is_allow() {
4246 warn!(
4247 correlation_id = %ctx.trace_id,
4248 agent_id = decision.decided_by.as_deref().unwrap_or("unknown"),
4249 action = ?decision.action,
4250 "Agent blocked request body"
4251 );
4252 self.metrics.record_blocked_request("agent_body_inspection");
4253
4254 let (status, message) = match &decision.action {
4255 crate::agents::AgentAction::Block { status, body, .. } => (
4256 *status,
4257 body.clone().unwrap_or_else(|| "Blocked".to_string()),
4258 ),
4259 _ => (403, "Forbidden".to_string()),
4260 };
4261
4262 return Err(Error::explain(ErrorType::HTTPStatus(status), message));
4263 }
4264
4265 trace!(
4266 correlation_id = %ctx.trace_id,
4267 "Agent allowed request body"
4268 );
4269 }
4270 Err(e) => {
4271 let fail_closed = ctx
4272 .route_config
4273 .as_ref()
4274 .map(|r| r.policies.failure_mode == zentinel_config::FailureMode::Closed)
4275 .unwrap_or(false);
4276
4277 if fail_closed {
4278 error!(
4279 correlation_id = %ctx.trace_id,
4280 error = %e,
4281 "Agent body inspection failed, blocking (fail-closed)"
4282 );
4283 self.log_manager.log_request_error(
4284 "error",
4285 "Agent body inspection failed, blocking (fail-closed)",
4286 &ctx.trace_id,
4287 ctx.route_id.as_deref(),
4288 ctx.upstream.as_deref(),
4289 Some(format!("error={}", e)),
4290 );
4291 return Err(Error::explain(
4292 ErrorType::HTTPStatus(503),
4293 "Service unavailable",
4294 ));
4295 } else {
4296 warn!(
4297 correlation_id = %ctx.trace_id,
4298 error = %e,
4299 "Agent body inspection failed, allowing (fail-open)"
4300 );
4301 self.log_manager.log_request_error(
4302 "warn",
4303 "Agent body inspection failed, allowing (fail-open)",
4304 &ctx.trace_id,
4305 ctx.route_id.as_deref(),
4306 ctx.upstream.as_deref(),
4307 Some(format!("error={}", e)),
4308 );
4309 }
4310 }
4311 }
4312
4313 Ok(())
4314 }
4315}
4316
4317#[cfg(test)]
4318mod cache_status_tests {
4319 use super::{apply_cache_status, cache_status_member};
4320 use crate::proxy::context::CacheStatus;
4321 use pingora::http::ResponseHeader;
4322
4323 #[test]
4324 fn members_describe_each_outcome() {
4325 assert_eq!(cache_status_member("edge", &CacheStatus::Hit), "edge; hit");
4326 assert_eq!(
4327 cache_status_member("edge", &CacheStatus::HitMemory),
4328 "edge; hit; detail=memory"
4329 );
4330 assert_eq!(
4331 cache_status_member("edge", &CacheStatus::HitDisk),
4332 "edge; hit; detail=disk"
4333 );
4334 assert_eq!(
4335 cache_status_member("edge", &CacheStatus::HitStale),
4336 "edge; fwd=stale"
4337 );
4338 assert_eq!(
4339 cache_status_member("edge", &CacheStatus::Miss),
4340 "edge; fwd=miss"
4341 );
4342 assert_eq!(
4343 cache_status_member("edge", &CacheStatus::Bypass("method")),
4344 "edge; fwd=bypass; detail=method"
4345 );
4346 }
4347
4348 #[test]
4355 fn appending_preserves_an_upstream_caches_member() {
4356 let mut response = ResponseHeader::build(200, None).expect("response builds");
4357 response
4359 .append_header("Cache-Status", "origin-shield; hit")
4360 .expect("shield member set");
4361
4362 apply_cache_status(&mut response, "edge", &CacheStatus::Miss);
4363
4364 let members: Vec<&str> = response
4365 .headers
4366 .get_all("Cache-Status")
4367 .iter()
4368 .map(|v| v.to_str().expect("ascii"))
4369 .collect();
4370
4371 assert_eq!(
4372 members,
4373 vec!["origin-shield; hit", "edge; fwd=miss"],
4374 "both caches must appear, origin-closest first (RFC 9211 s.2)"
4375 );
4376 }
4377
4378 #[test]
4380 fn inserting_would_destroy_the_chain() {
4381 let mut response = ResponseHeader::build(200, None).expect("response builds");
4382 response
4383 .append_header("Cache-Status", "origin-shield; hit")
4384 .expect("shield member set");
4385 response
4386 .insert_header("Cache-Status", "edge; fwd=miss")
4387 .expect("insert");
4388
4389 assert_eq!(
4390 response.headers.get_all("Cache-Status").iter().count(),
4391 1,
4392 "insert_header replaces; this test documents why append is required"
4393 );
4394 }
4395}