praxis_protocol/http/pingora/handler/
with_body.rs

1// SPDX-License-Identifier: MIT
2// Copyright (c) 2024 Praxis Contributors
3
4//! Pingora HTTP handler with body filter hooks enabled.
5
6use std::{sync::Arc, time::Duration};
7
8use arc_swap::ArcSwap;
9use async_trait::async_trait;
10use bytes::Bytes;
11use pingora_core::{
12    Result,
13    modules::http::{HttpModules, compression::ResponseCompressionBuilder},
14    upstreams::peer::HttpPeer,
15};
16use pingora_proxy::{ProxyHttp, Session};
17use praxis_filter::{CompressionConfig, FilterPipeline};
18use tokio::sync::Semaphore;
19use tracing::debug;
20
21use super::{
22    adjust_compression, emit_request_metrics, handle_connect_failure, logging_cleanup, record_passive_health,
23    request_body_filter, request_filter, response_body_filter, response_filter, upstream_peer, upstream_request, via,
24};
25use crate::http::pingora::context::PingoraRequestCtx;
26
27// -----------------------------------------------------------------------------
28// PingoraHttpHandler
29// -----------------------------------------------------------------------------
30
31/// Pingora HTTP handler that overrides body filter hooks.
32///
33/// Used when the pipeline contains filters that declare
34/// body access via [`BodyAccess`].
35///
36/// The pipeline is held behind [`ArcSwap`] so it can be
37/// atomically replaced by hot config reload without
38/// disrupting in-flight requests.
39///
40/// ```ignore
41/// // Requires a `FilterPipeline` and Pingora server runtime.
42/// use std::sync::Arc;
43///
44/// use arc_swap::ArcSwap;
45/// use praxis_protocol::http::pingora::handler::PingoraHttpHandler;
46///
47/// let handler = PingoraHttpHandler::new(
48///     Arc::new(ArcSwap::from_pointee(pipeline)),
49///     None,
50///     None,
51/// );
52/// ```
53///
54/// [`BodyAccess`]: praxis_filter::BodyAccess
55/// [`ArcSwap`]: arc_swap::ArcSwap
56pub struct PingoraHttpHandler {
57    /// Compression configuration snapshot for module registration.
58    ///
59    /// Used only by [`init_downstream_modules`] to register the
60    /// compression module at startup. Per-request compression
61    /// levels are read from the live pipeline via [`ArcSwap`]
62    /// so that hot-reload updates take effect immediately.
63    ///
64    /// Module registration itself is one-shot in Pingora;
65    /// adding compression to a listener that had none at
66    /// startup requires a restart.
67    ///
68    /// [`init_downstream_modules`]: Self::init_downstream_modules
69    /// [`ArcSwap`]: arc_swap::ArcSwap
70    compression: Option<CompressionConfig>,
71
72    /// Per-listener connection semaphore for max connections.
73    connection_semaphore: Option<Arc<Semaphore>>,
74
75    /// Per-listener downstream read timeout.
76    downstream_read_timeout: Option<Duration>,
77
78    /// Swappable filter pipeline.
79    pipeline: Arc<ArcSwap<FilterPipeline>>,
80}
81
82impl PingoraHttpHandler {
83    /// Create a handler with body filter support.
84    pub(super) fn new(
85        pipeline: Arc<ArcSwap<FilterPipeline>>,
86        downstream_read_timeout: Option<Duration>,
87        connection_semaphore: Option<Arc<Semaphore>>,
88    ) -> Self {
89        let compression = pipeline.load().compression_config().cloned();
90        Self {
91            compression,
92            connection_semaphore,
93            downstream_read_timeout,
94            pipeline,
95        }
96    }
97}
98
99#[async_trait]
100impl ProxyHttp for PingoraHttpHandler {
101    type CTX = PingoraRequestCtx;
102
103    fn new_ctx(&self) -> Self::CTX {
104        PingoraRequestCtx::default()
105    }
106
107    /// Registers Pingora's compression module when compression is
108    /// configured. Otherwise skips module registration to avoid
109    /// per-request `Box` allocation overhead.
110    fn init_downstream_modules(&self, modules: &mut HttpModules) {
111        if let Some(cfg) = &self.compression {
112            debug!(level = cfg.default_level, "registering compression module");
113            modules.add_module(ResponseCompressionBuilder::enable(cfg.default_level));
114        }
115    }
116
117    #[expect(clippy::cast_possible_truncation, reason = "millis fit u64")]
118    async fn early_request_filter(&self, session: &mut Session, ctx: &mut Self::CTX) -> Result<()>
119    where
120        Self::CTX: Send + Sync,
121    {
122        if praxis_core::memory::is_exceeded() {
123            return reject_503(session, "5", "memory pressure exceeded").await;
124        }
125
126        let (exceeded, permit) = crate::connections::try_acquire_global();
127        ctx._global_connection_permit = permit;
128        if exceeded {
129            return reject_503(session, "1", "global max connections exceeded").await;
130        }
131
132        if let Some(sem) = &self.connection_semaphore {
133            if let Ok(permit) = Arc::clone(sem).try_acquire_owned() {
134                ctx._connection_permit = Some(permit);
135            } else {
136                return reject_503(session, "1", "max connections exceeded").await;
137            }
138        }
139
140        if let Some(timeout) = self.downstream_read_timeout {
141            debug!(
142                timeout_ms = timeout.as_millis() as u64,
143                "applying downstream read timeout"
144            );
145            session.set_read_timeout(Some(timeout));
146        }
147        Ok(())
148    }
149
150    async fn request_filter(&self, session: &mut Session, ctx: &mut Self::CTX) -> Result<bool> {
151        let pipeline = ctx.pin_pipeline(&self.pipeline);
152        request_filter::execute(&pipeline, session, ctx).await
153    }
154
155    async fn request_body_filter(
156        &self,
157        session: &mut Session,
158        body: &mut Option<Bytes>,
159        end_of_stream: bool,
160        ctx: &mut Self::CTX,
161    ) -> Result<()>
162    where
163        Self::CTX: Send + Sync,
164    {
165        let pipeline = ctx.pipeline(&self.pipeline);
166        request_body_filter::execute(&pipeline, session, body, end_of_stream, ctx).await
167    }
168
169    fn response_body_filter(
170        &self,
171        _session: &mut Session,
172        body: &mut Option<Bytes>,
173        end_of_stream: bool,
174        ctx: &mut Self::CTX,
175    ) -> Result<Option<Duration>>
176    where
177        Self::CTX: Send + Sync,
178    {
179        let pipeline = ctx.pipeline(&self.pipeline);
180        response_body_filter::execute(&pipeline, body, end_of_stream, ctx)
181    }
182
183    fn fail_to_connect(
184        &self,
185        _session: &mut Session,
186        _peer: &HttpPeer,
187        ctx: &mut Self::CTX,
188        e: Box<pingora_core::Error>,
189    ) -> Box<pingora_core::Error> {
190        handle_connect_failure(ctx, e)
191    }
192
193    async fn upstream_request_filter(
194        &self,
195        session: &mut Session,
196        upstream_request: &mut pingora_http::RequestHeader,
197        ctx: &mut Self::CTX,
198    ) -> Result<()>
199    where
200        Self::CTX: Send + Sync,
201    {
202        let is_upgrade = session.is_upgrade_req();
203        upstream_request::strip_hop_by_hop(upstream_request, is_upgrade);
204        upstream_request::strip_reserved_internal(upstream_request);
205        upstream_request::apply_rewritten_path(upstream_request, ctx)?;
206        upstream_request::apply_mutated_content_length(upstream_request, ctx);
207        let client_ver = ctx.client_http_version.unwrap_or(http::Version::HTTP_11);
208        via::append_request_via(upstream_request, client_ver);
209        Ok(())
210    }
211
212    async fn response_filter(
213        &self,
214        session: &mut Session,
215        upstream_response: &mut pingora_http::ResponseHeader,
216        ctx: &mut Self::CTX,
217    ) -> Result<()>
218    where
219        Self::CTX: Send + Sync,
220    {
221        let pipeline = ctx.pipeline(&self.pipeline);
222        let result = response_filter::execute(&pipeline, upstream_response, ctx).await;
223        if result.is_ok() {
224            let client_ver = ctx.client_http_version.unwrap_or(http::Version::HTTP_11);
225            via::append_response_via(upstream_response, client_ver);
226            adjust_compression(session, upstream_response, pipeline.compression_config());
227        }
228        result
229    }
230
231    async fn upstream_peer(&self, _session: &mut Session, ctx: &mut Self::CTX) -> Result<Box<HttpPeer>> {
232        upstream_peer::execute(ctx).await
233    }
234
235    async fn logging(&self, session: &mut Session, e: Option<&pingora_core::Error>, ctx: &mut Self::CTX) {
236        let pipeline = ctx.pipeline(&self.pipeline);
237        emit_request_metrics(session, ctx);
238        record_passive_health(&pipeline, e, ctx);
239        logging_cleanup(&pipeline, ctx).await;
240    }
241}
242
243// ---------------------------------------------------------------------------
244// Utilities
245// ---------------------------------------------------------------------------
246
247/// Write a 503 response with `Retry-After` and return the corresponding error.
248async fn reject_503(session: &mut Session, retry_after: &'static str, reason: &'static str) -> Result<()> {
249    tracing::warn!(reason, "rejecting request");
250    let mut header = pingora_http::ResponseHeader::build(503, None)?;
251    header.append_header("Retry-After", retry_after)?;
252    session.write_response_header(Box::new(header), true).await?;
253    Err(pingora_core::Error::explain(
254        pingora_core::ErrorType::HTTPStatus(503),
255        reason,
256    ))
257}
praxis_protocol/http/pingora/handler/with_body.rs

praxis_protocol/http/pingora/handler/
with_body.rs
