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