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::model_routing;
35use super::model_routing_metrics::get_model_routing_metrics;
36use super::ZentinelProxy;
37
38struct NoHeaderAccessor;
40impl HeaderAccessor for NoHeaderAccessor {
41 fn get_header(&self, _name: &str) -> Option<String> {
42 None
43 }
44}
45
46#[async_trait]
47impl ProxyHttp for ZentinelProxy {
48 type CTX = RequestContext;
49
50 fn new_ctx(&self) -> Self::CTX {
51 RequestContext::new()
52 }
53
54 fn fail_to_connect(
55 &self,
56 _session: &mut Session,
57 peer: &HttpPeer,
58 ctx: &mut Self::CTX,
59 e: Box<Error>,
60 ) -> Box<Error> {
61 error!(
62 correlation_id = %ctx.trace_id,
63 route_id = ctx.route_id.as_deref().unwrap_or("unknown"),
64 upstream = ctx.upstream.as_deref().unwrap_or("unknown"),
65 peer_address = %peer.address(),
66 error = %e,
67 "Failed to connect to upstream peer"
68 );
69 self.log_manager.log_request_error(
70 "error",
71 "Failed to connect to upstream peer",
72 &ctx.trace_id,
73 ctx.route_id.as_deref(),
74 ctx.upstream.as_deref(),
75 Some(format!("peer={} error={}", peer.address(), e)),
76 );
77 e
79 }
80
81 async fn early_request_filter(
84 &self,
85 session: &mut Session,
86 ctx: &mut Self::CTX,
87 ) -> Result<(), Box<Error>> {
88 let req_header = session.req_header();
90 let method = req_header.method.as_str();
91 let path = req_header.uri.path();
92 let host = crate::http_helpers::extract_request_host(req_header);
93
94 if let Some(ref challenge_manager) = self.acme_challenges {
96 if let Some(token) = crate::acme::ChallengeManager::extract_token(path) {
97 if let Some(key_authorization) = challenge_manager.get_response(token) {
98 debug!(
99 token = %token,
100 "Serving ACME HTTP-01 challenge response"
101 );
102
103 let mut resp = ResponseHeader::build(200, None)?;
105 resp.insert_header("Content-Type", "text/plain")?;
106 resp.insert_header("Content-Length", key_authorization.len().to_string())?;
107
108 session.write_response_header(Box::new(resp), false).await?;
110 session
111 .write_response_body(Some(Bytes::from(key_authorization)), true)
112 .await?;
113
114 return Err(Error::explain(
116 ErrorType::InternalError,
117 "ACME challenge served",
118 ));
119 } else {
120 warn!(
122 token = %token,
123 "ACME challenge token not found"
124 );
125 }
126 }
127 }
128
129 ctx.method = method.to_string();
130 ctx.path = path.to_string();
131 ctx.host = Some(host.to_string());
132
133 let route_match = {
135 let route_matcher = self.route_matcher.read();
136 let mut request_info = RequestInfo::new(method, path, host);
137
138 if route_matcher.needs_headers() {
140 request_info = request_info
141 .with_headers(RequestInfo::build_headers(req_header.headers.iter()));
142 }
143
144 match route_matcher.match_request(&request_info) {
145 Some(m) => m,
146 None => return Ok(()), }
148 };
149
150 ctx.trace_id = self.get_trace_id(session);
151 ctx.route_id = Some(route_match.route_id.to_string());
152 ctx.route_config = Some(route_match.config.clone());
153
154 if let Some(traceparent) = req_header.headers.get(crate::otel::TRACEPARENT_HEADER) {
156 if let Ok(s) = traceparent.to_str() {
157 ctx.trace_context = crate::otel::TraceContext::parse_traceparent(s);
158 }
159 }
160
161 if let Some(tracer) = crate::otel::get_tracer() {
163 ctx.otel_span = Some(tracer.start_span(method, path, ctx.trace_context.as_ref()));
164 }
165
166 if route_match.config.service_type == zentinel_config::ServiceType::Builtin {
168 trace!(
169 correlation_id = %ctx.trace_id,
170 route_id = %route_match.route_id,
171 builtin_handler = ?route_match.config.builtin_handler,
172 "Handling builtin route in early_request_filter"
173 );
174
175 let handled = self
177 .handle_builtin_route(session, ctx, &route_match)
178 .await?;
179
180 if handled {
181 return Err(Error::explain(
183 ErrorType::InternalError,
184 "Builtin handler complete",
185 ));
186 }
187 }
188
189 Ok(())
190 }
191
192 async fn upstream_peer(
193 &self,
194 session: &mut Session,
195 ctx: &mut Self::CTX,
196 ) -> Result<Box<HttpPeer>, Box<Error>> {
197 self.reload_coordinator.inc_requests();
199
200 if ctx.config.is_none() {
202 ctx.config = Some(self.config_manager.current());
203 }
204
205 if ctx.client_ip.is_empty() {
207 ctx.client_ip = session
208 .client_addr()
209 .map(|a| a.to_string())
210 .unwrap_or_else(|| "unknown".to_string());
211 }
212
213 let req_header = session.req_header();
214
215 if ctx.method.is_empty() {
217 ctx.method = req_header.method.to_string();
218 ctx.path = req_header.uri.path().to_string();
219 ctx.query = req_header.uri.query().map(|q| q.to_string());
220 ctx.host = Some(crate::http_helpers::extract_request_host(req_header).to_string());
221 }
222 ctx.user_agent = req_header
223 .headers
224 .get("user-agent")
225 .and_then(|v| v.to_str().ok())
226 .map(|s| s.to_string());
227 ctx.referer = req_header
228 .headers
229 .get("referer")
230 .and_then(|v| v.to_str().ok())
231 .map(|s| s.to_string());
232
233 trace!(
234 correlation_id = %ctx.trace_id,
235 client_ip = %ctx.client_ip,
236 "Request received, initializing context"
237 );
238
239 let route_match = if let Some(ref route_config) = ctx.route_config {
241 let route_id = ctx.route_id.as_deref().unwrap_or("");
242 crate::routing::RouteMatch {
243 route_id: zentinel_common::RouteId::new(route_id),
244 config: route_config.clone(),
245 }
246 } else {
247 let (match_result, route_duration) = {
249 let route_matcher = self.route_matcher.read();
250 let host = ctx.host.as_deref().unwrap_or("");
251
252 let mut request_info = RequestInfo::new(&ctx.method, &ctx.path, host);
254
255 if route_matcher.needs_headers() {
257 request_info = request_info
258 .with_headers(RequestInfo::build_headers(req_header.headers.iter()));
259 }
260
261 if route_matcher.needs_query_params() {
263 request_info =
264 request_info.with_query_params(RequestInfo::parse_query_params(&ctx.path));
265 }
266
267 trace!(
268 correlation_id = %ctx.trace_id,
269 method = %request_info.method,
270 path = %request_info.path,
271 host = %request_info.host,
272 "Built request info for route matching"
273 );
274
275 let route_start = std::time::Instant::now();
276 let route_match = route_matcher.match_request(&request_info).ok_or_else(|| {
277 warn!(
278 correlation_id = %ctx.trace_id,
279 method = %request_info.method,
280 path = %request_info.path,
281 host = %request_info.host,
282 "No matching route found for request"
283 );
284 self.log_manager.log_request_error(
285 "warn",
286 "No matching route found for request",
287 &ctx.trace_id,
288 None,
289 None,
290 Some(format!(
291 "method={} path={} host={}",
292 request_info.method, request_info.path, request_info.host
293 )),
294 );
295 Error::explain(ErrorType::HTTPStatus(404), "No matching route found")
296 })?;
297 let route_duration = route_start.elapsed();
298 (route_match, route_duration)
300 };
301
302 ctx.route_id = Some(match_result.route_id.to_string());
303 ctx.route_config = Some(match_result.config.clone());
304
305 if ctx.trace_id.is_empty() {
307 ctx.trace_id = self.get_trace_id(session);
308
309 if let Some(traceparent) = req_header.headers.get(crate::otel::TRACEPARENT_HEADER) {
311 if let Ok(s) = traceparent.to_str() {
312 ctx.trace_context = crate::otel::TraceContext::parse_traceparent(s);
313 }
314 }
315
316 if let Some(tracer) = crate::otel::get_tracer() {
318 ctx.otel_span =
319 Some(tracer.start_span(&ctx.method, &ctx.path, ctx.trace_context.as_ref()));
320 }
321 }
322
323 trace!(
324 correlation_id = %ctx.trace_id,
325 route_id = %match_result.route_id,
326 route_duration_us = route_duration.as_micros(),
327 service_type = ?match_result.config.service_type,
328 "Route matched"
329 );
330 match_result
331 };
332
333 if route_match.config.service_type == zentinel_config::ServiceType::Builtin {
335 trace!(
336 correlation_id = %ctx.trace_id,
337 route_id = %route_match.route_id,
338 builtin_handler = ?route_match.config.builtin_handler,
339 "Route type is builtin, skipping upstream"
340 );
341 ctx.upstream = Some(format!("_builtin_{}", route_match.route_id));
343 return Err(Error::explain(
345 ErrorType::InternalError,
346 "Builtin handler handled in request_filter",
347 ));
348 }
349
350 if route_match.config.service_type == zentinel_config::ServiceType::Static {
352 trace!(
353 correlation_id = %ctx.trace_id,
354 route_id = %route_match.route_id,
355 "Route type is static, checking for static server"
356 );
357 if self
359 .static_servers
360 .get(route_match.route_id.as_str())
361 .await
362 .is_some()
363 {
364 ctx.upstream = Some(format!("_static_{}", route_match.route_id));
366 info!(
367 correlation_id = %ctx.trace_id,
368 route_id = %route_match.route_id,
369 path = %ctx.path,
370 "Serving static file"
371 );
372 return Err(Error::explain(
374 ErrorType::InternalError,
375 "Static file serving handled in request_filter",
376 ));
377 }
378 }
379
380 let mut model_routing_applied = false;
383 if let Some(ref inference) = route_match.config.inference {
384 if let Some(ref model_routing) = inference.model_routing {
385 let model = model_routing::extract_model_from_headers(&req_header.headers);
387
388 if let Some(ref model_name) = model {
389 if let Some(routing_result) =
391 model_routing::find_upstream_for_model(model_routing, model_name)
392 {
393 debug!(
394 correlation_id = %ctx.trace_id,
395 route_id = %route_match.route_id,
396 model = %model_name,
397 upstream = %routing_result.upstream,
398 is_default = routing_result.is_default,
399 provider_override = ?routing_result.provider,
400 "Model-based routing selected upstream"
401 );
402
403 ctx.record_model_routing(
404 &routing_result.upstream,
405 Some(model_name.clone()),
406 routing_result.provider,
407 );
408 model_routing_applied = true;
409
410 if let Some(metrics) = get_model_routing_metrics() {
412 metrics.record_model_routed(
413 route_match.route_id.as_str(),
414 model_name,
415 &routing_result.upstream,
416 );
417 if routing_result.is_default {
418 metrics.record_default_upstream(route_match.route_id.as_str());
419 }
420 if let Some(provider) = routing_result.provider {
421 metrics.record_provider_override(
422 route_match.route_id.as_str(),
423 &routing_result.upstream,
424 provider.as_str(),
425 );
426 }
427 }
428 }
429 } else if let Some(ref default_upstream) = model_routing.default_upstream {
430 debug!(
432 correlation_id = %ctx.trace_id,
433 route_id = %route_match.route_id,
434 upstream = %default_upstream,
435 "Model-based routing using default upstream (no model header)"
436 );
437 ctx.record_model_routing(default_upstream, None, None);
438 model_routing_applied = true;
439
440 if let Some(metrics) = get_model_routing_metrics() {
442 metrics.record_no_model_header(route_match.route_id.as_str());
443 }
444 }
445 }
446 }
447
448 if !model_routing_applied {
450 if let Some(ref upstream) = route_match.config.upstream {
451 ctx.upstream = Some(upstream.clone());
452 trace!(
453 correlation_id = %ctx.trace_id,
454 route_id = %route_match.route_id,
455 upstream = %upstream,
456 "Upstream configured for route"
457 );
458 } else {
459 warn!(
463 correlation_id = %ctx.trace_id,
464 route_id = %route_match.route_id,
465 "Route has no upstream configured, returning 500"
466 );
467 crate::http_helpers::write_error(
468 session,
469 500,
470 "Internal Server Error",
471 "text/plain",
472 )
473 .await?;
474 return Err(Error::explain(
475 ErrorType::HTTPStatus(500),
476 "Route has no valid upstream",
477 ));
478 }
479 }
480
481 if let Some(ref fallback_config) = route_match.config.fallback {
484 let upstream_name = ctx.upstream.as_ref().unwrap();
485
486 let is_healthy = if let Some(pool) = self.upstream_pools.get(upstream_name).await {
488 pool.has_healthy_targets().await
489 } else {
490 false };
492
493 let is_budget_exhausted = ctx.inference_budget_exhausted;
495
496 let current_model = ctx.inference_model.as_deref();
498
499 let evaluator = FallbackEvaluator::new(
501 fallback_config,
502 ctx.tried_upstreams(),
503 ctx.fallback_attempt,
504 );
505
506 if let Some(decision) = evaluator.should_fallback_before_request(
508 upstream_name,
509 is_healthy,
510 is_budget_exhausted,
511 current_model,
512 ) {
513 info!(
514 correlation_id = %ctx.trace_id,
515 route_id = %route_match.route_id,
516 from_upstream = %upstream_name,
517 to_upstream = %decision.next_upstream,
518 reason = %decision.reason,
519 fallback_attempt = ctx.fallback_attempt + 1,
520 "Triggering fallback routing"
521 );
522
523 if let Some(metrics) = get_fallback_metrics() {
525 metrics.record_fallback_attempt(
526 route_match.route_id.as_str(),
527 upstream_name,
528 &decision.next_upstream,
529 &decision.reason,
530 );
531 }
532
533 ctx.record_fallback(decision.reason, &decision.next_upstream);
535
536 if let Some((original, mapped)) = decision.model_mapping {
538 if let Some(metrics) = get_fallback_metrics() {
540 metrics.record_model_mapping(
541 route_match.route_id.as_str(),
542 &original,
543 &mapped,
544 );
545 }
546
547 ctx.record_model_mapping(original, mapped);
548 trace!(
549 correlation_id = %ctx.trace_id,
550 original_model = ?ctx.model_mapping_applied().map(|m| &m.0),
551 mapped_model = ?ctx.model_mapping_applied().map(|m| &m.1),
552 "Applied model mapping for fallback"
553 );
554 }
555 }
556 }
557
558 debug!(
559 correlation_id = %ctx.trace_id,
560 route_id = %route_match.route_id,
561 upstream = ?ctx.upstream,
562 method = %req_header.method,
563 path = %req_header.uri.path(),
564 host = ctx.host.as_deref().unwrap_or("-"),
565 client_ip = %ctx.client_ip,
566 "Processing request"
567 );
568
569 if ctx
571 .upstream
572 .as_ref()
573 .is_some_and(|u| u.starts_with("_static_"))
574 {
575 return Err(Error::explain(
577 ErrorType::InternalError,
578 "Static route should be handled in request_filter",
579 ));
580 }
581
582 let upstream_name = ctx
583 .upstream
584 .as_ref()
585 .ok_or_else(|| Error::explain(ErrorType::InternalError, "No upstream configured"))?;
586
587 trace!(
588 correlation_id = %ctx.trace_id,
589 upstream = %upstream_name,
590 "Looking up upstream pool"
591 );
592
593 let pool = self
594 .upstream_pools
595 .get(upstream_name)
596 .await
597 .ok_or_else(|| {
598 error!(
599 correlation_id = %ctx.trace_id,
600 upstream = %upstream_name,
601 "Upstream pool not found"
602 );
603 self.log_manager.log_request_error(
604 "error",
605 "Upstream pool not found",
606 &ctx.trace_id,
607 ctx.route_id.as_deref(),
608 Some(upstream_name),
609 None,
610 );
611 Error::explain(
612 ErrorType::InternalError,
613 format!("Upstream pool '{}' not found", upstream_name),
614 )
615 })?;
616
617 let max_retries = route_match
619 .config
620 .retry_policy
621 .as_ref()
622 .map(|r| r.max_attempts)
623 .unwrap_or(1);
624
625 trace!(
626 correlation_id = %ctx.trace_id,
627 upstream = %upstream_name,
628 max_retries = max_retries,
629 "Starting upstream peer selection"
630 );
631
632 let mut last_error = None;
633 let selection_start = std::time::Instant::now();
634
635 for attempt in 1..=max_retries {
636 ctx.upstream_attempts = attempt;
637
638 trace!(
639 correlation_id = %ctx.trace_id,
640 upstream = %upstream_name,
641 attempt = attempt,
642 max_retries = max_retries,
643 "Attempting to select upstream peer"
644 );
645
646 match pool.select_peer_with_metadata(None).await {
647 Ok((mut peer, metadata)) => {
648 let selection_duration = selection_start.elapsed();
649 pool.increment_active();
651 let peer_addr = peer.address().to_string();
653 ctx.selected_upstream_address = Some(peer_addr.clone());
654
655 if metadata.contains_key("sticky_session_new") {
657 ctx.sticky_session_new_assignment = true;
658 ctx.sticky_session_set_cookie =
659 metadata.get("sticky_set_cookie_header").cloned();
660 ctx.sticky_target_index = metadata
661 .get("sticky_target_index")
662 .and_then(|s| s.parse().ok());
663
664 trace!(
665 correlation_id = %ctx.trace_id,
666 sticky_target_index = ?ctx.sticky_target_index,
667 "New sticky session assignment, will set cookie"
668 );
669 }
670
671 debug!(
672 correlation_id = %ctx.trace_id,
673 upstream = %upstream_name,
674 peer_address = %peer_addr,
675 attempt = attempt,
676 selection_duration_us = selection_duration.as_micros(),
677 sticky_session_hit = metadata.contains_key("sticky_session_hit"),
678 sticky_session_new = ctx.sticky_session_new_assignment,
679 "Selected upstream peer"
680 );
681 if let Some(ref rc) = ctx.route_config {
683 if let Some(timeout_secs) = rc.policies.timeout_secs {
684 peer.options.read_timeout = Some(Duration::from_secs(timeout_secs));
685 }
686 }
687
688 if let Some(connect_secs) = ctx.filter_connect_timeout_secs {
690 peer.options.connection_timeout = Some(Duration::from_secs(connect_secs));
691 }
692 if let Some(upstream_secs) = ctx.filter_upstream_timeout_secs {
693 peer.options.read_timeout = Some(Duration::from_secs(upstream_secs));
694 }
695
696 return Ok(Box::new(peer));
697 }
698 Err(e) => {
699 warn!(
700 correlation_id = %ctx.trace_id,
701 upstream = %upstream_name,
702 attempt = attempt,
703 max_retries = max_retries,
704 error = %e,
705 "Failed to select upstream peer"
706 );
707 last_error = Some(e);
708
709 if attempt < max_retries {
710 let backoff = Duration::from_millis(100 * 2_u64.pow(attempt - 1));
712 trace!(
713 correlation_id = %ctx.trace_id,
714 backoff_ms = backoff.as_millis(),
715 "Backing off before retry"
716 );
717 sleep(backoff).await;
718 }
719 }
720 }
721 }
722
723 let selection_duration = selection_start.elapsed();
724 error!(
725 correlation_id = %ctx.trace_id,
726 upstream = %upstream_name,
727 attempts = max_retries,
728 selection_duration_ms = selection_duration.as_millis(),
729 last_error = ?last_error,
730 "All upstream selection attempts failed"
731 );
732 self.log_manager.log_request_error(
733 "error",
734 "All upstream selection attempts failed",
735 &ctx.trace_id,
736 ctx.route_id.as_deref(),
737 Some(upstream_name),
738 Some(format!("attempts={} error={:?}", max_retries, last_error)),
739 );
740
741 if ctx.used_fallback() {
743 if let Some(metrics) = get_fallback_metrics() {
744 metrics.record_fallback_exhausted(ctx.route_id.as_deref().unwrap_or("unknown"));
745 }
746 }
747
748 Err(Error::explain(
749 ErrorType::InternalError,
750 format!("All upstream attempts failed: {:?}", last_error),
751 ))
752 }
753
754 async fn request_filter(
755 &self,
756 session: &mut Session,
757 ctx: &mut Self::CTX,
758 ) -> Result<bool, Box<Error>> {
759 trace!(
760 correlation_id = %ctx.trace_id,
761 route_id = ctx.route_id.as_deref().unwrap_or("unknown"),
762 "Starting request filter phase"
763 );
764
765 if let Some(server_addr) = session.downstream_session.server_addr() {
767 let server_addr_str = server_addr.to_string();
768 let config = ctx
769 .config
770 .get_or_insert_with(|| self.config_manager.current());
771 for listener in &config.listeners {
772 if listener.address == server_addr_str {
773 session.downstream_session.set_read_timeout(Some(
775 std::time::Duration::from_secs(listener.request_timeout_secs),
776 ));
777 ctx.listener_keepalive_timeout_secs = Some(listener.keepalive_timeout_secs);
779 break;
780 }
781 }
782 }
783
784 if let Some(route_id) = ctx.route_id.as_deref() {
787 if self.rate_limit_manager.has_route_limiter(route_id) {
788 let rate_result = self.rate_limit_manager.check(
789 route_id,
790 &ctx.client_ip,
791 &ctx.path,
792 Option::<&NoHeaderAccessor>::None,
793 );
794
795 if rate_result.limit > 0 {
797 ctx.rate_limit_info = Some(super::context::RateLimitHeaderInfo {
798 limit: rate_result.limit,
799 remaining: rate_result.remaining,
800 reset_at: rate_result.reset_at,
801 });
802 }
803
804 if !rate_result.allowed {
805 use zentinel_config::RateLimitAction;
806
807 match rate_result.action {
808 RateLimitAction::Reject => {
809 warn!(
810 correlation_id = %ctx.trace_id,
811 route_id = route_id,
812 client_ip = %ctx.client_ip,
813 limiter = %rate_result.limiter,
814 limit = rate_result.limit,
815 remaining = rate_result.remaining,
816 "Request rate limited"
817 );
818 self.metrics.record_blocked_request("rate_limited");
819
820 let audit_entry = AuditLogEntry::rate_limited(
822 &ctx.trace_id,
823 &ctx.method,
824 &ctx.path,
825 &ctx.client_ip,
826 &rate_result.limiter,
827 )
828 .with_route_id(route_id)
829 .with_status_code(rate_result.status_code);
830 self.log_manager.log_audit(&audit_entry);
831
832 let body = rate_result
834 .message
835 .unwrap_or_else(|| "Rate limit exceeded".to_string());
836
837 let retry_after = rate_result.reset_at.saturating_sub(
839 std::time::SystemTime::now()
840 .duration_since(std::time::UNIX_EPOCH)
841 .unwrap_or_default()
842 .as_secs(),
843 );
844 crate::http_helpers::write_rate_limit_error(
845 session,
846 rate_result.status_code,
847 &body,
848 rate_result.limit,
849 rate_result.remaining,
850 rate_result.reset_at,
851 retry_after,
852 )
853 .await?;
854 return Ok(true); }
856 RateLimitAction::LogOnly => {
857 debug!(
858 correlation_id = %ctx.trace_id,
859 route_id = route_id,
860 "Rate limit exceeded (log only mode)"
861 );
862 }
864 RateLimitAction::Delay => {
865 if let Some(delay_ms) = rate_result.suggested_delay_ms {
867 let actual_delay = delay_ms.min(rate_result.max_delay_ms);
869
870 if actual_delay > 0 {
871 debug!(
872 correlation_id = %ctx.trace_id,
873 route_id = route_id,
874 suggested_delay_ms = delay_ms,
875 max_delay_ms = rate_result.max_delay_ms,
876 actual_delay_ms = actual_delay,
877 "Applying rate limit delay"
878 );
879
880 tokio::time::sleep(std::time::Duration::from_millis(
881 actual_delay,
882 ))
883 .await;
884 }
885 }
886 }
888 }
889 }
890 }
891 }
892
893 if let Some(route_id) = ctx.route_id.as_deref() {
896 if let Some(ref route_config) = ctx.route_config {
897 if route_config.service_type == zentinel_config::ServiceType::Inference
898 && self.inference_rate_limit_manager.has_route(route_id)
899 {
900 let headers = &session.req_header().headers;
903
904 let body = ctx.body_buffer.as_slice();
906
907 let rate_limit_key = &ctx.client_ip;
909
910 if let Some(check_result) = self.inference_rate_limit_manager.check(
911 route_id,
912 rate_limit_key,
913 headers,
914 body,
915 ) {
916 ctx.inference_rate_limit_enabled = true;
918 ctx.inference_estimated_tokens = check_result.estimated_tokens;
919 ctx.inference_rate_limit_key = Some(rate_limit_key.to_string());
920 ctx.inference_model = check_result.model.clone();
921
922 if !check_result.is_allowed() {
923 let retry_after_ms = check_result.retry_after_ms();
924 let retry_after_secs = retry_after_ms.div_ceil(1000);
925
926 warn!(
927 correlation_id = %ctx.trace_id,
928 route_id = route_id,
929 client_ip = %ctx.client_ip,
930 estimated_tokens = check_result.estimated_tokens,
931 model = ?check_result.model,
932 retry_after_ms = retry_after_ms,
933 "Inference rate limit exceeded (tokens)"
934 );
935 self.metrics
936 .record_blocked_request("inference_rate_limited");
937
938 let audit_entry = AuditLogEntry::new(
940 &ctx.trace_id,
941 AuditEventType::RateLimitExceeded,
942 &ctx.method,
943 &ctx.path,
944 &ctx.client_ip,
945 )
946 .with_route_id(route_id)
947 .with_status_code(429)
948 .with_reason(format!(
949 "Token rate limit exceeded: estimated {} tokens, model={:?}",
950 check_result.estimated_tokens, check_result.model
951 ));
952 self.log_manager.log_audit(&audit_entry);
953
954 let body = "Token rate limit exceeded";
956 let reset_at = std::time::SystemTime::now()
957 .duration_since(std::time::UNIX_EPOCH)
958 .unwrap_or_default()
959 .as_secs()
960 + retry_after_secs;
961
962 crate::http_helpers::write_rate_limit_error(
964 session,
965 429,
966 body,
967 0, 0, reset_at,
970 retry_after_secs,
971 )
972 .await?;
973 return Ok(true); }
975
976 trace!(
977 correlation_id = %ctx.trace_id,
978 route_id = route_id,
979 estimated_tokens = check_result.estimated_tokens,
980 model = ?check_result.model,
981 "Inference rate limit check passed"
982 );
983
984 if self.inference_rate_limit_manager.has_budget(route_id) {
986 ctx.inference_budget_enabled = true;
987
988 if let Some(budget_result) =
989 self.inference_rate_limit_manager.check_budget(
990 route_id,
991 rate_limit_key,
992 check_result.estimated_tokens,
993 )
994 {
995 if !budget_result.is_allowed() {
996 let retry_after_secs = budget_result.retry_after_secs();
997
998 warn!(
999 correlation_id = %ctx.trace_id,
1000 route_id = route_id,
1001 client_ip = %ctx.client_ip,
1002 estimated_tokens = check_result.estimated_tokens,
1003 retry_after_secs = retry_after_secs,
1004 "Token budget exhausted"
1005 );
1006
1007 ctx.inference_budget_exhausted = true;
1008 self.metrics.record_blocked_request("budget_exhausted");
1009
1010 let audit_entry = AuditLogEntry::new(
1012 &ctx.trace_id,
1013 AuditEventType::RateLimitExceeded,
1014 &ctx.method,
1015 &ctx.path,
1016 &ctx.client_ip,
1017 )
1018 .with_route_id(route_id)
1019 .with_status_code(429)
1020 .with_reason("Token budget exhausted".to_string());
1021 self.log_manager.log_audit(&audit_entry);
1022
1023 let body = "Token budget exhausted";
1025 let reset_at = std::time::SystemTime::now()
1026 .duration_since(std::time::UNIX_EPOCH)
1027 .unwrap_or_default()
1028 .as_secs()
1029 + retry_after_secs;
1030
1031 crate::http_helpers::write_rate_limit_error(
1032 session,
1033 429,
1034 body,
1035 0,
1036 0,
1037 reset_at,
1038 retry_after_secs,
1039 )
1040 .await?;
1041 return Ok(true);
1042 }
1043
1044 let remaining = match &budget_result {
1046 zentinel_common::budget::BudgetCheckResult::Allowed {
1047 remaining,
1048 } => *remaining as i64,
1049 zentinel_common::budget::BudgetCheckResult::Soft {
1050 remaining,
1051 ..
1052 } => *remaining,
1053 _ => 0,
1054 };
1055 ctx.inference_budget_remaining = Some(remaining);
1056
1057 if let Some(status) = self
1059 .inference_rate_limit_manager
1060 .budget_status(route_id, rate_limit_key)
1061 {
1062 ctx.inference_budget_period_reset = Some(status.period_end);
1063 }
1064
1065 trace!(
1066 correlation_id = %ctx.trace_id,
1067 route_id = route_id,
1068 budget_remaining = remaining,
1069 "Token budget check passed"
1070 );
1071 }
1072 }
1073
1074 if self
1076 .inference_rate_limit_manager
1077 .has_cost_attribution(route_id)
1078 {
1079 ctx.inference_cost_enabled = true;
1080 }
1081 }
1082 }
1083 }
1084 }
1085
1086 if let Some(ref route_config) = ctx.route_config {
1088 if let Some(ref inference) = route_config.inference {
1089 if let Some(ref guardrails) = inference.guardrails {
1090 if let Some(ref pi_config) = guardrails.prompt_injection {
1091 if pi_config.enabled && !ctx.body_buffer.is_empty() {
1092 ctx.guardrails_enabled = true;
1093
1094 if let Some(content) = extract_inference_content(&ctx.body_buffer) {
1096 let result = self
1097 .guardrail_processor
1098 .check_prompt_injection(
1099 pi_config,
1100 &content,
1101 ctx.inference_model.as_deref(),
1102 ctx.route_id.as_deref(),
1103 &ctx.trace_id,
1104 )
1105 .await;
1106
1107 match result {
1108 PromptInjectionResult::Blocked {
1109 status,
1110 message,
1111 detections,
1112 } => {
1113 warn!(
1114 correlation_id = %ctx.trace_id,
1115 route_id = ctx.route_id.as_deref().unwrap_or("unknown"),
1116 detection_count = detections.len(),
1117 "Prompt injection detected, blocking request"
1118 );
1119
1120 self.metrics.record_blocked_request("prompt_injection");
1121
1122 ctx.guardrail_detection_categories =
1124 detections.iter().map(|d| d.category.clone()).collect();
1125
1126 let audit_entry = AuditLogEntry::new(
1128 &ctx.trace_id,
1129 AuditEventType::Blocked,
1130 &ctx.method,
1131 &ctx.path,
1132 &ctx.client_ip,
1133 )
1134 .with_route_id(ctx.route_id.as_deref().unwrap_or("unknown"))
1135 .with_status_code(status)
1136 .with_reason("Prompt injection detected".to_string());
1137 self.log_manager.log_audit(&audit_entry);
1138
1139 crate::http_helpers::write_json_error(
1141 session,
1142 status,
1143 "prompt_injection_blocked",
1144 Some(&message),
1145 )
1146 .await?;
1147 return Ok(true);
1148 }
1149 PromptInjectionResult::Detected { detections } => {
1150 warn!(
1152 correlation_id = %ctx.trace_id,
1153 route_id = ctx.route_id.as_deref().unwrap_or("unknown"),
1154 detection_count = detections.len(),
1155 "Prompt injection detected (logged only)"
1156 );
1157 ctx.guardrail_detection_categories =
1158 detections.iter().map(|d| d.category.clone()).collect();
1159 }
1160 PromptInjectionResult::Warning { detections } => {
1161 ctx.guardrail_warning = true;
1163 ctx.guardrail_detection_categories =
1164 detections.iter().map(|d| d.category.clone()).collect();
1165 debug!(
1166 correlation_id = %ctx.trace_id,
1167 "Prompt injection warning set"
1168 );
1169 }
1170 PromptInjectionResult::Clean => {
1171 trace!(
1172 correlation_id = %ctx.trace_id,
1173 "No prompt injection detected"
1174 );
1175 }
1176 PromptInjectionResult::Error { message } => {
1177 trace!(
1179 correlation_id = %ctx.trace_id,
1180 error = %message,
1181 "Prompt injection check error (failure mode applied)"
1182 );
1183 }
1184 }
1185 }
1186 }
1187 }
1188 }
1189 }
1190 }
1191
1192 if let Some(route_id) = ctx.route_id.as_deref() {
1194 if let Some(ref route_config) = ctx.route_config {
1195 for filter_id in &route_config.filters {
1196 if let Some(result) = self.geo_filter_manager.check(filter_id, &ctx.client_ip) {
1197 ctx.geo_country_code = result.country_code.clone();
1199 ctx.geo_lookup_performed = true;
1200
1201 if !result.allowed {
1202 warn!(
1203 correlation_id = %ctx.trace_id,
1204 route_id = route_id,
1205 client_ip = %ctx.client_ip,
1206 country = ?result.country_code,
1207 filter_id = %filter_id,
1208 "Request blocked by geo filter"
1209 );
1210 self.metrics.record_blocked_request("geo_blocked");
1211
1212 let audit_entry = AuditLogEntry::new(
1214 &ctx.trace_id,
1215 AuditEventType::Blocked,
1216 &ctx.method,
1217 &ctx.path,
1218 &ctx.client_ip,
1219 )
1220 .with_route_id(route_id)
1221 .with_status_code(result.status_code)
1222 .with_reason(format!(
1223 "Geo blocked: country={}, filter={}",
1224 result.country_code.as_deref().unwrap_or("unknown"),
1225 filter_id
1226 ));
1227 self.log_manager.log_audit(&audit_entry);
1228
1229 let body = result
1231 .block_message
1232 .unwrap_or_else(|| "Access denied".to_string());
1233
1234 crate::http_helpers::write_error(
1235 session,
1236 result.status_code,
1237 &body,
1238 "text/plain",
1239 )
1240 .await?;
1241 return Ok(true); }
1243
1244 break;
1246 }
1247 }
1248 }
1249 }
1250
1251 let config_for_filters = std::sync::Arc::clone(
1254 ctx.config
1255 .get_or_insert_with(|| self.config_manager.current()),
1256 );
1257 if super::filters::apply_request_filters(session, ctx, &config_for_filters).await? {
1258 return Ok(true); }
1260
1261 let is_websocket_upgrade = session
1263 .req_header()
1264 .headers
1265 .get(http::header::UPGRADE)
1266 .map(|v| v.as_bytes().eq_ignore_ascii_case(b"websocket"))
1267 .unwrap_or(false);
1268
1269 if is_websocket_upgrade {
1270 ctx.is_websocket_upgrade = true;
1271
1272 if let Some(ref route_config) = ctx.route_config {
1274 if !route_config.websocket {
1275 warn!(
1276 correlation_id = %ctx.trace_id,
1277 route_id = ctx.route_id.as_deref().unwrap_or("unknown"),
1278 client_ip = %ctx.client_ip,
1279 "WebSocket upgrade rejected: not enabled for route"
1280 );
1281
1282 self.metrics.record_blocked_request("websocket_not_enabled");
1283
1284 let audit_entry = AuditLogEntry::new(
1286 &ctx.trace_id,
1287 AuditEventType::Blocked,
1288 &ctx.method,
1289 &ctx.path,
1290 &ctx.client_ip,
1291 )
1292 .with_route_id(ctx.route_id.as_deref().unwrap_or("unknown"))
1293 .with_action("websocket_rejected")
1294 .with_reason("WebSocket not enabled for route");
1295 self.log_manager.log_audit(&audit_entry);
1296
1297 crate::http_helpers::write_error(
1299 session,
1300 403,
1301 "WebSocket not enabled for this route",
1302 "text/plain",
1303 )
1304 .await?;
1305 return Ok(true); }
1307
1308 debug!(
1309 correlation_id = %ctx.trace_id,
1310 route_id = ctx.route_id.as_deref().unwrap_or("unknown"),
1311 "WebSocket upgrade request allowed"
1312 );
1313
1314 if route_config.websocket_inspection {
1316 let has_compression = session
1318 .req_header()
1319 .headers
1320 .get("Sec-WebSocket-Extensions")
1321 .and_then(|v| v.to_str().ok())
1322 .map(|s| s.contains("permessage-deflate"))
1323 .unwrap_or(false);
1324
1325 if has_compression {
1326 debug!(
1327 correlation_id = %ctx.trace_id,
1328 route_id = ctx.route_id.as_deref().unwrap_or("unknown"),
1329 "WebSocket inspection skipped: permessage-deflate negotiated"
1330 );
1331 ctx.websocket_skip_inspection = true;
1332 } else {
1333 ctx.websocket_inspection_enabled = true;
1334
1335 ctx.websocket_inspection_agents = self.agent_manager.get_agents_for_event(
1337 zentinel_agent_protocol::EventType::WebSocketFrame,
1338 );
1339
1340 debug!(
1341 correlation_id = %ctx.trace_id,
1342 route_id = ctx.route_id.as_deref().unwrap_or("unknown"),
1343 agent_count = ctx.websocket_inspection_agents.len(),
1344 "WebSocket frame inspection enabled"
1345 );
1346 }
1347 }
1348 }
1349 }
1350
1351 if let Some(route_config) = ctx.route_config.clone() {
1354 if route_config.service_type == zentinel_config::ServiceType::Static {
1355 trace!(
1356 correlation_id = %ctx.trace_id,
1357 route_id = ctx.route_id.as_deref().unwrap_or("unknown"),
1358 "Handling static file route"
1359 );
1360 let route_match = crate::routing::RouteMatch {
1362 route_id: zentinel_common::RouteId::new(ctx.route_id.as_deref().unwrap_or("")),
1363 config: route_config.clone(),
1364 };
1365 return self.handle_static_route(session, ctx, &route_match).await;
1366 } else if route_config.service_type == zentinel_config::ServiceType::Builtin {
1367 trace!(
1368 correlation_id = %ctx.trace_id,
1369 route_id = ctx.route_id.as_deref().unwrap_or("unknown"),
1370 builtin_handler = ?route_config.builtin_handler,
1371 "Handling builtin route"
1372 );
1373 let route_match = crate::routing::RouteMatch {
1375 route_id: zentinel_common::RouteId::new(ctx.route_id.as_deref().unwrap_or("")),
1376 config: route_config.clone(),
1377 };
1378 return self.handle_builtin_route(session, ctx, &route_match).await;
1379 }
1380 }
1381
1382 if let Some(route_id) = ctx.route_id.clone() {
1384 if let Some(validator) = self.validators.get(&route_id).await {
1385 trace!(
1386 correlation_id = %ctx.trace_id,
1387 route_id = %route_id,
1388 "Running API schema validation"
1389 );
1390 if let Some(result) = self
1391 .validate_api_request(session, ctx, &route_id, &validator)
1392 .await?
1393 {
1394 debug!(
1395 correlation_id = %ctx.trace_id,
1396 route_id = %route_id,
1397 validation_passed = result,
1398 "API validation complete"
1399 );
1400 return Ok(result);
1401 }
1402 }
1403 }
1404
1405 let client_addr = session
1407 .client_addr()
1408 .map(|a| format!("{}", a))
1409 .unwrap_or_else(|| "unknown".to_string());
1410 let client_port = session.client_addr().map(|_| 0).unwrap_or(0);
1411
1412 let req_header = session.req_header_mut();
1413
1414 req_header
1416 .insert_header("X-Correlation-Id", &ctx.trace_id)
1417 .ok();
1418 req_header.insert_header("X-Forwarded-By", "Zentinel").ok();
1419
1420 let config = ctx
1422 .config
1423 .get_or_insert_with(|| self.config_manager.current());
1424
1425 const HEADER_LIMIT_THRESHOLD: usize = 1024 * 1024; let header_count = req_header.headers.len();
1430 if config.limits.max_header_count < HEADER_LIMIT_THRESHOLD
1431 && header_count > config.limits.max_header_count
1432 {
1433 warn!(
1434 correlation_id = %ctx.trace_id,
1435 header_count = header_count,
1436 limit = config.limits.max_header_count,
1437 "Request blocked: exceeds header count limit"
1438 );
1439
1440 self.metrics.record_blocked_request("header_count_exceeded");
1441 return Err(Error::explain(ErrorType::InternalError, "Too many headers"));
1442 }
1443
1444 if config.limits.max_header_size_bytes < HEADER_LIMIT_THRESHOLD {
1446 let total_header_size: usize = req_header
1447 .headers
1448 .iter()
1449 .map(|(k, v)| k.as_str().len() + v.len())
1450 .sum();
1451
1452 if total_header_size > config.limits.max_header_size_bytes {
1453 warn!(
1454 correlation_id = %ctx.trace_id,
1455 header_size = total_header_size,
1456 limit = config.limits.max_header_size_bytes,
1457 "Request blocked: exceeds header size limit"
1458 );
1459
1460 self.metrics.record_blocked_request("header_size_exceeded");
1461 return Err(Error::explain(
1462 ErrorType::InternalError,
1463 "Headers too large",
1464 ));
1465 }
1466 }
1467
1468 trace!(
1470 correlation_id = %ctx.trace_id,
1471 "Processing request through agents"
1472 );
1473 if let Err(e) = self
1474 .process_agents(session, ctx, &client_addr, client_port)
1475 .await
1476 {
1477 if let ErrorType::HTTPStatus(status) = e.etype() {
1480 let error_msg = e.to_string();
1482 let body = error_msg
1483 .split("context:")
1484 .nth(1)
1485 .map(|s| s.trim())
1486 .unwrap_or("Request blocked");
1487 debug!(
1488 correlation_id = %ctx.trace_id,
1489 status = status,
1490 body = %body,
1491 "Sending HTTP error response for agent block"
1492 );
1493 crate::http_helpers::write_error(session, *status, body, "text/plain").await?;
1494 return Ok(true); }
1496 return Err(e);
1498 }
1499
1500 trace!(
1501 correlation_id = %ctx.trace_id,
1502 "Request filter phase complete, forwarding to upstream"
1503 );
1504
1505 Ok(false) }
1507
1508 async fn request_body_filter(
1515 &self,
1516 _session: &mut Session,
1517 body: &mut Option<Bytes>,
1518 end_of_stream: bool,
1519 ctx: &mut Self::CTX,
1520 ) -> Result<(), Box<Error>> {
1521 use zentinel_config::BodyStreamingMode;
1522
1523 if ctx.is_websocket_upgrade {
1525 if let Some(ref handler) = ctx.websocket_handler {
1526 let result = handler.process_client_data(body.take()).await;
1527 match result {
1528 crate::websocket::ProcessResult::Forward(data) => {
1529 *body = data;
1530 }
1531 crate::websocket::ProcessResult::Close(reason) => {
1532 warn!(
1533 correlation_id = %ctx.trace_id,
1534 code = reason.code,
1535 reason = %reason.reason,
1536 "WebSocket connection closed by agent (client->server)"
1537 );
1538 return Err(Error::explain(
1540 ErrorType::InternalError,
1541 format!("WebSocket closed: {} {}", reason.code, reason.reason),
1542 ));
1543 }
1544 }
1545 }
1546 return Ok(());
1548 }
1549
1550 let chunk_len = body.as_ref().map(|b| b.len()).unwrap_or(0);
1552 if chunk_len > 0 {
1553 ctx.request_body_bytes += chunk_len as u64;
1554
1555 trace!(
1556 correlation_id = %ctx.trace_id,
1557 chunk_size = chunk_len,
1558 total_body_bytes = ctx.request_body_bytes,
1559 end_of_stream = end_of_stream,
1560 streaming_mode = ?ctx.request_body_streaming_mode,
1561 "Processing request body chunk"
1562 );
1563
1564 let config = ctx
1566 .config
1567 .get_or_insert_with(|| self.config_manager.current());
1568 if ctx.request_body_bytes > config.limits.max_body_size_bytes as u64 {
1569 warn!(
1570 correlation_id = %ctx.trace_id,
1571 body_bytes = ctx.request_body_bytes,
1572 limit = config.limits.max_body_size_bytes,
1573 "Request body size limit exceeded"
1574 );
1575 self.metrics.record_blocked_request("body_size_exceeded");
1576 return Err(Error::explain(
1577 ErrorType::InternalError,
1578 "Request body too large",
1579 ));
1580 }
1581 }
1582
1583 if ctx.body_inspection_enabled && !ctx.body_inspection_agents.is_empty() {
1585 let config = ctx
1586 .config
1587 .get_or_insert_with(|| self.config_manager.current());
1588 let max_inspection_bytes = config
1589 .waf
1590 .as_ref()
1591 .map(|w| w.body_inspection.max_inspection_bytes as u64)
1592 .unwrap_or(1024 * 1024);
1593
1594 match ctx.request_body_streaming_mode {
1595 BodyStreamingMode::Stream => {
1596 if body.is_some() {
1598 self.process_body_chunk_streaming(body, end_of_stream, ctx)
1599 .await?;
1600 } else if end_of_stream && ctx.agent_needs_more {
1601 self.process_body_chunk_streaming(body, end_of_stream, ctx)
1603 .await?;
1604 }
1605 }
1606 BodyStreamingMode::Hybrid { buffer_threshold } => {
1607 if ctx.body_bytes_inspected < buffer_threshold as u64 {
1609 if let Some(ref chunk) = body {
1611 let bytes_to_buffer = std::cmp::min(
1612 chunk.len(),
1613 (buffer_threshold as u64 - ctx.body_bytes_inspected) as usize,
1614 );
1615 ctx.body_buffer.extend_from_slice(&chunk[..bytes_to_buffer]);
1616 ctx.body_bytes_inspected += bytes_to_buffer as u64;
1617
1618 if ctx.body_bytes_inspected >= buffer_threshold as u64 || end_of_stream
1620 {
1621 self.send_buffered_body_to_agents(
1623 end_of_stream && chunk.len() == bytes_to_buffer,
1624 ctx,
1625 )
1626 .await?;
1627 ctx.body_buffer.clear();
1628
1629 if bytes_to_buffer < chunk.len() {
1631 let remaining = chunk.slice(bytes_to_buffer..);
1632 let mut remaining_body = Some(remaining);
1633 self.process_body_chunk_streaming(
1634 &mut remaining_body,
1635 end_of_stream,
1636 ctx,
1637 )
1638 .await?;
1639 }
1640 }
1641 }
1642 } else {
1643 self.process_body_chunk_streaming(body, end_of_stream, ctx)
1645 .await?;
1646 }
1647 }
1648 BodyStreamingMode::Buffer => {
1649 if let Some(ref chunk) = body {
1651 if ctx.body_bytes_inspected < max_inspection_bytes {
1652 let bytes_to_inspect = std::cmp::min(
1653 chunk.len() as u64,
1654 max_inspection_bytes - ctx.body_bytes_inspected,
1655 ) as usize;
1656
1657 ctx.body_buffer
1658 .extend_from_slice(&chunk[..bytes_to_inspect]);
1659 ctx.body_bytes_inspected += bytes_to_inspect as u64;
1660
1661 trace!(
1662 correlation_id = %ctx.trace_id,
1663 bytes_inspected = ctx.body_bytes_inspected,
1664 max_inspection_bytes = max_inspection_bytes,
1665 buffer_size = ctx.body_buffer.len(),
1666 "Buffering body for agent inspection"
1667 );
1668 }
1669 }
1670
1671 let should_send =
1673 end_of_stream || ctx.body_bytes_inspected >= max_inspection_bytes;
1674 if should_send && !ctx.body_buffer.is_empty() {
1675 self.send_buffered_body_to_agents(end_of_stream, ctx)
1676 .await?;
1677 ctx.body_buffer.clear();
1678 }
1679 }
1680 }
1681 }
1682
1683 if end_of_stream {
1684 trace!(
1685 correlation_id = %ctx.trace_id,
1686 total_body_bytes = ctx.request_body_bytes,
1687 bytes_inspected = ctx.body_bytes_inspected,
1688 "Request body complete"
1689 );
1690 }
1691
1692 Ok(())
1693 }
1694
1695 async fn response_filter(
1696 &self,
1697 session: &mut Session,
1698 upstream_response: &mut ResponseHeader,
1699 ctx: &mut Self::CTX,
1700 ) -> Result<(), Box<Error>> {
1701 let status = upstream_response.status.as_u16();
1702 let duration = ctx.elapsed();
1703
1704 trace!(
1705 correlation_id = %ctx.trace_id,
1706 status = status,
1707 "Starting response filter phase"
1708 );
1709
1710 if status == 101 && ctx.is_websocket_upgrade {
1712 if ctx.websocket_inspection_enabled && !ctx.websocket_skip_inspection {
1713 let inspector = crate::websocket::WebSocketInspector::with_metrics(
1715 self.agent_manager.clone(),
1716 ctx.route_id
1717 .clone()
1718 .unwrap_or_else(|| "unknown".to_string()),
1719 ctx.trace_id.clone(),
1720 ctx.client_ip.clone(),
1721 100, Some(self.metrics.clone()),
1723 );
1724
1725 let handler = crate::websocket::WebSocketHandler::new(
1726 std::sync::Arc::new(inspector),
1727 1024 * 1024, );
1729
1730 ctx.websocket_handler = Some(std::sync::Arc::new(handler));
1731
1732 info!(
1733 correlation_id = %ctx.trace_id,
1734 route_id = ctx.route_id.as_deref().unwrap_or("unknown"),
1735 agent_count = ctx.websocket_inspection_agents.len(),
1736 "WebSocket upgrade successful, frame inspection enabled"
1737 );
1738 } else if ctx.websocket_skip_inspection {
1739 debug!(
1740 correlation_id = %ctx.trace_id,
1741 route_id = ctx.route_id.as_deref().unwrap_or("unknown"),
1742 "WebSocket upgrade successful, inspection skipped (compression negotiated)"
1743 );
1744 } else {
1745 debug!(
1746 correlation_id = %ctx.trace_id,
1747 route_id = ctx.route_id.as_deref().unwrap_or("unknown"),
1748 "WebSocket upgrade successful"
1749 );
1750 }
1751 }
1752
1753 upstream_response.insert_header("X-Correlation-Id", &ctx.trace_id)?;
1755
1756 if let Some(ref rate_info) = ctx.rate_limit_info {
1758 upstream_response.insert_header("X-RateLimit-Limit", rate_info.limit.to_string())?;
1759 upstream_response
1760 .insert_header("X-RateLimit-Remaining", rate_info.remaining.to_string())?;
1761 upstream_response.insert_header("X-RateLimit-Reset", rate_info.reset_at.to_string())?;
1762 }
1763
1764 if ctx.inference_budget_enabled {
1766 if let Some(remaining) = ctx.inference_budget_remaining {
1767 upstream_response.insert_header("X-Budget-Remaining", remaining.to_string())?;
1768 }
1769 if let Some(period_reset) = ctx.inference_budget_period_reset {
1770 let reset_datetime = chrono::DateTime::from_timestamp(period_reset as i64, 0)
1772 .map(|dt| dt.format("%Y-%m-%dT%H:%M:%SZ").to_string())
1773 .unwrap_or_else(|| period_reset.to_string());
1774 upstream_response.insert_header("X-Budget-Period-Reset", reset_datetime)?;
1775 }
1776 }
1777
1778 if let Some(ref country_code) = ctx.geo_country_code {
1780 upstream_response.insert_header("X-GeoIP-Country", country_code)?;
1781 }
1782
1783 if let Some(ref route_config) = ctx.route_config {
1785 let mods = &route_config.policies.response_headers;
1786 for (old_name, new_name) in &mods.rename {
1788 if let Some(value) = upstream_response
1789 .headers
1790 .get(old_name)
1791 .and_then(|v| v.to_str().ok())
1792 {
1793 let owned = value.to_string();
1794 upstream_response
1795 .insert_header(new_name.clone(), &owned)
1796 .ok();
1797 upstream_response.remove_header(old_name);
1798 }
1799 }
1800 for (name, value) in &mods.set {
1801 upstream_response
1802 .insert_header(name.clone(), value.as_str())
1803 .ok();
1804 }
1805 for (name, value) in &mods.add {
1806 upstream_response
1807 .append_header(name.clone(), value.as_str())
1808 .ok();
1809 }
1810 for name in &mods.remove {
1811 upstream_response.remove_header(name);
1812 }
1813 }
1814
1815 if let Some(ref cache_status) = ctx.cache_status {
1817 let status_header_enabled = ctx
1818 .config
1819 .as_ref()
1820 .and_then(|c| c.cache.as_ref())
1821 .map(|c| c.status_header)
1822 .unwrap_or(false);
1823
1824 if status_header_enabled {
1825 let cache_name = ctx
1826 .config
1827 .as_ref()
1828 .and_then(|c| c.cache.as_ref())
1829 .map(|c| c.status_header_name.as_str())
1830 .unwrap_or("zentinel");
1831
1832 let value = match cache_status {
1833 super::context::CacheStatus::HitMemory => {
1834 format!("{cache_name}; hit; detail=memory")
1835 }
1836 super::context::CacheStatus::HitDisk => {
1837 format!("{cache_name}; hit; detail=disk")
1838 }
1839 super::context::CacheStatus::Hit => format!("{cache_name}; hit"),
1840 super::context::CacheStatus::HitStale => format!("{cache_name}; fwd=stale"),
1841 super::context::CacheStatus::Miss => format!("{cache_name}; fwd=miss"),
1842 super::context::CacheStatus::Bypass(reason) => match *reason {
1843 "method" => format!("{cache_name}; fwd=bypass; detail=method"),
1844 "disabled" => format!("{cache_name}; fwd=bypass; detail=disabled"),
1845 "no-route" => format!("{cache_name}; fwd=bypass; detail=no-route"),
1846 _ => format!("{cache_name}; fwd=bypass"),
1847 },
1848 };
1849 upstream_response.insert_header("Cache-Status", &value).ok();
1850 }
1851 }
1852
1853 if let Some(config) = ctx.config.as_ref().map(std::sync::Arc::clone) {
1855 super::filters::apply_response_filters(upstream_response, ctx, &config);
1856 }
1857
1858 if ctx.compress_enabled {
1860 session.upstream_compression.adjust_level(6);
1861 }
1862
1863 if let Some(keepalive_secs) = ctx.listener_keepalive_timeout_secs {
1865 session
1866 .downstream_session
1867 .set_keepalive(Some(keepalive_secs));
1868 }
1869
1870 if ctx.sticky_session_new_assignment {
1872 if let Some(ref set_cookie_header) = ctx.sticky_session_set_cookie {
1873 upstream_response.insert_header("Set-Cookie", set_cookie_header)?;
1874 trace!(
1875 correlation_id = %ctx.trace_id,
1876 sticky_target_index = ?ctx.sticky_target_index,
1877 "Added sticky session Set-Cookie header"
1878 );
1879 }
1880 }
1881
1882 if ctx.guardrail_warning {
1884 upstream_response.insert_header("X-Guardrail-Warning", "prompt-injection-detected")?;
1885 }
1886
1887 if ctx.used_fallback() {
1889 upstream_response.insert_header("X-Fallback-Used", "true")?;
1890
1891 if let Some(ref upstream) = ctx.upstream {
1892 upstream_response.insert_header("X-Fallback-Upstream", upstream)?;
1893 }
1894
1895 if let Some(ref reason) = ctx.fallback_reason {
1896 upstream_response.insert_header("X-Fallback-Reason", reason.to_string())?;
1897 }
1898
1899 if let Some(ref original) = ctx.original_upstream {
1900 upstream_response.insert_header("X-Original-Upstream", original)?;
1901 }
1902
1903 if let Some(ref mapping) = ctx.model_mapping_applied {
1904 upstream_response
1905 .insert_header("X-Model-Mapping", format!("{} -> {}", mapping.0, mapping.1))?;
1906 }
1907
1908 trace!(
1909 correlation_id = %ctx.trace_id,
1910 fallback_attempt = ctx.fallback_attempt,
1911 fallback_upstream = ctx.upstream.as_deref().unwrap_or("unknown"),
1912 original_upstream = ctx.original_upstream.as_deref().unwrap_or("unknown"),
1913 "Added fallback response headers"
1914 );
1915
1916 if status < 400 {
1918 if let Some(metrics) = get_fallback_metrics() {
1919 metrics.record_fallback_success(
1920 ctx.route_id.as_deref().unwrap_or("unknown"),
1921 ctx.upstream.as_deref().unwrap_or("unknown"),
1922 );
1923 }
1924 }
1925 }
1926
1927 if ctx.inference_rate_limit_enabled {
1929 let content_type = upstream_response
1931 .headers
1932 .get("content-type")
1933 .and_then(|ct| ct.to_str().ok());
1934
1935 if is_sse_response(content_type) {
1936 let provider = ctx
1938 .route_config
1939 .as_ref()
1940 .and_then(|r| r.inference.as_ref())
1941 .map(|i| i.provider)
1942 .unwrap_or_default();
1943
1944 ctx.inference_streaming_response = true;
1945 ctx.inference_streaming_counter = Some(StreamingTokenCounter::new(
1946 provider,
1947 ctx.inference_model.clone(),
1948 ));
1949
1950 trace!(
1951 correlation_id = %ctx.trace_id,
1952 content_type = ?content_type,
1953 model = ?ctx.inference_model,
1954 "Initialized streaming token counter for SSE response"
1955 );
1956 }
1957 }
1958
1959 if !ctx.route_agent_ids.is_empty() {
1961 let agent_ids = ctx.route_agent_ids.clone();
1962 let mut resp_headers_map: std::collections::HashMap<String, Vec<String>> =
1963 std::collections::HashMap::with_capacity(upstream_response.headers.len());
1964 for (name, value) in upstream_response.headers.iter() {
1965 resp_headers_map
1966 .entry(name.as_str().to_string())
1967 .or_default()
1968 .push(value.to_str().unwrap_or("").to_string());
1969 }
1970
1971 let agent_ctx = crate::agents::AgentCallContext {
1972 correlation_id: zentinel_common::CorrelationId::from_string(&ctx.trace_id),
1973 metadata: zentinel_agent_protocol::RequestMetadata {
1974 correlation_id: ctx.trace_id.clone(),
1975 request_id: uuid::Uuid::new_v4().to_string(),
1976 client_ip: ctx.client_ip.clone(),
1977 client_port: 0,
1978 server_name: ctx.host.clone(),
1979 protocol: "HTTP/1.1".to_string(),
1980 tls_version: None,
1981 tls_cipher: None,
1982 route_id: ctx.route_id.clone(),
1983 upstream_id: ctx.upstream.clone(),
1984 timestamp: chrono::Utc::now().to_rfc3339(),
1985 traceparent: ctx.traceparent(),
1986 },
1987 route_id: ctx.route_id.clone(),
1988 upstream_id: ctx.upstream.clone(),
1989 request_body: None,
1990 response_body: None,
1991 };
1992
1993 match self
1994 .agent_manager
1995 .process_response_headers(&agent_ctx, status, &resp_headers_map, &agent_ids)
1996 .await
1997 {
1998 Ok(decision) => {
1999 for op in &decision.response_headers {
2001 match op {
2002 zentinel_agent_protocol::HeaderOp::Set { name, value } => {
2003 upstream_response
2004 .insert_header(name.clone(), value.as_str())
2005 .ok();
2006 }
2007 zentinel_agent_protocol::HeaderOp::Add { name, value } => {
2008 upstream_response
2009 .append_header(name.clone(), value.as_str())
2010 .ok();
2011 }
2012 zentinel_agent_protocol::HeaderOp::Remove { name } => {
2013 upstream_response.remove_header(name);
2014 }
2015 }
2016 }
2017
2018 let has_body_agents = self
2020 .agent_manager
2021 .any_agent_handles_event(
2022 &agent_ids,
2023 zentinel_agent_protocol::EventType::ResponseBodyChunk,
2024 )
2025 .await;
2026 if has_body_agents {
2027 ctx.response_agent_processing_enabled = true;
2028 upstream_response.insert_header("Connection", "close").ok();
2031 session.downstream_session.set_keepalive(None);
2032 debug!(
2033 correlation_id = %ctx.trace_id,
2034 "Enabling response body agent processing (agent subscribes to ResponseBody)"
2035 );
2036 }
2037
2038 debug!(
2039 correlation_id = %ctx.trace_id,
2040 response_headers_modified = !decision.response_headers.is_empty(),
2041 needs_body = ctx.response_agent_processing_enabled,
2042 "Response headers processed through agents"
2043 );
2044 }
2045 Err(e) => {
2046 warn!(
2047 correlation_id = %ctx.trace_id,
2048 error = %e,
2049 "Agent response header processing failed, continuing without agent"
2050 );
2051 }
2052 }
2053 }
2054
2055 if status >= 400 {
2057 trace!(
2058 correlation_id = %ctx.trace_id,
2059 status = status,
2060 "Handling error response"
2061 );
2062 self.handle_error_response(upstream_response, ctx).await?;
2063 }
2064
2065 self.metrics.record_request(
2067 ctx.route_id.as_deref().unwrap_or("unknown"),
2068 &ctx.method,
2069 status,
2070 duration,
2071 );
2072
2073 if let Some(ref mut span) = ctx.otel_span {
2075 span.set_status(status);
2076 if let Some(ref upstream) = ctx.upstream {
2077 span.set_upstream(upstream, "");
2078 }
2079 if status >= 500 {
2080 span.record_error(&format!("HTTP {}", status));
2081 }
2082 }
2083
2084 if let Some(ref upstream) = ctx.upstream {
2086 let success = status < 500;
2087
2088 trace!(
2089 correlation_id = %ctx.trace_id,
2090 upstream = %upstream,
2091 success = success,
2092 status = status,
2093 "Recording passive health check result"
2094 );
2095
2096 let error_msg = if !success {
2097 Some(format!("HTTP {}", status))
2098 } else {
2099 None
2100 };
2101 self.passive_health
2102 .record_outcome(upstream, success, error_msg.as_deref())
2103 .await;
2104
2105 if let Some(pool) = self.upstream_pools.get(upstream).await {
2107 pool.report_result(upstream, success).await;
2108 }
2109
2110 if !success {
2111 warn!(
2112 correlation_id = %ctx.trace_id,
2113 upstream = %upstream,
2114 status = status,
2115 "Upstream returned error status"
2116 );
2117 }
2118 }
2119
2120 if status >= 500 {
2122 error!(
2123 correlation_id = %ctx.trace_id,
2124 route_id = ctx.route_id.as_deref().unwrap_or("unknown"),
2125 upstream = ctx.upstream.as_deref().unwrap_or("none"),
2126 method = %ctx.method,
2127 path = %ctx.path,
2128 status = status,
2129 duration_ms = duration.as_millis(),
2130 attempts = ctx.upstream_attempts,
2131 "Request completed with server error"
2132 );
2133 self.log_manager.log_request_error(
2134 "error",
2135 "Request completed with server error",
2136 &ctx.trace_id,
2137 ctx.route_id.as_deref(),
2138 ctx.upstream.as_deref(),
2139 Some(format!(
2140 "status={} method={} path={} duration_ms={}",
2141 status,
2142 ctx.method,
2143 ctx.path,
2144 duration.as_millis()
2145 )),
2146 );
2147 } else if status >= 400 {
2148 warn!(
2149 correlation_id = %ctx.trace_id,
2150 route_id = ctx.route_id.as_deref().unwrap_or("unknown"),
2151 upstream = ctx.upstream.as_deref().unwrap_or("none"),
2152 method = %ctx.method,
2153 path = %ctx.path,
2154 status = status,
2155 duration_ms = duration.as_millis(),
2156 "Request completed with client error"
2157 );
2158 self.log_manager.log_request_error(
2159 "warn",
2160 "Request completed with client error",
2161 &ctx.trace_id,
2162 ctx.route_id.as_deref(),
2163 ctx.upstream.as_deref(),
2164 Some(format!(
2165 "status={} method={} path={}",
2166 status, ctx.method, ctx.path
2167 )),
2168 );
2169 } else {
2170 debug!(
2171 correlation_id = %ctx.trace_id,
2172 route_id = ctx.route_id.as_deref().unwrap_or("unknown"),
2173 upstream = ctx.upstream.as_deref().unwrap_or("none"),
2174 method = %ctx.method,
2175 path = %ctx.path,
2176 status = status,
2177 duration_ms = duration.as_millis(),
2178 attempts = ctx.upstream_attempts,
2179 "Request completed"
2180 );
2181 }
2182
2183 Ok(())
2184 }
2185
2186 async fn upstream_request_filter(
2189 &self,
2190 _session: &mut Session,
2191 upstream_request: &mut pingora::http::RequestHeader,
2192 ctx: &mut Self::CTX,
2193 ) -> Result<()>
2194 where
2195 Self::CTX: Send + Sync,
2196 {
2197 trace!(
2198 correlation_id = %ctx.trace_id,
2199 route_id = ctx.route_id.as_deref().unwrap_or("unknown"),
2200 "Applying upstream request modifications"
2201 );
2202
2203 upstream_request
2205 .insert_header("X-Trace-Id", &ctx.trace_id)
2206 .ok();
2207
2208 if let Some(ref span) = ctx.otel_span {
2210 let sampled = ctx
2211 .trace_context
2212 .as_ref()
2213 .map(|c| c.sampled)
2214 .unwrap_or(true);
2215 let traceparent =
2216 crate::otel::create_traceparent(&span.trace_id, &span.span_id, sampled);
2217 upstream_request
2218 .insert_header(crate::otel::TRACEPARENT_HEADER, &traceparent)
2219 .ok();
2220 }
2221
2222 upstream_request
2224 .insert_header("X-Forwarded-By", "Zentinel")
2225 .ok();
2226
2227 if let Some(ref route_config) = ctx.route_config {
2231 let mods = &route_config.policies.request_headers;
2232
2233 for (old_name, new_name) in &mods.rename {
2235 if let Some(value) = upstream_request
2236 .headers
2237 .get(old_name)
2238 .and_then(|v| v.to_str().ok())
2239 {
2240 let owned = value.to_string();
2241 upstream_request
2242 .insert_header(new_name.clone(), &owned)
2243 .ok();
2244 upstream_request.remove_header(old_name);
2245 }
2246 }
2247
2248 for (name, value) in &mods.set {
2250 upstream_request
2251 .insert_header(name.clone(), value.as_str())
2252 .ok();
2253 }
2254
2255 for (name, value) in &mods.add {
2257 upstream_request
2258 .append_header(name.clone(), value.as_str())
2259 .ok();
2260 }
2261
2262 for name in &mods.remove {
2264 upstream_request.remove_header(name);
2265 }
2266
2267 trace!(
2268 correlation_id = %ctx.trace_id,
2269 "Applied request header modifications"
2270 );
2271 }
2272
2273 if let Some(ref config) = ctx.config {
2275 super::filters::apply_request_headers_filters(upstream_request, ctx, config);
2276 }
2277
2278 upstream_request.remove_header("X-Internal-Token");
2280 upstream_request.remove_header("Authorization-Internal");
2281
2282 if let Some(ref route_config) = ctx.route_config {
2285 if let Some(ref shadow_config) = route_config.shadow {
2286 let pools_snapshot = self.upstream_pools.snapshot().await;
2288 let upstream_pools = std::sync::Arc::new(pools_snapshot);
2289
2290 let route_id = ctx
2292 .route_id
2293 .clone()
2294 .unwrap_or_else(|| "unknown".to_string());
2295
2296 let shadow_manager = crate::shadow::ShadowManager::new(
2298 upstream_pools,
2299 shadow_config.clone(),
2300 Some(std::sync::Arc::clone(&self.metrics)),
2301 route_id,
2302 );
2303
2304 if shadow_manager.should_shadow(upstream_request) {
2306 trace!(
2307 correlation_id = %ctx.trace_id,
2308 shadow_upstream = %shadow_config.upstream,
2309 percentage = shadow_config.percentage,
2310 "Shadowing request"
2311 );
2312
2313 let shadow_headers = upstream_request.clone();
2315
2316 let shadow_ctx = crate::upstream::RequestContext {
2318 client_ip: ctx.client_ip.parse().ok(),
2319 headers: std::collections::HashMap::new(), path: ctx.path.clone(),
2321 method: ctx.method.clone(),
2322 };
2323
2324 let buffer_body = shadow_config.buffer_body
2326 && crate::shadow::should_buffer_method(&ctx.method);
2327
2328 if buffer_body {
2329 trace!(
2333 correlation_id = %ctx.trace_id,
2334 "Deferring shadow request until body is buffered"
2335 );
2336 ctx.shadow_pending = Some(crate::proxy::context::ShadowPendingRequest {
2337 headers: shadow_headers,
2338 manager: std::sync::Arc::new(shadow_manager),
2339 request_ctx: shadow_ctx,
2340 include_body: true,
2341 });
2342 if !ctx.body_inspection_enabled {
2345 ctx.body_inspection_enabled = true;
2346 }
2349 } else {
2350 shadow_manager.shadow_request(shadow_headers, None, shadow_ctx);
2352 ctx.shadow_sent = true;
2353 }
2354 }
2355 }
2356 }
2357
2358 Ok(())
2359 }
2360
2361 fn response_body_filter(
2367 &self,
2368 _session: &mut Session,
2369 body: &mut Option<Bytes>,
2370 end_of_stream: bool,
2371 ctx: &mut Self::CTX,
2372 ) -> Result<Option<Duration>, Box<Error>> {
2373 if ctx.is_websocket_upgrade {
2376 if let Some(ref handler) = ctx.websocket_handler {
2377 let handler = handler.clone();
2378 let data = body.take();
2379
2380 let result = tokio::task::block_in_place(|| {
2383 tokio::runtime::Handle::current()
2384 .block_on(async { handler.process_server_data(data).await })
2385 });
2386
2387 match result {
2388 crate::websocket::ProcessResult::Forward(data) => {
2389 *body = data;
2390 }
2391 crate::websocket::ProcessResult::Close(reason) => {
2392 warn!(
2393 correlation_id = %ctx.trace_id,
2394 code = reason.code,
2395 reason = %reason.reason,
2396 "WebSocket connection closed by agent (server->client)"
2397 );
2398 let close_frame =
2401 crate::websocket::WebSocketFrame::close(reason.code, &reason.reason);
2402 let codec = crate::websocket::WebSocketCodec::new(1024 * 1024);
2403 if let Ok(encoded) = codec.encode_frame(&close_frame, false) {
2404 *body = Some(Bytes::from(encoded));
2405 }
2406 }
2407 }
2408 }
2409 return Ok(None);
2411 }
2412
2413 if ctx.response_agent_processing_enabled && !ctx.route_agent_ids.is_empty() {
2415 if let Some(ref chunk) = body {
2416 ctx.response_agent_body_buffer.extend_from_slice(chunk);
2417 }
2418
2419 if end_of_stream {
2420 let agent_ids = ctx.route_agent_ids.clone();
2421 let buffer = std::mem::take(&mut ctx.response_agent_body_buffer);
2422 let chunk_index = 0u32;
2423 let total_size = Some(buffer.len());
2424 let trace_id = ctx.trace_id.clone();
2425 let client_ip = ctx.client_ip.clone();
2426 let host = ctx.host.clone();
2427 let route_id = ctx.route_id.clone();
2428 let upstream_id = ctx.upstream.clone();
2429 let traceparent = ctx.traceparent();
2430 let agent_mgr = self.agent_manager.clone();
2431
2432 let result = tokio::task::block_in_place(|| {
2435 tokio::runtime::Handle::current().block_on(async {
2436 let agent_ctx = crate::agents::AgentCallContext {
2437 correlation_id: zentinel_common::CorrelationId::from_string(&trace_id),
2438 metadata: zentinel_agent_protocol::RequestMetadata {
2439 correlation_id: trace_id.clone(),
2440 request_id: uuid::Uuid::new_v4().to_string(),
2441 client_ip,
2442 client_port: 0,
2443 server_name: host,
2444 protocol: "HTTP/1.1".to_string(),
2445 tls_version: None,
2446 tls_cipher: None,
2447 route_id: route_id.clone(),
2448 upstream_id: upstream_id.clone(),
2449 timestamp: chrono::Utc::now().to_rfc3339(),
2450 traceparent,
2451 },
2452 route_id,
2453 upstream_id,
2454 request_body: None,
2455 response_body: None,
2456 };
2457
2458 agent_mgr
2459 .process_response_body_streaming(
2460 &agent_ctx,
2461 &buffer,
2462 true, chunk_index,
2464 buffer.len(),
2465 total_size,
2466 &agent_ids,
2467 )
2468 .await
2469 })
2470 });
2471
2472 match result {
2473 Ok(decision) => {
2474 if let Some(mutation) = decision.response_body_mutation {
2476 if let Some(ref data) = mutation.data {
2477 if !data.is_empty() {
2478 if let Ok(decoded) = base64::Engine::decode(
2480 &base64::engine::general_purpose::STANDARD,
2481 data,
2482 ) {
2483 debug!(
2484 correlation_id = %ctx.trace_id,
2485 original_size = buffer.len(),
2486 new_size = decoded.len(),
2487 "Agent replaced response body"
2488 );
2489 *body = Some(Bytes::from(decoded));
2490 ctx.response_agent_body_complete = true;
2491 } else {
2492 warn!(
2493 correlation_id = %ctx.trace_id,
2494 "Failed to decode agent response body mutation (invalid base64)"
2495 );
2496 }
2497 }
2498 }
2500 }
2502
2503 }
2507 Err(e) => {
2508 warn!(
2509 correlation_id = %ctx.trace_id,
2510 error = %e,
2511 "Agent response body processing failed, passing through original"
2512 );
2513 }
2514 }
2515 } else if !end_of_stream {
2516 *body = None;
2518 return Ok(None);
2519 }
2520 }
2521
2522 if let Some(ref chunk) = body {
2524 ctx.response_bytes += chunk.len() as u64;
2525
2526 trace!(
2527 correlation_id = %ctx.trace_id,
2528 chunk_size = chunk.len(),
2529 total_response_bytes = ctx.response_bytes,
2530 end_of_stream = end_of_stream,
2531 "Processing response body chunk"
2532 );
2533
2534 if let Some(ref mut counter) = ctx.inference_streaming_counter {
2536 let result = counter.process_chunk(chunk);
2537
2538 if result.content.is_some() || result.is_done {
2539 trace!(
2540 correlation_id = %ctx.trace_id,
2541 has_content = result.content.is_some(),
2542 is_done = result.is_done,
2543 chunks_processed = counter.chunks_processed(),
2544 accumulated_content_len = counter.content().len(),
2545 "Processed SSE chunk for token counting"
2546 );
2547 }
2548 }
2549
2550 if ctx.response_body_inspection_enabled
2554 && !ctx.response_body_inspection_agents.is_empty()
2555 {
2556 let config = ctx
2557 .config
2558 .get_or_insert_with(|| self.config_manager.current());
2559 let max_inspection_bytes = config
2560 .waf
2561 .as_ref()
2562 .map(|w| w.body_inspection.max_inspection_bytes as u64)
2563 .unwrap_or(1024 * 1024);
2564
2565 if ctx.response_body_bytes_inspected < max_inspection_bytes {
2566 let bytes_to_inspect = std::cmp::min(
2567 chunk.len() as u64,
2568 max_inspection_bytes - ctx.response_body_bytes_inspected,
2569 ) as usize;
2570
2571 ctx.response_body_bytes_inspected += bytes_to_inspect as u64;
2575 ctx.response_body_chunk_index += 1;
2576
2577 trace!(
2578 correlation_id = %ctx.trace_id,
2579 bytes_inspected = ctx.response_body_bytes_inspected,
2580 max_inspection_bytes = max_inspection_bytes,
2581 chunk_index = ctx.response_body_chunk_index,
2582 "Tracking response body for inspection"
2583 );
2584 }
2585 }
2586 }
2587
2588 if end_of_stream {
2589 trace!(
2590 correlation_id = %ctx.trace_id,
2591 total_response_bytes = ctx.response_bytes,
2592 response_bytes_inspected = ctx.response_body_bytes_inspected,
2593 "Response body complete"
2594 );
2595 }
2596
2597 Ok(None)
2599 }
2600
2601 async fn connected_to_upstream(
2604 &self,
2605 _session: &mut Session,
2606 reused: bool,
2607 peer: &HttpPeer,
2608 #[cfg(unix)] _fd: RawFd,
2609 #[cfg(windows)] _sock: std::os::windows::io::RawSocket,
2610 digest: Option<&Digest>,
2611 ctx: &mut Self::CTX,
2612 ) -> Result<(), Box<Error>> {
2613 ctx.connection_reused = reused;
2615
2616 if reused {
2618 trace!(
2619 correlation_id = %ctx.trace_id,
2620 upstream = ctx.upstream.as_deref().unwrap_or("unknown"),
2621 peer_address = %peer.address(),
2622 "Reusing existing upstream connection"
2623 );
2624 } else {
2625 debug!(
2626 correlation_id = %ctx.trace_id,
2627 upstream = ctx.upstream.as_deref().unwrap_or("unknown"),
2628 peer_address = %peer.address(),
2629 ssl = digest.as_ref().map(|d| d.ssl_digest.is_some()).unwrap_or(false),
2630 "Established new upstream connection"
2631 );
2632 }
2633
2634 Ok(())
2635 }
2636
2637 fn request_cache_filter(&self, session: &mut Session, ctx: &mut Self::CTX) -> Result<()> {
2647 let route_id = match ctx.route_id.as_deref() {
2649 Some(id) => id,
2650 None => {
2651 trace!(
2652 correlation_id = %ctx.trace_id,
2653 "Cache filter: no route ID, skipping cache"
2654 );
2655 return Ok(());
2656 }
2657 };
2658
2659 if !self.cache_manager.is_enabled(route_id) {
2661 ctx.cache_status = Some(super::context::CacheStatus::Bypass("disabled"));
2662 trace!(
2663 correlation_id = %ctx.trace_id,
2664 route_id = %route_id,
2665 "Cache disabled for route"
2666 );
2667 return Ok(());
2668 }
2669
2670 if !self
2672 .cache_manager
2673 .is_method_cacheable(route_id, &ctx.method)
2674 {
2675 ctx.cache_status = Some(super::context::CacheStatus::Bypass("method"));
2676 trace!(
2677 correlation_id = %ctx.trace_id,
2678 route_id = %route_id,
2679 method = %ctx.method,
2680 "Method not cacheable"
2681 );
2682 return Ok(());
2683 }
2684
2685 if !self.cache_manager.is_path_cacheable(route_id, &ctx.path) {
2687 ctx.cache_status = Some(super::context::CacheStatus::Bypass("excluded"));
2688 trace!(
2689 correlation_id = %ctx.trace_id,
2690 route_id = %route_id,
2691 path = %ctx.path,
2692 "Path excluded from caching"
2693 );
2694 return Ok(());
2695 }
2696
2697 debug!(
2699 correlation_id = %ctx.trace_id,
2700 route_id = %route_id,
2701 method = %ctx.method,
2702 path = %ctx.path,
2703 "Enabling HTTP caching for request"
2704 );
2705
2706 let storage = get_cache_storage();
2708 let eviction = get_cache_eviction();
2709 let cache_lock = get_cache_lock();
2710
2711 session.cache.enable(
2713 storage,
2714 Some(eviction),
2715 None, Some(cache_lock),
2717 None, );
2719
2720 ctx.cache_eligible = true;
2722
2723 trace!(
2724 correlation_id = %ctx.trace_id,
2725 route_id = %route_id,
2726 cache_enabled = session.cache.enabled(),
2727 "Cache enabled for request"
2728 );
2729
2730 Ok(())
2731 }
2732
2733 fn cache_key_callback(&self, session: &Session, ctx: &mut Self::CTX) -> Result<CacheKey> {
2738 let req_header = session.req_header();
2739 let method = req_header.method.as_str();
2740 let path = req_header.uri.path();
2741 let host = ctx.host.as_deref().unwrap_or("unknown");
2742 let query = req_header.uri.query();
2743
2744 let key_string = crate::cache::CacheManager::generate_cache_key(method, host, path, query);
2746
2747 trace!(
2748 correlation_id = %ctx.trace_id,
2749 cache_key = %key_string,
2750 "Generated cache key"
2751 );
2752
2753 Ok(CacheKey::new("", format!("{}", req_header.uri), ""))
2756 }
2757
2758 fn cache_miss(&self, session: &mut Session, ctx: &mut Self::CTX) {
2763 session.cache.cache_miss();
2765
2766 ctx.cache_status = Some(super::context::CacheStatus::Miss);
2767
2768 if let Some(route_id) = ctx.route_id.as_deref() {
2770 self.cache_manager.stats().record_miss();
2771
2772 trace!(
2773 correlation_id = %ctx.trace_id,
2774 route_id = %route_id,
2775 path = %ctx.path,
2776 "Cache miss"
2777 );
2778 }
2779 }
2780
2781 async fn cache_hit_filter(
2787 &self,
2788 session: &mut Session,
2789 meta: &CacheMeta,
2790 hit_handler: &mut HitHandler,
2791 is_fresh: bool,
2792 ctx: &mut Self::CTX,
2793 ) -> Result<Option<ForcedFreshness>>
2794 where
2795 Self::CTX: Send + Sync,
2796 {
2797 let req_header = session.req_header();
2799 let method = req_header.method.as_str();
2800 let path = req_header.uri.path();
2801 let host = req_header.uri.host().unwrap_or("localhost");
2802 let query = req_header.uri.query();
2803
2804 let cache_key = crate::cache::CacheManager::generate_cache_key(method, host, path, query);
2806
2807 if self.cache_manager.should_invalidate(&cache_key) {
2809 info!(
2810 correlation_id = %ctx.trace_id,
2811 route_id = ctx.route_id.as_deref().unwrap_or("unknown"),
2812 cache_key = %cache_key,
2813 "Cache entry invalidated by purge request"
2814 );
2815 return Ok(Some(ForcedFreshness::ForceExpired));
2817 }
2818
2819 if is_fresh {
2821 let is_disk_hit = hit_handler
2823 .as_any()
2824 .downcast_ref::<HybridHitHandler>()
2825 .is_some()
2826 || hit_handler
2827 .as_any()
2828 .downcast_ref::<DiskHitHandler>()
2829 .is_some();
2830
2831 let stats = self.cache_manager.stats();
2832 if is_disk_hit {
2833 ctx.cache_status = Some(super::context::CacheStatus::HitDisk);
2834 stats.record_disk_hit();
2835 } else {
2836 ctx.cache_status = Some(super::context::CacheStatus::HitMemory);
2837 stats.record_memory_hit();
2838 }
2839 stats.record_hit();
2840
2841 debug!(
2842 correlation_id = %ctx.trace_id,
2843 route_id = ctx.route_id.as_deref().unwrap_or("unknown"),
2844 is_fresh = is_fresh,
2845 tier = if is_disk_hit { "disk" } else { "memory" },
2846 "Cache hit (fresh)"
2847 );
2848 } else {
2849 ctx.cache_status = Some(super::context::CacheStatus::HitStale);
2850
2851 trace!(
2852 correlation_id = %ctx.trace_id,
2853 route_id = ctx.route_id.as_deref().unwrap_or("unknown"),
2854 is_fresh = is_fresh,
2855 "Cache hit (stale)"
2856 );
2857 }
2858
2859 Ok(None)
2861 }
2862
2863 fn response_cache_filter(
2868 &self,
2869 _session: &Session,
2870 resp: &ResponseHeader,
2871 ctx: &mut Self::CTX,
2872 ) -> Result<RespCacheable> {
2873 let route_id = match ctx.route_id.as_deref() {
2874 Some(id) => id,
2875 None => {
2876 return Ok(RespCacheable::Uncacheable(NoCacheReason::Custom(
2877 "no_route",
2878 )));
2879 }
2880 };
2881
2882 if !self.cache_manager.is_enabled(route_id) {
2884 return Ok(RespCacheable::Uncacheable(NoCacheReason::Custom(
2885 "disabled",
2886 )));
2887 }
2888
2889 let status = resp.status.as_u16();
2890
2891 if !self.cache_manager.is_status_cacheable(route_id, status) {
2893 trace!(
2894 correlation_id = %ctx.trace_id,
2895 route_id = %route_id,
2896 status = status,
2897 "Status code not cacheable"
2898 );
2899 return Ok(RespCacheable::Uncacheable(NoCacheReason::OriginNotCache));
2900 }
2901
2902 if let Some(cache_control) = resp.headers.get("cache-control") {
2904 if let Ok(cc_str) = cache_control.to_str() {
2905 if crate::cache::CacheManager::is_no_cache(cc_str) {
2906 trace!(
2907 correlation_id = %ctx.trace_id,
2908 route_id = %route_id,
2909 cache_control = %cc_str,
2910 "Response has no-cache directive"
2911 );
2912 return Ok(RespCacheable::Uncacheable(NoCacheReason::OriginNotCache));
2913 }
2914 }
2915 }
2916
2917 let cache_control = resp
2919 .headers
2920 .get("cache-control")
2921 .and_then(|v| v.to_str().ok());
2922 let ttl = self.cache_manager.calculate_ttl(route_id, cache_control);
2923
2924 if ttl.is_zero() {
2925 trace!(
2926 correlation_id = %ctx.trace_id,
2927 route_id = %route_id,
2928 "TTL is zero, not caching"
2929 );
2930 return Ok(RespCacheable::Uncacheable(NoCacheReason::OriginNotCache));
2931 }
2932
2933 let config = self
2935 .cache_manager
2936 .get_route_config(route_id)
2937 .unwrap_or_default();
2938
2939 let now = std::time::SystemTime::now();
2941 let fresh_until = now + ttl;
2942
2943 let header = resp.clone();
2945
2946 let cache_meta = CacheMeta::new(
2948 fresh_until,
2949 now,
2950 config.stale_while_revalidate_secs as u32,
2951 config.stale_if_error_secs as u32,
2952 header,
2953 );
2954
2955 self.cache_manager.stats().record_store();
2957
2958 debug!(
2959 correlation_id = %ctx.trace_id,
2960 route_id = %route_id,
2961 status = status,
2962 ttl_secs = ttl.as_secs(),
2963 stale_while_revalidate_secs = config.stale_while_revalidate_secs,
2964 stale_if_error_secs = config.stale_if_error_secs,
2965 "Caching response"
2966 );
2967
2968 Ok(RespCacheable::Cacheable(cache_meta))
2969 }
2970
2971 fn should_serve_stale(
2975 &self,
2976 _session: &mut Session,
2977 ctx: &mut Self::CTX,
2978 error: Option<&Error>,
2979 ) -> bool {
2980 let route_id = match ctx.route_id.as_deref() {
2981 Some(id) => id,
2982 None => return false,
2983 };
2984
2985 let config = match self.cache_manager.get_route_config(route_id) {
2987 Some(c) => c,
2988 None => return false,
2989 };
2990
2991 if let Some(e) = error {
2993 if e.esource() == &pingora::ErrorSource::Upstream {
2995 debug!(
2996 correlation_id = %ctx.trace_id,
2997 route_id = %route_id,
2998 error = %e,
2999 stale_if_error_secs = config.stale_if_error_secs,
3000 "Considering stale-if-error"
3001 );
3002 return config.stale_if_error_secs > 0;
3003 }
3004 }
3005
3006 if error.is_none() && config.stale_while_revalidate_secs > 0 {
3008 trace!(
3009 correlation_id = %ctx.trace_id,
3010 route_id = %route_id,
3011 stale_while_revalidate_secs = config.stale_while_revalidate_secs,
3012 "Allowing stale-while-revalidate"
3013 );
3014 return true;
3015 }
3016
3017 false
3018 }
3019
3020 fn range_header_filter(
3030 &self,
3031 session: &mut Session,
3032 response: &mut ResponseHeader,
3033 ctx: &mut Self::CTX,
3034 ) -> pingora_proxy::RangeType
3035 where
3036 Self::CTX: Send + Sync,
3037 {
3038 let supports_range = ctx.route_config.as_ref().is_none_or(|config| {
3040 matches!(
3042 config.service_type,
3043 zentinel_config::ServiceType::Static | zentinel_config::ServiceType::Web
3044 )
3045 });
3046
3047 if !supports_range {
3048 trace!(
3049 correlation_id = %ctx.trace_id,
3050 route_id = ctx.route_id.as_deref().unwrap_or("unknown"),
3051 "Range request not supported for this route type"
3052 );
3053 return pingora_proxy::RangeType::None;
3054 }
3055
3056 let range_type = pingora_proxy::range_header_filter(session.req_header(), response, None);
3058
3059 match &range_type {
3060 pingora_proxy::RangeType::None => {
3061 trace!(
3062 correlation_id = %ctx.trace_id,
3063 route_id = ctx.route_id.as_deref().unwrap_or("unknown"),
3064 "No range request or not applicable"
3065 );
3066 }
3067 pingora_proxy::RangeType::Single(range) => {
3068 trace!(
3069 correlation_id = %ctx.trace_id,
3070 route_id = ctx.route_id.as_deref().unwrap_or("unknown"),
3071 range_start = range.start,
3072 range_end = range.end,
3073 "Processing single-range request"
3074 );
3075 }
3076 pingora_proxy::RangeType::Multi(multi) => {
3077 trace!(
3078 correlation_id = %ctx.trace_id,
3079 route_id = ctx.route_id.as_deref().unwrap_or("unknown"),
3080 range_count = multi.ranges.len(),
3081 "Processing multi-range request"
3082 );
3083 }
3084 pingora_proxy::RangeType::Invalid => {
3085 debug!(
3086 correlation_id = %ctx.trace_id,
3087 route_id = ctx.route_id.as_deref().unwrap_or("unknown"),
3088 "Invalid range header"
3089 );
3090 }
3091 }
3092
3093 range_type
3094 }
3095
3096 async fn fail_to_proxy(
3099 &self,
3100 session: &mut Session,
3101 e: &Error,
3102 ctx: &mut Self::CTX,
3103 ) -> pingora_proxy::FailToProxy
3104 where
3105 Self::CTX: Send + Sync,
3106 {
3107 let error_code = match e.etype() {
3108 ErrorType::ConnectRefused => 503,
3110 ErrorType::ConnectTimedout => 504,
3111 ErrorType::ConnectNoRoute => 502,
3112
3113 ErrorType::ReadTimedout => 504,
3115 ErrorType::WriteTimedout => 504,
3116
3117 ErrorType::TLSHandshakeFailure => 502,
3119 ErrorType::InvalidCert => 502,
3120
3121 ErrorType::InvalidHTTPHeader => 400,
3123 ErrorType::H2Error => 502,
3124
3125 ErrorType::ConnectProxyFailure => 502,
3127 ErrorType::ConnectionClosed => 502,
3128
3129 ErrorType::HTTPStatus(status) => *status,
3131
3132 ErrorType::InternalError => 500,
3135
3136 _ => 502,
3138 };
3139
3140 error!(
3141 correlation_id = %ctx.trace_id,
3142 route_id = ctx.route_id.as_deref().unwrap_or("unknown"),
3143 upstream = ctx.upstream.as_deref().unwrap_or("unknown"),
3144 error_type = ?e.etype(),
3145 error = %e,
3146 error_code = error_code,
3147 "Proxy error occurred"
3148 );
3149
3150 self.metrics
3152 .record_blocked_request(&format!("proxy_error_{}", error_code));
3153
3154 let error_message = match error_code {
3158 400 => "Bad Request",
3159 502 => "Bad Gateway",
3160 503 => "Service Unavailable",
3161 504 => "Gateway Timeout",
3162 _ => "Internal Server Error",
3163 };
3164
3165 let body = format!(
3167 r#"{{"error":"{} {}","trace_id":"{}"}}"#,
3168 error_code, error_message, ctx.trace_id
3169 );
3170
3171 let mut header = pingora::http::ResponseHeader::build(error_code, None).unwrap();
3173 header
3174 .insert_header("Content-Type", "application/json")
3175 .ok();
3176 header
3177 .insert_header("Content-Length", body.len().to_string())
3178 .ok();
3179 header
3180 .insert_header("X-Correlation-Id", ctx.trace_id.as_str())
3181 .ok();
3182 header.insert_header("Connection", "close").ok();
3183
3184 if let Err(write_err) = session.write_response_header(Box::new(header), false).await {
3186 warn!(
3187 correlation_id = %ctx.trace_id,
3188 error = %write_err,
3189 "Failed to write error response header"
3190 );
3191 } else {
3192 if let Err(write_err) = session
3194 .write_response_body(Some(bytes::Bytes::from(body)), true)
3195 .await
3196 {
3197 warn!(
3198 correlation_id = %ctx.trace_id,
3199 error = %write_err,
3200 "Failed to write error response body"
3201 );
3202 }
3203 }
3204
3205 pingora_proxy::FailToProxy {
3208 error_code,
3209 can_reuse_downstream: false,
3210 }
3211 }
3212
3213 fn error_while_proxy(
3219 &self,
3220 peer: &HttpPeer,
3221 session: &mut Session,
3222 e: Box<Error>,
3223 ctx: &mut Self::CTX,
3224 client_reused: bool,
3225 ) -> Box<Error> {
3226 let error_type = e.etype().clone();
3227 let upstream_id = ctx.upstream.as_deref().unwrap_or("unknown");
3228
3229 let is_retryable = matches!(
3231 error_type,
3232 ErrorType::ConnectTimedout
3233 | ErrorType::ReadTimedout
3234 | ErrorType::WriteTimedout
3235 | ErrorType::ConnectionClosed
3236 | ErrorType::ConnectRefused
3237 );
3238
3239 warn!(
3241 correlation_id = %ctx.trace_id,
3242 route_id = ctx.route_id.as_deref().unwrap_or("unknown"),
3243 upstream = %upstream_id,
3244 peer_address = %peer.address(),
3245 error_type = ?error_type,
3246 error = %e,
3247 client_reused = client_reused,
3248 is_retryable = is_retryable,
3249 "Error during proxy operation"
3250 );
3251
3252 let peer_address = peer.address().to_string();
3255 let upstream_pools = self.upstream_pools.clone();
3256 let upstream_id_owned = upstream_id.to_string();
3257 tokio::spawn(async move {
3258 if let Some(pool) = upstream_pools.get(&upstream_id_owned).await {
3259 pool.report_result(&peer_address, false).await;
3260 }
3261 });
3262
3263 self.metrics
3265 .record_blocked_request(&format!("proxy_error_{:?}", error_type));
3266
3267 let mut enhanced_error = e.more_context(format!(
3269 "Upstream: {}, Peer: {}, Attempts: {}",
3270 upstream_id,
3271 peer.address(),
3272 ctx.upstream_attempts
3273 ));
3274
3275 if is_retryable {
3280 let can_retry = if client_reused {
3281 !session.as_ref().retry_buffer_truncated()
3283 } else {
3284 true
3286 };
3287
3288 enhanced_error.retry.decide_reuse(can_retry);
3289
3290 if can_retry {
3291 debug!(
3292 correlation_id = %ctx.trace_id,
3293 upstream = %upstream_id,
3294 error_type = ?error_type,
3295 "Error is retryable, will attempt retry"
3296 );
3297 }
3298 } else {
3299 enhanced_error.retry.decide_reuse(false);
3301 }
3302
3303 enhanced_error
3304 }
3305
3306 async fn logging(&self, session: &mut Session, _error: Option<&Error>, ctx: &mut Self::CTX) {
3307 self.reload_coordinator.dec_requests();
3309
3310 if !ctx.shadow_sent {
3312 if let Some(shadow_pending) = ctx.shadow_pending.take() {
3313 let body = if shadow_pending.include_body && !ctx.body_buffer.is_empty() {
3314 Some(ctx.body_buffer.clone())
3316 } else {
3317 None
3318 };
3319
3320 trace!(
3321 correlation_id = %ctx.trace_id,
3322 body_size = body.as_ref().map(|b| b.len()).unwrap_or(0),
3323 "Firing deferred shadow request with buffered body"
3324 );
3325
3326 shadow_pending.manager.shadow_request(
3327 shadow_pending.headers,
3328 body,
3329 shadow_pending.request_ctx,
3330 );
3331 ctx.shadow_sent = true;
3332 }
3333 }
3334
3335 let duration = ctx.elapsed();
3336
3337 let status = session
3339 .response_written()
3340 .map(|r| r.status.as_u16())
3341 .unwrap_or(0);
3342
3343 if let (Some(ref peer_addr), Some(ref upstream_id)) =
3346 (&ctx.selected_upstream_address, &ctx.upstream)
3347 {
3348 let success = status > 0 && status < 500;
3350
3351 if let Some(pool) = self.upstream_pools.get(upstream_id).await {
3352 pool.report_result_with_latency(peer_addr, success, Some(duration))
3353 .await;
3354 pool.decrement_active();
3355 trace!(
3356 correlation_id = %ctx.trace_id,
3357 upstream = %upstream_id,
3358 peer_address = %peer_addr,
3359 success = success,
3360 duration_ms = duration.as_millis(),
3361 status = status,
3362 "Reported result to adaptive load balancer"
3363 );
3364 }
3365
3366 if ctx.inference_rate_limit_enabled && success {
3368 let cold_detected = self.warmth_tracker.record_request(peer_addr, duration);
3369 if cold_detected {
3370 debug!(
3371 correlation_id = %ctx.trace_id,
3372 upstream = %upstream_id,
3373 peer_address = %peer_addr,
3374 duration_ms = duration.as_millis(),
3375 "Cold model detected on inference upstream"
3376 );
3377 }
3378 }
3379 }
3380
3381 if ctx.inference_rate_limit_enabled {
3384 if let (Some(route_id), Some(ref rate_limit_key)) =
3385 (ctx.route_id.as_deref(), &ctx.inference_rate_limit_key)
3386 {
3387 let response_headers = session
3389 .response_written()
3390 .map(|r| &r.headers)
3391 .cloned()
3392 .unwrap_or_default();
3393
3394 let streaming_result = if ctx.inference_streaming_response {
3396 ctx.inference_streaming_counter
3397 .as_ref()
3398 .map(|counter| counter.finalize())
3399 } else {
3400 None
3401 };
3402
3403 if let Some(ref result) = streaming_result {
3405 debug!(
3406 correlation_id = %ctx.trace_id,
3407 output_tokens = result.output_tokens,
3408 input_tokens = ?result.input_tokens,
3409 source = ?result.source,
3410 content_length = result.content_length,
3411 "Finalized streaming token count"
3412 );
3413 }
3414
3415 if ctx.inference_streaming_response {
3417 if let Some(ref route_config) = ctx.route_config {
3418 if let Some(ref inference) = route_config.inference {
3419 if let Some(ref guardrails) = inference.guardrails {
3420 if let Some(ref pii_config) = guardrails.pii_detection {
3421 if pii_config.enabled {
3422 if let Some(ref counter) = ctx.inference_streaming_counter {
3424 let response_content = counter.content();
3425 if !response_content.is_empty() {
3426 let pii_result = self
3427 .guardrail_processor
3428 .check_pii(
3429 pii_config,
3430 response_content,
3431 ctx.route_id.as_deref(),
3432 &ctx.trace_id,
3433 )
3434 .await;
3435
3436 match pii_result {
3437 crate::inference::PiiCheckResult::Detected {
3438 detections,
3439 redacted_content: _,
3440 } => {
3441 warn!(
3442 correlation_id = %ctx.trace_id,
3443 route_id = ctx.route_id.as_deref().unwrap_or("unknown"),
3444 detection_count = detections.len(),
3445 "PII detected in inference response"
3446 );
3447
3448 ctx.pii_detection_categories = detections
3450 .iter()
3451 .map(|d| d.category.clone())
3452 .collect();
3453
3454 for detection in &detections {
3456 self.metrics.record_pii_detected(
3457 ctx.route_id.as_deref().unwrap_or("unknown"),
3458 &detection.category,
3459 );
3460 }
3461 }
3462 crate::inference::PiiCheckResult::Clean => {
3463 trace!(
3464 correlation_id = %ctx.trace_id,
3465 "No PII detected in response"
3466 );
3467 }
3468 crate::inference::PiiCheckResult::Error { message } => {
3469 debug!(
3470 correlation_id = %ctx.trace_id,
3471 error = %message,
3472 "PII detection check failed"
3473 );
3474 }
3475 }
3476 }
3477 }
3478 }
3479 }
3480 }
3481 }
3482 }
3483 }
3484
3485 let empty_body: &[u8] = &[];
3489
3490 if let Some(actual_estimate) = self.inference_rate_limit_manager.record_actual(
3491 route_id,
3492 rate_limit_key,
3493 &response_headers,
3494 empty_body,
3495 ctx.inference_estimated_tokens,
3496 ) {
3497 let (actual_tokens, source_info) = if let Some(ref streaming) = streaming_result
3499 {
3500 if let Some(total_tokens) = streaming.total_tokens {
3502 (total_tokens, "streaming_api")
3503 } else if actual_estimate.source == crate::inference::TokenSource::Estimated
3504 {
3505 let total = ctx.inference_input_tokens + streaming.output_tokens;
3508 (total, "streaming_tiktoken")
3509 } else {
3510 (actual_estimate.tokens, "headers")
3511 }
3512 } else {
3513 (actual_estimate.tokens, "headers")
3514 };
3515
3516 ctx.inference_actual_tokens = Some(actual_tokens);
3517
3518 debug!(
3519 correlation_id = %ctx.trace_id,
3520 route_id = route_id,
3521 estimated_tokens = ctx.inference_estimated_tokens,
3522 actual_tokens = actual_tokens,
3523 source = source_info,
3524 streaming_response = ctx.inference_streaming_response,
3525 model = ?ctx.inference_model,
3526 "Recorded actual inference tokens"
3527 );
3528
3529 if ctx.inference_budget_enabled {
3531 let alerts = self.inference_rate_limit_manager.record_budget(
3532 route_id,
3533 rate_limit_key,
3534 actual_tokens,
3535 );
3536
3537 for alert in alerts.iter() {
3539 warn!(
3540 correlation_id = %ctx.trace_id,
3541 route_id = route_id,
3542 tenant = %alert.tenant,
3543 threshold_pct = alert.threshold * 100.0,
3544 tokens_used = alert.tokens_used,
3545 tokens_limit = alert.tokens_limit,
3546 "Token budget alert threshold crossed"
3547 );
3548 }
3549
3550 if let Some(status) = self
3552 .inference_rate_limit_manager
3553 .budget_status(route_id, rate_limit_key)
3554 {
3555 ctx.inference_budget_remaining = Some(status.tokens_remaining as i64);
3556 }
3557 }
3558
3559 if ctx.inference_cost_enabled {
3561 if let Some(model) = ctx.inference_model.as_deref() {
3562 let (input_tokens, output_tokens) = if let Some(ref streaming) =
3564 streaming_result
3565 {
3566 let input =
3568 streaming.input_tokens.unwrap_or(ctx.inference_input_tokens);
3569 let output = streaming.output_tokens;
3570 (input, output)
3571 } else {
3572 let input = ctx.inference_input_tokens;
3574 let output = actual_tokens.saturating_sub(input);
3575 (input, output)
3576 };
3577 ctx.inference_output_tokens = output_tokens;
3578
3579 if let Some(cost_result) = self
3580 .inference_rate_limit_manager
3581 .calculate_cost(route_id, model, input_tokens, output_tokens)
3582 {
3583 ctx.inference_request_cost = Some(cost_result.total_cost);
3584
3585 trace!(
3586 correlation_id = %ctx.trace_id,
3587 route_id = route_id,
3588 model = model,
3589 input_tokens = input_tokens,
3590 output_tokens = output_tokens,
3591 total_cost = cost_result.total_cost,
3592 currency = %cost_result.currency,
3593 "Calculated inference request cost"
3594 );
3595 }
3596 }
3597 }
3598 }
3599 }
3600 }
3601
3602 if self.log_manager.should_log_access(status) {
3604 let access_entry = AccessLogEntry {
3605 timestamp: chrono::Utc::now().to_rfc3339(),
3606 trace_id: ctx.trace_id.clone(),
3607 method: ctx.method.clone(),
3608 path: ctx.path.clone(),
3609 query: ctx.query.clone(),
3610 protocol: "HTTP/1.1".to_string(),
3611 status,
3612 body_bytes: ctx.response_bytes,
3613 duration_ms: duration.as_millis() as u64,
3614 client_ip: ctx.client_ip.clone(),
3615 user_agent: ctx.user_agent.clone(),
3616 referer: ctx.referer.clone(),
3617 host: ctx.host.clone(),
3618 route_id: ctx.route_id.clone(),
3619 upstream: ctx.upstream.clone(),
3620 upstream_attempts: ctx.upstream_attempts,
3621 instance_id: self.app_state.instance_id.clone(),
3622 namespace: ctx.namespace.clone(),
3623 service: ctx.service.clone(),
3624 body_bytes_sent: ctx.response_bytes,
3626 upstream_addr: ctx.selected_upstream_address.clone(),
3627 connection_reused: ctx.connection_reused,
3628 rate_limit_hit: status == 429,
3629 geo_country: ctx.geo_country_code.clone(),
3630 };
3631 self.log_manager.log_access(&access_entry);
3632 }
3633
3634 if tracing::enabled!(tracing::Level::DEBUG) {
3636 let write_pending_ms = session.upstream_write_pending_time().as_millis() as u64;
3638 debug!(
3639 trace_id = %ctx.trace_id,
3640 method = %ctx.method,
3641 path = %ctx.path,
3642 route_id = ?ctx.route_id,
3643 upstream = ?ctx.upstream,
3644 status = status,
3645 duration_ms = duration.as_millis() as u64,
3646 upstream_write_pending_ms = write_pending_ms,
3647 upstream_attempts = ctx.upstream_attempts,
3648 error = ?_error.map(|e| e.to_string()),
3649 "Request completed"
3650 );
3651 }
3652
3653 if ctx.is_websocket_upgrade && status == 101 {
3655 info!(
3656 trace_id = %ctx.trace_id,
3657 route_id = ?ctx.route_id,
3658 upstream = ?ctx.upstream,
3659 client_ip = %ctx.client_ip,
3660 "WebSocket connection established"
3661 );
3662 }
3663
3664 if let Some(span) = ctx.otel_span.take() {
3666 span.end();
3667 }
3668 }
3669}
3670
3671impl ZentinelProxy {
3676 async fn process_body_chunk_streaming(
3678 &self,
3679 body: &mut Option<Bytes>,
3680 end_of_stream: bool,
3681 ctx: &mut RequestContext,
3682 ) -> Result<(), Box<Error>> {
3683 let chunk_data: Vec<u8> = body.as_ref().map(|b| b.to_vec()).unwrap_or_default();
3685 let chunk_index = ctx.request_body_chunk_index;
3686 ctx.request_body_chunk_index += 1;
3687 ctx.body_bytes_inspected += chunk_data.len() as u64;
3688
3689 debug!(
3690 correlation_id = %ctx.trace_id,
3691 chunk_index = chunk_index,
3692 chunk_size = chunk_data.len(),
3693 end_of_stream = end_of_stream,
3694 "Streaming body chunk to agents"
3695 );
3696
3697 let agent_ctx = crate::agents::AgentCallContext {
3699 correlation_id: zentinel_common::CorrelationId::from_string(&ctx.trace_id),
3700 metadata: zentinel_agent_protocol::RequestMetadata {
3701 correlation_id: ctx.trace_id.clone(),
3702 request_id: ctx.trace_id.clone(),
3703 client_ip: ctx.client_ip.clone(),
3704 client_port: 0,
3705 server_name: ctx.host.clone(),
3706 protocol: "HTTP/1.1".to_string(),
3707 tls_version: None,
3708 tls_cipher: None,
3709 route_id: ctx.route_id.clone(),
3710 upstream_id: ctx.upstream.clone(),
3711 timestamp: chrono::Utc::now().to_rfc3339(),
3712 traceparent: ctx.traceparent(),
3713 },
3714 route_id: ctx.route_id.clone(),
3715 upstream_id: ctx.upstream.clone(),
3716 request_body: None, response_body: None,
3718 };
3719
3720 let agent_ids = ctx.body_inspection_agents.clone();
3721 let total_size = None; match self
3724 .agent_manager
3725 .process_request_body_streaming(
3726 &agent_ctx,
3727 &chunk_data,
3728 end_of_stream,
3729 chunk_index,
3730 ctx.body_bytes_inspected as usize,
3731 total_size,
3732 &agent_ids,
3733 )
3734 .await
3735 {
3736 Ok(decision) => {
3737 ctx.agent_needs_more = decision.needs_more;
3739
3740 if let Some(ref mutation) = decision.request_body_mutation {
3742 if !mutation.is_pass_through() {
3743 if mutation.is_drop() {
3744 *body = None;
3746 trace!(
3747 correlation_id = %ctx.trace_id,
3748 chunk_index = chunk_index,
3749 "Agent dropped body chunk"
3750 );
3751 } else if let Some(ref new_data) = mutation.data {
3752 *body = Some(Bytes::from(new_data.clone()));
3754 trace!(
3755 correlation_id = %ctx.trace_id,
3756 chunk_index = chunk_index,
3757 original_size = chunk_data.len(),
3758 new_size = new_data.len(),
3759 "Agent mutated body chunk"
3760 );
3761 }
3762 }
3763 }
3764
3765 if !decision.needs_more && !decision.is_allow() {
3767 warn!(
3768 correlation_id = %ctx.trace_id,
3769 action = ?decision.action,
3770 "Agent blocked request body"
3771 );
3772 self.metrics.record_blocked_request("agent_body_inspection");
3773
3774 let (status, message) = match &decision.action {
3775 crate::agents::AgentAction::Block { status, body, .. } => (
3776 *status,
3777 body.clone().unwrap_or_else(|| "Blocked".to_string()),
3778 ),
3779 _ => (403, "Forbidden".to_string()),
3780 };
3781
3782 return Err(Error::explain(ErrorType::HTTPStatus(status), message));
3783 }
3784
3785 trace!(
3786 correlation_id = %ctx.trace_id,
3787 needs_more = decision.needs_more,
3788 "Agent processed body chunk"
3789 );
3790 }
3791 Err(e) => {
3792 let fail_closed = ctx
3793 .route_config
3794 .as_ref()
3795 .map(|r| r.policies.failure_mode == zentinel_config::FailureMode::Closed)
3796 .unwrap_or(false);
3797
3798 if fail_closed {
3799 error!(
3800 correlation_id = %ctx.trace_id,
3801 error = %e,
3802 "Agent streaming body inspection failed, blocking (fail-closed)"
3803 );
3804 self.log_manager.log_request_error(
3805 "error",
3806 "Agent streaming body inspection failed, blocking (fail-closed)",
3807 &ctx.trace_id,
3808 ctx.route_id.as_deref(),
3809 ctx.upstream.as_deref(),
3810 Some(format!("error={}", e)),
3811 );
3812 return Err(Error::explain(
3813 ErrorType::HTTPStatus(503),
3814 "Service unavailable",
3815 ));
3816 } else {
3817 warn!(
3818 correlation_id = %ctx.trace_id,
3819 error = %e,
3820 "Agent streaming body inspection failed, allowing (fail-open)"
3821 );
3822 self.log_manager.log_request_error(
3823 "warn",
3824 "Agent streaming body inspection failed, allowing (fail-open)",
3825 &ctx.trace_id,
3826 ctx.route_id.as_deref(),
3827 ctx.upstream.as_deref(),
3828 Some(format!("error={}", e)),
3829 );
3830 }
3831 }
3832 }
3833
3834 Ok(())
3835 }
3836
3837 async fn send_buffered_body_to_agents(
3839 &self,
3840 end_of_stream: bool,
3841 ctx: &mut RequestContext,
3842 ) -> Result<(), Box<Error>> {
3843 debug!(
3844 correlation_id = %ctx.trace_id,
3845 buffer_size = ctx.body_buffer.len(),
3846 end_of_stream = end_of_stream,
3847 agent_count = ctx.body_inspection_agents.len(),
3848 decompression_enabled = ctx.decompression_enabled,
3849 "Sending buffered body to agents for inspection"
3850 );
3851
3852 let body_for_inspection = if ctx.decompression_enabled {
3854 if let Some(ref encoding) = ctx.body_content_encoding {
3855 let config = crate::decompression::DecompressionConfig {
3856 max_ratio: ctx.max_decompression_ratio,
3857 max_output_bytes: ctx.max_decompression_bytes,
3858 };
3859
3860 match crate::decompression::decompress_body(&ctx.body_buffer, encoding, &config) {
3861 Ok(result) => {
3862 ctx.body_was_decompressed = true;
3863 self.metrics
3864 .record_decompression_success(encoding, result.ratio);
3865 debug!(
3866 correlation_id = %ctx.trace_id,
3867 encoding = %encoding,
3868 compressed_size = result.compressed_size,
3869 decompressed_size = result.decompressed_size,
3870 ratio = result.ratio,
3871 "Body decompressed for agent inspection"
3872 );
3873 result.data
3874 }
3875 Err(e) => {
3876 let failure_reason = match &e {
3878 crate::decompression::DecompressionError::RatioExceeded { .. } => {
3879 "ratio_exceeded"
3880 }
3881 crate::decompression::DecompressionError::SizeExceeded { .. } => {
3882 "size_exceeded"
3883 }
3884 crate::decompression::DecompressionError::InvalidData { .. } => {
3885 "invalid_data"
3886 }
3887 crate::decompression::DecompressionError::UnsupportedEncoding {
3888 ..
3889 } => "unsupported",
3890 crate::decompression::DecompressionError::IoError(_) => "io_error",
3891 };
3892 self.metrics
3893 .record_decompression_failure(encoding, failure_reason);
3894
3895 let fail_closed = ctx
3897 .route_config
3898 .as_ref()
3899 .map(|r| {
3900 r.policies.failure_mode == zentinel_config::FailureMode::Closed
3901 })
3902 .unwrap_or(false);
3903
3904 if fail_closed {
3905 error!(
3906 correlation_id = %ctx.trace_id,
3907 error = %e,
3908 encoding = %encoding,
3909 "Decompression failed, blocking (fail-closed)"
3910 );
3911 return Err(Error::explain(
3912 ErrorType::HTTPStatus(400),
3913 "Invalid compressed body",
3914 ));
3915 } else {
3916 warn!(
3917 correlation_id = %ctx.trace_id,
3918 error = %e,
3919 encoding = %encoding,
3920 "Decompression failed, sending compressed body (fail-open)"
3921 );
3922 ctx.body_buffer.clone()
3923 }
3924 }
3925 }
3926 } else {
3927 ctx.body_buffer.clone()
3928 }
3929 } else {
3930 ctx.body_buffer.clone()
3931 };
3932
3933 let agent_ctx = crate::agents::AgentCallContext {
3934 correlation_id: zentinel_common::CorrelationId::from_string(&ctx.trace_id),
3935 metadata: zentinel_agent_protocol::RequestMetadata {
3936 correlation_id: ctx.trace_id.clone(),
3937 request_id: ctx.trace_id.clone(),
3938 client_ip: ctx.client_ip.clone(),
3939 client_port: 0,
3940 server_name: ctx.host.clone(),
3941 protocol: "HTTP/1.1".to_string(),
3942 tls_version: None,
3943 tls_cipher: None,
3944 route_id: ctx.route_id.clone(),
3945 upstream_id: ctx.upstream.clone(),
3946 timestamp: chrono::Utc::now().to_rfc3339(),
3947 traceparent: ctx.traceparent(),
3948 },
3949 route_id: ctx.route_id.clone(),
3950 upstream_id: ctx.upstream.clone(),
3951 request_body: Some(body_for_inspection.clone()),
3952 response_body: None,
3953 };
3954
3955 let agent_ids = ctx.body_inspection_agents.clone();
3956 match self
3957 .agent_manager
3958 .process_request_body(&agent_ctx, &body_for_inspection, end_of_stream, &agent_ids)
3959 .await
3960 {
3961 Ok(decision) => {
3962 if !decision.is_allow() {
3963 warn!(
3964 correlation_id = %ctx.trace_id,
3965 action = ?decision.action,
3966 "Agent blocked request body"
3967 );
3968 self.metrics.record_blocked_request("agent_body_inspection");
3969
3970 let (status, message) = match &decision.action {
3971 crate::agents::AgentAction::Block { status, body, .. } => (
3972 *status,
3973 body.clone().unwrap_or_else(|| "Blocked".to_string()),
3974 ),
3975 _ => (403, "Forbidden".to_string()),
3976 };
3977
3978 return Err(Error::explain(ErrorType::HTTPStatus(status), message));
3979 }
3980
3981 trace!(
3982 correlation_id = %ctx.trace_id,
3983 "Agent allowed request body"
3984 );
3985 }
3986 Err(e) => {
3987 let fail_closed = ctx
3988 .route_config
3989 .as_ref()
3990 .map(|r| r.policies.failure_mode == zentinel_config::FailureMode::Closed)
3991 .unwrap_or(false);
3992
3993 if fail_closed {
3994 error!(
3995 correlation_id = %ctx.trace_id,
3996 error = %e,
3997 "Agent body inspection failed, blocking (fail-closed)"
3998 );
3999 self.log_manager.log_request_error(
4000 "error",
4001 "Agent body inspection failed, blocking (fail-closed)",
4002 &ctx.trace_id,
4003 ctx.route_id.as_deref(),
4004 ctx.upstream.as_deref(),
4005 Some(format!("error={}", e)),
4006 );
4007 return Err(Error::explain(
4008 ErrorType::HTTPStatus(503),
4009 "Service unavailable",
4010 ));
4011 } else {
4012 warn!(
4013 correlation_id = %ctx.trace_id,
4014 error = %e,
4015 "Agent body inspection failed, allowing (fail-open)"
4016 );
4017 self.log_manager.log_request_error(
4018 "warn",
4019 "Agent body inspection failed, allowing (fail-open)",
4020 &ctx.trace_id,
4021 ctx.route_id.as_deref(),
4022 ctx.upstream.as_deref(),
4023 Some(format!("error={}", e)),
4024 );
4025 }
4026 }
4027 }
4028
4029 Ok(())
4030 }
4031}