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

praxis_protocol/http/pingora/
context.rs
