// Copyright 2026 Cloudflare, Inc. // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. // You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, software // distributed under the License is distributed on an "AS IS" BASIS, // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // See the License for the specific language governing permissions and // limitations under the License. //! # pingora-proxy //! //! Programmable HTTP proxy built on top of [pingora_core]. //! //! # Features //! - HTTP/1.x and HTTP/2 for both downstream and upstream //! - Connection pooling //! - TLSv1.3, mutual TLS, customizable CA //! - Request/Response scanning, modification or rejection //! - Dynamic upstream selection //! - Configurable retry and failover //! - Fully programmable and customizable at any stage of a HTTP request //! //! # How to use //! //! Users of this crate defines their proxy by implementing [ProxyHttp] trait, which contains the //! callbacks to be invoked at each stage of a HTTP request. //! //! Then the service can be passed into [`http_proxy_service()`] for a [pingora_core::server::Server] to //! run it. //! //! See `examples/load_balancer.rs` for a detailed example. use async_trait::async_trait; use bytes::Bytes; use futures::future::BoxFuture; use futures::future::FutureExt; use http::{header, version::Version, Method}; use log::{debug, error, trace, warn}; use once_cell::sync::Lazy; use pingora_http::{RequestHeader, ResponseHeader}; use std::fmt::Debug; use std::str; use std::sync::{ atomic::{AtomicBool, AtomicU64, AtomicU8, Ordering}, Arc, }; use std::time::Duration; use tokio::sync::{mpsc, Notify}; use tokio::time; use pingora_cache::NoCacheReason; use pingora_core::apps::{ HttpPersistentSettings, HttpServerApp, HttpServerOptions, ReusedHttpStream, }; use pingora_core::connectors::http::custom; use pingora_core::connectors::{http::Connector, ConnectorOptions}; use pingora_core::modules::http::compression::ResponseCompressionBuilder; use pingora_core::modules::http::{HttpModuleCtx, HttpModules}; use pingora_core::protocols::http::client::HttpSession as ClientSession; use pingora_core::protocols::http::custom::CustomMessageWrite; use pingora_core::protocols::http::subrequest::server::SubrequestHandle; use pingora_core::protocols::http::v1::client::HttpSession as HttpSessionV1; use pingora_core::protocols::http::v2::server::H2Options; use pingora_core::protocols::http::HttpTask; use pingora_core::protocols::http::ServerSession as HttpSession; use pingora_core::protocols::http::SERVER_NAME; use pingora_core::protocols::Stream; use pingora_core::protocols::{Digest, UniqueID}; use pingora_core::server::configuration::ServerConf; use pingora_core::server::{RuntimeOpts, ShutdownWatch}; use pingora_core::upstreams::peer::{HttpPeer, Peer}; use pingora_error::{Error, ErrorSource, ErrorType::*, OrErr, Result}; const TASK_BUFFER_SIZE: usize = 4; type DownstreamCustomMessageReader = Box> + Unpin + Send + Sync + 'static>; mod proxy_cache; mod proxy_common; mod proxy_custom; mod proxy_h1; mod proxy_h2; mod proxy_purge; mod proxy_trait; pub mod subrequest; use subrequest::{BodyMode, Ctx as SubrequestCtx}; pub use proxy_cache::range_filter::{range_header_filter, MultiRangeInfo, RangeType}; pub use proxy_purge::PurgeStatus; pub use proxy_trait::{FailToProxy, ProxyHttp, ProxyWarnLogContext}; pub mod prelude { pub use crate::{http_proxy, http_proxy_service, ProxyHttp, ProxyWarnLogContext, Session}; } pub type ProcessCustomSession = Arc< dyn Fn(Arc>, Stream, &ShutdownWatch) -> BoxFuture<'static, Option> + Send + Sync + Unpin + 'static, >; /// The concrete type that holds the user defined HTTP proxy. /// /// Users don't need to interact with this object directly. pub struct HttpProxy where C: custom::Connector, // Upstream custom connector { inner: SV, // TODO: name it better than inner client_upstream: Connector, shutdown: Notify, shutdown_flag: Arc, pub server_options: Option, pub h2_options: Option, pub downstream_modules: HttpModules, #[cfg(feature = "upstream_modules")] pub upstream_modules: HttpModules, max_retries: usize, process_custom_session: Option>, } impl HttpProxy { /// Create a new [`HttpProxy`] with the given [`ProxyHttp`] implementation and [`ServerConf`]. /// /// After creating an `HttpProxy`, you should call [`HttpProxy::handle_init_modules()`] to /// initialize the downstream modules before processing requests. /// /// For most use cases, prefer using [`http_proxy_service()`] which wraps the `HttpProxy` in a /// [`Service`]. This constructor is useful when you need to integrate `HttpProxy` into a custom /// accept loop (e.g., for SNI-based routing decisions before TLS termination). /// /// # Example /// /// ```ignore /// use pingora_proxy::HttpProxy; /// use std::sync::Arc; /// /// let mut proxy = HttpProxy::new(my_proxy_app, server_conf); /// proxy.handle_init_modules(); /// let proxy = Arc::new(proxy); /// // Use proxy.process_new_http() in your custom accept loop /// ``` pub fn new(inner: SV, conf: Arc) -> Self { HttpProxy { inner, client_upstream: Connector::new(Some(ConnectorOptions::from_server_conf(&conf))), shutdown: Notify::new(), shutdown_flag: Arc::new(AtomicBool::new(false)), server_options: None, h2_options: None, downstream_modules: HttpModules::new(), #[cfg(feature = "upstream_modules")] upstream_modules: HttpModules::new(), max_retries: conf.max_retries, process_custom_session: None, } } } impl HttpProxy where C: custom::Connector, { fn new_custom( inner: SV, conf: Arc, connector: C, on_custom: Option>, server_options: Option, client_options: Option, ) -> Self where SV: ProxyHttp + Send + Sync + 'static, SV::CTX: Send + Sync, { let client_options = client_options.unwrap_or_else(|| ConnectorOptions::from_server_conf(&conf)); let client_upstream = Connector::new_custom(Some(client_options), connector); HttpProxy { inner, client_upstream, shutdown: Notify::new(), shutdown_flag: Arc::new(AtomicBool::new(false)), server_options, downstream_modules: HttpModules::new(), #[cfg(feature = "upstream_modules")] upstream_modules: HttpModules::new(), max_retries: conf.max_retries, process_custom_session: on_custom, h2_options: None, } } /// Return the number of times a pooled upstream connection was found to contain /// unexpected data from the server. pub fn unexpected_data_connection_count(&self) -> u64 { self.client_upstream.unexpected_data_connection_count() } /// Return a shared reference to the unexpected data connection counter for periodic metric reporting. pub fn unexpected_data_connection_counter(&self) -> Arc { self.client_upstream.unexpected_data_connection_counter() } /// Initialize the downstream modules for this proxy. /// /// This method must be called after creating an [`HttpProxy`] with [`HttpProxy::new()`] /// and before processing any requests. It invokes [`ProxyHttp::init_downstream_modules()`] /// to set up any HTTP modules configured by the user's proxy implementation. /// /// Note: When using [`http_proxy_service()`] or [`http_proxy_service_with_name()`], /// this method is called automatically. pub fn handle_init_modules(&mut self) where SV: ProxyHttp, { self.inner .init_downstream_modules(&mut self.downstream_modules); #[cfg(feature = "upstream_modules")] self.inner.init_upstream_modules(&mut self.upstream_modules); } async fn handle_new_request( &self, mut downstream_session: Box, ) -> Option> where SV: ProxyHttp + Send + Sync, SV::CTX: Send + Sync, { // phase 1 read request header let res = tokio::select! { biased; // biased select is cheaper, and we don't want to drop already buffered requests res = downstream_session.read_request() => { res } _ = self.shutdown.notified() => { // service shutting down, dropping the connection to stop more req from coming in return None; } }; match res { Ok(true) => { // TODO: check n==0 debug!("Successfully get a new request"); } Ok(false) => { return None; // TODO: close connection? } Err(mut e) => { e.as_down(); error!( "Fail to proxy: {e}, downstream session type: {}", downstream_session.session_type() ); if matches!(e.etype, InvalidHTTPHeader) { downstream_session .respond_error(400) .await .unwrap_or_else(|e| { error!("failed to send error response to downstream: {e}"); }); } // otherwise the connection must be broken, no need to send anything downstream_session.shutdown().await; return None; } } trace!( "Request header: {:?}", downstream_session.req_header().as_ref() ); // CONNECT method proxying is not default supported by the proxy http logic itself, // since the tunneling process changes the request-response flow. // https://datatracker.ietf.org/doc/html/rfc9110#name-connect // Also because the method impacts message framing in a way is currently unaccounted for // (https://datatracker.ietf.org/doc/html/rfc9112#section-6.3-2.2) // it is safest to disallow use of the method by default. if !self .server_options .as_ref() .is_some_and(|opts| opts.allow_connect_method_proxying) && downstream_session.req_header().method == Method::CONNECT { downstream_session .respond_error(405) .await .unwrap_or_else(|e| { error!("failed to send error response to downstream: {e}"); }); downstream_session.shutdown().await; return None; } Some(downstream_session) } // return bool: server_session can be reused, and error if any async fn proxy_to_upstream( &self, session: &mut Session, ctx: &mut SV::CTX, ) -> (bool, Option>) where SV: ProxyHttp + Send + Sync, SV::CTX: Send + Sync, { let peer = match self.inner.upstream_peer(session, ctx).await { Ok(p) => p, Err(e) => return (false, Some(e)), }; let client_session = self.client_upstream.get_http_session(&*peer).await; match client_session { Ok((client_session, client_reused)) => { let (server_reused, error) = match client_session { ClientSession::H1(mut h1) => { let (server_reused, client_reuse, error) = self .proxy_to_h1_upstream(session, &mut h1, client_reused, &peer, ctx) .await; if client_reuse { let session = ClientSession::H1(h1); self.client_upstream .release_http_session(session, &*peer, peer.idle_timeout()) .await; } (server_reused, error) } ClientSession::H2(mut h2) => { let (server_reused, mut error) = self .proxy_to_h2_upstream(session, &mut h2, client_reused, &peer, ctx) .await; let session = ClientSession::H2(h2); self.client_upstream .release_http_session(session, &*peer, peer.idle_timeout()) .await; if let Some(e) = error.as_mut() { // try to downgrade if A. origin says so or B. origin sends an invalid // response, which usually means origin h2 is not production ready if matches!(e.etype, H2Downgrade | InvalidH2) { if peer .get_alpn() .is_none_or(|alpn| alpn.get_min_http_version() == 1) { // Add the peer to prefer h1 so that all following requests // will use h1 self.client_upstream.prefer_h1(&*peer); } else { // the peer doesn't allow downgrading to h1 (e.g. gRPC) e.retry = false.into(); } } } (server_reused, error) } ClientSession::Custom(mut c) => { let (server_reused, error) = self .proxy_to_custom_upstream(session, &mut c, client_reused, &peer, ctx) .await; let session = ClientSession::Custom(c); self.client_upstream .release_http_session(session, &*peer, peer.idle_timeout()) .await; (server_reused, error) } }; ( server_reused, error.map(|e| { self.inner .error_while_proxy(&peer, session, e, ctx, client_reused) }), ) } Err(mut e) => { e.as_up(); let new_err = self.inner.fail_to_connect(session, &peer, ctx, e); (false, Some(new_err.into_up())) } } } async fn upstream_filter( &self, session: &mut Session, task: &mut HttpTask, ctx: &mut SV::CTX, ) -> Result> where SV: ProxyHttp + Send + Sync, SV::CTX: Send + Sync, { let duration = match task { HttpTask::Header(header, _eos) => { self.inner .upstream_response_filter(session, header, ctx) .await?; None } HttpTask::Body(data, eos) | HttpTask::UpgradedBody(data, eos) => self .inner .upstream_response_body_filter(session, data, *eos, ctx)?, HttpTask::Trailer(Some(trailers)) => { self.inner .upstream_response_trailer_filter(session, trailers, ctx)?; None } _ => { // task does not support a filter None } }; Ok(duration) } async fn finish( &self, mut session: Session, ctx: &mut SV::CTX, reuse: bool, error: Option>, ) -> Option where SV: ProxyHttp + Send + Sync, SV::CTX: Send + Sync, { self.inner .logging(&mut session, error.as_deref(), ctx) .await; if let Some(e) = error { session.downstream_session.on_proxy_failure(e); } if reuse { // TODO: log error let mut persistent_settings = HttpPersistentSettings::for_session(&session); if let Some(uc) = self.inner.persist_connection_context(&session, ctx) { persistent_settings.set_user_context(uc); } session .downstream_session .finish() .await .ok() .flatten() .map(|s| ReusedHttpStream::from_reusable_stream(s, persistent_settings)) } else { None } } fn cleanup_sub_req(&self, session: &mut Session) { if let Some(ctx) = session.subrequest_ctx.as_mut() { ctx.release_write_lock(); } } } use pingora_cache::HttpCache; use pingora_core::protocols::http::compression::ResponseCompressionCtx; /// The established HTTP session /// /// This object is what users interact with in order to access the request itself or change the proxy /// behavior. pub struct Session { /// the HTTP session to downstream (the client) pub downstream_session: Box, /// The interface to control HTTP caching pub cache: HttpCache, /// (de)compress responses coming into the proxy (from upstream) pub upstream_compression: ResponseCompressionCtx, /// ignore downstream range (skip downstream range filters) pub ignore_downstream_range: bool, /// Were the upstream request headers modified? pub upstream_headers_mutated_for_cache: bool, /// Upstream predicate for whether this HTTP/1 request is an upgrade. h1_upgrade_request_status: H1UpgradeRequestStatus, /// The context from parent request, if this is a subrequest. pub subrequest_ctx: Option>, /// Handle to allow spawning subrequests, assigned by the `Subrequest` app logic. pub subrequest_spawner: Option, // Downstream filter modules pub downstream_modules_ctx: HttpModuleCtx, /// Upstream filter modules. These run before `upstream_compression` and see the raw /// (pre-compression) upstream response body. #[cfg(feature = "upstream_modules")] pub upstream_modules_ctx: HttpModuleCtx, /// Upstream response body bytes received (payload only). Set by proxy layer. /// TODO: move this into an upstream session digest for future fields. upstream_body_bytes_received: usize, /// Whether proxy task filtering has seen a downstream 101 upgrade header. downstream_task_seen_upgraded: bool, /// Upstream write pending time. Set by proxy layer (HTTP/1.x only). upstream_write_pending_time: Duration, /// Flag that is set when the shutdown process has begun. shutdown_flag: Arc, } impl Session { fn new( downstream_session: impl Into>, downstream_modules: &HttpModules, #[cfg(feature = "upstream_modules")] upstream_modules: &HttpModules, shutdown_flag: Arc, ) -> Self { Session { downstream_session: downstream_session.into(), cache: HttpCache::new(), // disable both upstream and downstream compression upstream_compression: ResponseCompressionCtx::new(0, false, false), ignore_downstream_range: false, upstream_headers_mutated_for_cache: false, h1_upgrade_request_status: H1UpgradeRequestStatus::default(), subrequest_ctx: None, subrequest_spawner: None, // optionally set later on downstream_modules_ctx: downstream_modules.build_ctx(), #[cfg(feature = "upstream_modules")] upstream_modules_ctx: upstream_modules.build_ctx(), upstream_body_bytes_received: 0, downstream_task_seen_upgraded: false, upstream_write_pending_time: Duration::ZERO, shutdown_flag, } } /// Create a new [Session] from the given [Stream] /// /// This function is mostly used for testing and mocking, given the downstream modules and /// shutdown flags will never be set. pub fn new_h1(stream: Stream) -> Self { let modules = HttpModules::new(); Self::new( Box::new(HttpSession::new_http1(stream)), &modules, #[cfg(feature = "upstream_modules")] &HttpModules::new(), Arc::new(AtomicBool::new(false)), ) } /// Create a new [Session] from the given [Stream] with modules /// /// This function is mostly used for testing and mocking, given the shutdown flag will never be /// set. pub fn new_h1_with_modules(stream: Stream, downstream_modules: &HttpModules) -> Self { Self::new( Box::new(HttpSession::new_http1(stream)), downstream_modules, #[cfg(feature = "upstream_modules")] &HttpModules::new(), Arc::new(AtomicBool::new(false)), ) } /// Run upstream module filters on the given [`HttpTask`]. /// /// Upstream modules process each task **before** `upstream_compression` and /// see the raw (pre-compression) upstream response. Like the downstream /// module path, `response_trailer_filter` and `response_done_filter` return /// values are converted to body tasks when present. #[cfg(feature = "upstream_modules")] pub async fn upstream_modules_filter_task(&mut self, t: &mut HttpTask) -> Result<()> { match t { HttpTask::Header(header, eos) => { self.upstream_modules_ctx .response_header_filter(header, *eos) .await?; } HttpTask::Body(body, eos) | HttpTask::UpgradedBody(body, eos) => { self.upstream_modules_ctx.response_body_filter(body, *eos)?; } HttpTask::Trailer(trailers) => { if let Some(buf) = self .upstream_modules_ctx .response_trailer_filter(trailers)? { *t = HttpTask::Body(Some(buf), true); } } HttpTask::Done => { if let Some(buf) = self.upstream_modules_ctx.response_done_filter()? { *t = HttpTask::Body(Some(buf), true); } } HttpTask::Failed(_) => {} } Ok(()) } pub fn as_downstream_mut(&mut self) -> &mut HttpSession { &mut self.downstream_session } pub fn as_downstream(&self) -> &HttpSession { &self.downstream_session } /// Write HTTP response with the given error code to the downstream. pub async fn respond_error(&mut self, error: u16) -> Result<()> { self.as_downstream_mut().respond_error(error).await } /// Write HTTP response with the given error code to the downstream with a body. pub async fn respond_error_with_body(&mut self, error: u16, body: Bytes) -> Result<()> { self.as_downstream_mut() .respond_error_with_body(error, body) .await } /// Write the given HTTP response header to the downstream /// /// Different from directly calling [HttpSession::write_response_header], this function also /// invokes the filter modules. pub async fn write_response_header( &mut self, mut resp: Box, end_of_stream: bool, ) -> Result<()> { self.downstream_modules_ctx .response_header_filter(&mut resp, end_of_stream) .await?; self.downstream_session.write_response_header(resp).await } /// Similar to `write_response_header()`, this fn will clone the `resp` internally pub async fn write_response_header_ref( &mut self, resp: &ResponseHeader, end_of_stream: bool, ) -> Result<(), Box> { self.write_response_header(Box::new(resp.clone()), end_of_stream) .await } /// Write the given HTTP response body chunk to the downstream /// /// Different from directly calling [HttpSession::write_response_body], this function also /// invokes the filter modules. pub async fn write_response_body( &mut self, mut body: Option, end_of_stream: bool, ) -> Result<()> { self.downstream_modules_ctx .response_body_filter(&mut body, end_of_stream)?; if body.is_none() && !end_of_stream { return Ok(()); } let data = body.unwrap_or_default(); self.downstream_session .write_response_body(data, end_of_stream) .await } // Run downstream module response filters on a single task, updating // `seen_upgraded` to track whether an upgrade has been seen. Used by both // `send_downstream_proxy_task` and `write_response_tasks`. async fn downstream_response_task_filter( &mut self, task: &mut HttpTask, seen_upgraded: &mut bool, ) -> Result<()> { match task { HttpTask::Header(resp, end) => { if *seen_upgraded { return reject_unexpected_task_after_h1_upgrade(self, "header", *seen_upgraded); } self.downstream_modules_ctx .response_header_filter(resp, *end) .await?; reject_mismatched_h1_upgrade_101(self, resp, "downstream_module_header_filter") .map_err(|e| e.into_in())?; if resp.status == http::StatusCode::SWITCHING_PROTOCOLS && self.downstream_session.is_upgrade(resp) == Some(true) { *seen_upgraded = true; } } HttpTask::Body(data, end) => { if *seen_upgraded { return reject_unexpected_task_after_h1_upgrade(self, "body", *seen_upgraded); } self.downstream_modules_ctx .response_body_filter(data, *end)?; } HttpTask::UpgradedBody(data, end) => { if !*seen_upgraded { return reject_unexpected_upgraded_body_before_h1_upgrade(self, *seen_upgraded); } self.downstream_modules_ctx .response_body_filter(data, *end)?; } HttpTask::Trailer(trailers) => { if *seen_upgraded { return reject_unexpected_task_after_h1_upgrade( self, "trailer", *seen_upgraded, ); } if let Some(buf) = self .downstream_modules_ctx .response_trailer_filter(trailers)? { // Write the trailers into the body if the filter // returns a buffer. // // Note, this will not work if end of stream has already // been seen or we've written content-length bytes. // (Trailers should never come after upgraded body) *task = HttpTask::Body(Some(buf), true); } } HttpTask::Done => { // `Done` can be sent in certain response paths to mark end // of response if not already done via trailers or body with // end flag set. // If the filter returns body bytes on Done, // write them into the response. After a 101, those bytes are // already in the upgraded protocol and must not be HTTP-framed. // // Note, this will not work if end of stream has already // been seen or we've written content-length bytes. if let Some(buf) = self.downstream_modules_ctx.response_done_filter()? { *task = if *seen_upgraded { HttpTask::UpgradedBody(Some(buf), true) } else { HttpTask::Body(Some(buf), true) }; } } _ => { /* Failed */ } } Ok(()) } /// Queue a downstream proxy task for cancel-safe writing after running /// downstream module filters. This allows decoupling cache writes from /// downstream writes. /// /// Only works with sessions that support the proxy task API. /// /// # Panics /// Panics if the session doesn't support the proxy task API. /// Use `write_response_tasks()` for sessions that don't support the proxy task API. pub async fn send_downstream_proxy_task(&mut self, mut task: HttpTask) -> Result<()> { let mut seen_upgraded = self.downstream_task_seen_upgraded || self.was_upgraded(); self.downstream_response_task_filter(&mut task, &mut seen_upgraded) .await?; self.downstream_task_seen_upgraded = seen_upgraded; self.downstream_session.send_downstream_proxy_task(task); Ok(()) } /// Enable or disable the cancel-safe proxy task API for this session. /// /// When disabled, the proxy falls back to the blocking `write_response_tasks` /// path. This can be called from request filters to opt out on a per-request /// basis. pub fn set_proxy_tasks_enabled(&mut self, enabled: bool) { self.downstream_session.set_proxy_tasks_enabled(enabled); } /// Check if there are pending downstream tasks queued for writing. /// Used for backpressure - don't queue more cache tasks if we have pending writes. /// Returns false for sessions that don't support the proxy task API. pub fn has_pending_downstream_tasks(&self) -> bool { self.downstream_session.supports_proxy_task_api() && self.downstream_session.has_pending_downstream_proxy_tasks() } /// Write all queued downstream proxy tasks. This is cancel-safe and can be called /// in a select! loop while waiting for upstream tasks. /// For sessions that don't support the proxy task API, this is a no-op. pub async fn write_downstream_proxy_tasks(&mut self) -> Result { if self.downstream_session.supports_proxy_task_api() { self.downstream_session.write_downstream_proxy_tasks().await } else { Ok(false) } } pub async fn write_response_tasks(&mut self, mut tasks: Vec) -> Result { let mut seen_upgraded = self.downstream_task_seen_upgraded || self.was_upgraded(); for task in tasks.iter_mut() { self.downstream_response_task_filter(task, &mut seen_upgraded) .await?; } self.downstream_task_seen_upgraded = seen_upgraded; self.downstream_session.response_duplex_vec(tasks).await } /// Mark the upstream headers as modified by caching. This should lead to range filters being /// skipped when responding to the downstream. pub fn mark_upstream_headers_mutated_for_cache(&mut self) { self.upstream_headers_mutated_for_cache = true; } /// Check whether the upstream headers were marked as mutated during the request. pub fn upstream_headers_mutated_for_cache(&self) -> bool { self.upstream_headers_mutated_for_cache } fn set_upstream_h1_upgrade_request_status(&mut self, upstream_is_upgrade_req: bool) { self.h1_upgrade_request_status = H1UpgradeRequestStatus { upstream: Some(upstream_is_upgrade_req), }; } fn h1_upgrade_request_snapshot(&self) -> H1UpgradeRequestSnapshot { H1UpgradeRequestSnapshot { downstream: self.downstream_session.is_upgrade_req(), upstream: self.h1_upgrade_request_status.upstream, } } /// Get the total upstream response body bytes received (payload only) recorded by the proxy layer. pub fn upstream_body_bytes_received(&self) -> usize { self.upstream_body_bytes_received } /// Set the total upstream response body bytes received (payload only). Intended for internal use by proxy layer. pub(crate) fn set_upstream_body_bytes_received(&mut self, n: usize) { self.upstream_body_bytes_received = n; } /// Get the upstream write pending time recorded by the proxy layer. Returns [`Duration::ZERO`] for HTTP/2. pub fn upstream_write_pending_time(&self) -> Duration { self.upstream_write_pending_time } /// Set the upstream write pending time. Intended for internal use by proxy layer. pub(crate) fn set_upstream_write_pending_time(&mut self, d: Duration) { self.upstream_write_pending_time = d; } /// Is the proxy process in the process of shutting down (e.g. due to graceful upgrade)? pub fn is_process_shutting_down(&self) -> bool { self.shutdown_flag.load(Ordering::Acquire) } pub fn downstream_custom_message(&mut self) -> Result> { if let Some(custom_session) = self.downstream_session.as_custom_mut() { custom_session .take_custom_message_reader() .map(Some) .ok_or(Error::explain( ReadError, "can't extract custom reader from downstream", )) } else { Ok(None) } } fn take_downstream_custom_message_reader( &mut self, downstream_custom_message_writer: &mut Option>, ) -> Result> { if downstream_custom_message_writer.is_none() { return Ok(None); } let Some(custom_session) = self.downstream_session.as_custom_mut() else { return Ok(None); }; let Some(reader) = custom_session.take_custom_message_reader() else { if let Some(writer) = downstream_custom_message_writer.take() { custom_session.restore_custom_message_writer(writer)?; } return Err(Error::explain( ReadError, "can't extract custom reader from downstream", )); }; Ok(Some(reader)) } } #[derive(Clone, Copy, Debug, Default)] struct H1UpgradeRequestStatus { upstream: Option, } #[derive(Clone, Copy, Debug)] struct H1UpgradeRequestSnapshot { downstream: bool, upstream: Option, } impl H1UpgradeRequestSnapshot { fn mismatch(self) -> bool { // No upstream predicate means this helper cannot prove a mismatch. The // current proxy paths record it before upstream responses can be handled. matches!(self.upstream, Some(upstream) if self.downstream != upstream) } } /// Rejects a 101 response when the downstream and upstream H1 upgrade state differs. /// /// Upstream and downstream must agree that this request is an upgrade before a /// 101 can establish a tunnel. Otherwise one side changes protocol while the /// other stays in HTTP handling, allowing tunneled traffic to bypass request /// processing or corrupt the connection state. fn reject_mismatched_h1_upgrade_101( session: &Session, header: &ResponseHeader, stage: &'static str, ) -> Result<()> { if header.status != http::StatusCode::SWITCHING_PROTOCOLS { return Ok(()); } let status = session.h1_upgrade_request_snapshot(); if status.mismatch() { return Error::e_explain( InvalidHTTPHeader, format!( "received 101 response with mismatched upstream/downstream upgrade status: stage={stage}, downstream_upgrade_req={}, upstream_upgrade_req={:?}, downstream_was_upgraded={}, downstream_task_seen_upgraded={}, response_version={:?}, response_upgrade_header_present={}, response_connection_header_present={}", status.downstream, status.upstream, session.was_upgraded(), session.downstream_task_seen_upgraded, header.version, header.headers.get(http::header::UPGRADE).is_some(), header.headers.get(http::header::CONNECTION).is_some(), ), ); } Ok(()) } fn reject_unexpected_task_after_h1_upgrade( session: &Session, task: &'static str, task_filter_seen_upgraded: bool, ) -> Result<()> { let status = session.h1_upgrade_request_snapshot(); Error::e_explain( InvalidHTTPHeader, format!( "received {task} task after downstream 101 upgrade: downstream_upgrade_req={}, upstream_upgrade_req={:?}, downstream_was_upgraded={}, downstream_task_seen_upgraded={}, task_filter_seen_upgraded={}", status.downstream, status.upstream, session.was_upgraded(), session.downstream_task_seen_upgraded, task_filter_seen_upgraded ), ) .map_err(|e| e.into_in()) } fn reject_unexpected_upgraded_body_before_h1_upgrade( session: &Session, task_filter_seen_upgraded: bool, ) -> Result<()> { let status = session.h1_upgrade_request_snapshot(); Error::e_explain( InvalidHTTPHeader, format!( "received upgraded body task before downstream 101 upgrade: downstream_upgrade_req={}, upstream_upgrade_req={:?}, downstream_was_upgraded={}, downstream_task_seen_upgraded={}, task_filter_seen_upgraded={}", status.downstream, status.upstream, session.was_upgraded(), session.downstream_task_seen_upgraded, task_filter_seen_upgraded ), ) .map_err(|e| e.into_in()) } impl AsRef for Session { fn as_ref(&self) -> &HttpSession { &self.downstream_session } } impl AsMut for Session { fn as_mut(&mut self) -> &mut HttpSession { &mut self.downstream_session } } use std::ops::{Deref, DerefMut}; impl Deref for Session { type Target = HttpSession; fn deref(&self) -> &Self::Target { &self.downstream_session } } impl DerefMut for Session { fn deref_mut(&mut self) -> &mut Self::Target { &mut self.downstream_session } } // generic HTTP 502 response sent when proxy_upstream_filter refuses to connect to upstream static BAD_GATEWAY: Lazy = Lazy::new(|| { let mut resp = ResponseHeader::build(http::StatusCode::BAD_GATEWAY, Some(3)).unwrap(); resp.insert_header(header::SERVER, &SERVER_NAME[..]) .unwrap(); resp.insert_header(header::CONTENT_LENGTH, 0).unwrap(); resp.insert_header(header::CACHE_CONTROL, "private, no-store") .unwrap(); resp }); impl HttpProxy where C: custom::Connector, { async fn process_request( self: &Arc, mut session: Session, mut ctx: ::CTX, ) -> Option where SV: ProxyHttp + Send + Sync + 'static, ::CTX: Send + Sync, { if let Err(e) = self .inner .early_request_filter(&mut session, &mut ctx) .await { return self .handle_error(session, &mut ctx, e, "Fail to early filter request:") .await; } if self.inner.allow_spawning_subrequest(&session, &ctx) { session.subrequest_spawner = Some(SubrequestSpawner::new(self.clone())); } let req = session.downstream_session.req_header_mut(); // Built-in downstream request filters go first if let Err(e) = session .downstream_modules_ctx .request_header_filter(req) .await { return self .handle_error( session, &mut ctx, e, "Failed in downstream modules request filter:", ) .await; } match self.inner.request_filter(&mut session, &mut ctx).await { Ok(response_sent) => { if response_sent { // TODO: log error self.inner.logging(&mut session, None, &mut ctx).await; self.cleanup_sub_req(&mut session); let mut persistent_settings = HttpPersistentSettings::for_session(&session); if let Some(uc) = self.inner.persist_connection_context(&session, &ctx) { persistent_settings.set_user_context(uc); } return session .downstream_session .finish() .await .ok() .flatten() .map(|s| ReusedHttpStream::from_reusable_stream(s, persistent_settings)); } /* else continue */ } Err(e) => { return self .handle_error(session, &mut ctx, e, "Fail to filter request:") .await; } } if let Some((reuse, err)) = self.proxy_cache(&mut session, &mut ctx).await { // cache hit return self.finish(session, &mut ctx, reuse, err).await; } // either uncacheable, or cache miss // there should not be a write lock in the sub req ctx after this point self.cleanup_sub_req(&mut session); // decide if the request is allowed to go to upstream match self .inner .proxy_upstream_filter(&mut session, &mut ctx) .await { Ok(proxy_to_upstream) => { if !proxy_to_upstream { // The hook can choose to write its own response, but if it doesn't, we respond // with a generic 502 if session.cache.enabled() { // drop the cache lock that this request may be holding onto session.cache.disable(NoCacheReason::DeclinedToUpstream); } if session.response_written().is_none() { match session.write_response_header_ref(&BAD_GATEWAY, true).await { Ok(()) => {} Err(e) => { return self .handle_error( session, &mut ctx, e, "Error responding with Bad Gateway:", ) .await; } } } return self.finish(session, &mut ctx, true, None).await; } /* else continue */ } Err(e) => { if session.cache.enabled() { session.cache.disable(NoCacheReason::InternalError); } return self .handle_error( session, &mut ctx, e, "Error deciding if we should proxy to upstream:", ) .await; } } let mut retries: usize = 0; let mut server_reuse = false; let mut proxy_error: Option> = None; while retries < self.max_retries { retries += 1; let (reuse, e) = self.proxy_to_upstream(&mut session, &mut ctx).await; server_reuse = reuse; match e { Some(error) => { let retry = error.retry(); // only log error that will be retried here, the final error will be logged below if retry && !self.inner.suppress_proxy_warn_log( &session, &ctx, &error, ProxyWarnLogContext::UpstreamRetry, ) { warn!( "Fail to proxy: {}, tries: {}, retry: {}, {}", error, retries, retry, self.inner.request_summary(&session, &ctx) ); } proxy_error = Some(error); if !retry { break; } } None => { proxy_error = None; break; } }; } // serve stale if error // Check both error and cache before calling the function because await is not cheap // allow unwrap until if let chains #[allow(clippy::unnecessary_unwrap)] let serve_stale_result = if proxy_error.is_some() && session.cache.can_serve_stale_error() { self.handle_stale_if_error(&mut session, &mut ctx, proxy_error.as_ref().unwrap()) .await } else { None }; let final_error = if let Some((reuse, stale_cache_error)) = serve_stale_result { // don't reuse server conn if serve stale polluted it server_reuse = server_reuse && reuse; stale_cache_error } else { proxy_error }; if let Some(e) = final_error.as_ref() { // If we have errored and are still holding a cache lock, release it. if session.cache.enabled() { let reason = if *e.esource() == ErrorSource::Upstream { NoCacheReason::UpstreamError } else { NoCacheReason::InternalError }; session.cache.disable(reason); } let res = self.inner.fail_to_proxy(&mut session, e, &mut ctx).await; // final error will have > 0 status unless downstream connection is dead if !self.inner.suppress_error_log(&session, &ctx, e) { error!( "Fail to proxy: {}, status: {}, tries: {}, retry: {}, {}", e, res.error_code, retries, false, // we never retry here self.inner.request_summary(&session, &ctx), ); } } // logging() will be called in finish() self.finish(session, &mut ctx, server_reuse, final_error) .await } async fn handle_error( &self, mut session: Session, ctx: &mut ::CTX, e: Box, context: &str, ) -> Option where SV: ProxyHttp + Send + Sync + 'static, ::CTX: Send + Sync, { let res = self.inner.fail_to_proxy(&mut session, &e, ctx).await; if !self.inner.suppress_error_log(&session, ctx, &e) { error!( "{context} {}, status: {}, {}", e, res.error_code, self.inner.request_summary(&session, ctx) ); } self.inner.logging(&mut session, Some(&e), ctx).await; self.cleanup_sub_req(&mut session); session.downstream_session.on_proxy_failure(e); if res.can_reuse_downstream { let mut persistent_settings = HttpPersistentSettings::for_session(&session); if let Some(uc) = self.inner.persist_connection_context(&session, ctx) { persistent_settings.set_user_context(uc); } session .downstream_session .finish() .await .ok() .flatten() .map(|s| ReusedHttpStream::from_reusable_stream(s, persistent_settings)) } else { None } } } /* Make process_subrequest() a trait to workaround https://github.com/rust-lang/rust/issues/78649 if process_subrequest() is implemented as a member of HttpProxy, rust complains error[E0391]: cycle detected when computing type of `proxy_cache::::proxy_cache::{opaque#0}` --> pingora-proxy/src/proxy_cache.rs:13:10 | 13 | ) -> Option<(bool, Option>)> */ #[async_trait] pub trait Subrequest { async fn process_subrequest( self: Arc, session: Box, sub_req_ctx: Box, ); } #[async_trait] impl Subrequest for HttpProxy where SV: ProxyHttp + Send + Sync + 'static, ::CTX: Send + Sync, C: custom::Connector, { async fn process_subrequest( self: Arc, session: Box, sub_req_ctx: Box, ) { debug!("starting subrequest"); let mut session = match self.handle_new_request(session).await { Some(downstream_session) => Session::new( downstream_session, &self.downstream_modules, #[cfg(feature = "upstream_modules")] &self.upstream_modules, self.shutdown_flag.clone(), ), None => return, // bad request }; // no real downstream to keepalive, but it doesn't matter what is set here because at the end // of this fn the dummy connection will be dropped session.set_keepalive(None); session.subrequest_ctx.replace(sub_req_ctx); trace!("processing subrequest"); let ctx = self.inner.new_ctx(); self.process_request(session, ctx).await; trace!("subrequest done"); } } /// A handle to the underlying HTTP proxy app that allows spawning subrequests. pub struct SubrequestSpawner { app: Arc, } /// A [`PreparedSubrequest`] that is ready to run. pub struct PreparedSubrequest { app: Arc, session: Box, sub_req_ctx: Box, } impl PreparedSubrequest { pub async fn run(self) { self.app .process_subrequest(self.session, self.sub_req_ctx) .await } pub fn session(&self) -> &HttpSession { self.session.as_ref() } pub fn session_mut(&mut self) -> &mut HttpSession { self.session.deref_mut() } } impl SubrequestSpawner { /// Create a new [`SubrequestSpawner`]. pub fn new(app: Arc) -> SubrequestSpawner { SubrequestSpawner { app } } /// Spawn a background subrequest and return a join handle. // TODO: allow configuring the subrequest session before use pub fn spawn_background_subrequest( &self, session: &HttpSession, ctx: SubrequestCtx, ) -> tokio::task::JoinHandle<()> { let new_app = self.app.clone(); // Clone the Arc let (mut session, handle) = subrequest::create_session(session); if ctx.body_mode() == BodyMode::NoBody { session .as_subrequest_mut() .expect("created subrequest session") .clear_request_body_headers(); } let sub_req_ctx = Box::new(ctx); handle.drain_tasks(); tokio::spawn(async move { new_app .process_subrequest(Box::new(session), sub_req_ctx) .await; }) } /// Create a subrequest that listens to `HttpTask`s sent from the returned `Sender` /// and sends `HttpTask`s to the returned `Receiver`. /// /// To run that subrequest, call `run()`. // TODO: allow configuring the subrequest session before use pub fn create_subrequest( &self, session: &HttpSession, ctx: SubrequestCtx, ) -> (PreparedSubrequest, SubrequestHandle) { let new_app = self.app.clone(); // Clone the Arc let (mut session, handle) = subrequest::create_session(session); if ctx.body_mode() == BodyMode::NoBody { session .as_subrequest_mut() .expect("created subrequest session") .clear_request_body_headers(); } let sub_req_ctx = Box::new(ctx); ( PreparedSubrequest { app: new_app, session: Box::new(session), sub_req_ctx, }, handle, ) } } #[async_trait] impl HttpServerApp for HttpProxy where SV: ProxyHttp + Send + Sync + 'static, ::CTX: Send + Sync, C: custom::Connector, { async fn process_new_http( self: &Arc, mut session: HttpSession, shutdown: &ShutdownWatch, ) -> Option { // Extract user context from the previous request before the session is moved into the Box let prev_user_ctx = session.take_connection_user_context(); let session = Box::new(session); // TODO: keepalive pool, use stack let mut session = match self.handle_new_request(session).await { Some(downstream_session) => Session::new( downstream_session, &self.downstream_modules, #[cfg(feature = "upstream_modules")] &self.upstream_modules, self.shutdown_flag.clone(), ), None => return None, // bad request }; if *shutdown.borrow() { // stop downstream from reusing if this service is shutting down soon session.set_keepalive(None); } let mut ctx = self.inner.new_ctx(); // Deliver user context from the previous request on this reused connection if let Some(prev_ctx) = prev_user_ctx { self.inner .on_connection_reuse(&mut session, &mut ctx, prev_ctx); } self.process_request(session, ctx).await } async fn http_cleanup(&self) { self.shutdown_flag.store(true, Ordering::Release); // Notify all keepalived requests blocking on read_request() to abort self.shutdown.notify_waiters(); } fn server_options(&self) -> Option<&HttpServerOptions> { self.server_options.as_ref() } fn h2_options(&self) -> Option { self.h2_options.clone() } async fn process_custom_session( self: Arc, stream: Stream, shutdown: &ShutdownWatch, ) -> Option { let app = self.clone(); let Some(process_custom_session) = app.process_custom_session.as_ref() else { warn!("custom was called on an empty on_custom"); return None; }; process_custom_session(self.clone(), stream, shutdown).await } // TODO implement h2_options } use pingora_core::services::listening::{RuntimeOptsOverride, Service}; /// Create an [`HttpProxy`] without wrapping it in a [`Service`]. /// /// This is useful when you need to integrate `HttpProxy` into a custom accept loop, /// for example when implementing SNI-based routing that decides between TLS passthrough /// and TLS termination on a single port. /// /// The returned `HttpProxy` is fully initialized and ready to process requests via /// [`HttpServerApp::process_new_http()`]. /// /// # Example /// /// ```ignore /// use pingora_proxy::http_proxy; /// use std::sync::Arc; /// /// // Create the proxy /// let proxy = Arc::new(http_proxy(&server_conf, my_proxy_app)); /// /// // In your custom accept loop: /// loop { /// let (stream, addr) = listener.accept().await?; /// /// // Peek SNI, decide routing... /// if should_terminate_tls { /// let tls_stream = my_acceptor.accept(stream).await?; /// let session = HttpSession::new_http1(Box::new(tls_stream)); /// proxy.process_new_http(session, &shutdown).await; /// } /// } /// ``` pub fn http_proxy(conf: &Arc, inner: SV) -> HttpProxy where SV: ProxyHttp, { let mut proxy = HttpProxy::new(inner, conf.clone()); proxy.handle_init_modules(); proxy } /// Create a [Service] from the user implemented [ProxyHttp]. /// /// The returned [Service] can be hosted by a [pingora_core::server::Server] directly. pub fn http_proxy_service(conf: &Arc, inner: SV) -> Service> where SV: ProxyHttp, { http_proxy_service_with_name(conf, inner, "Pingora HTTP Proxy Service") } /// Create a [Service] from the user implemented [ProxyHttp]. /// /// The returned [Service] can be hosted by a [pingora_core::server::Server] directly. pub fn http_proxy_service_with_name( conf: &Arc, inner: SV, name: &str, ) -> Service> where SV: ProxyHttp, { let mut proxy = HttpProxy::new(inner, conf.clone()); proxy.handle_init_modules(); Service::new(name.to_string(), proxy) } /// Create a [Service] from the user implemented [ProxyHttp]. /// /// The returned [Service] can be hosted by a [pingora_core::server::Server] directly. pub fn http_proxy_service_with_name_custom( conf: &Arc, inner: SV, name: &str, connector: C, on_custom: ProcessCustomSession, ) -> Service> where SV: ProxyHttp + Send + Sync + 'static, SV::CTX: Send + Sync + 'static, C: custom::Connector, { let mut proxy = HttpProxy::new_custom(inner, conf.clone(), connector, Some(on_custom), None, None); proxy.handle_init_modules(); Service::new(name.to_string(), proxy) } /// A builder for a [Service] that can be used to create a [HttpProxy] instance /// /// The [ProxyServiceBuilder] can be used to construct a [HttpProxy] service with a custom name, /// connector, and custom session handler. /// pub struct ProxyServiceBuilder where SV: ProxyHttp + Send + Sync + 'static, SV::CTX: Send + Sync + 'static, C: custom::Connector, { conf: Arc, inner: SV, name: String, connector: C, custom: Option>, server_options: Option, client_options: Option, runtime_opts_override: Option, } impl ProxyServiceBuilder where SV: ProxyHttp + Send + Sync + 'static, SV::CTX: Send + Sync + 'static, { /// Create a new [ProxyServiceBuilder] with the given [ServerConf] and [ProxyHttp] /// implementation. /// /// The returned builder can be used to construct a [HttpProxy] service with a custom name, /// connector, and custom session handler. /// /// The [ProxyServiceBuilder] will default to using the [ProxyHttp] implementation and no custom /// session handler. /// pub fn new(conf: &Arc, inner: SV) -> Self { ProxyServiceBuilder { conf: conf.clone(), inner, name: "Pingora HTTP Proxy Service".into(), connector: (), custom: None, server_options: None, client_options: None, runtime_opts_override: None, } } } impl ProxyServiceBuilder where SV: ProxyHttp + Send + Sync + 'static, SV::CTX: Send + Sync + 'static, C: custom::Connector, { /// Sets the name of the [HttpProxy] service. pub fn name(mut self, name: impl AsRef) -> Self { self.name = name.as_ref().to_owned(); self } /// Set a custom connector and custom session handler for the [ProxyServiceBuilder]. /// /// The custom connector is used to establish a connection to the upstream server. /// /// The custom session handler is used to handle custom protocol specific logic /// between the proxy and the upstream server. /// /// Returns a new [ProxyServiceBuilder] with the custom connector and session handler. pub fn custom( self, connector: C2, on_custom: ProcessCustomSession, ) -> ProxyServiceBuilder { let Self { conf, inner, name, server_options, client_options, runtime_opts_override, .. } = self; ProxyServiceBuilder { conf, inner, name, connector, custom: Some(on_custom), server_options, client_options, runtime_opts_override, } } /// Set the upstream client connector options for the [ProxyServiceBuilder]. /// /// Returns a new [ProxyServiceBuilder] with the upstream client connector options set. pub fn client_options(mut self, options: ConnectorOptions) -> Self { self.client_options = Some(options); self } /// Set the server options for the [ProxyServiceBuilder]. /// /// Returns a new [ProxyServiceBuilder] with the server options set. pub fn server_options(mut self, options: HttpServerOptions) -> Self { self.server_options = Some(options); self } /// Set a runtime options override for the [Service] built by this builder. /// /// Returning [`None`] from the override uses the global runtime options. pub fn runtime_opts_override(mut self, override_fn: F) -> Self where F: Fn(&RuntimeOpts) -> Option + Send + Sync + 'static, { self.runtime_opts_override = Some(Arc::new(override_fn)); self } /// Builds a new [Service] from the [ProxyServiceBuilder]. /// /// This function takes ownership of the [ProxyServiceBuilder] and returns a new [Service] with /// a fully initialized [HttpProxy]. /// /// The returned [Service] is ready to be used by a [pingora_core::server::Server]. pub fn build(self) -> Service> { let Self { conf, inner, name, connector, custom, server_options, client_options, runtime_opts_override, } = self; let mut proxy = HttpProxy::new_custom( inner, conf, connector, custom, server_options, client_options, ); proxy.handle_init_modules(); let mut service = Service::new(name, proxy); if let Some(runtime_opts_override) = runtime_opts_override { service.set_runtime_opts_override(runtime_opts_override); } service } } #[cfg(test)] mod tests { use super::*; use pingora_core::modules::http::{HttpModule, HttpModuleBuilder}; use pingora_core::protocols::l4::stream::Stream as L4Stream; use pingora_core::protocols::l4::virt::{VirtualSockOpt, VirtualSocket, VirtualSocketStream}; use std::pin::Pin; use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::Mutex; use std::task::{Context, Poll}; use tokio::io::{AsyncRead, AsyncWrite, ReadBuf}; #[derive(Debug)] struct StaticVirtualSocket { read_buf: &'static [u8], read_pos: usize, write_buf: Arc>>, } impl StaticVirtualSocket { fn new(read_buf: &'static [u8], write_buf: Arc>>) -> Self { Self { read_buf, read_pos: 0, write_buf, } } } impl AsyncRead for StaticVirtualSocket { fn poll_read( mut self: Pin<&mut Self>, _cx: &mut Context<'_>, buf: &mut ReadBuf<'_>, ) -> Poll> { let remaining = self.read_buf.len() - self.read_pos; let to_read = remaining.min(buf.remaining()); if to_read > 0 { buf.put_slice(&self.read_buf[self.read_pos..self.read_pos + to_read]); self.read_pos += to_read; } Poll::Ready(Ok(())) } } impl AsyncWrite for StaticVirtualSocket { fn poll_write( self: Pin<&mut Self>, _cx: &mut Context<'_>, buf: &[u8], ) -> Poll> { self.write_buf.lock().unwrap().extend_from_slice(buf); Poll::Ready(Ok(buf.len())) } fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { Poll::Ready(Ok(())) } fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { Poll::Ready(Ok(())) } } impl VirtualSocket for StaticVirtualSocket { fn set_socket_option(&self, _opt: VirtualSockOpt) -> std::io::Result<()> { Ok(()) } } async fn new_upgrade_request_session(written: Arc>>) -> Session { let socket = StaticVirtualSocket::new( b"GET / HTTP/1.1\r\nHost: example.com\r\nUpgrade: websocket\r\nConnection: Upgrade\r\n\r\n", written, ); let stream = L4Stream::from(VirtualSocketStream::new(Box::new(socket))); let mut session = Session::new_h1(Box::new(stream)); session.read_request().await.unwrap(); session } fn upgrade_response_header() -> ResponseHeader { let mut header = ResponseHeader::build(http::StatusCode::SWITCHING_PROTOCOLS, Some(2)).unwrap(); header .insert_header(http::header::UPGRADE, "websocket") .unwrap(); header .insert_header(http::header::CONNECTION, "Upgrade") .unwrap(); header } struct SwitchTo101Module; #[async_trait] impl HttpModule for SwitchTo101Module { async fn response_header_filter( &mut self, resp: &mut ResponseHeader, _end_of_stream: bool, ) -> Result<()> { resp.set_status(http::StatusCode::SWITCHING_PROTOCOLS)?; resp.set_version(Version::HTTP_11); Ok(()) } fn as_any(&self) -> &dyn std::any::Any { self } fn as_any_mut(&mut self) -> &mut dyn std::any::Any { self } } struct SwitchTo101ModuleBuilder; impl HttpModuleBuilder for SwitchTo101ModuleBuilder { fn init(&self) -> pingora_core::modules::http::Module { Box::new(SwitchTo101Module) } } struct DoneBytesModule { called: Arc, } #[async_trait] impl HttpModule for DoneBytesModule { fn response_done_filter(&mut self) -> Result> { self.called.store(true, Ordering::Release); Ok(Some(Bytes::from_static(b"hello"))) } fn as_any(&self) -> &dyn std::any::Any { self } fn as_any_mut(&mut self) -> &mut dyn std::any::Any { self } } struct DoneBytesModuleBuilder { called: Arc, } impl HttpModuleBuilder for DoneBytesModuleBuilder { fn init(&self) -> pingora_core::modules::http::Module { Box::new(DoneBytesModule { called: self.called.clone(), }) } } struct DoneEmptyModule { called: Arc, } impl HttpModule for DoneEmptyModule { fn response_done_filter(&mut self) -> Result> { self.called.store(true, Ordering::Release); Ok(None) } fn as_any(&self) -> &dyn std::any::Any { self } fn as_any_mut(&mut self) -> &mut dyn std::any::Any { self } } struct DoneEmptyModuleBuilder { called: Arc, } impl HttpModuleBuilder for DoneEmptyModuleBuilder { fn init(&self) -> pingora_core::modules::http::Module { Box::new(DoneEmptyModule { called: self.called.clone(), }) } } fn assert_raw_upgrade_payload(written: &[u8]) { assert!( written.starts_with(b"HTTP/1.1 101 Switching Protocols\r\n"), "unexpected response: {:?}", String::from_utf8_lossy(written) ); assert!( written.ends_with(b"\r\n\r\nhello"), "upgrade payload should be written as raw tunneled bytes: {:?}", String::from_utf8_lossy(written) ); assert!( !written .windows(b"\r\n5\r\nhello".len()) .any(|w| w == b"\r\n5\r\nhello"), "upgrade payload must not be chunk framed: {:?}", String::from_utf8_lossy(written) ); } #[tokio::test] async fn write_response_tasks_rejects_body_after_101() { let written = Arc::new(Mutex::new(Vec::new())); let mut session = new_upgrade_request_session(written.clone()).await; let err = session .write_response_tasks(vec![ HttpTask::Header(Box::new(upgrade_response_header()), false), HttpTask::Body(Some(Bytes::from_static(b"hello")), true), ]) .await .unwrap_err(); assert_eq!(err.etype(), &InvalidHTTPHeader); assert_eq!(err.esource(), &ErrorSource::Internal); assert!(written.lock().unwrap().is_empty()); } #[tokio::test] async fn write_response_tasks_allows_upgraded_body_after_101() { let written = Arc::new(Mutex::new(Vec::new())); let mut session = new_upgrade_request_session(written.clone()).await; let response_done = session .write_response_tasks(vec![ HttpTask::Header(Box::new(upgrade_response_header()), false), HttpTask::UpgradedBody(Some(Bytes::from_static(b"hello")), true), ]) .await .unwrap(); assert!(response_done); let written = written.lock().unwrap().clone(); assert_raw_upgrade_payload(&written); } #[tokio::test] async fn write_response_tasks_rejects_upgraded_body_before_101() { let written = Arc::new(Mutex::new(Vec::new())); let mut session = new_upgrade_request_session(written.clone()).await; session.set_upstream_h1_upgrade_request_status(true); let err = session .write_response_tasks(vec![HttpTask::UpgradedBody( Some(Bytes::from_static(b"hello")), true, )]) .await .unwrap_err(); assert_eq!(err.etype(), &InvalidHTTPHeader); assert_eq!(err.esource(), &ErrorSource::Internal); assert!(written.lock().unwrap().is_empty()); } #[tokio::test] async fn write_response_tasks_rejects_trailer_after_101() { let written = Arc::new(Mutex::new(Vec::new())); let mut session = new_upgrade_request_session(written.clone()).await; let err = session .write_response_tasks(vec![ HttpTask::Header(Box::new(upgrade_response_header()), false), HttpTask::Trailer(Some(Box::new(http::HeaderMap::new()))), ]) .await .unwrap_err(); assert_eq!(err.etype(), &InvalidHTTPHeader); assert_eq!(err.esource(), &ErrorSource::Internal); assert!(written.lock().unwrap().is_empty()); } #[tokio::test] async fn write_response_tasks_runs_done_filter_after_101_as_upgraded_body() { let written = Arc::new(Mutex::new(Vec::new())); let called = Arc::new(AtomicBool::new(false)); let socket = StaticVirtualSocket::new( b"GET / HTTP/1.1\r\nHost: example.com\r\nUpgrade: websocket\r\nConnection: Upgrade\r\n\r\n", written.clone(), ); let stream = L4Stream::from(VirtualSocketStream::new(Box::new(socket))); let mut modules = HttpModules::new(); modules.add_module(Box::new(DoneBytesModuleBuilder { called: called.clone(), })); let mut session = Session::new_h1_with_modules(Box::new(stream), &modules); session.read_request().await.unwrap(); let response_done = session .write_response_tasks(vec![ HttpTask::Header(Box::new(upgrade_response_header()), false), HttpTask::Done, ]) .await .unwrap(); assert!(response_done); assert!(called.load(Ordering::Acquire)); let written = written.lock().unwrap().clone(); assert_raw_upgrade_payload(&written); } #[tokio::test] async fn write_response_tasks_allows_empty_done_after_101() { let written = Arc::new(Mutex::new(Vec::new())); let called = Arc::new(AtomicBool::new(false)); let socket = StaticVirtualSocket::new( b"GET / HTTP/1.1\r\nHost: example.com\r\nUpgrade: websocket\r\nConnection: Upgrade\r\n\r\n", written.clone(), ); let stream = L4Stream::from(VirtualSocketStream::new(Box::new(socket))); let mut modules = HttpModules::new(); modules.add_module(Box::new(DoneEmptyModuleBuilder { called: called.clone(), })); let mut session = Session::new_h1_with_modules(Box::new(stream), &modules); session.read_request().await.unwrap(); let response_done = session .write_response_tasks(vec![ HttpTask::Header(Box::new(upgrade_response_header()), false), HttpTask::Done, ]) .await .unwrap(); assert!(response_done); assert!(called.load(Ordering::Acquire)); let written = written.lock().unwrap().clone(); assert!( written.starts_with(b"HTTP/1.1 101 Switching Protocols\r\n"), "unexpected response: {:?}", String::from_utf8_lossy(&written) ); assert!( written.ends_with(b"\r\n\r\n"), "empty Done filter should only finish the upgraded response: {:?}", String::from_utf8_lossy(&written) ); } #[tokio::test] async fn write_response_tasks_rejects_module_created_101_with_upgrade_mismatch() { let written = Arc::new(Mutex::new(Vec::new())); let socket = StaticVirtualSocket::new( b"GET / HTTP/1.1\r\nHost: example.com\r\nUpgrade: websocket\r\nConnection: Upgrade\r\n\r\n", written.clone(), ); let stream = L4Stream::from(VirtualSocketStream::new(Box::new(socket))); let mut modules = HttpModules::new(); modules.add_module(Box::new(SwitchTo101ModuleBuilder)); let mut session = Session::new_h1_with_modules(Box::new(stream), &modules); session.read_request().await.unwrap(); session.h1_upgrade_request_status = H1UpgradeRequestStatus { upstream: Some(false), }; let err = session .write_response_tasks(vec![ HttpTask::Header( Box::new(ResponseHeader::build(200, Some(0)).unwrap()), false, ), HttpTask::Body(Some(Bytes::from_static(b"hello")), true), ]) .await .unwrap_err(); assert_eq!(err.etype(), &InvalidHTTPHeader); assert_eq!(err.esource(), &ErrorSource::Internal); assert!(written.lock().unwrap().is_empty()); } #[tokio::test] async fn write_response_tasks_rejects_module_created_101_before_body() { let written = Arc::new(Mutex::new(Vec::new())); let socket = StaticVirtualSocket::new( b"GET / HTTP/1.1\r\nHost: example.com\r\nUpgrade: websocket\r\nConnection: Upgrade\r\n\r\n", written.clone(), ); let stream = L4Stream::from(VirtualSocketStream::new(Box::new(socket))); let mut modules = HttpModules::new(); modules.add_module(Box::new(SwitchTo101ModuleBuilder)); let mut session = Session::new_h1_with_modules(Box::new(stream), &modules); session.read_request().await.unwrap(); let err = session .write_response_tasks(vec![ HttpTask::Header( Box::new(ResponseHeader::build(200, Some(0)).unwrap()), false, ), HttpTask::Body(Some(Bytes::from_static(b"hello")), true), ]) .await .unwrap_err(); assert_eq!(err.etype(), &InvalidHTTPHeader); assert_eq!(err.esource(), &ErrorSource::Internal); assert!(written.lock().unwrap().is_empty()); } #[tokio::test] async fn send_downstream_proxy_task_rejects_body_after_101() { let written = Arc::new(Mutex::new(Vec::new())); let mut session = new_upgrade_request_session(written.clone()).await; session.set_proxy_tasks_enabled(true); session .send_downstream_proxy_task(HttpTask::Header( Box::new(upgrade_response_header()), false, )) .await .unwrap(); let err = session .send_downstream_proxy_task(HttpTask::Body(Some(Bytes::from_static(b"hello")), true)) .await .unwrap_err(); assert_eq!(err.etype(), &InvalidHTTPHeader); assert_eq!(err.esource(), &ErrorSource::Internal); assert!(written.lock().unwrap().is_empty()); } #[tokio::test] async fn send_downstream_proxy_task_allows_upgraded_body_after_101() { let written = Arc::new(Mutex::new(Vec::new())); let mut session = new_upgrade_request_session(written.clone()).await; session.set_proxy_tasks_enabled(true); session .send_downstream_proxy_task(HttpTask::Header( Box::new(upgrade_response_header()), false, )) .await .unwrap(); session .send_downstream_proxy_task(HttpTask::UpgradedBody( Some(Bytes::from_static(b"hello")), true, )) .await .unwrap(); let response_done = session.write_downstream_proxy_tasks().await.unwrap(); assert!(response_done); let written = written.lock().unwrap().clone(); assert_raw_upgrade_payload(&written); } #[tokio::test] async fn send_downstream_proxy_task_rejects_upgraded_body_before_101() { let written = Arc::new(Mutex::new(Vec::new())); let mut session = new_upgrade_request_session(written.clone()).await; session.set_upstream_h1_upgrade_request_status(true); session.set_proxy_tasks_enabled(true); let err = session .send_downstream_proxy_task(HttpTask::UpgradedBody( Some(Bytes::from_static(b"hello")), true, )) .await .unwrap_err(); assert_eq!(err.etype(), &InvalidHTTPHeader); assert_eq!(err.esource(), &ErrorSource::Internal); assert!(!session.has_pending_downstream_tasks()); assert!(written.lock().unwrap().is_empty()); } #[tokio::test] async fn send_downstream_proxy_task_runs_done_filter_after_101_as_upgraded_body() { let written = Arc::new(Mutex::new(Vec::new())); let called = Arc::new(AtomicBool::new(false)); let socket = StaticVirtualSocket::new( b"GET / HTTP/1.1\r\nHost: example.com\r\nUpgrade: websocket\r\nConnection: Upgrade\r\n\r\n", written.clone(), ); let stream = L4Stream::from(VirtualSocketStream::new(Box::new(socket))); let mut modules = HttpModules::new(); modules.add_module(Box::new(DoneBytesModuleBuilder { called: called.clone(), })); let mut session = Session::new_h1_with_modules(Box::new(stream), &modules); session.read_request().await.unwrap(); session.set_proxy_tasks_enabled(true); session .send_downstream_proxy_task(HttpTask::Header( Box::new(upgrade_response_header()), false, )) .await .unwrap(); session .send_downstream_proxy_task(HttpTask::Done) .await .unwrap(); let response_done = session.write_downstream_proxy_tasks().await.unwrap(); assert!(response_done); assert!(called.load(Ordering::Acquire)); let written = written.lock().unwrap().clone(); assert_raw_upgrade_payload(&written); } }