diff --git a/pingora-proxy/src/lib.rs b/pingora-proxy/src/lib.rs index c6b256bc..4f7fd65b 100644 --- a/pingora-proxy/src/lib.rs +++ b/pingora-proxy/src/lib.rs @@ -352,8 +352,8 @@ where "Fail to proxy: {e}, downstream session type: {}", downstream_session.session_type() ); - downstream_session - .respond_error(400) + self.inner + .request_error_filter(&mut downstream_session, &e) .await .unwrap_or_else(|e| { error!("failed to send error response to downstream: {e}"); @@ -1891,6 +1891,7 @@ mod tests { read_buf: Vec, read_pos: usize, write_buf: Arc>>, + shutdown: Arc, } impl StaticVirtualSocket { @@ -1899,6 +1900,7 @@ mod tests { read_buf: read_buf.to_vec(), read_pos: 0, write_buf, + shutdown: Arc::new(AtomicBool::new(false)), } } } @@ -1934,6 +1936,7 @@ mod tests { } fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { + self.shutdown.store(true, Ordering::Relaxed); Poll::Ready(Ok(())) } } @@ -1977,6 +1980,229 @@ mod tests { } } + fn unread_request_session( + request: &[u8], + ) -> (HttpSession, Arc>>, Arc) { + let written = Arc::new(Mutex::new(Vec::new())); + let socket = StaticVirtualSocket::new(request, written.clone()); + let shutdown = socket.shutdown.clone(); + let stream = L4Stream::from(VirtualSocketStream::new(Box::new(socket))); + (HttpSession::new_http1(Box::new(stream)), written, shutdown) + } + + const MALFORMED_REQUEST: &[u8] = b"GET / HTTP/1.1\r\nBad Header: value\r\n\r\n"; + const INVALID_REQUEST: &[u8] = + b"GET / HTTP/1.1\r\nHost: example.com\r\nContent-Length: 1\r\nContent-Length: 2\r\n\r\n"; + const CUSTOM_ERROR_BODY: &[u8] = b"bad request\n"; + + enum RequestErrorAction { + Respond, + Close, + Fail, + RespondThenFail, + } + + struct RequestErrorProxy { + calls: AtomicUsize, + action: RequestErrorAction, + } + + #[async_trait] + impl ProxyHttp for RequestErrorProxy { + type CTX = (); + + fn new_ctx(&self) -> Self::CTX { + panic!("rejected requests must not create a proxy context"); + } + + async fn upstream_peer( + &self, + _session: &mut Session, + _ctx: &mut Self::CTX, + ) -> Result> { + unreachable!("rejected requests must not reach upstream selection"); + } + + async fn request_filter( + &self, + _session: &mut Session, + _ctx: &mut Self::CTX, + ) -> Result { + unreachable!("rejected requests must not reach request filters"); + } + + async fn logging(&self, _session: &mut Session, _e: Option<&Error>, _ctx: &mut Self::CTX) { + unreachable!("rejected requests must not reach normal logging"); + } + + async fn request_error_filter(&self, session: &mut HttpSession, e: &Error) -> Result<()> { + assert_eq!(e.etype(), &InvalidHTTPHeader); + assert_eq!(e.esource(), &ErrorSource::Downstream); + self.calls.fetch_add(1, Ordering::Relaxed); + // These accessors are usable even if parsing failed before a header was created. + let _ = session.client_addr(); + let _ = session.digest(); + if matches!(self.action, RequestErrorAction::Close) { + return Ok(()); + } + if matches!(self.action, RequestErrorAction::Fail) { + return Err(Error::new(WriteError)); + } + let mut response = ResponseHeader::build(422, Some(2))?; + response.set_content_length(CUSTOM_ERROR_BODY.len())?; + response.insert_header(header::CONTENT_TYPE, "text/plain")?; + session + .write_error_response(response, Bytes::from_static(CUSTOM_ERROR_BODY)) + .await?; + if matches!(self.action, RequestErrorAction::RespondThenFail) { + return Err(Error::new(WriteError)); + } + Ok(()) + } + } + + fn request_error_proxy(action: RequestErrorAction) -> Arc> { + Arc::new(HttpProxy::new( + RequestErrorProxy { + calls: AtomicUsize::new(0), + action, + }, + Arc::new(ServerConf::default()), + )) + } + + fn assert_rejection_response( + written: &Mutex>, + status: u16, + body: &[u8], + content_length: usize, + ) { + let written = written.lock().unwrap(); + let mut headers = [httparse::EMPTY_HEADER; 16]; + let mut response = httparse::Response::new(&mut headers); + let httparse::Status::Complete(header_len) = response.parse(&written).unwrap() else { + panic!("incomplete error response"); + }; + assert_eq!(response.code, Some(status)); + let header = |name: &str| { + response + .headers + .iter() + .find(|h| h.name.eq_ignore_ascii_case(name)) + .map(|h| h.value) + }; + assert_eq!(header("connection"), Some(b"close".as_slice())); + assert_eq!( + header("content-length"), + Some(content_length.to_string().as_bytes()) + ); + if status == 422 { + assert_eq!(header("content-type"), Some(b"text/plain".as_slice())); + assert!(header("server").is_none()); + } else { + assert!(header("server").is_some()); + } + // Also rules out an appended fallback response or a response to pipelined input. + assert_eq!(&written[header_len..], body); + } + + #[tokio::test] + async fn request_error_filter_defaults_to_400() { + let proxy = HttpProxy::new(DefaultRetryProxy, Arc::new(ServerConf::default())); + for request in [MALFORMED_REQUEST, INVALID_REQUEST] { + let (session, written, shutdown) = unread_request_session(request); + assert!(proxy.handle_new_request(Box::new(session)).await.is_none()); + assert!(shutdown.load(Ordering::Relaxed)); + assert_rejection_response(&written, 400, b"", 0); + } + } + + #[tokio::test] + async fn request_error_filter_customizes_rejected_requests_before_context_creation() { + for request in [MALFORMED_REQUEST, INVALID_REQUEST] { + let proxy = request_error_proxy(RequestErrorAction::Respond); + let (session, written, closed) = unread_request_session(request); + let (_shutdown_tx, shutdown) = tokio::sync::watch::channel(false); + assert!(proxy.process_new_http(session, &shutdown).await.is_none()); + assert_eq!(proxy.inner.calls.load(Ordering::Relaxed), 1); + assert!(closed.load(Ordering::Relaxed)); + assert_rejection_response(&written, 422, CUSTOM_ERROR_BODY, CUSTOM_ERROR_BODY.len()); + } + } + + #[tokio::test] + async fn request_error_filter_preserves_head_body_suppression() { + let request = b"HEAD / HTTP/1.1\r\nHost: example.com\r\nContent-Length: 1\r\nContent-Length: 2\r\n\r\n"; + let proxy = request_error_proxy(RequestErrorAction::Respond); + let (session, written, shutdown) = unread_request_session(request); + assert!(proxy.handle_new_request(Box::new(session)).await.is_none()); + assert_eq!(proxy.inner.calls.load(Ordering::Relaxed), 1); + assert!(shutdown.load(Ordering::Relaxed)); + assert_rejection_response(&written, 422, b"", CUSTOM_ERROR_BODY.len()); + } + + #[tokio::test] + async fn request_error_filter_closes_without_a_fallback_response() { + for action in [RequestErrorAction::Close, RequestErrorAction::Fail] { + let proxy = request_error_proxy(action); + let (session, written, shutdown) = unread_request_session(MALFORMED_REQUEST); + assert!(proxy.handle_new_request(Box::new(session)).await.is_none()); + assert_eq!(proxy.inner.calls.load(Ordering::Relaxed), 1); + assert!(shutdown.load(Ordering::Relaxed)); + assert!(written.lock().unwrap().is_empty()); + } + } + + #[tokio::test] + async fn request_error_filter_does_not_append_a_response_after_callback_failure() { + let proxy = request_error_proxy(RequestErrorAction::RespondThenFail); + let (session, written, shutdown) = unread_request_session(MALFORMED_REQUEST); + assert!(proxy.handle_new_request(Box::new(session)).await.is_none()); + assert_eq!(proxy.inner.calls.load(Ordering::Relaxed), 1); + assert!(shutdown.load(Ordering::Relaxed)); + assert_rejection_response(&written, 422, CUSTOM_ERROR_BODY, CUSTOM_ERROR_BODY.len()); + } + + #[tokio::test] + async fn request_error_filter_does_not_process_pipelined_input_after_rejection() { + let mut request = MALFORMED_REQUEST.to_vec(); + request.extend_from_slice(b"GET / HTTP/1.1\r\nHost: example.com\r\n\r\n"); + let proxy = request_error_proxy(RequestErrorAction::Respond); + let (session, written, closed) = unread_request_session(&request); + let (_shutdown_tx, shutdown) = tokio::sync::watch::channel(false); + assert!(proxy.process_new_http(session, &shutdown).await.is_none()); + assert_eq!(proxy.inner.calls.load(Ordering::Relaxed), 1); + assert!(closed.load(Ordering::Relaxed)); + assert_rejection_response(&written, 422, CUSTOM_ERROR_BODY, CUSTOM_ERROR_BODY.len()); + } + + #[tokio::test] + async fn request_error_filter_skips_valid_requests_and_connection_errors() { + for request in [ + b"GET / HTTP/1.1\r\nHost: example.com\r\n\r\n".as_slice(), + b"", + b"GET / HTTP/1.1\r\nHost:", + ] { + let proxy = request_error_proxy(RequestErrorAction::Respond); + let (session, written, shutdown) = unread_request_session(request); + let result = proxy.handle_new_request(Box::new(session)).await; + assert_eq!(result.is_some(), request.ends_with(b"\r\n\r\n")); + assert_eq!(proxy.inner.calls.load(Ordering::Relaxed), 0); + assert!(written.lock().unwrap().is_empty()); + if request.ends_with(b"Host:") { + assert!(shutdown.load(Ordering::Relaxed)); + } + } + } + + #[tokio::test] + async fn request_error_filter_skips_shutdown() { + let proxy = request_error_proxy(RequestErrorAction::Respond); + proxy.http_cleanup().await; + assert!(proxy.handle_new_request(pending_session()).await.is_none()); + assert_eq!(proxy.inner.calls.load(Ordering::Relaxed), 0); + } + fn default_policy_would_retry_for_session( session: &mut Session, retry: RetryType, diff --git a/pingora-proxy/src/proxy_trait.rs b/pingora-proxy/src/proxy_trait.rs index d675c33c..3f31ffcd 100644 --- a/pingora-proxy/src/proxy_trait.rs +++ b/pingora-proxy/src/proxy_trait.rs @@ -620,6 +620,43 @@ pub trait ProxyHttp { e } + /// Handle a request header rejected with [`InvalidHTTPHeader`]. + /// + /// This runs before a proxy [`Session`] or [`Self::CTX`] is created. It covers + /// the malformed-header rejection in the HTTP/1 request-reading path, not all + /// protocol errors. The default implementation sends a 400 response. + /// + /// The request header may be unavailable or invalid. Request-dependent methods, + /// including `req_header()` and body-reading methods, may panic and must not be + /// called here. Connection metadata such as `client_addr()` and `digest()` can + /// be used with `e` for logging. Normal request filters and [`Self::logging()`] + /// are not called for this rejected request. + /// + /// Override this callback to write a custom response with + /// [`HttpSession::write_error_response()`], setting `Content-Length` to match + /// the body, or return `Ok(())` without writing to close silently. The connection + /// is closed after this callback returns, including on error. A returned error + /// is logged without attempting a fallback response. + /// + /// For example, a callback can replace the default response headers and body: + /// + /// ``` + /// # use bytes::Bytes; + /// # use pingora_core::protocols::http::ServerSession; + /// # use pingora_error::Result; + /// # use pingora_http::ResponseHeader; + /// # async fn custom_error(session: &mut ServerSession) -> Result<()> { + /// let body = Bytes::from_static(b"Invalid request\n"); + /// let mut response = ResponseHeader::build(400, Some(2))?; + /// response.set_content_length(body.len())?; + /// response.insert_header("Content-Type", "text/plain")?; + /// session.write_error_response(response, body).await + /// # } + /// ``` + async fn request_error_filter(&self, session: &mut HttpSession, _e: &Error) -> Result<()> { + session.respond_error(400).await + } + /// This filter is called when the request encounters a fatal error. /// /// Users may write an error response to the downstream if the downstream is still writable.