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