praxis_protocol/http/pingora/
context.rs

1// SPDX-License-Identifier: MIT
2// Copyright (c) 2024 Praxis Contributors
3
4//! Per-request context that carries filter pipeline results through Pingora's request/response lifecycle hooks.
5
6use std::{collections::VecDeque, net::IpAddr, sync::Arc, time::Instant};
7
8use bytes::Bytes;
9use praxis_core::connectivity::Upstream;
10use praxis_filter::{BodyBuffer, BodyMode, FilterPipeline, Request, Response, TrustedHeaderMutation};
11use tokio::sync::OwnedSemaphorePermit;
12
13// -----------------------------------------------------------------------------
14// PingoraRequestCtx
15// -----------------------------------------------------------------------------
16
17/// Per-request context carrying filter pipeline results through Pingora hooks.
18///
19/// ```
20/// use std::sync::Arc;
21///
22/// use praxis_protocol::http::pingora::context::PingoraRequestCtx;
23///
24/// let mut ctx = PingoraRequestCtx::default();
25/// ctx.cluster = Some(Arc::from("api-cluster"));
26/// assert_eq!(ctx.cluster.as_deref(), Some("api-cluster"));
27/// ```
28#[expect(clippy::struct_excessive_bools, reason = "lifecycle flags")]
29pub struct PingoraRequestCtx {
30    /// Connection permit from the per-listener semaphore.
31    ///
32    /// Held for the lifetime of the request. RAII drop
33    /// releases the permit when the context is dropped,
34    /// including error and timeout paths.
35    pub _connection_permit: Option<OwnedSemaphorePermit>,
36
37    /// Permit from the process-wide connection semaphore.
38    ///
39    /// Present only when `runtime.max_connections` is configured.
40    pub _global_connection_permit: Option<OwnedSemaphorePermit>,
41
42    /// Downstream client IP address.
43    pub client_addr: Option<IpAddr>,
44
45    /// HTTP version of the downstream client request.
46    ///
47    /// Captured during `request_filter` so the response-phase Via
48    /// header can reflect the protocol the client used.
49    pub client_http_version: Option<http::Version>,
50
51    /// Name of the cluster selected by a cluster-selecting filter.
52    pub cluster: Option<Arc<str>>,
53
54    /// Cached per-filter body-done indices. Swapped into each
55    /// [`HttpFilterContext`] and written back after execution so
56    /// that the heap allocation is reused across pipeline phases.
57    ///
58    /// [`HttpFilterContext`]: praxis_filter::HttpFilterContext
59    pub cached_body_done_indices: Vec<bool>,
60
61    /// Cached per-filter execution indices. Same lifecycle as
62    /// [`cached_body_done_indices`].
63    ///
64    /// [`cached_body_done_indices`]: Self::cached_body_done_indices
65    pub cached_executed_filter_indices: Vec<bool>,
66
67    /// Whether the downstream connection uses TLS.
68    ///
69    /// Derived from the Pingora session's SSL digest during
70    /// `request_filter`. Used by the forwarded headers filter
71    /// to set `X-Forwarded-Proto` correctly for HTTP/1.1
72    /// connections where the URI lacks a scheme.
73    pub downstream_tls: bool,
74
75    /// Verified downstream TLS peer identity.
76    ///
77    /// Set once from the SSL digest in `request_filter` before
78    /// the first filter runs.  Cloned (not moved) into each
79    /// `HttpFilterContext` so it is available in both pre-read
80    /// body phases and the main filter pipeline.  `None` for
81    /// non-mTLS or no-client-cert connections.
82    pub peer_identity: Option<praxis_tls::TlsPeerIdentity>,
83
84    /// Whether the connection was upgraded via 101 Switching Protocols.
85    ///
86    /// Set during `response_filter` when the upstream returns 101.
87    /// Body filter hooks skip processing when true, since post-upgrade
88    /// bytes are raw protocol frames (e.g. `WebSocket`), not HTTP bodies.
89    pub connection_upgraded: bool,
90
91    /// Type-safe request-scoped extension container. Swapped into each
92    /// [`HttpFilterContext`] and written back after filter execution,
93    /// following the same lifecycle as [`filter_metadata`].
94    ///
95    /// [`HttpFilterContext`]: praxis_filter::HttpFilterContext
96    /// [`filter_metadata`]: PingoraRequestCtx::filter_metadata
97    pub extensions: praxis_filter::RequestExtensions,
98
99    /// Durable per-request metadata that persists across all lifecycle
100    /// phases. Swapped into each [`HttpFilterContext`] and written back
101    /// after filter execution.
102    ///
103    /// [`HttpFilterContext`]: praxis_filter::HttpFilterContext
104    pub filter_metadata: std::collections::HashMap<String, String>,
105
106    /// Ordered log of trusted header mutations from pre-read body
107    /// processing. Swapped into each [`HttpFilterContext`] and written
108    /// back after filter execution.
109    ///
110    /// [`HttpFilterContext`]: praxis_filter::HttpFilterContext
111    pub pre_read_mutations: Vec<TrustedHeaderMutation>,
112
113    /// Structured per-request metadata keyed by namespace.
114    /// Swapped into each [`HttpFilterContext`] and written back
115    /// after filter execution.
116    ///
117    /// [`HttpFilterContext`]: praxis_filter::HttpFilterContext
118    pub structured_metadata: std::collections::HashMap<String, serde_json::Value>,
119
120    /// Post-mutation request body length produced during `StreamBuffer`
121    /// pre-read.
122    ///
123    /// Stored so `upstream_request_filter` can repair request framing
124    /// before Pingora sends headers to the backend.
125    pub mutated_request_body_len: Option<usize>,
126
127    /// Pipeline pinned for this request's entire lifecycle.
128    ///
129    /// Set once during `request_filter` by cloning the [`Arc`] from the
130    /// listener's [`ArcSwap`]. All subsequent hooks (request body, response,
131    /// response body, logging) use this reference instead of re-loading
132    /// from the [`ArcSwap`], ensuring that a hot configuration reload
133    /// cannot change the pipeline mid-request.
134    ///
135    /// [`ArcSwap`]: arc_swap::ArcSwap
136    pub pinned_pipeline: Option<Arc<FilterPipeline>>,
137
138    /// Filter results from body pre-read. Carried into the next
139    /// [`HttpFilterContext`] so that branch chains attached to the
140    /// first `on_request` filter can evaluate body-derived results.
141    ///
142    /// [`HttpFilterContext`]: praxis_filter::HttpFilterContext
143    pub filter_results: std::collections::HashMap<&'static str, praxis_filter::FilterResultSet>,
144
145    /// Typed per-filter state that persists across all lifecycle
146    /// phases. Keyed by stable filter invocation ID, unique within
147    /// the request's pinned [`FilterPipeline`]. Swapped into each
148    /// [`HttpFilterContext`] and written back after filter execution,
149    /// following the same pattern as [`filter_metadata`].
150    ///
151    /// [`FilterPipeline`]: praxis_filter::FilterPipeline
152    /// [`HttpFilterContext`]: praxis_filter::HttpFilterContext
153    /// [`filter_metadata`]: Self::filter_metadata
154    pub filter_state: std::collections::HashMap<usize, Box<dyn std::any::Any + Send + Sync>>,
155
156    /// Cluster name snapshot retained for metrics emission in the
157    /// `logging()` hook, after `cluster` has been consumed by filter
158    /// context construction.
159    pub metrics_cluster: Option<Arc<str>>,
160
161    /// Pre-built [`SharedString`] for the metrics cluster label.
162    ///
163    /// Cached when `metrics_cluster` is set so that
164    /// `emit_request_metrics` avoids an `Arc` clone per request.
165    ///
166    /// [`SharedString`]: ::metrics::SharedString
167    pub metrics_cluster_shared: Option<::metrics::SharedString>,
168
169    /// Pre-read body chunks (`StreamBuffer` mode). When `StreamBuffer` is
170    /// active, the body is read during `request_filter` (before upstream
171    /// selection) so that body-based routing can influence `upstream_peer`.
172    /// The `request_body_filter` hook then forwards these stored chunks
173    /// instead of reading from the session.
174    ///
175    /// Uses `VecDeque` so that draining from the front is O(1).
176    pub pre_read_body: Option<VecDeque<Bytes>>,
177
178    /// Buffer for request body accumulation in [`StreamBuffer`] mode.
179    ///
180    /// [`StreamBuffer`]: praxis_filter::BodyMode::StreamBuffer
181    pub request_body_buffer: Option<BodyBuffer>,
182
183    /// Accumulated request body bytes seen so far.
184    pub request_body_bytes: u64,
185
186    /// Per-request body delivery mode for the request direction.
187    /// Seeded from static pipeline capabilities, then potentially
188    /// upgraded by filters during `on_request`.
189    pub request_body_mode: BodyMode,
190
191    /// Whether the request body has been released (`StreamBuffer` mode).
192    /// Once true, remaining chunks bypass buffering and stream through.
193    pub request_body_released: bool,
194
195    /// Whether the request method is idempotent (GET, HEAD, OPTIONS).
196    pub request_is_idempotent: bool,
197
198    /// Snapshot of the original request for body/response body phases.
199    pub request_snapshot: Option<Request>,
200
201    /// When this request was received.
202    pub request_start: Instant,
203
204    /// Buffer for response body accumulation in [`StreamBuffer`] mode.
205    ///
206    /// [`StreamBuffer`]: praxis_filter::BodyMode::StreamBuffer
207    pub response_body_buffer: Option<BodyBuffer>,
208
209    /// Accumulated response body bytes seen so far.
210    pub response_body_bytes: u64,
211
212    /// Per-request body delivery mode for the response direction.
213    /// Seeded from static pipeline capabilities, then potentially
214    /// upgraded by filters during `on_response`.
215    pub response_body_mode: BodyMode,
216
217    /// Whether the response body has been released (`StreamBuffer` mode).
218    pub response_body_released: bool,
219
220    /// Snapshot of response headers after response-phase filters.
221    ///
222    /// Used to evaluate `response_conditions` consistently during
223    /// response-body hooks, where Pingora no longer exposes mutable
224    /// response headers.
225    pub response_header_snapshot: Option<Response>,
226
227    /// Upstream response status code, captured during `response_filter`
228    /// for passive health recording in the `logging` hook.
229    pub upstream_response_status: Option<u16>,
230
231    /// Whether the response phase has been executed. Used to ensure
232    /// cleanup (e.g. least-connections counter release) in the
233    /// `logging()` hook when errors bypass `response_filter`.
234    pub response_phase_done: bool,
235
236    /// Number of upstream connection retries attempted.
237    pub retries: u32,
238
239    /// Index of the selected endpoint in the cluster's
240    /// endpoint list. Set during load balancing; used
241    /// for passive health recording in the logging hook.
242    pub selected_endpoint_index: Option<usize>,
243
244    /// Rewritten URI path for the upstream request.
245    ///
246    /// Set by the `path_rewrite` filter via [`HttpFilterContext`] and
247    /// applied in `upstream_request_filter`.
248    ///
249    /// [`HttpFilterContext`]: praxis_filter::HttpFilterContext
250    pub rewritten_path: Option<String>,
251
252    /// Upstream endpoint selected by the load balancer filter.
253    pub upstream: Option<Upstream>,
254
255    /// Saved upstream for retry (cloned before first use).
256    pub upstream_for_retry: Option<Upstream>,
257}
258
259/// Build an [`HttpFilterContext`] from a `PingoraRequestCtx`.
260///
261/// Macro (not a function) so Rust's disjoint field borrowing
262/// works: `filter_context_for` borrows `self.request_snapshot`
263/// immutably while `cluster`, `upstream`, and `rewritten_path`
264/// are taken mutably. A function call with `&mut self` would
265/// collapse these into a single mutable borrow.
266///
267/// [`HttpFilterContext`]: praxis_filter::HttpFilterContext
268macro_rules! filter_context {
269    ($ctx:expr, $pipeline:expr, $request:expr, $response_header:expr) => {{
270        $pipeline.prepare_extensions(&mut $ctx.extensions);
271        praxis_filter::HttpFilterContext {
272            buffered_request_body: $ctx
273                .pre_read_body
274                .as_ref()
275                .map(|chunks| chunks.front().cloned().unwrap_or_default()),
276            body_done_indices: std::mem::take(&mut $ctx.cached_body_done_indices),
277            branch_iterations: std::collections::HashMap::new(),
278            client_addr: $ctx.client_addr,
279            cluster: $ctx.cluster.take(),
280            current_filter_id: None,
281            downstream_tls: $ctx.downstream_tls,
282            extensions: std::mem::take(&mut $ctx.extensions),
283            executed_filter_indices: std::mem::take(&mut $ctx.cached_executed_filter_indices),
284            extra_request_headers: Vec::new(),
285            request_headers_to_remove: Vec::new(),
286            request_headers_to_set: Vec::new(),
287            filter_metadata: std::mem::take(&mut $ctx.filter_metadata),
288            pre_read_mutations: std::mem::take(&mut $ctx.pre_read_mutations),
289            structured_metadata: std::mem::take(&mut $ctx.structured_metadata),
290            filter_results: std::mem::take(&mut $ctx.filter_results),
291            filter_state: std::mem::take(&mut $ctx.filter_state),
292            health_registry: $pipeline.health_registry(),
293            peer_identity: $ctx.peer_identity.clone(),
294            id_generator: $pipeline.id_generator(),
295            kv_stores: $pipeline.kv_stores(),
296            subrequest_connector: $pipeline.subrequest_connector(),
297            request: $request,
298            request_body_bytes: $ctx.request_body_bytes,
299            request_body_mode: $ctx.request_body_mode,
300            request_start: $ctx.request_start,
301            response_body_bytes: $ctx.response_body_bytes,
302            response_body_mode: $ctx.response_body_mode,
303            response_header: $response_header,
304            response_headers_modified: false,
305            rewritten_path: $ctx.rewritten_path.take(),
306            selected_endpoint_index: $ctx.selected_endpoint_index,
307            time_source: $pipeline.time_source(),
308            upstream: $ctx.upstream.take(),
309        }
310    }};
311}
312
313impl PingoraRequestCtx {
314    /// Build an [`HttpFilterContext`] using an external request reference.
315    ///
316    /// Takes `cluster` and `upstream` from `self` (leaving `None`
317    /// behind) so that filters can reassign them. The caller must
318    /// write those fields back after filter execution.
319    ///
320    /// ```
321    /// use praxis_filter::{FilterPipeline, FilterRegistry, Request};
322    /// use praxis_protocol::http::pingora::context::PingoraRequestCtx;
323    ///
324    /// let registry = FilterRegistry::with_builtins();
325    /// let pipeline = FilterPipeline::build(&mut [], &registry).unwrap();
326    /// let request = Request {
327    ///     method: http::Method::GET,
328    ///     uri: http::Uri::from_static("/"),
329    ///     headers: http::HeaderMap::new(),
330    /// };
331    /// let mut ctx = PingoraRequestCtx::default();
332    /// let filter_ctx = ctx.build_filter_context(&pipeline, &request, None);
333    /// assert!(filter_ctx.cluster.is_none());
334    /// ```
335    ///
336    /// [`HttpFilterContext`]: praxis_filter::HttpFilterContext
337    pub fn build_filter_context<'a>(
338        &mut self,
339        pipeline: &'a FilterPipeline,
340        request: &'a Request,
341        response_header: Option<&'a mut Response>,
342    ) -> praxis_filter::HttpFilterContext<'a> {
343        filter_context!(self, pipeline, request, response_header)
344    }
345
346    /// Build an [`HttpFilterContext`] from the stored [`request_snapshot`].
347    ///
348    /// Uses disjoint field borrowing so that `request_snapshot` is
349    /// borrowed immutably while `cluster` and `upstream` are taken
350    /// mutably.
351    ///
352    /// Returns `None` when `request_snapshot` is not set.
353    ///
354    /// ```
355    /// use praxis_filter::{FilterPipeline, FilterRegistry, Request};
356    /// use praxis_protocol::http::pingora::context::PingoraRequestCtx;
357    ///
358    /// let registry = FilterRegistry::with_builtins();
359    /// let pipeline = FilterPipeline::build(&mut [], &registry).unwrap();
360    /// let mut ctx = PingoraRequestCtx::default();
361    /// ctx.request_snapshot = Some(Request {
362    ///     method: http::Method::GET,
363    ///     uri: http::Uri::from_static("/"),
364    ///     headers: http::HeaderMap::new(),
365    /// });
366    /// let filter_ctx = ctx.filter_context_for(&pipeline, None);
367    /// assert!(filter_ctx.is_some());
368    /// ```
369    ///
370    /// [`HttpFilterContext`]: praxis_filter::HttpFilterContext
371    /// [`request_snapshot`]: PingoraRequestCtx::request_snapshot
372    pub fn filter_context_for<'a>(
373        &'a mut self,
374        pipeline: &'a FilterPipeline,
375        response_header: Option<&'a mut Response>,
376    ) -> Option<praxis_filter::HttpFilterContext<'a>> {
377        let request = self.request_snapshot.as_ref()?;
378        Some(filter_context!(self, pipeline, request, response_header))
379    }
380
381    /// Build an [`HttpFilterContext`] plus the saved response header for body conditions.
382    ///
383    /// Response body hooks do not receive mutable headers, but their
384    /// `response_conditions` still need the response status and headers
385    /// captured during the response phase.
386    ///
387    /// [`HttpFilterContext`]: praxis_filter::HttpFilterContext
388    pub fn response_body_context_for<'a>(
389        &'a mut self,
390        pipeline: &'a FilterPipeline,
391    ) -> Option<(praxis_filter::HttpFilterContext<'a>, Option<&'a Response>)> {
392        let request = self.request_snapshot.as_ref()?;
393        let response_header = self.response_header_snapshot.as_ref();
394        Some((filter_context!(self, pipeline, request, None), response_header))
395    }
396
397    /// Pin the current pipeline for this request's entire lifecycle.
398    ///
399    /// Clones the [`Arc`] from the [`ArcSwap`] and stores it in
400    /// [`pinned_pipeline`]. All subsequent hooks should call
401    /// [`pipeline`] instead of re-loading from the [`ArcSwap`].
402    ///
403    /// Called once by `request_filter` in both body-capable and
404    /// no-body handlers.
405    ///
406    /// [`ArcSwap`]: arc_swap::ArcSwap
407    /// [`pinned_pipeline`]: Self::pinned_pipeline
408    /// [`pipeline`]: Self::pipeline
409    pub fn pin_pipeline(&mut self, swap: &arc_swap::ArcSwap<FilterPipeline>) -> Arc<FilterPipeline> {
410        if let Some(existing) = &self.pinned_pipeline {
411            return Arc::clone(existing);
412        }
413        let pipeline = swap.load_full();
414        self.pinned_pipeline = Some(Arc::clone(&pipeline));
415        pipeline
416    }
417
418    /// Return the pinned pipeline, falling back to a fresh
419    /// [`ArcSwap`] load when no pipeline was pinned.
420    ///
421    /// The fallback covers early-failure paths where a lifecycle
422    /// hook runs before `request_filter` (e.g. after
423    /// `early_request_filter` rejection triggers `logging`).
424    ///
425    /// Per-body-chunk hooks (`request_body_filter`,
426    /// `response_body_filter`) call this on every chunk,
427    /// incurring one [`Arc::clone`] per invocation.
428    ///
429    /// [`ArcSwap`]: arc_swap::ArcSwap
430    pub fn pipeline(&self, swap: &arc_swap::ArcSwap<FilterPipeline>) -> Arc<FilterPipeline> {
431        self.pinned_pipeline
432            .as_ref()
433            .map_or_else(|| swap.load_full(), Arc::clone)
434    }
435}
436
437impl Default for PingoraRequestCtx {
438    #[expect(
439        clippy::too_many_lines,
440        reason = "context default enumerates all lifecycle fields explicitly"
441    )]
442    fn default() -> Self {
443        Self {
444            _connection_permit: None,
445            _global_connection_permit: None,
446            cached_body_done_indices: Vec::new(),
447            cached_executed_filter_indices: Vec::new(),
448            client_addr: None,
449            client_http_version: None,
450            cluster: None,
451            connection_upgraded: false,
452            downstream_tls: false,
453            peer_identity: None,
454            extensions: praxis_filter::RequestExtensions::new(),
455            filter_metadata: std::collections::HashMap::new(),
456            pre_read_mutations: Vec::new(),
457            structured_metadata: std::collections::HashMap::new(),
458            mutated_request_body_len: None,
459            pinned_pipeline: None,
460            filter_results: std::collections::HashMap::new(),
461            filter_state: std::collections::HashMap::new(),
462            metrics_cluster: None,
463            metrics_cluster_shared: None,
464            pre_read_body: None,
465            request_body_buffer: None,
466            request_body_bytes: 0,
467            request_body_mode: BodyMode::Stream,
468            request_body_released: false,
469            request_is_idempotent: false,
470            request_snapshot: None,
471            request_start: Instant::now(),
472            response_body_buffer: None,
473            response_body_bytes: 0,
474            response_body_mode: BodyMode::Stream,
475            response_body_released: false,
476            response_header_snapshot: None,
477            upstream_response_status: None,
478            response_phase_done: false,
479            retries: 0,
480            rewritten_path: None,
481            selected_endpoint_index: None,
482            upstream: None,
483            upstream_for_retry: None,
484        }
485    }
486}
487
488// -----------------------------------------------------------------------------
489// Tests
490// -----------------------------------------------------------------------------
491
492#[cfg(test)]
493#[expect(clippy::allow_attributes, reason = "blanket test suppressions")]
494#[allow(
495    clippy::unwrap_used,
496    clippy::expect_used,
497    clippy::indexing_slicing,
498    clippy::significant_drop_tightening,
499    clippy::too_many_lines,
500    reason = "tests"
501)]
502mod tests {
503    use std::{
504        collections::VecDeque,
505        net::{IpAddr, Ipv4Addr},
506        sync::Arc,
507    };
508
509    use bytes::Bytes;
510    use http::{HeaderMap, Method, Uri};
511    use praxis_core::connectivity::Upstream;
512    use praxis_filter::{BodyBuffer, BodyMode, FilterPipeline, FilterRegistry};
513
514    use super::*;
515
516    #[test]
517    fn default_state_has_no_client_addr() {
518        let ctx = default_ctx();
519        assert!(ctx.client_addr.is_none(), "default client_addr should be None");
520    }
521
522    #[test]
523    fn default_state_has_no_cluster() {
524        let ctx = default_ctx();
525        assert!(ctx.cluster.is_none(), "default cluster should be None");
526    }
527
528    #[test]
529    fn default_state_has_zero_retries() {
530        let ctx = default_ctx();
531        assert_eq!(ctx.retries, 0, "default retries should be zero");
532    }
533
534    #[test]
535    fn default_state_flags_are_false() {
536        let ctx = default_ctx();
537        assert!(
538            !ctx.request_body_released,
539            "default request_body_released should be false"
540        );
541        assert!(
542            !ctx.response_body_released,
543            "default response_body_released should be false"
544        );
545        assert!(
546            !ctx.request_is_idempotent,
547            "default request_is_idempotent should be false"
548        );
549        assert!(!ctx.response_phase_done, "default response_phase_done should be false");
550    }
551
552    #[test]
553    fn default_state_buffers_are_none() {
554        let ctx = default_ctx();
555        assert!(
556            ctx.request_body_buffer.is_none(),
557            "default request_body_buffer should be None"
558        );
559        assert!(
560            ctx.response_body_buffer.is_none(),
561            "default response_body_buffer should be None"
562        );
563        assert!(ctx.pre_read_body.is_none(), "default pre_read_body should be None");
564    }
565
566    #[test]
567    fn default_state_snapshots_are_none() {
568        let ctx = default_ctx();
569        assert!(
570            ctx.request_snapshot.is_none(),
571            "default request_snapshot should be None"
572        );
573        assert!(ctx.upstream.is_none(), "default upstream should be None");
574        assert!(
575            ctx.upstream_for_retry.is_none(),
576            "default upstream_for_retry should be None"
577        );
578    }
579
580    #[test]
581    fn set_client_addr() {
582        let mut ctx = default_ctx();
583        let addr = IpAddr::V4(Ipv4Addr::new(192, 168, 1, 100));
584        ctx.client_addr = Some(addr);
585        assert_eq!(
586            ctx.client_addr.unwrap(),
587            addr,
588            "client_addr should match assigned value"
589        );
590    }
591
592    #[test]
593    fn set_cluster() {
594        let mut ctx = default_ctx();
595        ctx.cluster = Some(Arc::from("api-cluster"));
596        assert_eq!(
597            ctx.cluster.as_deref(),
598            Some("api-cluster"),
599            "cluster should match assigned value"
600        );
601    }
602
603    #[test]
604    fn set_upstream() {
605        let mut ctx = default_ctx();
606        let upstream = Upstream {
607            address: Arc::from("10.0.0.1:80"),
608            tls: None,
609            connection: Arc::new(praxis_core::connectivity::ConnectionOptions::default()),
610        };
611        ctx.upstream = Some(upstream.clone());
612        assert_eq!(
613            &*ctx.upstream.as_ref().unwrap().address,
614            "10.0.0.1:80",
615            "upstream address should match assigned value"
616        );
617    }
618
619    #[test]
620    fn increment_retries() {
621        let mut ctx = default_ctx();
622        ctx.retries += 1;
623        ctx.retries += 1;
624        assert_eq!(ctx.retries, 2, "retries should be 2 after two increments");
625    }
626
627    #[test]
628    fn release_request_body_flag() {
629        let mut ctx = default_ctx();
630        assert!(!ctx.request_body_released, "request_body_released should start false");
631        ctx.request_body_released = true;
632        assert!(
633            ctx.request_body_released,
634            "request_body_released should be true after setting"
635        );
636    }
637
638    #[test]
639    fn release_response_body_flag() {
640        let mut ctx = default_ctx();
641        assert!(!ctx.response_body_released, "response_body_released should start false");
642        ctx.response_body_released = true;
643        assert!(
644            ctx.response_body_released,
645            "response_body_released should be true after setting"
646        );
647    }
648
649    #[test]
650    fn response_phase_done_flag() {
651        let mut ctx = default_ctx();
652        assert!(!ctx.response_phase_done, "response_phase_done should start false");
653        ctx.response_phase_done = true;
654        assert!(
655            ctx.response_phase_done,
656            "response_phase_done should be true after setting"
657        );
658    }
659
660    #[test]
661    fn set_pre_read_body() {
662        let mut ctx = default_ctx();
663        let chunks = VecDeque::from([Bytes::from_static(b"chunk1"), Bytes::from_static(b"chunk2")]);
664        ctx.pre_read_body = Some(chunks);
665        let body = ctx.pre_read_body.as_ref().unwrap();
666        assert_eq!(body.len(), 2, "pre_read_body should contain 2 chunks");
667        assert_eq!(body[0], Bytes::from_static(b"chunk1"), "first chunk should be 'chunk1'");
668        assert_eq!(
669            body[1],
670            Bytes::from_static(b"chunk2"),
671            "second chunk should be 'chunk2'"
672        );
673    }
674
675    #[test]
676    fn set_request_snapshot() {
677        let mut ctx = default_ctx();
678        let snapshot = Request {
679            method: Method::POST,
680            uri: "/api/data".parse::<Uri>().unwrap(),
681            headers: HeaderMap::new(),
682        };
683        ctx.request_snapshot = Some(snapshot);
684        let snap = ctx.request_snapshot.as_ref().unwrap();
685        assert_eq!(snap.method, Method::POST, "snapshot method should be POST");
686        assert_eq!(snap.uri.path(), "/api/data", "snapshot URI path should be /api/data");
687    }
688
689    #[test]
690    fn request_body_buffer_lifecycle() {
691        let mut ctx = default_ctx();
692        let mut buf = BodyBuffer::new(100);
693        buf.push(Bytes::from_static(b"data")).unwrap();
694        ctx.request_body_buffer = Some(buf);
695
696        assert!(
697            ctx.request_body_buffer.is_some(),
698            "buffer should be present after assignment"
699        );
700        let taken = ctx.request_body_buffer.take().unwrap();
701        assert_eq!(
702            taken.freeze(),
703            Bytes::from_static(b"data"),
704            "frozen buffer should contain pushed data"
705        );
706        assert!(ctx.request_body_buffer.is_none(), "buffer should be None after take");
707    }
708
709    #[test]
710    fn default_request_body_mode_is_stream() {
711        let ctx = default_ctx();
712        assert_eq!(
713            ctx.request_body_mode,
714            BodyMode::Stream,
715            "default request_body_mode should be Stream"
716        );
717    }
718
719    #[test]
720    fn default_response_body_mode_is_stream() {
721        let ctx = default_ctx();
722        assert_eq!(
723            ctx.response_body_mode,
724            BodyMode::Stream,
725            "default response_body_mode should be Stream"
726        );
727    }
728
729    #[test]
730    fn set_request_body_mode() {
731        let mut ctx = default_ctx();
732        ctx.request_body_mode = BodyMode::StreamBuffer { max_bytes: Some(4096) };
733        assert_eq!(
734            ctx.request_body_mode,
735            BodyMode::StreamBuffer { max_bytes: Some(4096) },
736            "request_body_mode should match assigned value"
737        );
738    }
739
740    #[test]
741    fn set_response_body_mode() {
742        let mut ctx = default_ctx();
743        ctx.response_body_mode = BodyMode::StreamBuffer { max_bytes: Some(8192) };
744        assert_eq!(
745            ctx.response_body_mode,
746            BodyMode::StreamBuffer { max_bytes: Some(8192) },
747            "response_body_mode should match assigned value"
748        );
749    }
750
751    // -------------------------------------------------------------------------
752    // Metadata Roundtrip Tests
753    // -------------------------------------------------------------------------
754
755    #[test]
756    fn metadata_roundtrip_through_filter_context() {
757        let registry = FilterRegistry::with_builtins();
758        let pipeline = FilterPipeline::build(&mut [], &registry).unwrap();
759        let request = Request {
760            method: Method::GET,
761            uri: "/".parse::<Uri>().unwrap(),
762            headers: HeaderMap::new(),
763        };
764
765        let mut ctx = default_ctx();
766        ctx.filter_metadata.insert("rpc.method".to_owned(), "echo".to_owned());
767        ctx.filter_metadata.insert("rpc.status".to_owned(), "ok".to_owned());
768
769        let fctx = ctx.build_filter_context(&pipeline, &request, None);
770        assert_eq!(
771            fctx.get_metadata("rpc.method"),
772            Some("echo"),
773            "metadata written before build should survive into filter context"
774        );
775        assert_eq!(
776            fctx.get_metadata("rpc.status"),
777            Some("ok"),
778            "multiple metadata keys should round-trip"
779        );
780    }
781
782    #[test]
783    fn metadata_written_in_filter_context_persists_back() {
784        let registry = FilterRegistry::with_builtins();
785        let pipeline = FilterPipeline::build(&mut [], &registry).unwrap();
786        let request = Request {
787            method: Method::GET,
788            uri: "/".parse::<Uri>().unwrap(),
789            headers: HeaderMap::new(),
790        };
791
792        let mut ctx = default_ctx();
793        let mut fctx = ctx.build_filter_context(&pipeline, &request, None);
794        fctx.set_metadata("trace.id", "abc-123");
795        fctx.set_metadata("trace.span", "42");
796
797        ctx.filter_metadata = fctx.filter_metadata;
798        assert_eq!(
799            ctx.filter_metadata.get("trace.id").map(String::as_str),
800            Some("abc-123"),
801            "metadata set in filter context should persist back to protocol context"
802        );
803        assert_eq!(
804            ctx.filter_metadata.get("trace.span").map(String::as_str),
805            Some("42"),
806            "multiple metadata keys should persist back"
807        );
808    }
809
810    // -------------------------------------------------------------------------
811    // Hot-Reload Pipeline Pinning (via production helpers)
812    // -------------------------------------------------------------------------
813
814    #[test]
815    fn pin_pipeline_captures_current_arc() {
816        let registry = FilterRegistry::with_builtins();
817        let pipeline_a = Arc::new(FilterPipeline::build(&mut [], &registry).unwrap());
818        let swap = arc_swap::ArcSwap::from(Arc::clone(&pipeline_a));
819
820        let mut ctx = default_ctx();
821        let pinned = ctx.pin_pipeline(&swap);
822        assert!(
823            Arc::ptr_eq(&pinned, &pipeline_a),
824            "pin_pipeline should return the current pipeline"
825        );
826        assert!(
827            Arc::ptr_eq(ctx.pinned_pipeline.as_ref().unwrap(), &pipeline_a),
828            "pinned_pipeline should be stored in ctx"
829        );
830    }
831
832    #[test]
833    fn pipeline_returns_pinned_after_reload() {
834        let registry = FilterRegistry::with_builtins();
835        let pipeline_a = Arc::new(FilterPipeline::build(&mut [], &registry).unwrap());
836        let swap = arc_swap::ArcSwap::from(Arc::clone(&pipeline_a));
837
838        let mut ctx = default_ctx();
839        ctx.pin_pipeline(&swap);
840
841        let pipeline_b = Arc::new(FilterPipeline::build(&mut [], &registry).unwrap());
842        swap.store(pipeline_b);
843
844        let later = ctx.pipeline(&swap);
845        assert!(
846            Arc::ptr_eq(&later, &pipeline_a),
847            "later hooks should still return pipeline A after reload"
848        );
849    }
850
851    #[test]
852    fn new_request_after_reload_pins_new_pipeline() {
853        let registry = FilterRegistry::with_builtins();
854        let pipeline_a = Arc::new(FilterPipeline::build(&mut [], &registry).unwrap());
855        let swap = arc_swap::ArcSwap::from(Arc::clone(&pipeline_a));
856
857        let mut ctx_a = default_ctx();
858        ctx_a.pin_pipeline(&swap);
859
860        let pipeline_b = Arc::new(FilterPipeline::build(&mut [], &registry).unwrap());
861        swap.store(Arc::clone(&pipeline_b));
862
863        let mut ctx_b = default_ctx();
864        ctx_b.pin_pipeline(&swap);
865
866        assert!(
867            Arc::ptr_eq(&ctx_a.pipeline(&swap), &pipeline_a),
868            "request A should use old pipeline"
869        );
870        assert!(
871            Arc::ptr_eq(&ctx_b.pipeline(&swap), &pipeline_b),
872            "request B should use new pipeline"
873        );
874    }
875
876    #[test]
877    fn old_pipeline_drops_after_ctx_drops() {
878        let registry = FilterRegistry::with_builtins();
879        let pipeline_a = Arc::new(FilterPipeline::build(&mut [], &registry).unwrap());
880        let weak_a = Arc::downgrade(&pipeline_a);
881        let swap = arc_swap::ArcSwap::from(pipeline_a);
882
883        let mut ctx = default_ctx();
884        ctx.pin_pipeline(&swap);
885
886        swap.store(Arc::new(FilterPipeline::build(&mut [], &registry).unwrap()));
887
888        assert!(weak_a.upgrade().is_some(), "old pipeline alive while ctx exists");
889
890        drop(ctx);
891
892        assert!(weak_a.upgrade().is_none(), "old pipeline drops with ctx");
893    }
894
895    #[test]
896    fn pipeline_helper_returns_pinned_for_every_phase() {
897        let registry = FilterRegistry::with_builtins();
898        let pipeline_a = Arc::new(FilterPipeline::build(&mut [], &registry).unwrap());
899        let swap = arc_swap::ArcSwap::from(Arc::clone(&pipeline_a));
900
901        let mut ctx = default_ctx();
902        ctx.pin_pipeline(&swap);
903
904        swap.store(Arc::new(FilterPipeline::build(&mut [], &registry).unwrap()));
905
906        for phase in ["request_body", "response", "response_body", "logging"] {
907            let p = ctx.pipeline(&swap);
908            assert!(
909                Arc::ptr_eq(&p, &pipeline_a),
910                "{phase}: should still return pinned pipeline A"
911            );
912        }
913    }
914
915    #[test]
916    fn pipeline_fallback_when_not_pinned() {
917        let registry = FilterRegistry::with_builtins();
918        let pipeline_a = Arc::new(FilterPipeline::build(&mut [], &registry).unwrap());
919        let swap = arc_swap::ArcSwap::from(Arc::clone(&pipeline_a));
920
921        let ctx = default_ctx();
922
923        let loaded = ctx.pipeline(&swap);
924        assert!(
925            Arc::ptr_eq(&loaded, &pipeline_a),
926            "unpinned ctx should fall back to current ArcSwap value"
927        );
928    }
929
930    #[test]
931    fn pin_pipeline_is_idempotent_after_reload() {
932        let registry = FilterRegistry::with_builtins();
933        let pipeline_a = Arc::new(FilterPipeline::build(&mut [], &registry).unwrap());
934        let swap = arc_swap::ArcSwap::from(Arc::clone(&pipeline_a));
935
936        let mut ctx = default_ctx();
937        ctx.pin_pipeline(&swap);
938
939        let pipeline_b = Arc::new(FilterPipeline::build(&mut [], &registry).unwrap());
940        swap.store(pipeline_b);
941
942        let second_pin = ctx.pin_pipeline(&swap);
943        assert!(
944            Arc::ptr_eq(&second_pin, &pipeline_a),
945            "repeated pin_pipeline after reload should return the original pin"
946        );
947    }
948
949    #[test]
950    fn filter_state_isolated_across_pipelines_with_same_ids() {
951        let registry = FilterRegistry::with_builtins();
952        let pipeline_a = Arc::new(FilterPipeline::build(&mut [], &registry).unwrap());
953        let pipeline_b = Arc::new(FilterPipeline::build(&mut [], &registry).unwrap());
954
955        let request = Request {
956            method: Method::GET,
957            uri: "/".parse::<Uri>().unwrap(),
958            headers: HeaderMap::new(),
959        };
960
961        let mut ctx_a = default_ctx();
962        ctx_a.pinned_pipeline = Some(Arc::clone(&pipeline_a));
963        ctx_a.request_snapshot = Some(request.clone());
964        let mut fctx_a = ctx_a.build_filter_context(&pipeline_a, &request, None);
965        fctx_a.current_filter_id = Some(0);
966        fctx_a.insert_filter_state(String::from("from_pipeline_a"));
967        ctx_a.filter_state = fctx_a.filter_state;
968
969        let mut ctx_b = default_ctx();
970        ctx_b.pinned_pipeline = Some(Arc::clone(&pipeline_b));
971        ctx_b.request_snapshot = Some(request.clone());
972        let fctx_b = ctx_b.build_filter_context(&pipeline_b, &request, None);
973
974        assert!(
975            fctx_b.filter_state.is_empty(),
976            "request B should have its own empty state map"
977        );
978
979        let fctx_a2 = ctx_a.filter_context_for(&pipeline_a, None).unwrap();
980        assert_eq!(
981            fctx_a2.filter_state.get(&0).and_then(|v| v.downcast_ref::<String>()),
982            Some(&String::from("from_pipeline_a")),
983            "request A should still see its own state in a later phase"
984        );
985    }
986
987    // -------------------------------------------------------------------------
988    // Test Utilities
989    // -------------------------------------------------------------------------
990
991    /// Create a default request context for tests.
992    fn default_ctx() -> PingoraRequestCtx {
993        PingoraRequestCtx::default()
994    }
995}
praxis_protocol/http/pingora/context.rs

praxis_protocol/http/pingora/
context.rs
